mbutrovich opened a new pull request, #4869: URL: https://github.com/apache/datafusion-comet/pull/4869
## Which issue does this PR close? Closes #4787. ## Rationale for this change A scalar subquery in a `RepartitionByExpression` (e.g. `DISTRIBUTE BY (_1 + (SELECT max(_2) FROM t))`) crashes native execution with `Subquery N not found for plan M`. The subquery lives in the exchange's partitioning expressions, not in the native child subtree. `CometNativeShuffleWriter.buildUnifiedPlan` serializes those partition expressions to the native plan, producing a `Subquery` proto keyed by `exprId.id`. But the subqueries registered against the native iterator come from `collectSubqueries(nativeChild)`, which only walks the child, so the partitioning subquery is serialized into the plan but never registered. The native lookup then fails. Columnar shuffle is unaffected: it computes partition keys on the JVM and resolves the subquery via `updateResult`, with no native serialization. ## What changes are included in this PR? `CometShuffleExchangeExec.prepareNativeShuffleDependency` now collects `ScalarSubquery` expressions from `outputPartitioning` and merges them into the spec's `NativeExecContext.subqueries`, so `CometNativeShuffleWriter` registers them against the iterator. This is the single point both the native-child path and the non-native-child (Spark-to-Arrow) path funnel through. ## How are these changes tested? Three new tests in `CometExecSuite`, each exercising a distinct path: - `scalar subquery in repartition`: native shuffle over a native child (the original repro). - `scalar subquery in repartition over non-native child`: native shuffle over `CometSparkToColumnarExec`, covering the convenience-overload path. - `scalar subquery in repartition (columnar shuffle)`: confirms the JVM shuffle path was never affected. The first two were verified to fail before the fix with `Subquery N not found`, and all three pass after. -- 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]
