From 9748e9cb80c57a9dea17dcc526f55918a700997e Mon Sep 17 00:00:00 2001 From: Matt Butrovich Date: Wed, 8 Jul 2026 17:46:31 -0400 Subject: [PATCH 1/3] fix: register partitioning scalar subqueries for native shuffle to avoid "Subquery N not found" (#4869) (cherry picked from commit 301e45964aa8524c343238624992dd8dc8c4225b) --- .../shuffle/CometShuffleExchangeExec.scala | 21 +++++++++- .../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 cb747d0375..770a7b5068 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 From cb070ddd11d534ed9a327ca9dfc78cd0a6d386a4 Mon Sep 17 00:00:00 2001 From: Matt Butrovich Date: Thu, 9 Jul 2026 10:22:15 -0400 Subject: [PATCH 2/3] fix clippy --- native/fs-hdfs/src/hdfs.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/native/fs-hdfs/src/hdfs.rs b/native/fs-hdfs/src/hdfs.rs index d0f4d63e1d..ba77f5bb82 100644 --- a/native/fs-hdfs/src/hdfs.rs +++ b/native/fs-hdfs/src/hdfs.rs @@ -100,7 +100,7 @@ impl HdfsManager { let hdfs_builder = hdfsNewBuilder(); let cstr_uri = CString::new(namenode_uri.as_bytes()).unwrap(); hdfsBuilderSetNameNode(hdfs_builder, cstr_uri.as_ptr()); - info!("Connecting to Namenode ({})", &namenode_uri); + info!("Connecting to Namenode ({})", namenode_uri); hdfsBuilderConnect(hdfs_builder) }; From d20493068c19d7c17487955760fb9d6a19c04383 Mon Sep 17 00:00:00 2001 From: Matt Butrovich Date: Thu, 9 Jul 2026 10:38:50 -0400 Subject: [PATCH 3/3] fix clippy --- native/spark-expr/src/conversion_funcs/cast.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/native/spark-expr/src/conversion_funcs/cast.rs b/native/spark-expr/src/conversion_funcs/cast.rs index 1f574f1231..939e3221b3 100644 --- a/native/spark-expr/src/conversion_funcs/cast.rs +++ b/native/spark-expr/src/conversion_funcs/cast.rs @@ -734,7 +734,7 @@ impl Display for Cast { write!( f, "Cast [data_type: {}, timezone: {}, child: {}, eval_mode: {:?}]", - self.data_type, self.cast_options.timezone, self.child, &self.cast_options.eval_mode + self.data_type, self.cast_options.timezone, self.child, self.cast_options.eval_mode ) } }