You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@gluten.apache.org by "minni31 (via GitHub)" <gi...@apache.org> on 2026/05/13 07:40:53 UTC
[PR] [VL] Add KeyGroupedPartitioning support to columnar shuffle [gluten]
minni31 opened a new pull request, #12084:
URL: https://github.com/apache/gluten/pull/12084
## CONTEXT
`KeyGroupedPartitioning` is a Spark partitioning scheme used by V2 data source connectors (e.g., Iceberg, Paimon) where data is partitioned by specific key expressions with known unique partition values. Currently, Gluten's columnar shuffle exchange does not handle this partitioning type, causing a fallback to vanilla Spark for any query involving V2 sources with key-grouped partitioning.
## WHAT
Adds `KeyGroupedPartitioning` support to the columnar shuffle exchange path in the Velox backend. The implementation reuses the existing JVM-side partition ID computation pattern (same mechanism as `RangePartitioning`):
- Adds `KeyGroupedPartitioning` to the validation whitelist in `ColumnarShuffleExchangeExecBase`, allowing the columnar shuffle to accept this partitioning type.
- Constructs a `KeyGroupedPartitioner` from the partitioning's `uniquePartitionValues`, mapping each partition key to its index.
- Computes partition IDs on the JVM side by evaluating partition key expressions against each row (via `BindReferences`) and looking up the result in the `KeyGroupedPartitioner`. The pid column is prepended to each batch so the native shuffle writer can read it directly.
- Reuses `RangePartitioningShortName` for the native partitioning descriptor since both Range and KeyGrouped use the same JVM-side pid prepend pattern — the native shuffle writer reads the prepended column rather than computing partition IDs natively.
- Each key extraction allocates a fresh array per row and converts to an immutable `ArraySeq` to avoid aliasing issues with mutable array reuse.
### Tests
| Suite | Tests | Status |
|-------|-------|--------|
| `VeloxShufflePartitioningSuite` | 22 tests: SxS tests for hash (A1-A4), range (B1-B3), round-robin (C1-C2), single (D1-D2), null semantics (E1-E2), data types (F1-F2), boundary cases (G1-G2), KeyGrouped unit tests (H1-H6) | Local pass |
Note: End-to-end KeyGroupedPartitioning tests require V2 data source connectors (Iceberg/Paimon) which are not available in this test module. The H-series unit tests validate key extraction, `KeyGroupedPartitioner` construction, and the full key-extraction-to-partition-lookup flow.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@gluten.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@gluten.apache.org
For additional commands, e-mail: commits-help@gluten.apache.org
Re: [PR] [VL] Add KeyGroupedPartitioning support to columnar shuffle [gluten]
Posted by "github-actions[bot] (via GitHub)" <gi...@apache.org>.
github-actions[bot] commented on PR #12084:
URL: https://github.com/apache/gluten/pull/12084#issuecomment-4438585997
Run Gluten Clickhouse CI on x86
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@gluten.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@gluten.apache.org
For additional commands, e-mail: commits-help@gluten.apache.org
Re: [PR] [VL] Add KeyGroupedPartitioning support to columnar shuffle [gluten]
Posted by "github-actions[bot] (via GitHub)" <gi...@apache.org>.
github-actions[bot] commented on PR #12084:
URL: https://github.com/apache/gluten/pull/12084#issuecomment-4439012497
Run Gluten Clickhouse CI on x86
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@gluten.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@gluten.apache.org
For additional commands, e-mail: commits-help@gluten.apache.org
Re: [PR] [VL] Add KeyGroupedPartitioning support to columnar shuffle [gluten]
Posted by "minni31 (via GitHub)" <gi...@apache.org>.
minni31 closed pull request #12084: [VL] Add KeyGroupedPartitioning support to columnar shuffle
URL: https://github.com/apache/gluten/pull/12084
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@gluten.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@gluten.apache.org
For additional commands, e-mail: commits-help@gluten.apache.org
Re: [PR] [VL] Add KeyGroupedPartitioning support to columnar shuffle [gluten]
Posted by "yaooqinn (via GitHub)" <gi...@apache.org>.
yaooqinn commented on code in PR #12084:
URL: https://github.com/apache/gluten/pull/12084#discussion_r3253171724
##########
backends-velox/src/main/scala/org/apache/spark/sql/execution/utils/ExecUtil.scala:
##########
@@ -128,10 +128,22 @@ object ExecUtil {
ascending = true,
samplePointsPerPartitionHint = SQLConf.get.rangeExchangeSampleSizePerPartition)
Some(part)
+ case k @ KeyGroupedPartitioning(_, n, _, _) =>
+ // Build a lookup map from partition key values to partition indices.
+ // KeyGroupedPartitioner.getPartition uses getOrElseUpdate: if a key is not found
+ // in the map, it falls back to hash-based assignment (nonNegativeMod). This is a
+ // best-effort fallback - all expected keys should be present in uniquePartitionValues.
+ val dataTypes = k.expressions.map(_.dataType)
+ val valueMap = scala.collection.mutable.Map.empty[Seq[Any], Int]
+ k.uniquePartitionValues.zipWithIndex.foreach {
+ case (partition, index) =>
+ valueMap.update(partition.toSeq(dataTypes), index)
Review Comment:
**[BLOCKING] `numPartitions` / `valueMap.size` mismatch when `isPartiallyClustered=true` or `partitionValues` contain duplicates**
`KeyGroupedPartitioner` is constructed with the original `n = KeyGroupedPartitioning.numPartitions` (i.e. `partitionValues.size`, including replicas), while `valueMap` is built from the deduped `uniquePartitionValues`. When `isPartiallyClustered=true` (the SPJ skew-handling path), `partitionValues` deliberately replicates entries → `numPartitions > uniquePartitionValues.size`. Result:
- pids `[uniqueValues.size, n)` are unreachable via the map lookup.
- The `getOrElseUpdate` hash fallback (`nonNegativeMod(key.hashCode, n)`) can land any unknown key into `[0, uniqueValues.size)`, **colliding with an already-occupied pid**.
- Two distinct partition keys then end up in the same output partition, breaking the KGP per-partition-uniqueness invariant that downstream SPJ-aware operators rely on.
For reference, vanilla Spark `master`'s `ShuffleExchangeExec` KGP branch ([`ShuffleExchangeExec.scala:378-383`](https://github.com/apache/spark/blob/master/sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/ShuffleExchangeExec.scala#L378-L383)) calls `k.toGrouped` first — which dedups `partitionKeys` and redefines `numPartitions := partitionKeys.length` — *before* constructing the `KeyGroupedPartitioner`. That canonicalization is what guarantees `valueMap.size == numPartitions`. (Spark 3.5's `ShuffleExchangeExec` has no KGP branch at all — KGP is consumed only by `BatchScanExec` and never re-shuffled, so the problem doesn't arise there either.)
Suggested fixes (any one):
1. Use `k.uniquePartitionValues.size` as the partitioner's `numPartitions` instead of `n` (matches the dedup invariant).
2. In `ColumnarShuffleExchangeExecBase.scala:107` reject `isPartiallyClustered=true` and any case where `partitionValues.size != uniquePartitionValues.size` via `ValidationResult.failed`, restricting the columnar path to the strictly-1:1 case.
3. Port the `.toGrouped`-equivalent canonicalization (dedup + renumber) before building `valueMap` and the partitioner.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@gluten.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@gluten.apache.org
For additional commands, e-mail: commits-help@gluten.apache.org