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..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 @@ -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( 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)]