chore: [branch-17] backport #4869

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

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

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))
   }
 
   /**

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