This is an automated email from the ASF dual-hosted git repository.

mbutrovich pushed a commit to branch branch-0.17
in repository https://gitbox.apache.org/repos/asf/datafusion-comet.git


The following commit(s) were added to refs/heads/branch-0.17 by this push:
     new 86ec9f2cec chore: [branch-17] backport #4869 (#4875)
86ec9f2cec is described below

commit 86ec9f2cec30cd8685dd2926241cb1c4641b6679
Author: Matt Butrovich <[email protected]>
AuthorDate: Thu Jul 9 15:00:56 2026 -0400

    chore: [branch-17] backport #4869 (#4875)
---
 native/fs-hdfs/src/hdfs.rs                         |  2 +-
 native/spark-expr/src/conversion_funcs/cast.rs     |  2 +-
 .../shuffle/CometShuffleExchangeExec.scala         | 21 ++++++++++-
 .../org/apache/comet/exec/CometExecSuite.scala     | 42 ++++++++++++++++++++++
 4 files changed, 64 insertions(+), 3 deletions(-)

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)
             };
 
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
         )
     }
 }
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


---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]

Reply via email to