Port SPARK-48949 scan-side partition filter to GpuBatchScanExec [fast-ut][databricks][reduced-it]

integration_tests/src/main/python/iceberg/iceberg_test.py

@@ -172,6 +172,88 @@ def distinct_after_spj(spark):
         gpu_plan_assertion=_assert_partial_clustering_spj_plan)
 
 
+@iceberg
+@ignore_order(local=True)
+@pytest.mark.skipif(
+    not is_spark_400_or_later(),
+    reason="spark.sql.sources.v2.bucketing.partition.filter.enabled was added in Spark 4.0.0")
+@pytest.mark.parametrize("partition_filter", [True, False], ids=["filtered", "unfiltered"])
+@pytest.mark.parametrize("partially_clustered", [True, False],
+                         ids=["partially_clustered", "clustered"])
+def test_iceberg_spj_partition_filter(spark_tmp_table_factory, partition_filter,
+                                      partially_clustered):
+    left_table = get_full_table_name(spark_tmp_table_factory)
+    right_table = get_full_table_name(spark_tmp_table_factory)
+    table_props = _build_tblprops({
+        # Keep separate INSERTs as separate scan splits so that id=1 is partially clustered.
+        "read.split.target-size": "1",
+        "read.split.open-file-cost": "1",
+    })
+    table_props_sql = ", ".join(f"'{k}' = '{v}'" for k, v in table_props.items())
+
+    def setup_iceberg_tables(spark):
+        spark.sql(
+            f"CREATE TABLE {left_table} (id INT, price DOUBLE) USING ICEBERG "
+            f"PARTITIONED BY (id) TBLPROPERTIES ({table_props_sql})")
+        spark.sql(
+            f"CREATE TABLE {right_table} (id INT, value STRING) USING ICEBERG "
+            f"PARTITIONED BY (id) TBLPROPERTIES ({table_props_sql})")
+
+        # Each side gets a key the other lacks (left id=3, right id=4), so partition filtering
+        # has something to prune on both scans. Spark decides which side pads and which
+        # replicates, so covering both branches of that decision means neither side may be a
+        # subset of the intersection. Each side also gets a key split across two files, inside
+        # the intersection, so whichever side pads reports numSplits=2 and partial clustering
+        # stays observable in the partition counts.
+        spark.sql(f"INSERT INTO {left_table} VALUES (1, 40.0), (2, 10.0), (3, 15.5)")
+        spark.sql(f"INSERT INTO {left_table} VALUES (1, 41.0)")
+        spark.sql(f"INSERT INTO {right_table} VALUES (1, 'a'), (2, 'b'), (4, 'd')")
+        spark.sql(f"INSERT INTO {right_table} VALUES (2, 'c')")
+
+    with_cpu_session(setup_iceberg_tables)
+
+    conf = {
+        "spark.sql.adaptive.enabled": "false",
+        "spark.sql.autoBroadcastJoinThreshold": "-1",
+        "spark.sql.sources.v2.bucketing.enabled": "true",
+        "spark.sql.sources.v2.bucketing.pushPartValues.enabled": "true",
+        "spark.sql.sources.v2.bucketing.partition.filter.enabled":
+            str(partition_filter).lower(),
+        "spark.sql.sources.v2.bucketing.partiallyClusteredDistribution.enabled":
+            str(partially_clustered).lower(),
+        "spark.sql.iceberg.planning.preserve-data-grouping": "true",
+    }
+
+    def join_after_spj(spark):
+        return spark.sql(
+            f"""
+            SELECT l.id, l.price, r.value
+            FROM {left_table} l
+            JOIN {right_table} r ON l.id = r.id
+            """)
+
+    # Both scans plan one partition per common partition value: the intersection {1, 2} when
+    # filtering, otherwise the union {1, 2, 3, 4}. Partial clustering adds one more, for the
+    # second file of the padding side's split key. KeyGroupedPartitioning.isPartiallyClustered
+    # would say this more directly but only exists on Spark 3.5.9+/4.0.3+/4.1.2+, which is the
+    # gate this test is deliberately avoiding.
+    expected_partitions = (2 if partition_filter else 4) + (1 if partially_clustered else 0)
+
+    def assert_plan(plan):
+        scans = _assert_spj_join_shape(plan, expect_spj=True)
+        counts = [scan.outputPartitioning().numPartitions() for scan in scans]
+        assert counts == [expected_partitions] * len(scans), \
+            f"Expected {expected_partitions} partitions per scan, found {counts}:\n{plan}"
+
+    # Asserting join output rather than SELECT DISTINCT keeps this off SPARK-55848, which is
+    # what forces the patch-level gate on test_iceberg_spj_partial_clustering_distinct.
+    assert_cpu_and_gpu_are_equal_collect_with_capture(
+        join_after_spj,
+        conf=conf,
+        require_non_empty=True,
+        gpu_plan_assertion=assert_plan)
+
+
 # Enough rows that every bucket of the wider bucket(4) side is populated, so reducing it to
 # gcd(4, 2) = 2 buckets moves rows into partition values the raw-keyed lookup cannot find.
 _SPJ_REDUCIBLE_ROWS = 64

sql-plugin/src/main/spark340/scala/com/nvidia/spark/rapids/shims/GpuBatchScanExec.scala

@@ -186,7 +186,17 @@ case class GpuBatchScanExec(
                 .get
                 .map(t => (InternalRowComparableWrapper(t._1, p.expressions), t._2))
                 .toMap
-              val nestGroupedPartitions = groupedPartitions.map {
+              // SPARK-48949. Inert here: the `...v2.bucketing.partition.filter.enabled` config
+              // that makes `commonPartitionValues` an intersection instead of a union does not
+              // exist before Spark 4.0, so this filter can never drop a group on the versions
+              // this shim serves. Kept identical to the spark350db143 copy so the next SPJ
+              // audit diff between the two stays cheap.
+              val filteredGroupedPartitions = groupedPartitions.filter {
+                case (partValues, _) =>
+                  commonPartValuesMap.keySet.contains(
+                    InternalRowComparableWrapper(partValues, p.expressions))
+              }
+              val nestGroupedPartitions = filteredGroupedPartitions.map {
                 case (partValue, splits) =>
                   // `commonPartValuesMap` should contain the part value since it's the super set.
                   val numSplits = commonPartValuesMap

sql-plugin/src/main/spark350db143/scala/com/nvidia/spark/rapids/shims/GpuBatchScanExec.scala

@@ -199,7 +199,16 @@ case class GpuBatchScanExec(
                 .get
                 .map(t => (InternalRowComparableWrapper(t._1, partExpressions), t._2))
                 .toMap
-            val nestGroupedPartitions = finalGroupedPartitions.map { case (partValue, splits) =>
+            // SPARK-48949: with `...v2.bucketing.partition.filter.enabled`,
+            // `commonPartitionValues` is the intersection of the two join sides rather than
+            // their union, so this scan can still enumerate groups the planner pruned away.
+            // Dropping them here is what keeps the assert below true.
+            val filteredGroupedPartitions = finalGroupedPartitions.filter {
+              case (partValues, _) =>
+                commonPartValuesMap.keySet.contains(
+                  InternalRowComparableWrapper(partValues, partExpressions))
+            }
+            val nestGroupedPartitions = filteredGroupedPartitions.map { case (partValue, splits) =>
               // `commonPartValuesMap` should contain the part value since it's the super set.
               val numSplits = commonPartValuesMap
                   .get(InternalRowComparableWrapper(partValue, partExpressions))