fix: stop gating columnar shuffle on native-serde checks it never uses

[Visorgood]

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, and my earlier rationale for keeping them was wrong: they are reachable. CometScanRule only rejects a stored collated column, so _2 COLLATE UTF8_LCASE over a plain-string Parquet column keeps CometNativeScan native and the collation arrives in a Project above it. VALUES reaches the gate too.

They also matter beyond the shuffle. Spark's collation-aware Murmur3Hash / LazilyGeneratedOrdering run on the JVM here, so no native step sees a collated partition key – but rejecting the exchange moves the whole stage off Comet, which is what keeps CometSort away from the collated key. supportedSortType only type-checks single-column sorts, so a multi-column collated sort slips through. With the checks removed, listagg DISTINCT ... COLLATE utf8_lcase returns aabb instead of ab and #1947 regresses.

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). Both also assert one CometShuffleExchangeExec in the Comet run, so a future fallback cannot leave both sides on plain Spark and pass while testing nothing.

One more test, collation introduced above the scan still falls back to Spark's shuffle, passes before and after: it pins the collation guard this PR keeps, which the suite did not cover.

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 102, failed 0

Both AQE configurations, since checkCometExchange strips the AQE plan. A full mvnw test run passed 3502 tests with no failures.

[andygrove]

Now that exchanges that used to fail the probe can go columnar in auto mode, could some TPC-DS/TPC-H plans pick up a Comet exchange where they previously fell back? CI hasn't run the plan-stability suites yet. Once it does, if any goldens move, please regenerate them in this PR with dev/regenerate-golden-files.sh.

Also, #5802 changes the same columnar shuffle on array/struct map key/value test in the other direction. It keeps 0 exchanges on 4.0+ and adds a flag-gated columnar test. Once this lands, that test passes without the flag. Could you and @sam-1112 coordinate which lands first, so the other can drop or adjust its columnar test?

[andygrove]

Thanks for tracing this through so carefully. The rationale really helped. One question about the collation checks you kept. The argument for dropping the exprToProto probes is that partitionIdExpression and the range ordering run on the JVM through UnsafeProjection and LazilyGeneratedOrdering. Doesn't that apply to collated string keys too? Spark's own collation-aware hash and ordering would run there. If there's a native step on the columnar path that a collated partition key reaches, could you point to it in a comment? If not, I think these checks should go too, so the function doesn't keep a guard the PR's own reasoning says is unnecessary.

[Visorgood]

You're right that nothing native sees the key: the partition id comes from Spark's Murmur3Hash / LazilyGeneratedOrdering on the JVM, and both are collation-aware.

But I believe the checks still can't go. Rejecting the exchange is what moves the whole stage off Comet, and that also keeps CometSort off the collated key. supportedSortType only type-checks single-column sorts (QueryPlanSerde.scala:1287), so a multi-column collated sort gets past it.

I removed both checks and ran CometCollationSuite: 4 failures. Three only change the reason string the test pins, since the query still falls back via the sort check. The fourth is a wrong answer:

• listagg DISTINCT with utf8_lcase collation (issue Fix listagg-collation.sql test in Spark 4.0.0 #1947) *** FAILED ***
  !== Spark Answer - 1 == == Comet Answer - 1 ==
  ![ab] [aabb]

The plan keeps CometColumnarExchange and CometSort over a two-column collated sort key, so Comet dedups a/A on raw bytes and #1947 is back.

Kept them, and rewrote the rationale in the description – my "unreachable" claim there was wrong. Also added a test for the fallback, which the suite didn't have. I'll file a separate issue for the single-column limit in supportedSortType; fixing that is what would make this check redundant.

[andygrove]

This helper only compares the Comet run with the Spark run. If a future change makes these queries fall back, both sides are plain Spark and the test still passes. Could the helper also assert one CometShuffleExchangeExec in the Comet run, for example with checkCometExchange(df, 1, false)? The map-key assignment test especially has no other test pinning that exact query to Comet.

[Visorgood]

Agreed, that would have passed while testing nothing. Added checkCometExchange(df, 1, false) at the top of the helper, so both assignment tests now pin the Comet run to one exchange before comparing partition ids.

[sam-1112]

@andygrove @Visorgood Thanks for flagging #5802.

For collation, I agree that JVM columnar shuffle has no native partition-id calculation. However, scan fallback does not make this guard unreachable: a non-default collation can be introduced above a normal scan or come from VALUES. The existing collation suite relies on the shuffle rule as a safety boundary against raw-byte Comet sort or aggregate behavior. I would keep the guard, but revise the rationale rather than describe it as unreachable.

For #5802, I suggest #6110 lands first. I will then rebase #5802 and update its columnar-shuffle tests and docs: with the dispatcher off, #6110 can use JVM columnar shuffle; with it on, #5802 can enable native shuffle.

[Visorgood]

Hey @sam-1112 ! You're right. CometScanRule only rejects a stored collated column, so _2 COLLATE UTF8_LCASE over a plain string column leaves CometNativeScan in place and the collation lands in a project above it – the exchange does reach the gate. VALUES gets there the same way, and CometCollationSuite already covers that.

Rewrote the rationale in the description and added a test for the fallback. It also turns out the check matters beyond the shuffle: with it removed, listagg DISTINCT under utf8_lcase returns aabb instead of ab, because CometSort stays on a two-column collated key. Details in my reply to Andy.

#6110 first works for me, thanks for offering to rebase #5802.

[andygrove]

The collation reasoning in your reply convinced me the two checks in columnarShuffleFailureReasons have to stay. The trouble is that the reason only lives in this thread, so the next person to read those checks will likely conclude they're dead code, as I did. Could you add a short comment at the checks saying they keep the stage off Comet so CometSort never sees a multi-column collated key that supportedSortType lets through?

Two of the new test comments also say something different from what the tests do. The collation test says Comet hashes raw bytes and would misroute rows, but on this path the partition id comes from Spark's collation-aware Murmur3Hash. And the comment above the UDF tests says they reproduce at default config, but they set spark.comet.exec.scalaUDF.codegen.enabled=false, which defaults to true. Could both say what's actually going on?