From 2057c3551e77bbc1420711f2b28534dfa7c470a7 Mon Sep 17 00:00:00 2001 From: huiyaoye Date: Mon, 12 Oct 2026 17:49:56 +0800 Subject: [PATCH 1/3] Fix zero-column round-robin shuffle presort COUNT(*) above an unkeyed REPARTITION can leave zero-column batches with positive row counts. Presorting them with an empty key list causes division by zero in ExternalSorter. Skip the default ascending presort when no keys are present. This helper has no limit or offset, so the input rows and repartitioning are preserved. Add regression coverage for zero-column batches, empty input and sorting with and without keys. --- .../datafusion-ext-plans/src/sort_exec.rs | 77 +++++++++++++++++++ 1 file changed, 77 insertions(+) diff --git a/native-engine/datafusion-ext-plans/src/sort_exec.rs b/native-engine/datafusion-ext-plans/src/sort_exec.rs index 6b7bb2b80..9dd30f8c6 100644 --- a/native-engine/datafusion-ext-plans/src/sort_exec.rs +++ b/native-engine/datafusion-ext-plans/src/sort_exec.rs @@ -124,6 +124,13 @@ pub fn create_default_ascending_sort_exec( execution_plan_metrics: Option, record_output: bool, ) -> Arc { + // This helper has no limit or offset, so an empty ordering is a no-op. + // Preserve zero-column batch row counts instead of passing them through + // RowConverter without any columns. + if key_exprs.is_empty() { + return input; + } + let mut sort_exec = SortExec::new( input, key_exprs @@ -1783,6 +1790,76 @@ mod test { Ok(()) } + + mod default_sort_tests { + use arrow::{compute::concat_batches, record_batch::RecordBatchOptions}; + use datafusion::{physical_expr::PhysicalExpr, physical_plan::ExecutionPlanProperties}; + + use super::*; + use crate::sort_exec::create_default_ascending_sort_exec; + + #[tokio::test] + async fn test_default_sort_zero_column_batches() -> Result<()> { + MemManager::init(100); + let schema = Arc::new(Schema::empty()); + let row_counts: &[&[usize]] = &[&[1], &[0, 4, 0, 100001], &[0], &[]]; + let mut partitions = vec![]; + for rows in row_counts { + let mut batches = vec![]; + for &num_rows in *rows { + batches.push(RecordBatch::try_new_with_options( + schema.clone(), + vec![], + &RecordBatchOptions::new().with_row_count(Some(num_rows)), + )?); + } + partitions.push(batches); + } + let input = Arc::new(TestMemoryExec::try_new(&partitions, schema.clone(), None)?); + let sort = create_default_ascending_sort_exec(input, &[], None, true); + assert_eq!(sort.schema(), schema); + assert_eq!( + sort.output_partitioning().partition_count(), + partitions.len() + ); + + let task_ctx = SessionContext::new().task_ctx(); + for (partition, rows) in row_counts.iter().enumerate() { + let batches = common::collect(sort.execute(partition, task_ctx.clone())?).await?; + assert!(batches.iter().all(|batch| batch.num_columns() == 0)); + assert_eq!( + batches.iter().map(RecordBatch::num_rows).sum::(), + rows.iter().sum::(), + ); + if rows.is_empty() { + assert!(batches.is_empty()); + } + } + Ok(()) + } + + #[tokio::test] + async fn test_default_sort_values() -> Result<()> { + MemManager::init(100); + for (with_key, expected) in [(false, [3, 1, 2]), (true, [1, 2, 3])] { + let input = build_single_column_table("id", vec![3, 1, 2])?; + let keys: Vec> = if with_key { + vec![Arc::new(Column::new("id", 0))] + } else { + vec![] + }; + let sort = create_default_ascending_sort_exec(input, &keys, None, true); + let batches = + common::collect(sort.execute(0, SessionContext::new().task_ctx())?).await?; + let batch = concat_batches(&sort.schema(), &batches)?; + assert_eq!( + batch.column(0).as_any().downcast_ref::(), + Some(&Int32Array::from(expected.to_vec())), + ); + } + Ok(()) + } + } } #[cfg(test)] From c3fc9530f5c964642ca1cd4e2f47a65746c1a945 Mon Sep 17 00:00:00 2001 From: huiyaoye Date: Wed, 14 Oct 2026 18:02:16 +0800 Subject: [PATCH 2/3] [AURON #2579] Add Spark regression test for zero-column round-robin shuffle --- .../spark/sql/AuronDataFrameSuite.scala | 25 +++++++++++++++++++ 1 file changed, 25 insertions(+) diff --git a/auron-spark-tests/spark35/src/test/scala/org/apache/spark/sql/AuronDataFrameSuite.scala b/auron-spark-tests/spark35/src/test/scala/org/apache/spark/sql/AuronDataFrameSuite.scala index b5374551d..5474ed613 100644 --- a/auron-spark-tests/spark35/src/test/scala/org/apache/spark/sql/AuronDataFrameSuite.scala +++ b/auron-spark-tests/spark35/src/test/scala/org/apache/spark/sql/AuronDataFrameSuite.scala @@ -21,6 +21,7 @@ import java.io.ByteArrayOutputStream import scala.util.Random import org.apache.spark.sql.catalyst.expressions.{EqualTo, Expression} +import org.apache.spark.sql.catalyst.plans.physical.RoundRobinPartitioning import org.apache.spark.sql.execution.auron.plan.{NativeAggExec, NativeShuffleExchangeExec} import org.apache.spark.sql.execution.exchange.ReusedExchangeExec import org.apache.spark.sql.functions._ @@ -29,6 +30,30 @@ import org.apache.spark.sql.test.SQLTestData.TestData2 class AuronDataFrameSuite extends DataFrameSuite with SparkQueryTestsBase { + testAuron("AURON-2579: count over round-robin repartition with zero-column input") { + withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") { + val df = sql(""" + |SELECT count(*) AS row_count + |FROM ( + | SELECT /*+ REPARTITION(2) */ * + | FROM ( + | SELECT explode(array( + | CAST(0 AS BIGINT), CAST(1 AS BIGINT), + | CAST(2 AS BIGINT), CAST(3 AS BIGINT))) AS id + | ) seed + |) inputs + |""".stripMargin) + + assert(getExecutedPlan(df).exists { + case shuffle: NativeShuffleExchangeExec => + shuffle.outputPartitioning == RoundRobinPartitioning(2) && + shuffle.child.output.isEmpty + case _ => false + }) + checkAnswer(df, Seq(Row(4L))) + } + } + testAuron("repartitionByRange") { val partitionNum = 10 withSQLConf( From e0660d90d6804252d4f07fb07f86803b0c349009 Mon Sep 17 00:00:00 2001 From: huiyaoye Date: Wed, 14 Oct 2026 19:53:38 +0800 Subject: [PATCH 3/3] [AURON #2579] Fix Spark regression test formatting --- .../test/scala/org/apache/spark/sql/AuronDataFrameSuite.scala | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/auron-spark-tests/spark35/src/test/scala/org/apache/spark/sql/AuronDataFrameSuite.scala b/auron-spark-tests/spark35/src/test/scala/org/apache/spark/sql/AuronDataFrameSuite.scala index 5474ed613..76f44b10c 100644 --- a/auron-spark-tests/spark35/src/test/scala/org/apache/spark/sql/AuronDataFrameSuite.scala +++ b/auron-spark-tests/spark35/src/test/scala/org/apache/spark/sql/AuronDataFrameSuite.scala @@ -47,7 +47,7 @@ class AuronDataFrameSuite extends DataFrameSuite with SparkQueryTestsBase { assert(getExecutedPlan(df).exists { case shuffle: NativeShuffleExchangeExec => shuffle.outputPartitioning == RoundRobinPartitioning(2) && - shuffle.child.output.isEmpty + shuffle.child.output.isEmpty case _ => false }) checkAnswer(df, Seq(Row(4L)))