-2
\$\begingroup\$

I have two dataframes: Budget and Forecast. For those dataframes, I'm trying to create snapshot record by joining with temp table snapshot_to_collect for loop. I'm joining the two dataframes together for each row in snapshot table and append into deltalake.

The code I have below works and gets me what I need, but it takes forever to process. I'm using databricks spark notebook to get the job done. What can I do to make the code run faster?

spark = SparkSession.builder.appName("optimize_code").getOrCreate()
df = spark.sql("""SELECT * FROM PERIOD.Snapshot""")
schema = StructType([
    StructField("PRJ_DSID", StringType(), True),
    StructField("DSID", StringType(), True),
    StructField("BID", StringType(), True),
    StructField("SNAP_DSID", StringType(), True),
    StructField("SNAP_EFF_PRD", StringType(), True),
    StructField("SNAP_DT", DateType(), True),
    StructField("PRJ_ELMT_DSID", StringType(), True),
    StructField("PRD_DSID", StringType(), True),
    StructField("COST_Raw_Cost", DoubleType(), True),
    StructField("FCST_Raw_Cost", DoubleType(), True),
    StructField("COST_PLN_VER_DSID", StringType(), True),
    StructField("FCST_PLN_VER_DSID", StringType(), True),
    StructField("FCST_PRJ_PLN_TYP_DSID", StringType(), True),
    StructField("COST_PRJ_PLN_TYP_DSID", StringType(), True)
])

# Create an empty DataFrame with the defined schema
combined = spark.createDataFrame([], schema)
# Loop through each row in the DataFrame
for row in df.rdd.toLocalIterator():
    # Convert the Row to a dictionary
    row_dict = row.asDict()
    snapshot_df = spark.createDataFrame([row_dict], schema)
    snapshots_to_collect.createOrReplaceTempView("snapshot_df")

    df_budget = spark.sql(f"""
        SELECT
            FACT_PRJ.PRJ_DSID,
            dim_prj.DSID,
            dim_prj.BID,
            snapshots_to_collect.DSID AS SNAP_DSID,
            snapshots_to_collect.SNAP_DT AS SNAP_EFF_PRD,
            snapshots_to_collect.SNAP_DT AS SNAP_DT,
            FACT_PRJ.PRJ_ELMT_DSID,
            FACT_PRJ.PRD_DSID,
            CASE
            WHEN DATE_TRUNC('MM', TO_DATE(dim_PRD.PRD_BEG_DT, 'MMM-yy')) > DATE_TRUNC('MM', TO_DATE(snapshots_to_collect.SNAP_DT, 'MMM-yy')) THEN NULL
            ELSE CAST(FACT_PRJ.RAW_COST AS numeric(14,2))
             END AS RAW_COST
        FROM Test.FR.FACT_PRJ
        INNER JOIN snapshots_to_collect
            ON 1=1
        INNER JOIN Test.FR.dim_prj_pln_typ
            ON FACT_PRJ.PRJ_PLN_TYP_DSID = dim_prj_pln_typ.DSID
            AND dim_prj_pln_typ.NAME IN ('WM Approved Cost Budget')
        INNER JOIN Test.FR.dim_prd
            ON FACT_PRJ.PRD_DSID = dim_prd.NAME
        INNER JOIN Test.FR.dim_prj_pln_ver
            ON FACT_PRJ.PLN_VER_DSID = dim_prj_pln_ver.DSID
            AND DIM_PRJ_PLN_VER.PLN_STAT_CD = 'B'
            AND DIM_PRJ_PLN_VER.CURR_PLN_STAT_FLG = 'Y'
        LEFT JOIN Test.FR.dim_prj_elmt
            ON FACT_PRJ.PRJ_ELMT_DSID = dim_prj_elmt.DSID
        LEFT JOIN Test.FR.dim_prj
            ON dim_prj.DSID = FACT_PRJ.PRJ_DSID
    """)
    
    df_COSTF = df_budget
    
    df_forecast = spark.sql(f"""
        SELECT
            FACT_PRJ.PRJ_DSID,
            dim_prj.DSID,
            dim_prj.BID,
            snapshots_to_collect.DSID AS SNAP_DSID,
            snapshots_to_collect.SNAP_DT AS SNAP_EFF_PRD,
            snapshots_to_collect.SNAP_DT AS SNAP_DT,
            FACT_PRJ.PRJ_ELMT_DSID,
            FACT_PRJ.PRJ_PLN_TYP_DSID FCST_PRJ_PLN_TYP_DSID,
            FACT_PRJ.PLN_VER_DSID FCST_PLN_VER_DSID,
            FACT_PRJ.PRD_DSID,
           CASE
            WHEN DATE_TRUNC('MM', TO_DATE(dim_PRD.PRD_BEG_DT, 'MMM-yy')) < DATE_TRUNC('MM', TO_DATE(snapshots_to_collect.SNAP_DT, 'MMM-yy')) THEN NULL
            ELSE CAST(FACT_PRJ.RAW_COST AS numeric(14,2))
             END AS FCST_Raw_Cost
        FROM Test.FR.FACT_PRJ
        INNER JOIN snapshots_to_collect
            ON 1=1
        INNER JOIN Test.FR.dim_prd
            ON FACT_PRJ.PRD_DSID = dim_prd.NAME
        LEFT JOIN Test.FR.dim_prj_elmt
            ON FACT_PRJ.PRJ_ELMT_DSID = dim_prj_elmt.DSID
        LEFT JOIN Test.FR.dim_prj
            ON dim_prj.DSID = FACT_PRJ.PRJ_DSID
        
        ORDER BY PRD_DSID DESC
    """)
    
    df_budget = df_budget.filter(df_budget['PRJ_DSID'].isNotNull()).withColumnRenamed("RAW_COST", "COST_RAW_COST")
    df_forecast = df_forecast.filter(df_forecast['PRJ_DSID'].isNotNull()).withColumnRenamed("RAW_COST", "FCST_RAW_COST")
    
    
    spark = SparkSession.builder.appName("concatenate_dfs").getOrCreate()
    
    join_condition = [
        'PRJ_DSID',
        'DSID',
        'BID',
        'SNAP_DSID',
        'PRJ_ELMT_DSID',
        'PRD_DSID',
        'SNAP_EFF_PRD',
        'SNAP_DT'
    ]
    
    # Joining actuals_df with df_CSB
    result_df = result_df.join(df_forecast, join_condition, 'outer')
    result_df = result_df.join(df_forecast, join_condition, 'outer')
    combined = result_df
    
    combined = spark.sql(f"""
    SELECT 
        BDGT_FCST_ACT.SNAP_DSID + '~' + DIM_PRJ.DSID + '~' + DIM_PRJ_ELMT.DSID + '~' + DIM_PRD.DSID AS REF_DSID,
        current_timestamp() AS DW_INS_DTTM,
        'N' AS DEL_FLG,
        COALESCE(BDGT_FCST_ACT.PRJ_DSID, 0) AS PRJ_DSID,
        COALESCE(BDGT_FCST_ACT.DSID, 0) AS DSID,
        COALESCE(BDGT_FCST_ACT.BID, 0) AS BID,
        COALESCE(BDGT_FCST_ACT.SNAP_DSID, 0) AS SNAP_DSID,
        BDGT_FCST_ACT.SNAP_EFF_PRD AS SNAP_EFF_PRD,
        BDGT_FCST_ACT.SNAP_DT AS SNAP_DT,
        COALESCE(BDGT_FCST_ACT.PRJ_ELMT_DSID, 0) AS PRJ_ELMT_DSID,
        COALESCE(BDGT_FCST_ACT.PRD_DSID, 0) AS PRD_DSID,
        COALESCE(BDGT_FCST_ACT.COST_Raw_Cost, 0) AS Budget_Raw_Cost,
        COALESCE(BDGT_FCST_ACT.FCST_Raw_Cost, 0) AS FCST_Raw_Cost
    FROM 
        BDGT_FCST_ACT
    LEFT JOIN 
        Test.FR.DIM_PRJ ON BDGT_FCST_ACT.PRJ_DSID = DIM_PRJ.DSID
    LEFT JOIN 
        Test.FR.DIM_PRJ_ELMT ON BDGT_FCST_ACT.PRJ_ELMT_DSID = DIM_PRJ_ELMT.DSID
    LEFT JOIN 
        Test.FR.DIM_PRD ON BDGT_FCST_ACT.PRD_DSID = DIM_PRD.DSID
    LEFT JOIN 
        Test.FR.DIM_ORG ON DIM_PRJ.EXE_ORG_DSID = DIM_ORG.DSID
    
    
    delta_path = f"abfss://[email protected]/Snapshots/sink"
    try:
        if DeltaTable.isDeltaTable(spark, delta_path):
            delta_table = DeltaTable.forPath(spark, delta_path)
            merge_condition = (
                "source.DSID = target.DSID AND " +
                "source.SNAP_DSID = target.SNAP_DSID AND " +
                "source.PRJ_ELMT_DSID = target.PRJ_ELMT_DSID AND " +
                "source.PRD_DSID = target.PRD_DSID AND " +
                "source.SNAP_DT != target.SNAP_DT"  # Check if SNAP_DT is different
            )
            
            delta_table.alias("target").merge(
                source=combined.alias("source"),
                condition=merge_condition
            ).whenMatchedUpdateAll().whenNotMatchedInsertAll().execute()
            print("Data merged successfully.")
        else:
            # If not, write the DataFrame as a new Delta table
            combined.write.format("delta").mode("overwrite").save(delta_path)
            print("Data written to a new Delta table.")
    except Exception as e:
        print(f"Error during upsert operation: {e}")
\$\endgroup\$
2
  • 4
    \$\begingroup\$ The title of the post should be a very short explanation of what the code does, rather than your concerns about the code. \$\endgroup\$
    – pacmaninbw
    Commented Apr 2 at 16:29
  • 1
    \$\begingroup\$ Is this code correct? I get SyntaxError: unterminated triple-quoted string literal (detected at line 163) trying to run it. \$\endgroup\$
    – ggorlen
    Commented Apr 2 at 22:57

1 Answer 1

0
\$\begingroup\$

This submission is about performance. It includes no CPU profiling, nor performance measurements of any kind, and does not generate example data which someone else could query against. This seems similar to tackling "I lost my keys" with "I'll just put on a blindfold before beginning my search."

spurious ON clause

            INNER JOIN snapshots_to_collect
                ON 1=1

Appending ON TRUE doesn't help the JOIN. Simply elide it.

Or perhaps the table's PK, which you did not disclose to us, suggests a way to avoid forming a giant Cartesian cross product.

motivation for Spark

We are told nothing about the data volumes, host memory sizes, or elapsed times. It's unclear why someone would choose to use Spark over plain SQL for this problem.

Consider going back to basics, building this up from a traditional RDBMS query, until you hit some barrier which motivates Spark. Then write down the technical reasons behind the decision to go with Spark, as comments in the source code. Future maintenance engineers will appreciate being able to review the design rationale and being able to evaluate whether its assumptions still hold.

query plan

We hope that each equi-join ON clause is accompanied by a supporting index. Use EXPLAIN PLAN SELECT ... to verify that, and generally to identify parts of the query that turn out to process many more rows than you would have intuitively expected them to.

incremental approach to 7-way JOIN

JOINing seven tables doesn't have to be slow. But each one you add is an opportunity to break the query by accidentally disabling an index or introducing a Cartesian cross product.

Start by JOINing just a pair of tables, and verify that meets your performance expectations. We might plausibly produce around ten thousand result rows per second.

Now JOIN against a third table, and again verify. Add one table at a time, and stop to investigate performance details if you encounter a step where performance craters.

reporting tables

Consider creating temporary reporting table(s) to cache the output of JOINing the first three or four tables. This lets you decouple processing stages, and reveals their individual contributions to total elapsed time. Writing rows is not very expensive, often on par with time to send result rows over a TCP connection.

Such report tables give you an opportunity to think about what those rows mean and whether you believe they are correct. Impose an appropriate PK. A report table can also be an attractive target for adding a covering index or a UNIQUE constraint. It's a way to steer EXPLAIN PLAN output, if you feel the backend optimizer should have used a different access path to obtain the result rows.

comments lie!

        # Joining actuals_df with df_CSB

Ordinarily this wouldn't be a great comment. It describes the same how that the code does, without explaining the why.

But here, instead of actuals_df the Author apparently intended result_df. Similarly for df_CSB versus df_forecast.

spurious JOIN

        result_df = result_df.join(df_forecast, join_condition, 'outer')
        result_df = result_df.join(df_forecast, join_condition, 'outer')

It's hard to see how that second line is helping anything. Simply elide it.

table aliasing

The identifiers in a Public API have a heavy documentation burden to carry, and will often need to be on the long side in order to be self explanatory. Similarly for database table and column names.

Local variables can be mercifully brief, as they have limited scope and the context is clear when we define and soon after use them.

        LEFT JOIN 
            Test.FR.DIM_PRJ_ELMT ON BDGT_FCST_ACT.PRJ_ELMT_DSID = DIM_PRJ_ELMT.DSID

Consider aliasing some table names in your queries, so this ON clause might end with e.g. ... = dpe.DSID


This codebase achieves only a subset of its design goals.

I would not be willing to delegate or accept maintenance tasks on it.

\$\endgroup\$

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.

Not the answer you're looking for? Browse other questions tagged or ask your own question.