fix: fall back to Spark for shuffle/sort/aggregate on non-default collated strings [Spark 4] by andygrove · Pull Request #4035 · apache/datafusion-comet

andygrove · 2026-04-22T17:10:09Z

Which issue does this PR close?

Closes #1947.

Rationale for this change

The listagg-collations.sql test is currently disabled for Spark 4.0.1 because Comet returns aabb instead of ab for:

SELECT lower(listagg(DISTINCT c1 COLLATE utf8_lcase) WITHIN GROUP (ORDER BY c1 COLLATE utf8_lcase))
FROM (VALUES ('a'), ('B'), ('b'), ('A')) AS t(c1);

Root cause: Comet accepts a collated StringType as a supported type in shuffle, sort, and aggregate, but hashes/sorts/compares the bytes directly. For this query the plan includes CometColumnarExchange hashpartitioning(collate(c1, utf8_lcase)), so rows with 'a' and 'A' are routed to different partitions and the downstream DISTINCT dedup (which runs in Spark with collation-aware equality) never sees them together.

Two things were broken:

The Spark 4.0 isStringCollationType shim checked dt.isInstanceOf[StringTypeWithCollation], but StringTypeWithCollation is an AbstractDataType used only for expression type matching, not a concrete data type. The check never matched a real collated StringType.
Even if it had matched, none of the shuffle/sort/aggregate type predicates were consulting it.

What changes are included in this PR?

common/src/main/spark-4.0/.../CometTypeShim.scala: Detect non-default collations by comparing collationId against the default StringType object's id.
spark/.../CometShuffleExchangeExec.scala: Reject collated-string keys in hash and range partitioning for both native and columnar shuffle paths.
spark/.../QueryPlanSerde.scala (supportedScalarSortElementType): Reject collated strings as a sort key.
spark/.../operators.scala (CometBaseAggregate): Reject collated-string grouping expressions.
dev/diffs/4.0.1.diff: Remove the sections that disabled the failing listagg-collations.sql query and its golden output; the test now runs unmodified against upstream Spark 4.0.1.

How are these changes tested?

New suite spark/src/test/spark-4.0/.../CometCollationSuite.scala with four tests:

The exact repro from Fix listagg-collation.sql test in Spark 4.0.0 #1947.
DISTINCT on a utf8_lcase collated column.
GROUP BY a utf8_lcase collated column.
A sanity test that default UTF8_BINARY strings still accelerate through Comet (guards against over-blocking).

The listagg-collations.sql SQLQueryTestSuite will also cover the fix via the updated dev/diffs/4.0.1.diff.

…lt collated strings The Spark 4.0 `CometTypeShim.isStringCollationType` check compared against `StringTypeWithCollation`, which is an `AbstractDataType` used only for expression type matching, so it never matched a concrete `StringType` with a non-default collation. As a result, Comet's columnar shuffle would hash-partition a collated column using byte-level hashing, sending rows that compare equal under the collation to different partitions. A downstream DISTINCT or GROUP BY then failed to deduplicate, producing incorrect aggregate results, for example on `listagg-collations.sql` query #5 in the Spark 4.0 SQL test suite. Detect non-default collations by comparing `collationId` against the default `StringType` object, and reject collated string columns in shuffle hash/range partitioning, sort, and aggregate grouping so those operators fall back to Spark. The previously disabled `listagg-collations.sql` query is re-enabled. Closes apache#1947

parthchandra

Claude: 4. Follow-up: BroadcastHashJoin with collated keys - Broadcast hash joins don't go through the shuffle exchange, so the new collation checks don't protect them. The doConvert in CometHashJoin and supportedSortMergeJoinEqualType don't check for collated strings. This could produce incorrect results for broadcast joins on collated keys. Worth a follow-up issue.

parthchandra · 2026-04-23T15:55:15Z

+  }
+
+  test("DISTINCT on utf8_lcase collated string groups case-insensitively") {
+    checkSparkAnswer(


Should we use checkSparkAnswerAndOperator here (and in the other tests in this suite)?

I updated the tests to either use checkSparkAnswerAndOperator or checkSparkAnswerAndFallbackReason to make them more explicit

…ER BY case Address PR apache#4035 review feedback: - Replace checkSparkAnswer with checkSparkAnswerAndFallbackReason on the listagg/DISTINCT/GROUP BY tests so regressions that silently re-enable the native path are caught. The shuffle-exchange rule is the first layer that fires for collated keys, so pin its reason string rather than the defense-in-depth aggregate check. - Add an ORDER BY test covering the range-partitioning fallback. - Rework the default-UTF8_BINARY sanity test to use withParquetTable so checkSparkAnswerAndOperator has a real Comet scan to attach to.

parthchandra

lgtm.

andygrove · 2026-04-24T18:40:40Z

Merged. Thanks @parthchandra

`CometShuffleExchangeExec.supportedHashPartitioningDataType` rejected struct, array and map partitioning keys, so any query repartitioning on a nested column fell back to Spark for the whole shuffle. The comment said "Native code does not support hashing complex types, see hash_funcs/utils.rs", but that file hashes nested types recursively (struct fields, list elements, map keys and values), and shuffle partitioning shares that kernel, and Spark's seed, with the `hash` expression. Adds the recursive struct/array/map cases to the gate, behind `spark.comet.shuffle.native.partitioning.hash.nested.enabled` (default true). Nesting is checked recursively through the same predicate, so a leaf type that cannot be hashed natively disqualifies the whole key and the shuffle still falls back: - collated strings, which Comet hashes as raw bytes (see apache#1947 / apache#4035, where rows equal under the collation reached different partitions and a downstream collation-aware DISTINCT produced a wrong answer) - CalendarInterval, which the native hasher has no branch for (apache#5059) Map keys are additionally restricted to Spark 4.0+. Map entry order is not semantically meaningful, so two equal maps must hash alike, and Spark 4.0+ normalizes a map shuffle key by wrapping it in `mapsort(...)`. Earlier versions insert no such normalization, so Comet would hash physical entry order. When the `mapsort` itself is not convertible -- CometMapSort supports scalar map keys only -- the existing expression check fails and the shuffle falls back. Co-authored-by: Claude Code <noreply@anthropic.com>

`CometShuffleExchangeExec.supportedHashPartitioningDataType` rejected struct, array and map partitioning keys, so any query repartitioning on a nested column fell back to Spark for the whole shuffle. The comment said "Native code does not support hashing complex types, see hash_funcs/utils.rs", but that file hashes nested types recursively (struct fields, list elements, map keys and values), and shuffle partitioning shares that kernel, and Spark's seed, with the `hash` expression. Adds the recursive struct/array/map cases to the gate, behind `spark.comet.shuffle.native.partitioning.hash.nested.enabled` (default true). Nesting is checked recursively through the same predicate, so a leaf type that cannot be hashed natively disqualifies the whole key and the shuffle still falls back: - collated strings, which Comet hashes as raw bytes (see apache#1947 / apache#4035, where rows equal under the collation reached different partitions and a downstream collation-aware DISTINCT produced a wrong answer) - CalendarInterval, which the native hasher has no branch for (apache#5059) Map keys are additionally restricted to Spark 4.0+. Map entry order is not semantically meaningful, so two equal maps must hash alike, and Spark 4.0+ normalizes a map shuffle key by wrapping it in `mapsort(...)`. Earlier versions insert no such normalization, so Comet would hash physical entry order. When the `mapsort` itself is not convertible -- CometMapSort supports scalar map keys only -- the existing expression check fails and the shuffle falls back. Also updates `CometFuzzTestSuite`'s "distribute by single column (complex types)" test, which asserted that native shuffle produces zero Comet exchanges for a complex partitioning key. The fuzz file it uses is generated without maps and without nested complex types, so every struct and array key in it is now natively supported and both shuffle modes produce one Comet exchange. Co-authored-by: Claude Code <noreply@anthropic.com>

`CometShuffleExchangeExec.supportedHashPartitioningDataType` rejected struct, array and map partitioning keys, so any query repartitioning on a nested column fell back to Spark for the whole shuffle. The comment said "Native code does not support hashing complex types, see hash_funcs/utils.rs", but that file hashes nested types recursively (struct fields, list elements, map keys and values), and shuffle partitioning shares that kernel, and Spark's seed, with the `hash` expression. Adds the recursive struct/array/map cases to the gate, behind `spark.comet.shuffle.native.partitioning.hash.nested.enabled` (default true). Nesting is checked recursively through the same predicate, so a leaf type that cannot be hashed natively disqualifies the whole key and the shuffle still falls back: - collated strings, which Comet hashes as raw bytes (see apache#1947 / apache#4035, where rows equal under the collation reached different partitions and a downstream collation-aware DISTINCT produced a wrong answer) - CalendarInterval, which the native hasher has no branch for (apache#5059) Map keys are additionally restricted to Spark 4.0+. Map entry order is not semantically meaningful, so two equal maps must hash alike, and Spark 4.0+ normalizes a map shuffle key by wrapping it in `mapsort(...)`. Earlier versions insert no such normalization, so Comet would hash physical entry order. When the `mapsort` itself is not convertible -- CometMapSort supports scalar map keys only -- the existing expression check fails and the shuffle falls back. The config defaults to false. The native hasher only vectorizes nested shapes whose leaves are primitives; `array<struct<...>>` and a map inside a struct fall through to a per-element path that re-enters `create_murmur3_hashes` for every element, so enabling this by default before measuring could make these shuffles slower than letting Spark do them. `CometFuzzTestSuite`'s "distribute by single column (complex types)" keeps its existing expectation, since the keys still fall back by default, and additionally asserts that they are admitted with the config enabled. Also corrects the comment in `CometFuzzTestBase` claiming that file has no nested complex types -- `generateSchema` does add `struct<array<..>>` and `array<struct<..>>` when both array and struct generation are on; maps are the only thing missing. Co-authored-by: Claude Code <noreply@anthropic.com>

`CometShuffleExchangeExec.supportedHashPartitioningDataType` rejected struct, array and map partitioning keys, so any query repartitioning on a nested column fell back to Spark for the whole shuffle. The comment said "Native code does not support hashing complex types, see hash_funcs/utils.rs", but that file hashes nested types recursively (struct fields, list elements, map keys and values), and shuffle partitioning shares that kernel, and Spark's seed, with the `hash` expression. Adds the recursive struct/array/map cases to the gate, behind `spark.comet.shuffle.native.partitioning.hash.nested.enabled` (default true). Nesting is checked recursively through the same predicate, so a leaf type that cannot be hashed natively disqualifies the whole key and the shuffle still falls back: - collated strings, which Comet hashes as raw bytes (see apache#1947 / apache#4035, where rows equal under the collation reached different partitions and a downstream collation-aware DISTINCT produced a wrong answer) - CalendarInterval, which the native hasher has no branch for (apache#5059) Map keys are additionally restricted to Spark 4.0+. Map entry order is not semantically meaningful, so two equal maps must hash alike, and Spark 4.0+ normalizes a map shuffle key by wrapping it in `mapsort(...)`. Earlier versions insert no such normalization, so Comet would hash physical entry order. When the `mapsort` itself is not convertible -- CometMapSort supports scalar map keys only -- the existing expression check fails and the shuffle falls back. The config defaults to false. The native hasher only vectorizes nested shapes whose leaves are primitives; `array<struct<...>>` and a map inside a struct fall through to a per-element path that re-enters `create_murmur3_hashes` for every element, so enabling this by default before measuring could make these shuffles slower than letting Spark do them. `CometFuzzTestSuite`'s "distribute by single column (complex types)" keeps its existing expectation, since the keys still fall back by default, and additionally asserts that they are admitted with the config enabled. Also corrects the comment in `CometFuzzTestBase` claiming that file has no nested complex types -- `generateSchema` does add `struct<array<..>>` and `array<struct<..>>` when both array and struct generation are on; maps are the only thing missing. Also covers the repartition-on-map route into the sliced-map `mapsort` defect fixed in apache#5630: a native OFFSET below a map shuffle key, asserting that both the exchange and the offset stay native so the sliced map actually reaches the native mapsort. Co-authored-by: Claude Code <noreply@anthropic.com>

andygrove force-pushed the fix-1947-listagg-collation branch from 49ab010 to 9077fca Compare April 22, 2026 17:12

andygrove added the spark 4.0 label Apr 22, 2026

andygrove force-pushed the fix-1947-listagg-collation branch from 9077fca to 441fa06 Compare April 22, 2026 19:15

ci: add CometCollationSuite to pr_build workflows

4bfe4cb

andygrove marked this pull request as ready for review April 22, 2026 21:04

andygrove requested review from kazuyukitanimura and parthchandra April 22, 2026 21:04

andygrove changed the title ~~fix: fall back to Spark for shuffle/sort/aggregate on non-default collated strings~~ fix: fall back to Spark for shuffle/sort/aggregate on non-default collated strings [Spark 4] Apr 23, 2026

parthchandra reviewed Apr 23, 2026

View reviewed changes

andygrove mentioned this pull request Apr 23, 2026

Broadcast hash join and sort-merge join can produce incorrect results on non-default collated keys [Spark 4] #4051

Closed

parthchandra approved these changes Apr 24, 2026

View reviewed changes

andygrove merged commit f089295 into apache:main Apr 24, 2026
235 of 238 checks passed

andygrove deleted the fix-1947-listagg-collation branch April 24, 2026 18:40

This was referenced Aug 31, 2026

Native shuffle rejects nested types as hash partitioning keys although the native hasher supports them #5566

Open

feat: support nested types as native shuffle hash partitioning keys #5567

Open

Navigation Menu

Uh oh!

fix: fall back to Spark for shuffle/sort/aggregate on non-default collated strings [Spark 4] - #4035

fix: fall back to Spark for shuffle/sort/aggregate on non-default collated strings [Spark 4]#4035
andygrove merged 3 commits into
apache:mainfrom
andygrove:fix-1947-listagg-collation

andygrove commented Apr 22, 2026

Uh oh!

parthchandra left a comment

Uh oh!

parthchandra Apr 23, 2026

Uh oh!

andygrove Apr 23, 2026

Uh oh!

Uh oh!

parthchandra left a comment

Uh oh!

Uh oh!

andygrove commented Apr 24, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Uh oh!

Conversation

andygrove commented Apr 22, 2026

Which issue does this PR close?

Rationale for this change

What changes are included in this PR?

How are these changes tested?

Uh oh!

parthchandra left a comment

Choose a reason for hiding this comment

Uh oh!

parthchandra Apr 23, 2026

Choose a reason for hiding this comment

Uh oh!

andygrove Apr 23, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

parthchandra left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

andygrove commented Apr 24, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants