Visorgood opened a new pull request, #6110:
URL: https://github.com/apache/datafusion-comet/pull/6110
## Which issue does this PR close?
Closes #5971.
## Rationale for this change
`columnarShuffleFailureReasons` asked whether Comet could serialize the
partitioning expressions to protobuf, in both the `RangePartitioning` and
`HashPartitioning` branches. Nothing on the columnar path consumes that:
1. `prepareJVMShuffleDependency` partitions on the JVM – `UnsafeProjection`
over `h.partitionIdExpression` / the sort keys, `LazilyGeneratedOrdering`,
Spark's `RangePartitioner`.
2. `CometShuffleDependency.outputPartitioning` is a Catalyst `Partitioning`,
not a proto message.
3. `PartitioningOuterClass.RangePartition` is built only in
`CometNativeShuffleWriter`.
4. That writer is reached only through `CometNativeShuffleHandle`;
`CometShuffleManager` hands the columnar path a different handle.
5. `CometCelebornShuffleManager` also reads the partitioning, but only via
`nativeDependency`, which requires `shuffleType == CometNativeShuffle`;
`rejectCometHandle` throws for both columnar handles.
So both probes rejected exchanges the JVM would have partitioned correctly,
and those queries fell back to Spark's shuffle for no compatibility reason.
This is an unnecessary-fallback bug, not a correctness bug.
The structurally identical probes in the native branch are untouched – there
the serialized expressions really do go native.
## What changes are included in this PR?
Dropped the `exprToProto` probe from both branches of
`columnarShuffleFailureReasons`.
The collation checks stay. They are unreachable rather than merely dead:
`CometScanRule` rejects collated columns outright ("a convenient place to force
the whole query to fall back to Spark for now"), so no collated data reaches
columnar shuffle to begin with. Removing them would drop a guard that becomes
meaningful the day that scan-level fallback goes away, with no way to test the
replacement today.
`inputs` and the `QueryPlanSerde` import are both still used elsewhere in
the method.
## How are these changes tested?
Four new tests in `CometColumnarShuffleSuite`, each confirmed to fail before
the change:
- `range partitioning on a nested floating-point key` – a `struct<double,
int>` sort key, which `CometSortOrder` reports `Incompatible` for under
`strictFloatingPoint` because `strictFloatingPointReason` recurses through
`containsType`.
- `range partitioning on an unserializable expression` and `hash
partitioning on an unserializable expression` – a Scala UDF with
`spark.comet.exec.scalaUDF.codegen.enabled=false`, so `CometScalaUDF.convert`
returns `None`. These reproduce **at default config**, which is the shape that
bites in practice.
- two `partition assignment matches Spark` tests comparing
`spark_partition_id()` per row. `checkShuffleAnswer` only compares the answer,
which is order-insensitive and would pass even if Comet routed rows to
different partitions, so assignment needs its own check (the same reasoning as
`CometNativeShuffleSuite`).
One existing expectation changed: `columnar shuffle on array/struct map
key/value` expected 0 Comet exchanges on Spark 4.0+, because Spark wraps map
shuffle keys in `mapsort(...)` and Comet cannot serialize that for array or
struct map keys. The columnar path computes partition ids on the JVM from
`h.partitionIdExpression`, `mapsort` included, so that verdict never applied to
it. The expectation is now 1, and the new map-key assignment test pins that the
distribution still matches Spark.
`CometExpressionSuite`'s two nested floating-point sort tests asserted the
removed reason string; they now assert the `CometSortOrder` reason, which is
what those tests are actually about – the `Sort` still falls back, only the
exchange no longer does.
Verified on Spark 4.1 with scalastyle and spotless enabled:
```
CometShuffleSuite, DisableAQECometShuffleSuite, CometShuffleManagerSuite
Suites: completed 3, aborted 0
Tests: succeeded 100, failed 0
```
Both AQE configurations, since `checkCometExchange` strips the AQE plan. A
full `mvnw test` run passed 3502 tests with no failures.
--
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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
