Generic sql_solid for executing SQL against Delta Table with Spark

[sspaeti]

I successfully integrated a solid that can MERGE against a Delta Table passing a Spark DataFrame as input (thanks to the Airline Demo example).

This looks something like this (code below):

As I'm seeking feedback and or points where things were done wrong, I also wanted to share some code around Delta Lake.

Some specific questions that came up while developing:

• I wasn't sure if I should define target_delta_table and input_dataframe as fixed parameters in sql_solid. I decided not so I can dynamically ingest these inputs during the pipeline and not hardcoded in the merge_example-solid. Should I change that to avoid errors during runtime?
• As now it's dynamic but still, I want to check if these parameters provided correctly as otherwise, it will raise an error. I tried with if input_defs['target_delta_table'] is None although this is wrong as this itself raises the error if input_def is not defined. I wanted to check the input_defs already inside the sql_solid, but I wasn't able to catch these two input parameters as this is not a real solid.
• I also want to keep it dynamic that I can use for INSERT, UPDATE and maybe SELECT as well. But perhaps it's better to create separate solids for that. Or what would be your recommendation?
• Any other comments or recommendations?

Dagster Code

Types

DeltaCoordinate = dict_with_fields(
    name="DeltaCoordinate",
    fields={
        "database": Field(String, description="database or schema or delta table"),
        "table_name": Field(String, description="table name of the delta table"),
        "s3_coordinate_bucket": Field(String, description="s3 bucket"),
        "s3_coordinate_key": Field(String, description="s3 delta table path"),
    },
)

SqlTableName = create_string_type("SqlTableName", description="The name of a database table")
DeltaTableName = create_string_type(
    "DeltaTableName",
    description="The name of a Delta table name that is available in `SparkContext`",
)

sql_solid function which returns the solid

def sql_solid(
    name,
    sql_statement,
    materialization_strategy,
    # target_delta_table: DeltaCoordinate,
    # src_df: DataFrame,
    table_name=None,
    input_defs=None,
    # input_defs=[
    #     InputDefinition("target_delta_table", DeltaCoordinate),
    # ],
):
    """Return a new solid that executes and materializes a SQL select statement.

    Args:
        name (str): The name of the new solid.
        sql_statement (str): The sql statement to execute which can be MERGE, INSERT, UPDATE.
        materialization_strategy (str): Must be 'delta_table' for now.
    Kwargs:
        input_defs (list[InputDefinition]): 'target_delta_table' (DeltaCoordinate) must be provided, which is
            the table where the sql_statement is running against.
            'input_dataframe' (DataFrame) must be provided for providing input data for the sql_statement

    Returns:
        function:
            The new SQL solid.
    """
    input_defs = check.opt_list_param(input_defs, "input_defs", InputDefinition)

    materialization_strategy_output_types = {  # pylint:disable=C0103
        "table": SqlTableName,
        "delta_table": DeltaCoordinate,
        # 'view': String,
        # 'query': SqlAlchemyQueryType,
        # 'subquery': SqlAlchemySubqueryType,
        # 'result_proxy': SqlAlchemyResultProxyType,
        # could also materialize as a Pandas table, as a Spark table, as an intermediate file, etc.
    }

    if materialization_strategy not in materialization_strategy_output_types:
        raise Exception(
            "Invalid materialization strategy {materialization_strategy}, must "
            "be one of {materialization_strategies}".format(
                materialization_strategy=materialization_strategy,
                materialization_strategies=str(list(materialization_strategy_output_types.keys())),
            )
        )

    output_description = (
        "The string name of the new table created by the solid"
        if materialization_strategy == "table" or materialization_strategy == "delta_table"
        else "The materialized SQL statement. If the materialization_strategy is "
        "'table', this is the string name of the new table created by the solid."
    )

    # sql_statement.replace(
    #     "delta.``", "delta.{target_delta_path}".format(target_delta_path=target_delta_path)
    # )

    description = """This solid executes the following SQL statement:
    {sql_statement}""".format(
        sql_statement=sql_statement
    )

    sql_statement = ("{sql_statement};").format(sql_statement=sql_statement)

    # start_solids_marker_1
    @solid(
        name=name,
        input_defs=input_defs,
        output_defs=[
            OutputDefinition(
                materialization_strategy_output_types[materialization_strategy],
                description=output_description,
            )
        ],
        description=description,
        required_resource_keys={"spark"},
        tags={"kind": "sql", "sql": sql_statement},
    )
    def _sql_solid(context, **input_defs):  # pylint: disable=unused-argument
        """Inner function defining the new solid.

        Args:
            context (SolidExecutionContext): Must expose a `spark` resource with an `spark_session` method
                wich can execute SQL against Delta Lake Tables.

        Returns:
            DeltaCoordinate:
                The Delta Table Coordinates where the SQL statements were running against.
        """
        if input_defs['target_delta_table'] is None:
            raise Exception("Input `target_delta_table` not provided.")
        if input_defs['input_dataframe'] is None:
            raise Exception("Input `input_dataframe` not provided.")
        ##
        ## Handling delta-table
        ##
        target_delta_path = _get_s3a_path(
            input_defs['target_delta_table']["s3_coordinate_bucket"],
            input_defs['target_delta_table']["s3_coordinate_key"],
        )
        context.log.info("Target Delta table path: {}".format(target_delta_path))

        # Set Delta-table path
        repl_sql_statement = sql_statement.replace(
            "target_delta_table", "delta.`{}`".format(target_delta_path)
        )

        context.log.info(
            "Executing sql statement:\n{sql_statement}".format(sql_statement=repl_sql_statement)
        )

        ##
        ## Handling spark dataframe
        ##

        # register input df for spark to be available in spark.sql
        input_defs['input_dataframe'].createOrReplaceTempView("input_dataframe")

        context.resources.spark.spark_session.sql(
            repl_sql_statement
        )  # text() function removed here (this would validate string as valid SQL, but with Delta-Merge does not work)
        yield AssetMaterialization(
            asset_key=input_defs['target_delta_table']['table_name'],
            description="SmartAnalytics tags per session_id",
            metadata_entries=[EventMetadataEntry.path(target_delta_path, "delta_table_path")],
        )

        yield Output(value=input_defs['target_delta_table'], output_name="result")

    # end_solids_marker_1
    return _sql_solid

My merge solid

merge_example = sql_solid(
    name="merge_example",
    sql_statement="""
    MERGE INTO target_delta_table trg
    USING input_dataframe AS src
    ON trg.session_session_id = src.session_session_id
    WHEN MATCHED THEN
        UPDATE SET
            trg.tag_key  = src.tag_key,
            trg.tag_value = src.tag_value
    WHEN NOT MATCHED THEN
        INSERT (session_session_id, tag_key, tag_value)
        VALUES (src.session_session_id, src.tag_key, src.tag_value)
    """,
    materialization_strategy="delta_table",
    # table_name="tag",
    input_defs=[
        InputDefinition("target_delta_table", DeltaCoordinate),
        InputDefinition("input_dataframe", DataFrame),
    ],
)