This is an automated email from the ASF dual-hosted git repository.
mbutrovich pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datafusion-comet.git
The following commit(s) were added to refs/heads/main by this push:
new 301e45964a fix: register partitioning scalar subqueries for native
shuffle to avoid "Subquery N not found" (#4869)
301e45964a is described below
commit 301e45964aa8524c343238624992dd8dc8c4225b
Author: Matt Butrovich <[email protected]>
AuthorDate: Wed Jul 8 17:46:31 2026 -0400
fix: register partitioning scalar subqueries for native shuffle to avoid
"Subquery N not found" (#4869)
---
.../shuffle/CometShuffleExchangeExec.scala | 21 ++++++++++-
.../org/apache/comet/exec/CometExecSuite.scala | 42 ++++++++++++++++++++++
2 files changed, 62 insertions(+), 1 deletion(-)
diff --git
a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleExchangeExec.scala
b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleExchangeExec.scala
index 565183278d..7629c3e59f 100644
---
a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleExchangeExec.scala
+++
b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleExchangeExec.scala
@@ -763,6 +763,25 @@ object CometShuffleExchangeExec
spec: NativeShuffleSpec): ShuffleDependency[Int, ColumnarBatch,
ColumnarBatch] = {
val numParts = thinRDD.getNumPartitions
+ // Subqueries in the partitioning expressions (e.g. DISTRIBUTE BY over a
subquery) belong to
+ // this exchange, not the native child, so the child's collectSubqueries
misses them. The
+ // writer serializes them with their exprId, so they must be registered
against the iterator or
+ // the native lookup fails with "Subquery N not found". Both the
native-child and
+ // non-native-child native-shuffle paths funnel through here.
+ //
+ // Only ScalarSubquery is matched because it is the sole expression
QueryPlanSerde turns into a
+ // native Subquery proto (and the only id the native side looks up); this
mirrors the other
+ // registration sites (CometExec.collectSubqueries,
CometNativeExec.prepareSubqueries). The
+ // exprId is stable under reuse, so a ReusedSubqueryExec plan still
resolves to the same result.
+ // The `case _ => Nil` fallthrough is safe: partitionings that are not
Expressions
+ // (SinglePartition, RoundRobinPartitioning) carry no key expressions to
hold a subquery.
+ val partitioningSubqueries = outputPartitioning match {
+ case e: Expression => e.collect { case s: ScalarSubquery => s }
+ case _ => Nil
+ }
+ val augmentedSpec = spec.copy(execContext =
+ spec.execContext.copy(subqueries = spec.execContext.subqueries ++
partitioningSubqueries))
+
// The code block below is mostly brought over from
// ShuffleExchangeExec::prepareShuffleDependency
val (partitioner, rangePartitionBounds) = outputPartitioning match {
@@ -831,7 +850,7 @@ object CometShuffleExchangeExec
shuffleWriteMetrics = metrics,
numParts = numParts,
rangePartitionBounds = rangePartitionBounds,
- nativeShuffleSpec = Some(spec))
+ nativeShuffleSpec = Some(augmentedSpec))
}
/**
diff --git a/spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala
b/spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala
index 67a1920a6b..ce6071bf1a 100644
--- a/spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala
+++ b/spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala
@@ -2209,6 +2209,48 @@ class CometExecSuite extends CometTestBase {
}
}
+ // Regression test for https://github.com/apache/datafusion-comet/issues/4787
+ // A scalar subquery inside a RepartitionByExpression (DISTRIBUTE BY) lives
in the shuffle's
+ // partitioning expressions, not the native child subtree, so it must be
registered separately
+ // for the native shuffle writer to resolve it.
+ test("scalar subquery in repartition") {
+ withParquetTable((0 until 10).map(i => (i, i)), "t") {
+ val df = sql("SELECT * FROM t DISTRIBUTE BY (_1 + (SELECT max(_2) FROM
t))")
+ checkSparkAnswerAndOperator(df)
+ }
+ }
+
+ // Same as above but forces a non-native shuffle child
(CometSparkToColumnarExec) by disabling
+ // the native scan and routing the parquet read through Spark-to-Arrow
conversion. This exercises
+ // the prepareShuffleDependency convenience-overload path, which builds its
own NativeExecContext.
+ test("scalar subquery in repartition over non-native child") {
+ withSQLConf(
+ CometConf.COMET_NATIVE_SCAN_ENABLED.key -> "false",
+ CometConf.COMET_CONVERT_FROM_PARQUET_ENABLED.key -> "true",
+ CometConf.COMET_SPARK_TO_ARROW_ENABLED.key -> "true",
+ CometConf.COMET_EXEC_SHUFFLE_ENABLED.key -> "true",
+ CometConf.COMET_SHUFFLE_MODE.key -> "native") {
+ withParquetTable((0 until 10).map(i => (i, i)), "t") {
+ val df = sql("SELECT * FROM t DISTRIBUTE BY (_1 + (SELECT max(_2) FROM
t))")
+ checkSparkAnswer(df)
+ }
+ }
+ }
+
+ // Columnar shuffle computes partition keys on the JVM (UnsafeProjection /
partitionIdExpression),
+ // so the partitioning subquery resolves via updateResult with no native
Subquery serialization.
+ // This confirms the "Subquery N not found" crash is specific to the native
shuffle path.
+ test("scalar subquery in repartition (columnar shuffle)") {
+ withSQLConf(
+ CometConf.COMET_EXEC_SHUFFLE_ENABLED.key -> "true",
+ CometConf.COMET_SHUFFLE_MODE.key -> "jvm") {
+ withParquetTable((0 until 10).map(i => (i, i)), "t") {
+ val df = sql("SELECT * FROM t DISTRIBUTE BY (_1 + (SELECT max(_2) FROM
t))")
+ checkSparkAnswer(df)
+ }
+ }
+ }
+
// Regression test for https://github.com/apache/datafusion-comet/issues/4042
// SPARK-43402 (Spark 4.0+) pushes scalar subqueries into
FileSourceScanExec.dataFilters.
// CometReuseSubquery re-applies subquery deduplication after Comet node
conversions, and
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
