Running Spark UDF written in the pipeline code

[haydarai]

Hi all, I’m trying to run a Spark UDF defined in the pipeline code with dagster-pyspark and it can’t run because of ModuleNotFoundError. I think it’s expected because we have to pass the UDF first right? So I see that we have pyspark_step_launcher in dagster-aws and dagster-databricks, but because my team uses a self-hosted one I cannot use these two packages. Do I have to make my own step launcher or there are better alternatives to do this? Many thanks!

Here's how the pipeline looks like

def ymd_to_datetime(ymd) -> datetime:
    try:
        ymd_datetime = datetime.strptime(str(ymd), "%Y%m%d")
    except ValueError:
        ymd_datetime = None
    return ymd_datetime

def ymd_to_ym(ymd: int) -> int:
    ymd_datetime = ymd_to_datetime(ymd)
    ym = int(ymd_datetime.strftime('%Y%m'))
    return ym

class UdfFixer(object):
    def __init__(self, input_function, spark_type):
        self.function, self.spark_type = input_function, spark_type

    def __call__(self, *args):
        return udf(self.function, self.spark_type)(*args)

ymd_to_ym_udf = UdfFixer(ymd_to_ym, IntegerType())


@solid(required_resource_keys={"pyspark", "pyspark_step_launcher"})
def convert_to_ym(context):
    df = context.resources.pyspark.spark_session.sql("select 20200101 as ymd")
    df2 = df.withColumn("ym", ymd_to_ym_udf(col("ymd")))
    return df2

@pipeline(
    mode_defs=[lab_mode], 
    preset_defs=[test_preset]
)
def toy_udf_pipeline():
    convert_to_ym()

@repository
def toy_udf():
    return [toy_udf_pipeline]

[sryza]

Hi @haydarai - are you running your pipeline within the spark cluster? Step launchers are for when you're running your pipeline outside a spark cluster, but need to launch a step within a spark cluster.

[haydarai]

I'm running it not within the cluster. If that's the case, what I should do to make the UDF available to use within the Dagster solid?

[sryza]

In general, distributing with code with Spark can be tricky and depends a lot on the setup. However it basically boils down to two choices:

• Install the libraries on the cluster ahead of time. This can be done with pip or whatever system you use for managing packages on those nodes.
• Package the code into a zip file and ask Spark to distribute it for you.

For the latter, I believe including the following in your configuration should work:

resources:
  pyspark:
    config:
      spark_conf:
        spark.submit.pyFiles: path/to/your/zip/file

If that doesn't work, this might:

@resource
def my_spark_resource(_):
    spark = SparkSession.builder.getOrCreate()
    spark.addPyFile('path/to/your/zip/file')
    return spark

...
mode = ModeDefinition(resource_defs={"pyspark": my_spark_resource})