Assert RapidsDeltaWrite in integration tests [databricks]

integration_tests/src/main/python/asserts.py

@@ -263,7 +263,15 @@ def _assert_gpu_and_cpu_writes_are_equal(
     print('### GPU RUN ###')
     gpu_start = time.time()
     gpu_path = base_path + '/GPU'
-    with_gpu_session(lambda spark : write_func(spark, gpu_path), conf=conf)
+
+    # Check if current test has delta_lake marker
+    from conftest import current_test_has_delta_marker
+    if current_test_has_delta_marker():
+        print("Delta Lake test detected - applying Delta write validation")
+        from delta_lake_utils import assert_rapids_delta_write
+        assert_rapids_delta_write(lambda spark: write_func(spark, gpu_path), conf=conf)
+    else:
+        with_gpu_session(lambda spark : write_func(spark, gpu_path), conf=conf)
     gpu_end = time.time()
     print('### WRITE: GPU TOOK {} CPU TOOK {} ###'.format(
         gpu_end - gpu_start, cpu_end - cpu_start))
@@ -315,7 +323,15 @@ def assert_gpu_and_cpu_save_as_table_are_equal_collect(table_name_factory, write
     print('### GPU RUN ###')
     gpu_start = time.time()
     gpu_table = table_name_factory.get() + '_gpu'
-    with_gpu_session(lambda spark : write_func(spark, gpu_table), conf=conf)
+    # Check if current test has delta_lake marker
+    from conftest import current_test_has_delta_marker
+    if current_test_has_delta_marker():
+        print("✓ Delta Lake test detected - applying Delta write validation")
+        from delta_lake_utils import assert_rapids_delta_write
+        assert_rapids_delta_write(lambda spark : write_func(spark, gpu_table), conf=conf)
+    else:
+        with_gpu_session(lambda spark : write_func(spark, gpu_table), conf=conf)
+
     gpu_end = time.time()
     print('### WRITE: GPU TOOK {} CPU TOOK {} ###'.format(
         gpu_end - gpu_start, cpu_end - cpu_start))

integration_tests/src/main/python/conftest.py

@@ -56,6 +56,7 @@ def array_columns_to_sort_locally():
 _allow_any_non_gpu = False
 _non_gpu_allowed = []
 _per_test_ansi_mode_enabled = None
+_current_test_has_delta_marker = False
 
 def is_allowing_any_non_gpu():
     return _allow_any_non_gpu
@@ -68,6 +69,11 @@ def is_per_test_ansi_mode_enabled():
     return _per_test_ansi_mode_enabled
 
 
+def current_test_has_delta_marker():
+    """Check if the current test has the @delta_lake marker."""
+    return _current_test_has_delta_marker
+
+
 def get_validate_execs_in_gpu_plan():
     return _validate_execs_in_gpu_plan
 
@@ -303,7 +309,10 @@ def pytest_runtest_setup(item):
         elif is_databricks_runtime():
             pytest.skip('Iceberg tests skipped on Databricks')
 
-    if item.get_closest_marker('delta_lake'):
+    global _current_test_has_delta_marker
+    _current_test_has_delta_marker = item.get_closest_marker('delta_lake') is not None
+
+    if _current_test_has_delta_marker:
         if not item.config.getoption('delta_lake'):
             pytest.skip('delta lake tests not configured to run')
 

integration_tests/src/main/python/delta_lake_auto_compact_test.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 
 import pytest
-from asserts import assert_gpu_and_cpu_writes_are_equal_collect, with_cpu_session, with_gpu_session
+from asserts import assert_gpu_and_cpu_writes_are_equal_collect, with_gpu_session
 from data_gen import copy_and_update, idfn
 from delta_lake_utils import *
 from marks import allow_non_gpu, delta_lake

integration_tests/src/main/python/delta_lake_delete_test.py

@@ -14,7 +14,7 @@
 
 import pytest
 
-from asserts import assert_equal, assert_gpu_and_cpu_writes_are_equal_collect, assert_gpu_fallback_write, assert_gpu_and_cpu_are_equal_collect, assert_gpu_fallback_collect
+from asserts import assert_gpu_and_cpu_writes_are_equal_collect, assert_gpu_fallback_write, assert_gpu_and_cpu_are_equal_collect, assert_gpu_fallback_collect
 from data_gen import *
 from delta_lake_utils import *
 from marks import *
@@ -42,18 +42,22 @@ def assert_delta_sql_delete_collect(spark_tmp_path, use_cdf, dest_table_func, de
                                     enable_deletion_vectors,
                                     partition_columns=None,
                                     conf=delta_delete_enabled_conf,
-                                    skip_sql_result_check=False):
+                                    skip_sql_result_check=False, expect_write=True):
     def read_data(spark, path):
         read_func = read_delta_path_with_cdf if use_cdf else read_delta_path
         df = read_func(spark, path)
         return df.sort(df.columns)
+
     def checker(data_path, do_delete):
         cpu_path = data_path + "/CPU"
         gpu_path = data_path + "/GPU"
         if not skip_sql_result_check:
             # compare resulting dataframe from the delete operation (some older Spark versions return empty here)
             cpu_result = with_cpu_session(lambda spark: do_delete(spark, cpu_path).collect(), conf=conf)
-            gpu_result = with_gpu_session(lambda spark: do_delete(spark, gpu_path).collect(), conf=conf)
+            if expect_write:
+                gpu_result = assert_rapids_delta_write(lambda spark: do_delete(spark, gpu_path).collect(), conf=conf)
+            else:
+                gpu_result = with_gpu_session(lambda spark: do_delete(spark, gpu_path).collect(), conf=conf)
             assert_equal(cpu_result, gpu_result)
         # compare table data results, read both via CPU to make sure GPU write can be read by CPU
         cpu_result = with_cpu_session(lambda spark: read_data(spark, cpu_path).collect(), conf=conf)
@@ -296,7 +300,7 @@ def generate_dest_data(spark):
     skip_sql_result = is_databricks_runtime()
     assert_delta_sql_delete_collect(spark_tmp_path, use_cdf, generate_dest_data,
                                     delete_sql, enable_deletion_vectors, partition_columns,
-                                    skip_sql_result_check=skip_sql_result)
+                                    skip_sql_result_check=skip_sql_result, expect_write=False)
 
 @allow_non_gpu(*delta_meta_allow)
 @delta_lake
@@ -318,7 +322,7 @@ def generate_dest_data(spark):
     skip_sql_result = is_databricks_runtime()
     assert_delta_sql_delete_collect(spark_tmp_path, use_cdf, generate_dest_data,
                                     delete_sql, enable_deletion_vectors, partition_columns,
-                                    skip_sql_result_check=skip_sql_result)
+                                    skip_sql_result_check=skip_sql_result, expect_write=False)
 
 @allow_non_gpu(*delta_meta_allow)
 @delta_lake

integration_tests/src/main/python/delta_lake_merge_common.py

@@ -47,11 +47,21 @@ def do_merge(spark, path):
     with_cpu_session(setup_tables)
     check_func(data_path, do_merge)
 
-def assert_collect(do_merge, read_delta_path, data_path, conf):
+def assert_collect(do_merge, data_path, conf, expect_write=True):
+    """
+    Execute the do_merge function in both CPU and GPU sessions and compare the results.
+    :param do_merge: A function that takes a Spark session and a path, performs a merge operation, and returns the result.
+    :param data_path: The base path where CPU and GPU data are stored.
+    :param conf: Configuration settings for the Spark sessions.
+    :param expect_write: A boolean indicating whether to expect a write operation in the GPU session.
+    """
     cpu_path = data_path + "/CPU"
     gpu_path = data_path + "/GPU"
     cpu_result = with_cpu_session(lambda spark: do_merge(spark, cpu_path), conf=conf)
-    gpu_result = with_gpu_session(lambda spark: do_merge(spark, gpu_path), conf=conf)
+    if expect_write:
+        gpu_result = assert_rapids_delta_write(lambda spark: do_merge(spark, gpu_path), conf=conf)
+    else:
+        gpu_result = with_gpu_session(lambda spark: do_merge(spark, gpu_path), conf=conf)
     assert_equal(cpu_result, gpu_result)
 
 # This method is used for making sure ExecutedCommand fallsback for Spark 3.5.3
@@ -64,7 +74,8 @@ def do_assert(do_merge, read_delta_path, data_path, conf):
 
 def assert_delta_sql_merge_collect(spark_tmp_path, spark_tmp_table_factory, use_cdf, enable_deletion_vectors,
                                    src_table_func, dest_table_func, merge_sql,
-                                   compare_logs, assert_func=assert_collect, partition_columns=None, conf=None):
+                                   compare_logs, assert_func=assert_collect, partition_columns=None, conf=None,
+                                   expect_write=True):
     assert conf is not None, "conf must be set"
 
     def read_data(spark, path):
@@ -76,7 +87,7 @@ def checker(data_path, do_merge):
         cpu_path = data_path + "/CPU"
         gpu_path = data_path + "/GPU"
         # compare resulting dataframe from the merge operation (some older Spark versions return empty here)
-        assert_func(do_merge, read_delta_path, data_path, conf)
+        assert_func(do_merge, data_path, conf, expect_write=expect_write)
         # compare merged table data results, read both via CPU to make sure GPU write can be read by CPU
         cpu_result = with_cpu_session(lambda spark: read_data(spark, cpu_path).collect(), conf=conf)
         gpu_result = with_cpu_session(lambda spark: read_data(spark, gpu_path).collect(), conf=conf)
@@ -104,15 +115,15 @@ def do_test_delta_merge_not_match_insert_only(spark_tmp_path, spark_tmp_table_fa
 
 def do_test_delta_merge_match_delete_only(spark_tmp_path, spark_tmp_table_factory, table_ranges,
                                           use_cdf, enable_deletion_vectors, partition_columns, num_slices, compare_logs,
-                                          conf, assert_func=assert_collect):
+                                          conf, assert_func=assert_collect, expect_write=True):
     src_range, dest_range = table_ranges
     src_table_func = lambda spark: make_df(spark, SetValuesGen(IntegerType(), src_range), num_slices)
     dest_table_func = lambda spark: make_df(spark, SetValuesGen(IntegerType(), dest_range), num_slices)
     merge_sql = "MERGE INTO {dest_table} USING {src_table} ON {dest_table}.a == {src_table}.a" \
                 " WHEN MATCHED THEN DELETE"
     assert_delta_sql_merge_collect(spark_tmp_path, spark_tmp_table_factory, use_cdf, enable_deletion_vectors,
                                    src_table_func, dest_table_func, merge_sql, compare_logs, assert_func,
-                                   partition_columns, conf=conf)
+                                   partition_columns, conf=conf, expect_write=expect_write)
 
 
 def do_test_delta_merge_standard_upsert(spark_tmp_path, spark_tmp_table_factory, use_cdf, enable_deletion_vectors,
@@ -139,15 +150,16 @@ def do_test_delta_merge_upsert_with_condition(spark_tmp_path, spark_tmp_table_fa
 
 def do_test_delta_merge_upsert_with_unmatchable_match_condition(spark_tmp_path,
                                                                 spark_tmp_table_factory, use_cdf, enable_deletion_vectors,
-                                                                num_slices, compare_logs, conf, assert_func=assert_collect):
+                                                                num_slices, compare_logs, conf, assert_func=assert_collect,
+                                                                expect_write=True):
     # Need to eliminate duplicate keys in the source table otherwise update semantics are ambiguous
     src_table_func = lambda spark: two_col_df(spark, int_gen, string_gen, num_slices=num_slices).groupBy("a").agg(f.max("b").alias("b"))
     dest_table_func = lambda spark: two_col_df(spark, SetValuesGen(IntegerType(), range(100)), string_gen, seed=1, num_slices=num_slices)
     merge_sql = "MERGE INTO {dest_table} USING {src_table} ON {dest_table}.a == {src_table}.a" \
                 " WHEN MATCHED AND {dest_table}.a > 100 THEN UPDATE SET *"
     assert_delta_sql_merge_collect(spark_tmp_path, spark_tmp_table_factory, use_cdf, enable_deletion_vectors,
                                    src_table_func, dest_table_func, merge_sql, compare_logs, assert_func,
-                                   conf=conf)
+                                   conf=conf, expect_write=expect_write)
 
 
 def do_test_delta_merge_update_with_aggregation(spark_tmp_path, spark_tmp_table_factory, use_cdf, enable_deletion_vectors,

integration_tests/src/main/python/delta_lake_merge_test.py

@@ -172,7 +172,7 @@ def checker(data_path, do_merge):
         cpu_path = data_path + "/CPU"
         gpu_path = data_path + "/GPU"
         # compare resulting dataframe from the merge operation (some older Spark versions return empty here)
-        assert_collect(do_merge, read_delta_path, data_path, materialize_conf)
+        assert_collect(do_merge, data_path, materialize_conf)
         # compare merged table data results, read both via CPU to make sure GPU write can be read by CPU
         cpu_result = with_cpu_session(lambda spark: read_data(spark, cpu_path).rdd.isCheckpointed(), conf=materialize_conf)
         gpu_result = with_cpu_session(lambda spark: read_data(spark, gpu_path).rdd.isCheckpointed(), conf=materialize_conf)
@@ -193,20 +193,22 @@ def checker(data_path, do_merge):
 @delta_lake
 @ignore_order
 @pytest.mark.skipif(is_before_spark_320(), reason="Delta Lake writes are not supported before Spark 3.2.x")
-@pytest.mark.parametrize("table_ranges", [(range(10), range(20)),  # partial delete of target
-                                          (range(5), range(5)),  # full delete of target
-                                          (range(10), range(20, 30))  # no-op delete
-                                          ], ids=idfn)
+@pytest.mark.parametrize("table_ranges_expect_write", [
+    ((range(10), range(20)), True),  # partial delete of target
+    ((range(5), range(5)), True),  # full delete of target
+    ((range(10), range(20, 30)), False)  # no-op delete. gpu write is not expected
+], ids=idfn)
 @pytest.mark.parametrize("use_cdf", [True, False], ids=idfn)
 @pytest.mark.parametrize("partition_columns", [None, ["a"], ["b"], ["a", "b"]], ids=idfn)
 @pytest.mark.parametrize("num_slices", num_slices_to_test, ids=idfn)
 @pytest.mark.parametrize("enable_deletion_vector", deletion_vector_values_with_350DB143_xfail_reasons(
                             enabled_xfail_reason='https://github.com/NVIDIA/spark-rapids/issues/12042'), ids=idfn)
-def test_delta_merge_match_delete_only(spark_tmp_path, spark_tmp_table_factory, table_ranges,
+def test_delta_merge_match_delete_only(spark_tmp_path, spark_tmp_table_factory, table_ranges_expect_write,
                                        use_cdf, partition_columns, num_slices, enable_deletion_vector):
+    table_ranges, expect_write = table_ranges_expect_write
     do_test_delta_merge_match_delete_only(spark_tmp_path, spark_tmp_table_factory, table_ranges,
                                           use_cdf, enable_deletion_vector, partition_columns, num_slices,
-                                          num_slices == 1, delta_merge_enabled_conf)
+                                          num_slices == 1, delta_merge_enabled_conf, expect_write=expect_write)
 
 @allow_non_gpu(*delta_meta_allow)
 @delta_lake
@@ -259,7 +261,8 @@ def test_delta_merge_upsert_with_unmatchable_match_condition(spark_tmp_path, spa
     do_test_delta_merge_upsert_with_unmatchable_match_condition(spark_tmp_path,
                                                                 spark_tmp_table_factory, use_cdf, enable_deletion_vector,
                                                                 num_slices, num_slices == 1,
-                                                                delta_merge_enabled_conf)
+                                                                delta_merge_enabled_conf,
+                                                                expect_write=False)
 
 @allow_non_gpu(*delta_meta_allow)
 @delta_lake

integration_tests/src/main/python/delta_lake_update_test.py

@@ -14,12 +14,12 @@
 
 import pytest
 
-from asserts import assert_equal, assert_gpu_and_cpu_writes_are_equal_collect, assert_gpu_fallback_write
+from asserts import assert_gpu_and_cpu_writes_are_equal_collect, assert_gpu_fallback_write
 from data_gen import *
 from delta_lake_utils import *
 from marks import *
 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
+    supports_delta_lake_deletion_vectors, with_cpu_session, is_before_spark_353
 
 delta_update_enabled_conf = copy_and_update(delta_writes_enabled_conf,
                                             {"spark.rapids.sql.command.UpdateCommand": "true",
@@ -48,7 +48,7 @@ def checker(data_path, do_update):
         gpu_path = data_path + "/GPU"
         # compare resulting dataframe from the update operation (some older Spark versions return empty here)
         cpu_result = with_cpu_session(lambda spark: do_update(spark, cpu_path).collect(), conf=conf)
-        gpu_result = with_gpu_session(lambda spark: do_update(spark, gpu_path).collect(), conf=conf)
+        gpu_result = assert_rapids_delta_write(lambda spark: do_update(spark, gpu_path).collect(), conf=conf)
         assert_equal(cpu_result, gpu_result)
         # compare table data results, read both via CPU to make sure GPU write can be read by CPU
         cpu_result = with_cpu_session(lambda spark: read_data(spark, cpu_path).collect(), conf=conf)

integration_tests/src/main/python/delta_lake_utils.py

@@ -18,8 +18,9 @@
 import re
 
 from spark_session import is_databricks122_or_later, supports_delta_lake_deletion_vectors, is_databricks143_or_later, \
-    with_cpu_session
+    with_cpu_session, with_gpu_session
 from asserts import assert_equal
+from conftest import spark_jvm
 
 delta_meta_allow = [
     "DeserializeToObjectExec",
@@ -34,6 +35,8 @@
     "SortExec"
 ]
 
+delta_write = ["RapidsDeltaWrite"]
+
 # Disable Deletion Vectors except for Databricks 14.3
 def deletion_vector_values_with_350DB143_xfail_reasons(enabled_xfail_reason=None, disabled_xfail_reason=None):
     # We will always set the deletion vectors to False
@@ -222,3 +225,41 @@ def setup_delta_dest_tables(spark, data_path, dest_table_func, use_cdf, enable_d
     for name in ["CPU", "GPU"]:
         path = "{}/{}".format(data_path, name)
         setup_delta_dest_table(spark, path, dest_table_func, use_cdf, partition_columns, enable_deletion_vectors)
+
+def assert_rapids_delta_write(do_test, conf):
+    """
+    Validates that a Delta write operation executed on the GPU produces the expected execution plans.
+    This function starts a plan capture mechanism using the Spark JVM's ExecutionPlanCaptureCallback,
+    runs the provided test function (`do_test`) within a GPU session, and collects the execution plans
+    generated during the write operation. It then checks that each expected Delta write class is present
+    in at least one captured plan.
+
+    Parameters
+    ----------
+    do_test : callable
+        A function that performs the Delta write operation to be validated.
+    conf : dict
+        A dictionary of configuration options to be passed to the GPU session.
+
+    Returns
+    -------
+    result : Any
+        The result returned by the `do_test` function.
+    """
+    jvm = spark_jvm()
+    jvm.org.apache.spark.sql.rapids.ExecutionPlanCaptureCallback.startCapture()
+    try:
+        result = with_gpu_session(do_test, conf=conf)
+        captured_plans = jvm.org.apache.spark.sql.rapids.ExecutionPlanCaptureCallback.getResultsWithTimeout(10000)
+        # Some write functions are no-op. We may not capture any GPU plan.
+        if len(captured_plans) > 0:
+            for cls in delta_write:
+                found = False
+                for plan in captured_plans:
+                    found = jvm.org.apache.spark.sql.rapids.ExecutionPlanCaptureCallback.contains(plan, cls)
+                    if found:
+                        break
+                assert found, f"{cls} is not found in any captured plan"
+        return result
+    finally:
+        jvm.org.apache.spark.sql.rapids.ExecutionPlanCaptureCallback.endCapture()