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,81 @@ 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")
+@pytest.mark.skipif(is_databricks_runtime(),
+                    reason="The Databricks 17.3 GPU merge regenerates row ids until #15884 lands; that PR carries the 17.3 test")
+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 through the standard write helper. For a Delta test it runs the
+    # GPU side under assert_rapids_delta_write, which asserts the GPU Delta write in the plans it
+    # captures, and only the GPU merge command writes that way, so a fallback to the CPU command
+    # fails here instead of passing on identical results. The helper then reads both tables back
+    # on the CPU, tracked columns included, and compares them engine to engine; the inserted
+    # row's id is masked because the file layout decides it and the join-based GPU merge lays
+    # out files differently from the CPU. The merge's own result row (rows affected, updated,
+    # deleted, inserted) is compared between the engines as well.
+    results = {}
+
+    def write_func(spark, path):
+        results[path] = do_merge(spark, path)
+
+    def read_tracked(spark, path):
+        return spark.sql("SELECT a, b, c, CASE WHEN a = 9 THEN NULL ELSE _metadata.row_id END AS row_id, "
+                         "_metadata.row_commit_version AS row_commit_version FROM delta.`{}`".format(path))
+
+    assert_gpu_and_cpu_writes_are_equal_collect(write_func, read_tracked, data_path, conf=conf)
+    assert_equal(results[data_path + "/CPU"], results[data_path + "/GPU"])
+
+    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()