feat: support nested types as native shuffle hash partitioning keys by viirya · Pull Request #5567 · apache/datafusion-comet

viirya · 2026-08-31T04:21:01Z

Which issue does this PR close?

Closes #5566.

Rationale for this change

CometShuffleExchangeExec.supportedHashPartitioningDataType rejected struct,
array and map partitioning keys, so a query repartitioning on a nested column fell
back to Spark for the whole shuffle. The comment justified this with:

Native code does not support hashing complex types, see hash_funcs/utils.rs

but that file hashes nested types recursively (struct fields,
List/LargeList/FixedSizeList elements, map keys and values), and shuffle
partitioning calls the same create_murmur3_hashes entry point, with the same
seed, that the hash expression uses.

Note that the two gates are not equivalent, so this is not simply "what hash
already allows": HashUtils.unsupportedReasonFor rejects decimal(precision > 18)
and TimeType at any depth, while the shuffle gate admits both. Comet does not
guarantee that native shuffle assigns rows to the same partitions as Spark
(#3079), and for a high-precision decimal it will not in general -- the native
hasher uses 16 little-endian i128 bytes where Spark uses the minimal big-endian
BigDecimal bytes.

What changes are included in this PR?

CometShuffleExchangeExec: add recursive StructType / ArrayType / MapType
cases to the hash partitioning gate, and replace the stale comment. Recursion
goes through the same predicate, so a leaf type that cannot be hashed natively
disqualifies the whole key and the shuffle still falls back -- this covers
collated strings (Fix listagg-collation.sql test in Spark 4.0.0 #1947 / fix: fall back to Spark for shuffle/sort/aggregate on non-default collated strings [Spark 4] #4035) and CalendarInterval (Hashing a CalendarInterval value fails with "Unsupported data type in hasher: Interval(MonthDayNano)" #5059) without needing
separate checks.
Map keys are admitted only on Spark 4.0+. Map entry order is not semantically
meaningful, so equal maps must hash alike; Spark 4.0+ normalizes a map shuffle
key with mapsort(...) (Add support for MapSort expression in Spark 4.0.0 #1941) while earlier versions do not, and Comet would
otherwise hash physical entry order. When the mapsort is not convertible
(CometMapSort supports scalar map keys only) the existing
partitioning-expression check already forces a fallback.
New config spark.comet.shuffle.native.partitioning.hash.nested.enabled,
disabled by default. 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. Enabling this by default before measuring could make those shuffles
slower than letting Spark do them, so the default stays off until there are
numbers. I plan to do that benchmark as a follow-up issue and PR.
Corrects the comment in CometFuzzTestBase claiming its Parquet file has no
nested complex types: generateSchema adds struct<array<..>> and
array<struct<..>> when both array and struct generation are on. Maps are the
only thing missing.

How are these changes tested?

CometNativeShuffleSuite and CometFuzzTestSuite, on two profiles:

Spark 4.1 (Scala 2.13): 52 and 45 tests, all passing
Spark 3.5 (Scala 2.12): 47 passing + 5 canceled by assume(isSpark40Plus, ...),
and 45 passing

New tests, all opting in via the config:

struct, array, and two-level (struct<array<..>>, array<struct<..>>) keys
a four-level key mixing all three recursive branches:
struct<a: array<struct<m: map<string, array<int>>, s: string>>, i: int>
multi-entry maps written in opposite key orders (map('a',1,'b',2) vs
map('b',2,'a',1)), asserting both rows share a spark_partition_id(), which is
the property the Spark 4.0+ restriction exists to protect
null keys at four levels -- top level, inside a struct, as an array element, and
as a map value -- with pairs of equal keys asserted into the same partition
partition assignment compared against Spark per row for shallow and deep keys,
with checkCometExchange so the test cannot pass by comparing Spark to Spark

Fallback tests: a map key on Spark 3.x, a map whose own key is nested, a collated
string nested inside a struct, and a CalendarInterval leaf (via make_interval)
inside a struct. CometFuzzTestSuite keeps its existing expectation that these
keys fall back by default, and additionally asserts they are admitted with the
config on.

Two notes on verification beyond the assertions:

checkShuffleAnswer asserts shuffleType == CometNativeShuffle at the plan
level. To confirm a nested key really is hashed in native code, I also inspected
the executed plan and exchange metrics for a deep key, which report native
shuffle-writer counters (repart_time, encode_time, interleave_time,
input_batches) that the Spark and columnar paths do not produce.
I sanity-checked that the partition-assignment test can fail, by perturbing the
native struct hash and confirming it reports every partition id shifted by one.

andygrove

Thanks @viirya. I checked the two things this rests on and they both hold up. The native shuffle partitioner really does call create_murmur3_hashes with seed 42, the same entry point the hash expression uses, and InsertMapSortInRepartitionExpressions really does rewrite nested map keys recursively on 4.0+. Recursing through the same predicate so that a bad leaf disqualifies the whole key is the right shape for this, and the collation and interval fallbacks fall out of it for free.

Most of what I have is about the tests. The map cases do not exercise the normalization they are gated on, the partition assignment test can pass without Comet running the shuffle at all, and there are no null keys anywhere. Details inline. I would also like to see some numbers before we turn this on by default.

andygrove · 2026-08-31T14:36:04Z

        true
      case dt if isTimeType(dt) =>
        true
+      case StructType(fields) if nestedHashPartitioningEnabled =>


Do you have any numbers for this? What I am worried about is the shapes that miss the vectorized paths. array<struct<...>> and a map nested inside a struct both fall through hash_list_with_primitive_elements! into hash_list_array!, which slices a one element array and re-enters create_murmur3_hashes for every element, plus a .columns().to_vec() per struct element on top of that. The four level key in the new tests pays that at every level.

Since the config defaults to true we are opting everyone into this, and if it turns out slower than just letting Spark do the shuffle then the default is wrong. struct<int, string> is not the interesting case here.

andygrove · 2026-08-31T14:36:04Z

+    }
+  }
+
+  test("native shuffle on map hash partitioning key") {


Every map in the new tests has a single entry, so mapsort is a no-op and the property this case is gated on never actually gets exercised. The reason we only admit maps on 4.0+ is that two equal maps with different physical entry order have to hash alike, and nothing here would notice if that stopped being true.

Could you add multi-entry maps written in different key orders, something like map('a', 1, 'b', 2) and map('b', 2, 'a', 1), and assert the two rows get the same spark_partition_id()?

andygrove · 2026-08-31T14:36:04Z

+      }
+    }
+  }
+  test("native shuffle nested hash partitioning key matches Spark's partition assignment") {


This is the only test that checks routing rather than the answer, so I would like it to be hard to fool, and right now it never asserts that Comet ran the shuffle. If the gate ever stops admitting these keys this quietly becomes Spark compared against Spark and still passes. Can you add a checkCometExchange(df, 1, true) alongside the sparkRows.nonEmpty check?

The other gap is nulls. None of the 200 shallow rows or the 40 deep rows have a null key at any level, and that is the case I would most want covered. The struct branch in hash_funcs/utils.rs is the one nested branch that does not look at its own null mask. List and Map both guard on is_null(row_idx), but DataType::Struct(_) just takes struct_array.columns() and recurses, while Spark is case null => seed. So the two only agree when the children happen to be null under a null parent. Parquet gives us that, which is why CometHashExpressionSuite is green, but a struct built in the plan may not, and Spark's own rewrite for a nested map key is If(IsNull(k), null, named_struct(...)), which is exactly that shape.

If a null struct ever hashes off leftover child values then two rows with the same key can land in different partitions, and that breaks grouping and joins whatever Spark would have done. Could you add null keys at each level here?

andygrove · 2026-08-31T14:36:04Z

+            s"SELECT /*+ REPARTITION(10, $keys) */ * FROM tbl)"
+        val cometRows = sql(query).collect().map(r => (r.getInt(0), r.getInt(1))).sorted
+        // `withSQLConf` returns Unit, so capture the Spark-side rows via a var rather than
+        // relying on the block's value.


withSQLConf returns T, not Unit. SQLTestUtils declares protected override def withSQLConf[T](pairs: (String, String)*)(f: => T): T and CometTestBase mixes it in, so both of these can be val sparkRows = withSQLConf(...) { ... }. Also missing a blank line before this test and before the interval one below it.

andygrove · 2026-08-31T14:36:04Z

+    }
+  }
+  test("native shuffle on nested hash partitioning key with interval leaf falls back") {
+    // CalendarIntervalType is allowed as a shuffle DATA column but the native hasher has no


I do not think this is testing CalendarIntervalType. INTERVAL '1' MONTH goes through constructMultiUnitsIntervalLiteral and comes out as YearMonthIntervalType(MONTH, MONTH). The gate rejects both so the assertion still holds, but #5059 and the CalendarIntervalType case in supportedSerializableDataType are about the calendar type, and that one stays uncovered.

Could you use make_interval(...), or set spark.sql.legacy.interval.enabled=true, so the test matches its comment?

andygrove · 2026-08-31T14:36:04Z

-      }
-      assert(cometShuffleExchanges.length == expectedNumCometShuffles)
+      // Both shuffle modes support the struct and array partitioning keys in this file (it is
+      // generated without maps and without nested complex types), so one Comet shuffle either way.


This file does have nested complex types. FuzzDataGenerator.generateSchema adds StructType(arraysOfPrimitives) and createArrayType(StructType(primitives)) when both generateArray and generateStruct are set, so struct<array<...>> and array<struct<...>> are both in there. Maps are the only thing missing. The comment in CometFuzzTestBase says the same thing and is equally wrong, so this was inherited, but could you fix both while you are here?

The knock-on is that this now asserts a native shuffle for array<decimal(36,18)> and struct<... decimal(36,18) ...>, where we hash 16 LE i128 bytes and Spark hashes the minimal BE BigDecimal bytes. I closed #3079 on the grounds that we do not need to allocate the same partitions as Spark, so I am not asking you to change the gate. But the description says this matches the hash expression's coverage and it does not, since HashUtils.unsupportedReasonFor rejects decimal(precision > 18) and TimeType at any depth. Worth saying so explicitly.

Last thing, this test only asserts a shuffle count and then calls collect(). Now that these keys run natively, is it worth comparing spark_partition_id() against Spark here too? This is the widest nested type coverage we have, and it is the one place a leaf type difference would show up across all of them at once.

viirya · 2026-09-01T16:08:59Z

Thanks for the detailed review — three of these were straight mistakes on my part.

withSQLConf return type. You're right that Spark 4.x declares withSQLConf[T](...)(f: => T): T, and my comment claiming it returns Unit was wrong. But the val form does not compile on 3.x: SQLHelper.withSQLConf there is public default void withSQLConf(Seq, Function0<BoxedUnit>), and -Pspark-3.5 fails with value nonEmpty is not a member of Unit. Since we build against both, I kept the var and fixed the comment to say what is actually going on. Added the missing blank lines.

The interval test wasn't testing CalendarIntervalType. Correct — INTERVAL '1' MONTH comes out as YearMonthIntervalType. Switched to make_interval(...), so it now matches its comment and covers the type #5059 is about.

The fuzz file does have nested complex types. Also correct: generateSchema adds struct<array<..>> and array<struct<..>> when both flags are set, and maps are the only omission. Fixed my comment and the inherited one in CometFuzzTestBase.

Map entry order. Added a test with map('a',1,'b',2) vs map('b',2,'a',1) (and a three-entry pair) asserting the rows share a spark_partition_id() and match Spark.

Parity test could pass without Comet. Added checkCometExchange(..., 1, true) to both the shallow and deep cases.

Null keys. Added a test with null keys at four levels — top level, inside a struct, as an array element, and as a map value — with pairs of equal keys asserted into the same partition and compared against Spark.

On the DataType::Struct(_) branch: you are right that it recurses into columns() without consulting its own null mask while List and Map guard on is_null(row_idx). I tried to construct a misroute through a plan-built null struct (both if(true, NULL, named_struct(...)) and nulling a struct that has non-null children) and could not — Comet's null struct came out with null children in both cases, so the two agreed. So I cannot show it biting today, but the asymmetry with the other two branches is real, and the new tests pin the current behavior. Happy to file it separately if you would like it tracked.

Default and performance. Agreed, and I have flipped it to false. You are right that struct<int, string> is not the interesting case: the shapes that miss the vectorized element paths fall into hash_list_array!, which slices a one-element array and re-enters create_murmur3_hashes per element, and I have no numbers for that. Turning it on by default before measuring was the wrong call. I would like to do the benchmark as a follow-up issue and PR rather than hold this one; all the tests opt in explicitly now.

On the decimal point. You are right that the description overstated the parity with the hash expression — HashUtils.unsupportedReasonFor rejects decimal(precision > 18) and TimeType at any depth and the shuffle gate does not. I have said so explicitly in the description. Worth noting what I actually measured: with the config on, every complex column in that fuzz file has partition parity with Spark, including array<decimal(36,18)> and the structs containing it. FuzzDataGenerator builds decimals from nextDouble(), so the values stay in the range where the 16-byte LE i128 and the minimal BE BigDecimal encodings agree — the divergence you describe is real in the code, just not reachable with the values that generator produces.

spark_partition_id() in the fuzz test. I tried this and backed it out. Pairing rows needs a stable identifier and neither attempt held up: c0 has many duplicate values, so sorting by it does not pair rows deterministically, and CAST(<key> AS STRING) is not a faithful key for nested values — that version failed for jvm shuffle too, which this change does not touch, so it was my test construction rather than routing. Doing it properly looks like it needs a real row id, i.e. materializing the fuzz table with something like monotonically_increasing_id(), which means changing the shared fixture that all three variants (native + 2x jvm) run against. I would rather not do that inside this PR unless you want it here. Happy to add it now or file a follow-up — which would you prefer?

sunchao · 2026-09-02T02:25:22Z

+      case MapType(keyType, valueType, _) if nestedHashPartitioningEnabled =>
+        // Map entry order is not semantically meaningful, so two equal maps must hash alike.
+        // Spark 4.0+ normalizes a map shuffle key by wrapping it in `mapsort(...)`, which is
+        // gated separately by CometMapSort (scalar map keys only) and, when unsupported, fails
+        // the expression check below. Earlier Spark versions insert no such normalization, so
+        // Comet would hash physical entry order and could route equal maps differently.
+        isSpark40Plus &&
+        supportedHashPartitioningDataType(keyType) &&
+        supportedHashPartitioningDataType(valueType)


[P2] Rebase sliced map offsets before enabling map shuffle keys

Could we fix spark_map_sort for sliced maps before admitting this case? With the nested-hash config enabled on Spark 4.0, a native OFFSET below a map repartition can pass a MapArray whose first entry offset is nonzero. DataFusion's limit uses batch.slice(skip, ...), and Arrow preserves the map's original entry offsets. The unified native shuffle evaluates mapsort(m) on that slice without a JVM round trip.

map_sort.rs builds its take indices only for the visible maps, but passes the original offsets to MapArray::try_new. For two two-entry maps, skipping the first leaves offsets [2,4]. Sorting creates two entries, so Arrow rejects the result with Max offset of 4 exceeds length of entries 2. The query fails instead of repartitioning. The map-sort code predates this PR, but the base gate rejected map hash keys, so the newly admitted path exposes it.

Please rebase the output offsets, or retain fallback for this case, and cover OFFSET followed by map repartition while asserting both operators run natively. This finding is source-derived. I have not executed the query.

viirya · 2026-09-02T07:17:46Z

@sunchao thanks — confirmed, and it is a query failure rather than a fallback. I executed it.

Unit level, slicing two two-entry maps to drop the first:

sliced len=1 offsets=OffsetBuffer(ScalarBuffer([2, 4]))
Err = Arrow error: Invalid argument error: Max offset of 4 exceeds length of entries 2

End to end, with the nested-hash config on:

SELECT * FROM (SELECT * FROM tbl ORDER BY _1 LIMIT 10 OFFSET 5) DISTRIBUTE BY _2

CometExchange hashpartitioning(mapsort(_2#3383), 10), REPARTITION_BY_COL, CometNativeShuffle
+- CometTakeOrderedAndProjectExec(limit=15, offset=5, orderBy=[_1#3382 ASC NULLS FIRST])
   +- CometNativeScan parquet [_1#3382,_2#3383]

CometNativeException: Error inserting batch: Arrow error:
  Invalid argument error: Max offset of 30 exceeds length of entries 20

One correction to the framing, though: this is not only exposed by the newly admitted
path. mapsort is not user-callable, but Spark 4.0+ has two rules that insert it —
InsertMapSortInRepartitionExpressions and InsertMapSortInGroupingExpressions. The
second one needs no repartition at all, so with the nested-hash config off (its
default) and no map shuffle key anywhere:

SELECT _2, count(*) FROM (SELECT * FROM tbl ORDER BY _1 LIMIT 15 OFFSET 5) GROUP BY _2

HashAggregate(keys=[_groupingmapsort#3400], functions=[count(1)])
+- HashAggregate(keys=[_groupingmapsort#3400], functions=[partial_count(1)])
   +- CometTakeOrderedAndProjectExec(limit=20, offset=5, orderBy=[_1 ASC NULLS FIRST])
      +- CometNativeScan parquet [_1,_2]

CometNativeException: Invalid argument error: Max offset of 40 exceeds length of entries 30

So it reproduces on current main and is a pre-existing bug rather than a regression
from this PR. Since it stands on its own I have split the fix out into #5630 (issue
#5629) rather than carrying it here — it rebases the output offsets as you suggested,
with a Rust unit test and an end-to-end test on the group-by path, both verified to
fail without the fix.

I have taken the repartition-on-map coverage back out of this PR for now, because on
main that path is still gated off and the test would be dead code. Once #5630 merges
I will sync this branch and add it there, asserting both operators run natively as you
asked.

sunchao · 2026-09-02T08:22:05Z

@viirya Agreed: this is a pre-existing mapsort defect, and #5567 adds another path that exposes it. The existing map-grouping rule is independent of the nested-hash option. Thanks for isolating the fix in #5630; I’m reviewing that PR now. Once this branch incorporates it, I’ll recheck the OFFSET → map-repartition coverage here, including the native-operator assertions.

`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>

sunchao

Thanks for syncing the mapsort fix and adding the OFFSET-to-map-repartition regression. At e819c2f9, the kernel matches the reviewed #5630 fix, and the new test asserts both native operators. I found no remaining P1/P2 in this update. I reused the earlier Arrow-only probe evidence. The new Spark test was not run here, and CI is still running.

viirya force-pushed the feat-native-shuffle-nested-hash-key branch 3 times, most recently from 1e6006c to 53ff3bd Compare August 31, 2026 07:03

andygrove reviewed Aug 31, 2026

View reviewed changes

viirya force-pushed the feat-native-shuffle-nested-hash-key branch from 53ff3bd to e0ce9d1 Compare September 1, 2026 16:08

viirya force-pushed the feat-native-shuffle-nested-hash-key branch from e0ce9d1 to 0e4f49a Compare September 1, 2026 17:14

sunchao reviewed Sep 2, 2026

View reviewed changes

This was referenced Sep 2, 2026

mapsort on a sliced map fails with "Max offset exceeds length of entries" #5629

Closed

fix: rebase map offsets in mapsort so sliced maps do not overrun entries #5630

Merged

viirya force-pushed the feat-native-shuffle-nested-hash-key branch from 0e4f49a to e819c2f Compare September 2, 2026 18:49

sunchao approved these changes Sep 2, 2026

View reviewed changes

Navigation Menu

Uh oh!

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

feat: support nested types as native shuffle hash partitioning keys#5567
viirya wants to merge 1 commit into
apache:mainfrom
viirya:feat-native-shuffle-nested-hash-key

viirya commented Aug 31, 2026 •

edited

Loading

Uh oh!

andygrove left a comment

Uh oh!

andygrove Aug 31, 2026

Uh oh!

andygrove Aug 31, 2026

Uh oh!

andygrove Aug 31, 2026

Uh oh!

andygrove Aug 31, 2026

Uh oh!

andygrove Aug 31, 2026

Uh oh!

andygrove Aug 31, 2026

Uh oh!

viirya commented Sep 1, 2026

Uh oh!

sunchao Sep 2, 2026

Uh oh!

viirya commented Sep 2, 2026

Uh oh!

sunchao commented Sep 2, 2026

Uh oh!

sunchao left a comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Uh oh!

Conversation

viirya commented Aug 31, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Which issue does this PR close?

Rationale for this change

What changes are included in this PR?

How are these changes tested?

Uh oh!

andygrove left a comment

Choose a reason for hiding this comment

Uh oh!

andygrove Aug 31, 2026

Choose a reason for hiding this comment

Uh oh!

andygrove Aug 31, 2026

Choose a reason for hiding this comment

Uh oh!

andygrove Aug 31, 2026

Choose a reason for hiding this comment

Uh oh!

andygrove Aug 31, 2026

Choose a reason for hiding this comment

Uh oh!

andygrove Aug 31, 2026

Choose a reason for hiding this comment

Uh oh!

andygrove Aug 31, 2026

Choose a reason for hiding this comment

Uh oh!

viirya commented Sep 1, 2026

Uh oh!

sunchao Sep 2, 2026

Choose a reason for hiding this comment

Uh oh!

viirya commented Sep 2, 2026

Uh oh!

sunchao commented Sep 2, 2026

Uh oh!

sunchao left a comment

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

viirya commented Aug 31, 2026 •

edited

Loading