[AURON #2579] Fix zero-column round-robin shuffle presort

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(

native-engine/datafusion-ext-plans/src/sort_exec.rs

@@ -124,6 +124,13 @@ pub fn create_default_ascending_sort_exec(
     execution_plan_metrics: Option<ExecutionPlanMetricsSet>,
     record_output: bool,
 ) -> Arc<dyn ExecutionPlan> {
+    // 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::<usize>(),
+                    rows.iter().sum::<usize>(),
+                );
+                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<Arc<dyn PhysicalExpr>> = 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::<Int32Array>(),
+                    Some(&Int32Array::from(expected.to_vec())),
+                );
+            }
+            Ok(())
+        }
+    }
 }
 
 #[cfg(test)]