[SQL] dropDuplicates and Window dedup produce incorrect results with SPJ partiallyClusteredDistribution

[paryoja]

What type of issue is this?

Bug

Spark version

4.0.1 (with Iceberg 1.10.1)

Describe the bug

When using Storage-Partitioned Join (SPJ) with spark.sql.sources.v2.bucketing.partiallyClusteredDistribution.enabled=true, both dropDuplicates() and Window-based dedup (row_number()) produce incorrect results — duplicate rows that should have been removed survive in the output.

Root cause

Partial clustering splits a partition with many files across multiple tasks to improve parallelism. However, downstream dedup operations (dropDuplicates, row_number() OVER (PARTITION BY ...)) rely on the assumption that all rows with the same partition key are co-located in a single task. Since SPJ eliminates the Exchange (shuffle), each split independently deduplicates, and duplicate partition keys survive across splits.

Steps to reproduce

1. Create two Iceberg tables partitioned by part_key:
   • Big table: 20 separate appends to partition p1 (= 20 data files), plus 1 append to p2
   • Small table: 1 append containing both p1 and p2
2. Enable SPJ with partial clustering:

spark.sql.sources.v2.bucketing.enabled = true
spark.sql.sources.v2.bucketing.partiallyClusteredDistribution.enabled = true
spark.sql.iceberg.planning.preserve-data-grouping = true
spark.sql.autoBroadcastJoinThreshold = -1

3. Perform a leftsemi join on part_key, then dropDuplicates(["part_key"])
4. The physical plan contains no Exchange node, confirming SPJ is active

import shutil
import tempfile

from pyspark.sql import SparkSession, Window, functions as F
from pyspark.sql.types import IntegerType, StringType, StructField, StructType
import pytest


ICEBERG_PKG = "org.apache.iceberg:iceberg-spark-runtime-4.0_2.13:1.10.1"

SCHEMA = StructType([
    StructField("id", StringType()),
    StructField("part_key", StringType()),
    StructField("value", IntegerType()),
    StructField("padding", StringType()),
])


@pytest.fixture(scope="module")
def warehouse_dir():
    d = tempfile.mkdtemp(prefix="iceberg_spj_test_")
    yield d
    shutil.rmtree(d, ignore_errors=True)


@pytest.fixture(scope="module")
def spark_iceberg(warehouse_dir):
    spark = (
        SparkSession.builder.master("local[4]")
        .appName("spj_window_dedup_bug")
        .config("spark.jars.packages", ICEBERG_PKG)
        .config(
            "spark.sql.extensions",
            "org.apache.iceberg.spark.extensions.IcebergSparkSessionExtensions",
        )
        .config("spark.sql.catalog.local", "org.apache.iceberg.spark.SparkCatalog")
        .config("spark.sql.catalog.local.type", "hadoop")
        .config("spark.sql.catalog.local.warehouse", warehouse_dir)
        .config("spark.ui.enabled", "false")
        .config("spark.driver.memory", "2g")
        .config("spark.sql.shuffle.partitions", "4")
        .config("spark.sql.adaptive.enabled", "false")
        .getOrCreate()
    )
    spark.sparkContext.setLogLevel("WARN")
    yield spark
    spark.stop()


@pytest.fixture()
def asymmetric_tables(spark_iceberg):
    """Two Iceberg tables with asymmetric file counts to trigger partial clustering."""
    big, small = "local.db.big", "local.db.small"
    for t in (big, small):
        spark_iceberg.sql(f"DROP TABLE IF EXISTS {t}")
        spark_iceberg.sql(f"""
            CREATE TABLE {t} (id STRING, part_key STRING, value INT, padding STRING)
            USING iceberg PARTITIONED BY (part_key)
        """)

    # Big table: 20 appends = 20 files for partition p1
    for i in range(20):
        data = [(f"big_{i}_{j}", "p1", i, "X" * 1000) for j in range(200)]
        spark_iceberg.createDataFrame(data, SCHEMA).writeTo(big).append()
    spark_iceberg.createDataFrame([("big_other", "p2", 99, "Y")], SCHEMA).writeTo(big).append()

    # Small table: 1 append = 1 file
    spark_iceberg.createDataFrame(
        [("small_0", "p1", 0, "Z"), ("small_1", "p2", 1, "Z")], SCHEMA
    ).writeTo(small).append()

    spark_iceberg.conf.set("spark.sql.sources.v2.bucketing.enabled", "true")
    spark_iceberg.conf.set(
        "spark.sql.sources.v2.bucketing.partiallyClusteredDistribution.enabled", "true"
    )
    spark_iceberg.conf.set("spark.sql.iceberg.planning.preserve-data-grouping", "true")
    spark_iceberg.conf.set("spark.sql.autoBroadcastJoinThreshold", "-1")

    return spark_iceberg.table(big), spark_iceberg.table(small)


def _get_plan(df) -> str:
    return df._jdf.queryExecution().executedPlan().toString()


class TestSPJDedupBug:

    def test_drop_duplicates_after_join_produces_duplicates(self, asymmetric_tables):
        """dropDuplicates after SPJ join → duplicates survive (expected 2, gets >2)."""
        big, small = asymmetric_tables

        deduped = big.join(small, on="part_key", how="leftsemi").dropDuplicates(["part_key"])

        assert "Exchange" not in _get_plan(deduped)
        assert deduped.count() > 2

    def test_window_dedup_after_join_produces_duplicates(self, asymmetric_tables):
        """row_number() Window dedup after SPJ join → duplicates survive."""
        big, small = asymmetric_tables

        joined = big.join(small, on="part_key", how="leftsemi")
        w = Window.partitionBy("part_key").orderBy(F.col("value").desc())
        deduped = joined.withColumn("_r", F.row_number().over(w)).filter("_r = 1").drop("_r")

        assert "Exchange" not in _get_plan(deduped)
        assert deduped.count() > 2

    def test_disabling_partial_clustering_fixes_it(self, asymmetric_tables, spark_iceberg):
        """Setting partiallyClusteredDistribution.enabled=false → dedup works."""
        big, small = asymmetric_tables
        spark_iceberg.conf.set(
            "spark.sql.sources.v2.bucketing.partiallyClusteredDistribution.enabled", "false"
        )

        deduped = big.join(small, on="part_key", how="leftsemi").dropDuplicates(["part_key"])

        assert deduped.count() == 2

Related

SPARK-38166: Duplicates after task failure in dropDuplicates and repartition

[naveenp2708]

I've been investigating this issue and traced the root cause to KeyGroupedPartitioning.satisfies0() returning true for ClusteredDistribution even when partial clustering has split the partition across multiple tasks. SPARK-53074 fixed the pre-join case (operators between scan and join), but this post-join case remains unaddressed.
The core problem: after a partially-clustered SPJ join, the output inherits KeyGroupedPartitioning which EnsureRequirements considers as satisfying ClusteredDistribution. No Exchange is inserted before downstream dedup operators (dropDuplicates, Window with row_number), so each split task independently deduplicates its subset — producing inflated results.
I'm working on a fix that adds an isPartiallyClustered flag to KeyGroupedPartitioning so that satisfies0(ClusteredDistribution) returns false when the partition was split, triggering the necessary Exchange. Plain SPJ joins without dedup remain unaffected. Will submit a JIRA and PR shortly.
JIRA: https://issues.apache.org/jira/browse/SPARK-55848

[szehon-ho]

@peter-toth @chirag-s-db fyi

[peter-toth]

I believe my refactor: #54330 fixes this issue implicitly as it properly tracks when a KeyedPartitioning is grouped, so we know when we need to regroup partitions to satisfy a ClusteredDistribution.

But because this is a correctness bug we should fix it on 3.5, 4.0 and 4.1 branches as well.

@naveenp2708, do you have a compact test case? If not, then I can add one to my PR Monday.

[naveenp2708]

Thanks @peter-toth I looked at #54330 and agree the KeyedPartitioning refactor should address this by properly tracking grouped state.

Here's a compact test case for KeyGroupedPartitioningSuite:

test("SPARK-55848: dropDuplicates after SPJ with partial clustering produces correct results") {
  withSQLConf(
    SQLConf.V2_BUCKETING_ENABLED.key -> "true",
    SQLConf.V2_BUCKETING_PARTIALLY_CLUSTERED_DISTRIBUTION_ENABLED.key -> "true",
    SQLConf.V2_BUCKETING_PUSH_PART_VALUES_ENABLED.key -> "true",
    SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1"
  ) {
    val joined = sql(s"""
      SELECT i.id, i.name
      FROM testcat.ns.$items i
      JOIN testcat.ns.$purchases p ON i.id = p.item_id
    """)
    val deduped = joined.dropDuplicates("id")
    assert(deduped.count() == joined.select("id").distinct().count(),
      "Partial clustering should not inflate dedup results")
  }
}

I also have a Window-based dedup test case and a no-regression test (verifying plain SPJ joins stay Exchange-free) if those would be useful.

Happy to submit these as a separate test-only PR for the backport branches (3.5, 4.0, 4.1), or you can include them in #54330. Let me know how you'd like to coordinate.

[szehon-ho]

HI @naveenp2708 @peter-toth @cloud-fan , what do you think, if we make a separate fix on the old code in old versions of Spark (4.1, 4.0, 3.5), even if it's a bit hacky. We can debate 4.1, but I just think it's a bit risky to also backport the #54330 to the old Spark versions.. We can also help if needed, to see if we can fix the use case in the old Spark versions.

[naveenp2708]

Hi @szehon-ho That's exactly what PR #54751 does — a targeted fix for branch-4.1 using an isPartiallyClustered flag on KeyGroupedPartitioning, without requiring the full #54330 refactor. Happy to adapt the same approach for 4.0 and 3.5 branches as well if needed.