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]