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

@@ -46,7 +46,7 @@ import org.apache.spark.sql.catalyst.InternalRow
 import org.apache.spark.sql.catalyst.analysis.UnresolvedAttribute
 import org.apache.spark.sql.catalyst.catalog.CatalogTable
 import org.apache.spark.sql.catalyst.encoders.{ExpressionEncoder, RowEncoder}
-import org.apache.spark.sql.catalyst.expressions.{Alias, And, Attribute, AttributeReference, BasePredicate, EqualNullSafe, Expression, If, IsNull, Literal, NamedExpression, Not, Or, PredicateHelper, UnsafeProjection}
+import org.apache.spark.sql.catalyst.expressions.{Alias, And, Attribute, AttributeReference, BasePredicate, CreateNamedStruct, EqualNullSafe, Expression, If, IsNull, Literal, NamedExpression, Not, Or, PredicateHelper, UnsafeProjection}
 import org.apache.spark.sql.catalyst.expressions.codegen.GeneratePredicate
 import org.apache.spark.sql.catalyst.plans.logical.{DeltaMergeIntoClause, DeltaMergeIntoMatchedClause, DeltaMergeIntoMatchedDeleteClause, DeltaMergeIntoMatchedUpdateClause, DeltaMergeIntoNotMatchedBySourceClause, DeltaMergeIntoNotMatchedBySourceDeleteClause, DeltaMergeIntoNotMatchedBySourceUpdateClause, DeltaMergeIntoNotMatchedClause, LogicalPlan, Project}
 import org.apache.spark.sql.catalyst.types.DataTypeUtils.toAttributes
@@ -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.{DataType, DataTypes, LongType, NullType, StringType, StructField, StructType}
 
 case class GpuMergeDataSizes(
     @JsonDeserialize(contentAs = classOf[java.lang.Long])
@@ -869,15 +869,20 @@ case class GpuMergeIntoCommand(
     val materializedValues = mutable.ArrayBuffer[NamedExpression]()
     def materializeNonDeterministic(
         exprs: Seq[Expression],
-        takesClause: Expression): Seq[Expression] = exprs.map {
-      case e if !e.deterministic =>
+        takesClause: Expression,
+        targetFields: Seq[StructField] = Nil): Seq[Expression] = exprs.zipWithIndex.map {
+      case (e, i) if !e.deterministic =>
+        // The value is evaluated here, before the clause output is typed, so an untyped NULL
+        // for a schema-evolved nested field has to get its target type now.
         val resolved = resolveOnJoinedPlan(Seq(e)).head
+        val typed = targetFields.lift(i).map(f => typeNullLiterals(resolved, f.dataType))
+          .getOrElse(resolved)
         val existing = joinedPlan.output.map(_.name) ++ materializedValues.map(_.name)
-        val alias = Alias(If(takesClause, resolved, Literal(null, resolved.dataType)),
+        val alias = Alias(If(takesClause, typed, Literal(null, typed.dataType)),
           uniqueColumnName(NON_DETERMINISTIC_VALUE_COL, existing))()
         materializedValues += alias
         alias.toAttribute
-      case e => e
+      case (e, _) => e
     }
 
     // The rows a clause acts on: the row kind, none of the earlier clauses of that kind taken,
@@ -946,7 +951,13 @@ 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 (e, field) => typeNullLiterals(e, field.dataType)
+      } ++ 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
@@ -1017,7 +1028,8 @@ case class GpuMergeIntoCommand(
     def clauseOutput(clause: DeltaMergeIntoClause, routing: Expression): Seq[Seq[Expression]] =
       clause match {
       case u: DeltaMergeIntoMatchedUpdateClause =>
-        updateOutput(materializeNonDeterministic(u.resolvedActions.map(_.expr), routing) ++
+        updateOutput(materializeNonDeterministic(u.resolvedActions.map(_.expr), routing,
+            outputRowSchema.fields.toSeq) ++
             rowTrackingUpdateExprs,
           And(incrUpdatedCountExpr, incrUpdatedMatchedCountExpr))
       case _: DeltaMergeIntoMatchedDeleteClause =>
@@ -1027,7 +1039,8 @@ case class GpuMergeIntoCommand(
             rowTrackingInsertExprs,
           incrInsertedCountExpr)
       case u: DeltaMergeIntoNotMatchedBySourceUpdateClause =>
-        updateOutput(materializeNonDeterministic(u.resolvedActions.map(_.expr), routing) ++
+        updateOutput(materializeNonDeterministic(u.resolvedActions.map(_.expr), routing,
+            outputRowSchema.fields.toSeq) ++
             rowTrackingUpdateExprs,
           And(incrUpdatedCountExpr, incrUpdatedNotMatchedBySourceCountExpr))
       case _: DeltaMergeIntoNotMatchedBySourceDeleteClause =>
@@ -1315,6 +1328,25 @@ case class GpuMergeIntoCommand(
     makeMetricUpdateUDF(name, deterministic).expr
   }
 
+  /**
+   * Gives an untyped NULL literal the type of the target column it is written to. A partial
+   * update of a struct column that schema evolution extended carries the same untyped NULL for
+   * the new nested field, inside the struct constructor, so the struct's fields are typed too.
+   */
+  private def typeNullLiterals(expr: Expression, dataType: DataType): Expression =
+    (expr, dataType) match {
+      case (l: Literal, _) if l.value == null && l.dataType == NullType =>
+        Literal(null, dataType)
+      case (struct: CreateNamedStruct, structType: StructType) =>
+        val typedChildren = struct.children.grouped(2).flatMap {
+          case Seq(name, value) =>
+            val fieldType = structType.fields.find(_.name == name.toString).map(_.dataType)
+            Seq(name, fieldType.map(typeNullLiterals(value, _)).getOrElse(value))
+        }.toSeq
+        CreateNamedStruct(typedChildren)
+      case _ => expr
+    }
+
   private def getTargetOutputCols(txn: OptimisticTransaction): Seq[NamedExpression] = {
     txn.metadata.schema.map { col =>
       targetOutputAttributesMap

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, columns=("a", "b", "c")):
+    # 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(*columns, "_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,159 @@ 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(*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_struct_update_schema_evolution_db173(spark_tmp_path,
+                                                                  spark_tmp_table_factory,
+                                                                  use_cdf):
+    # The INSERT clause adds the nested field s.y to the target. The UPDATE clause sets only
+    # s.x, so the new nested field in its 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, (10, "y1")), (3, (30, "y3"))],
+                                     "a INT, s STRUCT<x: INT, y: STRING>")
+
+    def dest_table_func(spark):
+        return spark.createDataFrame([(1, (1,)), (2, (2,)), (4, (4,))], "a INT, s STRUCT<x: INT>")
+
+    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.s.x = src.s.x "
+            "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, 10, None), (2, 2, None), (3, 30, "y3"), (4, 4, 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", "s.x", "s.y").collect()),
+            conf=delta_merge_enabled_conf)
+        assert expected == actual, f"{run}: expected {expected}, got {actual}"
+    if use_cdf:
+        _assert_cdf_rows_equal(data_path, columns=("a", "s.x", "s.y"))
+
+
+@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_nondeterministic_struct_update_schema_evolution_db173(spark_tmp_path,
+                                                                           spark_tmp_table_factory,
+                                                                           use_cdf):
+    # Same shape as the partial struct update, but the updated nested value is non-deterministic.
+    # That value is evaluated in a projection ahead of the clause output, so the new nested
+    # field s.y has to be typed before that evaluation. rand(7) < 2 is always true, which keeps
+    # the expected values deterministic.
+    data_path = spark_tmp_path + "/DELTA_DATA"
+    src_table = spark_tmp_table_factory.get()
+
+    def src_table_func(spark):
+        return spark.createDataFrame([(1, (10, "y1")), (3, (30, "y3"))],
+                                     "a INT, s STRUCT<x: INT, y: STRING>")
+
+    def dest_table_func(spark):
+        return spark.createDataFrame([(1, (1,)), (2, (2,)), (4, (4,))], "a INT, s STRUCT<x: INT>")
+
+    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.s.x = src.s.x + CAST(rand(7) < 2 AS INT) "
+            "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, 11, None), (2, 2, None), (3, 30, "y3"), (4, 4, 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", "s.x", "s.y").collect()),
+            conf=delta_merge_enabled_conf)
+        assert expected == actual, f"{run}: expected {expected}, got {actual}"
+    if use_cdf:
+        _assert_cdf_rows_equal(data_path, columns=("a", "s.x", "s.y"))
 
 
 @allow_non_gpu("ColumnarToRowExec", *delta_meta_allow)