Columnar shuffle rejects range partitioning based on a native-serde check it never uses

[0lai0]

What is the problem?

CometShuffleExchangeExec decides whether columnar (JVM) shuffle can handle a RangePartitioning by asking whether Comet can serialize the sort orders to protobuf:

https://github.com/apache/datafusion-comet/blob/main/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleExchangeExec.scala#L654-L659

case RangePartitioning(orderings, _) =>
  for (o <- orderings) {
    if (QueryPlanSerde.exprToProto(o, inputs).isEmpty) {
      reasons += s"unsupported range partitioning sort order: $o"
    }
  }

But on the columnar path the range partitioning is performed entirely on the JVM by Spark's own RangePartitioner, over an UnsafeProjection of the sort keys and a LazilyGeneratedOrdering:

https://github.com/apache/datafusion-comet/blob/main/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleExchangeExec.scala#L1002-L1022

case RangePartitioning(sortingExpressions, numPartitions) =>
  val rddForSampling = rdd.mapPartitionsInternal { iter =>
    val projection = UnsafeProjection.create(sortingExpressions.map(_.child), outputAttributes)
    ...
  }
  implicit val ordering = new LazilyGeneratedOrdering(orderingAttributes)
  new RangePartitioner(numPartitions, rddForSampling, ascending = true, ...)

The serialized SortOrder produced by the check is never sent to native code on this path. So the gate rejects columnar shuffle for orderings that the JVM would have partitioned correctly, and the query falls back to Spark's shuffle for no compatibility reason.

Impact

Any ordering expression that Comet cannot serialize disables Comet's columnar shuffle for that exchange, even though nothing about the partitioning would have run natively. Examples:

• floating-point keys nested in arrays or structs while spark.comet.exec.strictFloatingPoint=true
• any ORDER BY / repartitionByRange expression Comet has no serde for

This is an unnecessary-fallback bug rather than a correctness bug.

History

Until #5506 this also fired for scalar float and double keys under spark.comet.exec.strictFloatingPoint=true, because CometSortOrder.getSupportLevel returned Incompatible for them and exprToProto therefore returned None. #5506 narrowed that verdict, so the scalar floating-point case no longer reaches this gate. The structural problem is unchanged.

Suggested fix

Drop the exprToProto probe from the columnar branch, keeping the collation check that follows it (collation genuinely affects JVM ordering semantics here).

Please verify before fixing — I have not confirmed this myself: that no part of the columnar shuffle write path consumes a serialized partitioning for RangePartitioning. If some writer does, the gate is load-bearing and the fix is instead to narrow it to whatever that writer actually needs.

Suggested test

Pick an ordering that Comet cannot serialize but the JVM can sort — a struct containing a float with spark.comet.exec.strictFloatingPoint=true is a ready-made one — and assert with spark.comet.exec.shuffle.mode=jvm that the exchange stays on Comet's columnar shuffle instead of falling back to Spark's.

[Visorgood]

Hello @0lai0!

Verified the part you flagged: nothing on the columnar path consumes a serialized partitioning.

prepareJVMShuffleDependency partitions entirely on the JVM – UnsafeProjection, LazilyGeneratedOrdering, Spark's RangePartitioner. CometShuffleDependency.outputPartitioning is a Catalyst Partitioning, not a proto message. PartitioningOuterClass.RangePartition is built only in CometNativeShuffleWriter (line 355). That writer is reached only via CometNativeShuffleHandle; the columnar path gets a different handle.

So your fix holds. The near-identical probe in the native branch (~line 556) must stay – there the serialized SortOrder really does go native, and it also rolls up fallback reasons from the expression tag.

I've confirmed it with a failing test. After #5506 scalar float/double are Compatible, but strictFloatingPointReason uses containsType recursively, so a struct<double, int> sort key still returns Incompatible and trips the gate:

- PROBE range partitioning on a struct containing a double *** FAILED ***
    has 0 CometShuffleExchangeExec node which doesn't match the expected: 1
    
    Exchange rangepartitioning(c#293 ASC NULLS FIRST, 4), REPARTITION_BY_NUM
    +- CometProject [c#293], [struct(a, _1#289, b, _2#290) AS c#293]
       +- CometNativeScan parquet [_1#289,_2#290]

The scan and projection stay on Comet; only the exchange drops to Spark. Nothing in the query is unsupported – the gate alone forces the transition. The #5506 test next to it still passes.

take. My plan: drop the probe from the columnar branch, keep the collation check, add this test beside the #5506 one at CometColumnarShuffleSuite.scala:795, and update that test's comment now both halves are covered.

[andygrove]

Thanks for verifying, @Visorgood. I walked the same chain and agree the probe is dead on the columnar path. A few additions before you open the PR.

One gap in the verification: CometNativeShuffleHandle isn't the only reader of the partitioning. CometCelebornShuffleManager.getWriter also pulls dependency.outputPartitioning and dependency.rangePartitionBounds (around line 171). It reaches the same conclusion — it's guarded on shuffleType == CometNativeShuffle, and rejectCometHandle bails out for both CometBypassMergeSortShuffleHandle and CometSerializedShuffleHandle — but since the ask in the description was to prove that no columnar write path consumes a serialized partitioning, that second manager should be part of the answer.

The bigger thing is that the range branch isn't the only dead probe in columnarShuffleFailureReasons. The HashPartitioning branch just above it runs the same exprToProto check, and the columnar path evaluates h.partitionIdExpression through UnsafeProjection on the JVM exactly the way it does the range keys — never serializing either. Same argument, same conclusion. Would you take both in one PR? Fixing only range leaves the identical bug one case up.

On keeping the collation check, I'm not sure that rationale survives. LazilyGeneratedOrdering is collation-aware, and on the hash side Spark's Murmur3Hash hashes the collation key, so JVM partitioning is correct for collated strings either way. supportedSerializableDataType also already accepts collated strings as data columns, so as things stand we reject a type as a partition key that we happily ship as payload. Either that check is dead too, or collation isn't surviving the Arrow round trip and we have a worse bug in the data path. I haven't checked the round trip — could you work out which before deciding to keep it?

Minor last point: spark.comet.exec.strictFloatingPoint defaults to false, so the struct-with-double repro needs an explicit opt-in. That's fine for a test, but the case that bites at default config is the second bullet in the description, an ordering expression with no serde at all — repartitionByRange($"a" + 1) and ORDER BY f(a) both put arbitrary expressions in RangePartitioning. Worth covering that shape too.

[Visorgood]

Thanks @andygrove! All four points addressed:

1. Celeborn. Confirmed, and the guard is stronger than dependency-level: getWriter reaches the partitioning only through nativeDependency, which requires shuffleType == CometNativeShuffle, and the None branch calls rejectCometHandle, which throws for both columnar handles. Comet columnar shuffle over Celeborn is unsupported outright. Added to the chain.

2. Hash branch. Agreed, taken in the same PR. Same argument, same conclusion – the columnar path evaluates h.partitionIdExpression through UnsafeProjection and never serializes it. Pinned with a failing-then-passing test.
   Removing it did surface one thing: columnar shuffle on array/struct map key/value expected 0 exchanges on 4.0+ because Spark wraps map keys in mapsort(...). Since the JVM computes partition ids including mapsort, that fallback was also unnecessary; expectation is now 1, with a spark_partition_id() test pinning that the distribution still matches Spark.

3. Collation. Neither half of your dichotomy, I think. I wrote probes for collated hash and range keys and both fell back before shuffle was consulted – CometScanRule rejects collated columns at the scan, with the comment "we don't need specific support for collation in scans, but this is a convenient place to force the whole query to fall back to Spark for now". So no collated data reaches columnar shuffle, the round trip never runs, and there's nothing to measure. The payload-vs-key inconsistency you spotted is real but unobservable for the same reason.
   I've left the checks. They're unreachable rather than dead, and they become meaningful the day that scan-level fallback goes away – at which point the Arrow round-trip question you raised is a real one that deserves its own answer. Happy to remove them if you'd rather, but I couldn't justify it with a test.

4. Default-config test. You're right that the struct-with-double case needs an opt-in. repartitionByRange($"a" + 1) won't do it either – Add has a serde, so the probe passes. I used a Scala UDF with spark.comet.exec.scalaUDF.codegen.enabled=false, where CometScalaUDF.convert returns None by config rather than by absence of support – so the test stays meaningful even if serde coverage grows. One for hash, one for range.

100 tests green across both AQE configurations; a full mvnw test passed 3502 with no failures. Here is the PR: #6110