[squanch-c]

Which issue does this PR close?

Closes #2579

Rationale for this change

COUNT(*) above an unkeyed REPARTITION can leave zero-column input batches
with positive row counts. Native round-robin presort then passes no columns
to RowConverter, which produces zero key rows and causes division by zero
in ExternalSorter.

Minimal reproduction with Auron native shuffle conversion enabled:

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;

Expected result: 4, as returned by the JVM engine and by the same query
without the REPARTITION hint. The original Native execution fails with:

output_with_sender[Shuffle] error:
output_with_sender[Sort] error:
attempt to divide by zero

The captured failing plan includes Project (Output: []) followed by
NativeShuffleExchange (Input: [], RoundRobinPartitioning(2)). This example
reproduced the failure on Spark 3.5.9 using the RSS path. SQL execution with
a rebuilt patched JAR has now been successfully validated on Spark 3.5.9
with Celeborn (RSS).

An empty-key guard existed in the former sorter implementation in #958 and
was removed with that implementation in #1078. This restores the behavior
in the shared presort helper, which has no limit or offset.

What changes are included in this PR?

• Return the input from create_default_ascending_sort_exec when no sort keys
  are present. Both local and RSS round-robin shuffle use this helper.
• Add two regression tests covering zero-column batches, multiple batches and
  partitions, empty input, and values with and without sort keys.

Are there any user-facing changes?

COUNT(*) over native round-robin repartitioning no longer fails when column
pruning leaves zero input columns. No configuration change is required.

How was this patch tested?

On macOS arm64 with nightly-2025-05-09, upstream vendor setup and unchanged
Cargo.lock:

cargo test --locked -p datafusion-ext-plans test_default_sort_ -- --nocapture
2 passed; 0 failed; 189 filtered out

Rust workspace formatting also passed:

cargo fmt --all -- --check
Exit code: 0

Pinned rustfmt and git diff --check passed. An earlier negative-control run
of the zero-column regression test reproduced division by zero with only the
helper guard removed.

Runtime SQL validation passed on Spark 3.5.9 with Celeborn (RSS), using a
rebuilt Auron 9.0.0-SNAPSHOT JAR with this fix backported. The previously
failing zero-column round-robin SQL now completes successfully.

The SQL reproduction was also added as the AURON-2579 regression test in
AuronDataFrameSuite and passed with Spark 3.5.9 and Scala 2.12 using:

./build/mvn \
  -Ppre,spark-3.5,scala-2.12,spark-tests \
  -DsparkVersion=3.5.9 \
  -pl auron-spark-tests/spark35 \
  test -DskipTests=false \
  -Dsuites=org.apache.spark.sql.AuronDataFrameSuite \
  -Dtests=AURON-2579

The test verifies a local native shuffle with zero input columns and
RoundRobinPartitioning(2), and checks that COUNT(*) returns 4.

Was this patch authored or co-authored using generative AI tooling?

• Yes
• No

Generated-by: OpenAI Codex (GPT-6; exact client build unavailable)

[squanch-c]

Runtime validation has now been completed on Spark 3.5.9 with Celeborn (RSS), using a rebuilt Auron 9.0.0-SNAPSHOT JAR with this fix backported.
The previously failing zero-column round-robin SQL now completes successfully.

[lyne7-sc]

The fix makes sense to me. could we also add the reproduction sql as a Spark regression test?

[squanch-c]

The fix makes sense to me. could we also add the reproduction sql as a Spark regression test?

Thanks for the review and the suggestion! I’ve added the SQL reproduction as a Spark regression test and verified that it passes with Spark 3.5.9 and Scala 2.12 using:

./build/mvn
-Ppre,spark-3.5,scala-2.12,spark-tests
-DsparkVersion=3.5.9
-pl auron-spark-tests/spark35
test -DskipTests=false
-Dsuites=org.apache.spark.sql.AuronDataFrameSuite
-Dtests=AURON-2579

Could you please take another look when you have a chance?

[slfan1989]

LGTM.

The fix preserves zero-column row counts, and the regression tests cover the failure case. Please update the PR description to reflect the completed SQL validation.

[squanch-c]

LGTM.

The fix preserves zero-column row counts, and the regression tests cover the failure case. Please update the PR description to reflect the completed SQL validation.

Thank you for the feedback. I have updated the PR description accordingly. Could you please review it again and let me know if it does not follow the project’s conventions.

[slfan1989]

@squanch-c Thanks for the contribution! @lyne7-sc Thanks for the review!