Add row tracking preservation tests for MERGE, UPDATE and DELETE on OSS Delta [databricks]

integration_tests/src/main/python/delta_lake_delete_test.py

@@ -23,7 +23,7 @@
 import pyarrow.parquet as pq
 from spark_session import is_before_spark_320, is_databricks_runtime, supports_delta_lake_deletion_vectors, \
     with_cpu_session, with_gpu_session, is_before_spark_353, is_spark_353_or_later, \
-    is_databricks173_or_later
+    is_databricks173_or_later, supports_delta_lake_row_tracking
 
 delta_delete_enabled_conf = copy_and_update(delta_writes_enabled_conf,
                                             {"spark.rapids.sql.command.DeleteCommand": "true",
@@ -391,9 +391,9 @@ def generate_dest_data(spark):
 @allow_non_gpu("ColumnarToRowExec", *delta_meta_allow)
 @delta_lake
 @ignore_order
-@pytest.mark.skipif(not is_databricks173_or_later(),
-                    reason="DBR 17.3 row tracking regression coverage")
-def test_delta_delete_preserves_row_tracking_db173(spark_tmp_path):
+@pytest.mark.skipif(not supports_delta_lake_row_tracking(),
+                    reason="Row tracking needs Delta Lake 3.3 or Databricks 17.3")
+def test_delta_delete_preserves_row_tracking(spark_tmp_path):
     conf = copy_and_update(delta_delete_enabled_conf, delta_row_tracking_dml_conf)
     assert_delta_row_tracking_dml(
         spark_tmp_path, "DELETE FROM delta.`{path}` WHERE a IN (2, 3)", conf)

integration_tests/src/main/python/delta_lake_merge_test.py

@@ -23,7 +23,8 @@
 from spark_session import (is_before_spark_320, is_databricks_runtime, spark_version,
                            supports_delta_lake_deletion_vectors, is_before_spark_353,
                            is_spark_400_or_later, is_databricks143,
-                           is_databricks173_or_later, is_spark_41x)
+                           is_databricks173_or_later, is_spark_41x,
+                           supports_delta_lake_row_tracking)
 
 delta_merge_enabled_conf = copy_and_update(delta_writes_enabled_conf,
                                            {"spark.rapids.sql.command.MergeIntoCommand": "true",
@@ -453,6 +454,85 @@ def test_delta_merge_standard_upsert(spark_tmp_path, spark_tmp_table_factory, us
                                         num_slices, num_slices == 1, delta_merge_enabled_conf)
 
 
+@allow_non_gpu("ColumnarToRowExec", *delta_meta_allow)
+@delta_lake
+@ignore_order
+@pytest.mark.skipif(not supports_delta_lake_row_tracking(),
+                    reason="Row tracking needs Delta Lake 3.3 or Databricks 17.3")
+def test_delta_merge_preserves_row_tracking(spark_tmp_path):
+    # A matched update and an insert touch a row-tracked target. (A NOT MATCHED BY SOURCE clause
+    # runs on the GPU only with Delta 4.1 and Databricks 17.3, so it is left out here.) The row
+    # ids of the rows that existed before the merge must survive on both engines, whether the row
+    # is updated or copied, and the commit version moves only for the updated row. Both tables
+    # start from the same single file, so those rows carry the same ids and versions on the CPU
+    # and on the GPU and are compared engine to engine, row tracking columns included. The
+    # inserted row gets a fresh id that the file layout decides, and the join-based GPU merge
+    # lays out files differently from the CPU, so it is checked per engine for freshness only,
+    # and the commit logs are not compared as the UPDATE and DELETE tests do.
+    conf = copy_and_update(delta_merge_enabled_conf, delta_row_tracking_dml_conf)
+    data_path = spark_tmp_path + "/DELTA_DATA"
+    with_cpu_session(lambda spark: setup_delta_row_tracking_dest_tables(
+        spark, data_path, row_tracking_dml_test_df), conf=conf)
+    merge_sql = ("MERGE INTO delta.`{path}` t "
+                 "USING (SELECT * FROM VALUES (2, 'B', 'y'), (9, 'I', 'y') AS s(a, b, c)) s "
+                 "ON t.a = s.a "
+                 "WHEN MATCHED THEN UPDATE SET t.c = s.c "
+                 "WHEN NOT MATCHED THEN INSERT *")
+    tracked_sql = ("SELECT a, b, c, _metadata.row_id AS row_id, "
+                   "_metadata.row_commit_version AS row_commit_version FROM delta.`{}`")
+
+    def tracked_rows(spark, path):
+        return {r["a"]: (r["b"], r["c"], r["row_id"], r["row_commit_version"])
+                for r in spark.sql(tracked_sql.format(path)).collect()}
+
+    before = {run: with_cpu_session(lambda spark, p=data_path + "/" + run: tracked_rows(spark, p),
+                                    conf=conf) for run in ["CPU", "GPU"]}
+
+    def do_merge(spark, path):
+        return spark.sql(merge_sql.format(path=path)).collect()
+
+    # The merge on both engines. The plans captured during the GPU run have to carry the GPU
+    # merge command and the GPU Delta write, so a fallback to the CPU command fails here instead
+    # of passing on identical results. (The command's own DataFrame plan only shows a
+    # CommandResult wrapper, so the check goes through the capture callback, as the REORG tests
+    # do.) The merge's own result row (rows affected, updated, deleted, inserted) is compared
+    # between the engines as well.
+    cpu_result = with_cpu_session(lambda spark: do_merge(spark, data_path + "/CPU"), conf=conf)
+    plan_callback = spark_jvm().org.apache.spark.sql.rapids.ExecutionPlanCaptureCallback
+    plan_callback.startCapture()
+    try:
+        gpu_result = with_gpu_session(lambda spark: do_merge(spark, data_path + "/GPU"), conf=conf)
+        captured_plans = plan_callback.getResultsWithTimeout(10000)
+    finally:
+        plan_callback.endCapture()
+    for class_name in ["GpuExecutedCommandExec"] + delta_write:
+        assert any(plan_callback.contains(plan, class_name) for plan in captured_plans), \
+            "{} is not found in the captured MERGE plans".format(class_name)
+    assert_equal(cpu_result, gpu_result)
+
+    # The rows that existed before, with their row id and commit version, engine to engine.
+    def existing_rows(spark, path):
+        return sorted(tuple(r) for r in
+                      spark.sql(tracked_sql.format(path) + " WHERE a <> 9").collect())
+    assert_equal(with_cpu_session(lambda spark: existing_rows(spark, data_path + "/CPU"), conf=conf),
+                 with_cpu_session(lambda spark: existing_rows(spark, data_path + "/GPU"), conf=conf))
+
+    for run in ["CPU", "GPU"]:
+        path = data_path + "/" + run
+        after = with_cpu_session(lambda spark: tracked_rows(spark, path), conf=conf)
+        assert sorted(after.keys()) == [1, 2, 3, 4, 9], "{}: {}".format(run, after)
+        for a in [1, 2, 3, 4]:
+            assert after[a][2] == before[run][a][2], \
+                "{}: row id of a={} changed: {} -> {}".format(run, a, before[run][a], after[a])
+        for a in [1, 3, 4]:  # copied unchanged
+            assert after[a][3] == before[run][a][3], \
+                "{}: commit version of copied a={} changed: {} -> {}".format(run, a, before[run][a], after[a])
+        assert after[2][3] > before[run][2][3], \
+            "{}: commit version of updated a=2 did not move: {} -> {}".format(run, before[run][2], after[2])
+        assert after[9][2] > max(v[2] for v in before[run].values()), \
+            "{}: inserted row id is not fresh: {}".format(run, after[9])
+
+
 @allow_non_gpu(*delta_meta_allow)
 @delta_lake
 @ignore_order

integration_tests/src/main/python/delta_lake_update_test.py

@@ -20,7 +20,7 @@
 from marks import *
 from spark_session import is_before_spark_320, is_databricks_runtime, \
     supports_delta_lake_deletion_vectors, with_cpu_session, is_before_spark_353, \
-    is_databricks173_or_later
+    is_databricks173_or_later, supports_delta_lake_row_tracking
 
 delta_update_enabled_conf = copy_and_update(delta_writes_enabled_conf,
                                             {"spark.rapids.sql.command.UpdateCommand": "true",
@@ -191,9 +191,9 @@ def generate_dest_data(spark):
 @allow_non_gpu("ColumnarToRowExec", *delta_meta_allow)
 @delta_lake
 @ignore_order
-@pytest.mark.skipif(not is_databricks173_or_later(),
-                    reason="DBR 17.3 row tracking regression coverage")
-def test_delta_update_preserves_row_tracking_db173(spark_tmp_path):
+@pytest.mark.skipif(not supports_delta_lake_row_tracking(),
+                    reason="Row tracking needs Delta Lake 3.3 or Databricks 17.3")
+def test_delta_update_preserves_row_tracking(spark_tmp_path):
     conf = copy_and_update(delta_update_enabled_conf, delta_row_tracking_dml_conf)
     assert_delta_row_tracking_dml(
         spark_tmp_path, "UPDATE delta.`{path}` SET c = b WHERE a IN (2, 3)", conf)

integration_tests/src/main/python/spark_session.py

@@ -354,6 +354,15 @@ def gpu_supports_delta_dv_scan():
     else:
         return is_spark_353_or_later()
 
+def supports_delta_lake_row_tracking():
+    """Whether the current Delta Lake runtime provides row tracking (delta.enableRowTracking and
+    the _metadata.row_id / row_commit_version fields). OSS: Delta Lake 3.3, which the plugin pairs
+    with Spark 3.5.x and later."""
+    if is_databricks_runtime():
+        return is_databricks173_or_later()
+    else:
+        return is_spark_350_or_later()
+
 def is_support_default_values_in_schema():
     # Spark 340 + and Databricks 330 + support
     return is_spark_340_or_later() or is_databricks113_or_later()