Fix DBR 18.3 GPU MERGE type mismatch with schema evolution [databricks]

delta-lake/common/src/main/db-400db173-410db183/scala/com/databricks/sql/transaction/tahoe/rapids/GpuMergeIntoCommand.scala

@@ -58,7 +58,7 @@ import org.apache.spark.sql.execution.metric.{SQLMetric, SQLMetrics}
 import org.apache.spark.sql.expressions.Window
 import org.apache.spark.sql.functions._
 import org.apache.spark.sql.nvidia.DFUDFShims
-import org.apache.spark.sql.types.{DataTypes, LongType, StringType, StructField, StructType}
+import org.apache.spark.sql.types.{DataTypes, LongType, NullType, StringType, StructField, StructType}
 
 case class GpuMergeDataSizes(
     @JsonDeserialize(contentAs = classOf[java.lang.Long])
@@ -946,7 +946,15 @@ case class GpuMergeIntoCommand(
         incrMetricExpr: Expression): Seq[Seq[Expression]] = {
       // Generate update expressions and set ROW_DELETED_COL = false and
       // CDC_TYPE_COLUMN_NAME = CDC_TYPE_NOT_CDC
-      val mainDataOutput = updateExprs :+ FalseLiteral :+ incrMetricExpr :+
+      // Schema evolution gives a column that the clause does not set an untyped NULL. The GPU
+      // processor concatenates this output with outputs that carry the column's real type, so
+      // the NULL is typed with the target column's type.
+      val typedUpdateExprs = updateExprs.zip(outputRowSchema.fields).map {
+        case (l: Literal, field) if l.value == null && l.dataType == NullType =>
+          Literal(null, field.dataType)
+        case (e, _) => e
+      } ++ updateExprs.drop(outputRowSchema.fields.length)
+      val mainDataOutput = typedUpdateExprs :+ FalseLiteral :+ incrMetricExpr :+
           CDC_TYPE_NOT_CDC_LITERAL
       val exprs = if (cdcEnabled) {
         // For update preimage, we have do a no-op copy with ROW_DELETED_COL = false and

integration_tests/src/main/python/delta_lake_merge_test.py

@@ -84,6 +84,19 @@ def _assert_gpu_merge_processor(do_merge, data_path, conf, expect_write=True):
         f"{class_name} was not found in the captured MERGE plans"
 
 
+def _assert_cdf_rows_equal(data_path):
+    # The change feed holds the update pre and post images, which the data comparison misses.
+    cdf_rows = {}
+    for run in ["CPU", "GPU"]:
+        cdf_rows[run] = with_cpu_session(
+            lambda spark: sorted((tuple(row) for row in
+                                  read_delta_path_with_cdf(spark, data_path + "/" + run)
+                                  .select("a", "b", "c", "_change_type", "_commit_version")
+                                  .collect()), key=str),
+            conf=delta_merge_enabled_conf)
+    assert_equal(cdf_rows["CPU"], cdf_rows["GPU"])
+
+
 @allow_non_gpu(*delta_meta_allow)
 @delta_lake
 @ignore_order
@@ -1445,7 +1458,8 @@ def read_func(spark, path):
 @ignore_order
 @pytest.mark.skipif(not is_databricks173_or_later(),
                     reason="NOT MATCHED BY SOURCE is supported on the GPU with Databricks 17.3+")
-def test_delta_merge_not_matched_by_source_schema_evolution_db173(spark_tmp_path, spark_tmp_table_factory):
+@pytest.mark.parametrize("use_cdf", [False, True], ids=idfn)
+def test_delta_merge_not_matched_by_source_schema_evolution_db173(spark_tmp_path, spark_tmp_table_factory, use_cdf):
     # Schema evolution adds a source column to the target. Target-only rows go through the
     # NOT MATCHED BY SOURCE clause and get NULL in the new column, like the copied rows do.
     data_path = spark_tmp_path + "/DELTA_DATA"
@@ -1460,7 +1474,7 @@ def dest_table_func(spark):
 
     def setup_tables(spark):
         setup_delta_dest_tables(spark, data_path, dest_table_func,
-                                use_cdf=False, enable_deletion_vectors=False)
+                                use_cdf=use_cdf, enable_deletion_vectors=False)
         src_table_func(spark).createOrReplaceTempView(src_table)
 
     def do_merge(spark, path):
@@ -1484,6 +1498,58 @@ def do_merge(spark, path):
                                  .select("a", "b", "c").collect()),
             conf=delta_merge_enabled_conf)
         assert expected == actual, f"{run}: expected {expected}, got {actual}"
+    if use_cdf:
+        _assert_cdf_rows_equal(data_path)
+
+
+@allow_non_gpu(*delta_meta_allow)
+@delta_lake
+@ignore_order
+@pytest.mark.skipif(not is_databricks173_or_later(),
+                    reason="Schema evolution MERGE processor typing is checked on Databricks 17.3+")
+@pytest.mark.parametrize("use_cdf", [False, True], ids=idfn)
+def test_delta_merge_partial_update_schema_evolution_db173(spark_tmp_path, spark_tmp_table_factory,
+                                                           use_cdf):
+    # The INSERT clause adds a source column to the target. The UPDATE clauses set only an
+    # existing column, so the new column in their output has to be typed like the one in the
+    # insert output for the GPU processor to concatenate them.
+    data_path = spark_tmp_path + "/DELTA_DATA"
+    src_table = spark_tmp_table_factory.get()
+
+    def src_table_func(spark):
+        return spark.createDataFrame([(1, "updated", "new-col"), (3, "inserted", "brand-new")],
+                                     "a INT, b STRING, c STRING")
+
+    def dest_table_func(spark):
+        return spark.createDataFrame([(1, "old"), (2, "keep"), (4, "stale")], "a INT, b STRING")
+
+    def setup_tables(spark):
+        setup_delta_dest_tables(spark, data_path, dest_table_func,
+                                use_cdf=use_cdf, enable_deletion_vectors=False)
+        src_table_func(spark).createOrReplaceTempView(src_table)
+
+    def do_merge(spark, path):
+        src_table_func(spark).createOrReplaceTempView(src_table)
+        return spark.sql(
+            "MERGE WITH SCHEMA EVOLUTION INTO delta.`{path}` AS dest USING {src_table} AS src "
+            "ON dest.a = src.a "
+            "WHEN MATCHED THEN UPDATE SET dest.b = concat(src.b, '-matched') "
+            "WHEN NOT MATCHED THEN INSERT * "
+            .format(path=path, src_table=src_table)).collect()
+
+    with_cpu_session(setup_tables)
+    _assert_gpu_merge_processor(do_merge, data_path, delta_merge_no_cpu_bridge_conf)
+    expected = [(1, "updated-matched", None), (2, "keep", None), (3, "inserted", "brand-new"),
+                (4, "stale", None)]
+    for run in ["CPU", "GPU"]:
+        actual = with_cpu_session(
+            lambda spark: sorted(tuple(row) for row in
+                                 read_delta_path(spark, data_path + "/" + run)
+                                 .select("a", "b", "c").collect()),
+            conf=delta_merge_enabled_conf)
+        assert expected == actual, f"{run}: expected {expected}, got {actual}"
+    if use_cdf:
+        _assert_cdf_rows_equal(data_path)
 
 
 @allow_non_gpu("ColumnarToRowExec", *delta_meta_allow)