DRAFT Bench 25358

benchmarks/bench.sh

@@ -105,6 +105,8 @@ spill_views:            Sort and GROUP BY queries that spill StringView/BinaryVi
 null_aware_join:        Null-aware (NOT IN) hash join micro-benchmarks: uncorrelated, non-equality-correlated and equality-correlated
                           NOT IN across NULL fractions, to measure the per-pair join-filter work the correlated cases do
                           (data generated inline by the suite's load SQL from range(); knobs: NAJ_ROWS, NAJ_LARGE_ROWS)
+staged_sort:            Automatic staged ORDER BY evaluation across ties, NULLs, expensive keys, and controls
+                          (inline data; knobs: STAGED_SORT_ROWS, STAGED_SORT_PAYLOAD_WIDTH)
 projection_subquery:    IN / NOT IN / EXISTS subqueries in the SELECT list (see https://github.com/apache/datafusion/issues/25341); each query projects the
                           boolean subquery result and aggregates it, so the cost is the decorrelation plan and not the output size
                           (q07 correlates on '<' instead of '=', so it keeps the nested-loop plan and acts as the control)
@@ -276,6 +278,9 @@ main() {
                     # Data is generated inline by the suite's load SQL from range().
                     echo "null_aware_join: no external data to generate"
                     ;;
+                staged_sort)
+                    echo "staged_sort: no external data to generate"
+                    ;;
                 projection_subquery)
                     # Data is generated inline by the suite's load SQL.
                     echo "projection_subquery: no external data to generate"
@@ -533,6 +538,9 @@ main() {
                 null_aware_join)
                     run_null_aware_join
                     ;;
+                staged_sort)
+                    run_staged_sort
+                    ;;
                 projection_subquery)
                     run_projection_subquery
                     ;;
@@ -992,6 +1000,17 @@ run_null_aware_join() {
       bash -c "$SQL_CARGO_COMMAND"
 }
 
+# Inline staged ORDER BY workloads; the suite sets its own sort thresholds.
+run_staged_sort() {
+    echo "Running staged_sort benchmark (rows=${STAGED_SORT_ROWS:-200000}, payload_width=${STAGED_SORT_PAYLOAD_WIDTH:-32})..."
+    debug_run env BENCH_NAME=staged_sort \
+      BENCH_RESULTS_FILE="$(sql_results_file staged_sort)" \
+      STAGED_SORT_ROWS="${STAGED_SORT_ROWS:-200000}" \
+      STAGED_SORT_PAYLOAD_WIDTH="${STAGED_SORT_PAYLOAD_WIDTH:-32}" \
+      ${QUERY:+BENCH_QUERY="${QUERY}"} \
+      bash -c "$SQL_CARGO_COMMAND"
+}
+
 # Runs the projection_subquery suite: IN / NOT IN / EXISTS subqueries that sit
 # in the SELECT list instead of a filter (see
 # https://github.com/apache/datafusion/issues/25341). The load SQL builds the

benchmarks/sql_benchmarks/README.md

@@ -28,6 +28,7 @@ in the community:
 | `push_down_topk`      | `ORDER BY ... LIMIT` over outer joins (TPC-H data); exercises pushing a TopK through a join |
 | `pwmj`                | Piecewise merge join: classic `INNER`/`LEFT`/`RIGHT`/`FULL` joins over `lhs.key < rhs.key`. Subgroups (`--subgroup`) are match regimes: `no_match`, `null_heavy`, `half_match`, `selective`, `all_match`. Generated with `range()`, no data needed. |
 | `smj`                 | Sort‑merge join benchmark                                          |
+| [`staged_sort`](staged_sort/README.md) | Automatic staged ORDER BY evaluation across ties, NULLs, expensive keys, and controls; inline data. |
 | [`sort_pushdown`](sort_pushdown/README.md) | SQL counterparts of the five sort pushdown variants, using existing Parquet data and query files. |
 | `sort tpch`           | Sorting benchmarks against the TPC-H lineitem table                |
 | `taxi`                | NYC taxi dataset benchmark                                         |

benchmarks/sql_benchmarks/staged_sort/README.md

@@ -0,0 +1,83 @@
+<!--
+SPDX-License-Identifier: Apache-2.0
+SPDX-FileCopyrightText: See the NOTICE file distributed with this work for additional information regarding copyright ownership
+SPDX-FileContributor: Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements
+-->
+
+# Staged sort
+
+Full `ORDER BY` can sort cheap prefix keys first and evaluate later, expensive
+keys only for tied rows. The sort automatically groups keys into stages; this
+suite does not set the staged-sort switch or prescribe stage sizes. The same SQL
+can run on `main`, which may not know that switch.
+
+| Query | Workload                                                                                          |
+| ----- | ------------------------------------------------------------------------------------------------- |
+| 01    | Unique `id` prefix; `lower` and `md5` can be skipped.                                             |
+| 02    | One sparse tie group (~10% of rows); only tied rows reach `md5`.                                 |
+| 03    | One prefix group; all rows reach `md5` (fallback).                                                |
+| 04    | Sparse ties reach `lower(payload)`, which resolves them and skips `md5`.                          |
+| 05    | Nullable prefix, descending with NULLS FIRST; NULLs tie.                                          |
+| 06    | Cheap keys only; control with no expensive stage.                                                 |
+| 07    | TopK (`LIMIT 100`); eager fetch-path control.                                                     |
+| 08    | Same keys as 02, but multiple sorted runs and merge.                                              |
+
+Each query builds a deterministic in-memory table during `load`, outside the
+timed region. Defaults: `STAGED_SORT_ROWS=200000` (minimum 100),
+`STAGED_SORT_PAYLOAD_WIDTH=32` filler characters appended to the decimal ID.
+IDs arrive in descending order rather than already sorted. Fixture assertions
+check row count, unique IDs, 100 prefix values, the all-tied key, NULL/non-NULL
+rows and minimum payload width. They also verify more than one zero sparse key
+and a majority of nonzero keys equal to unique IDs: exactly one tied group, with
+all other rows resolved. The zero group contains roughly 10% of the rows. A final unique ID resolves
+any remaining ties. `SELECT *` keeps computed sort keys in `SortExec` rather
+than projecting them first; untimed `expect_plan` checks this placement.
+
+Each query sets `target_partitions=1`. Queries 01–07 set the in-place sort
+threshold to 1 GiB, intended to give one run at defaults without a memory limit;
+08 sets it to zero to exercise multiple runs. Very large row/width overrides or
+memory limits can change this behavior. Check plans and execution metrics with
+`--debug` or `EXPLAIN ANALYZE`; debug output affects the first run's timing.
+The timed region includes scanning, sorting, merging where applicable, and
+collecting the full output, not table creation or fixture assertions.
+
+From `benchmarks/`:
+
+```bash
+./bench.sh data staged_sort  # no external data
+./bench.sh run staged_sort
+cargo run --release --bin benchmark_runner -- staged_sort --rows 200000 --payload-width 32 --dry-run
+
+# Same branch, explicit A/B; use matching knobs and machine conditions.
+DATAFUSION_EXECUTION_ENABLE_STAGED_SORT=false cargo run --release --bin benchmark_runner -- staged_sort
+DATAFUSION_EXECUTION_ENABLE_STAGED_SORT=true cargo run --release --bin benchmark_runner -- staged_sort
+```
+
+For a `main` baseline, copy the suite files into a main worktree and leave
+`DATAFUSION_EXECUTION_ENABLE_STAGED_SORT` unset; do not pass an unknown setting
+to main. Use identical fixture options and settings on both branches.
+
+## Correctness and plan checks
+
+Persist full ordered results from the baseline and validate the candidate against
+the same files. Results can be large; keep them outside the repository. The
+runner's persist/validate modes execute queries, so use small fixtures for a
+smoke check and do not treat those runs as performance measurements:
+
+```bash
+mkdir -p /tmp/staged-sort-results
+export STAGED_SORT_RESULTS_DIR=/tmp/staged-sort-results
+# Baseline branch (or staged sort disabled), from benchmarks/:
+DATAFUSION_EXECUTION_ENABLE_STAGED_SORT=false cargo run --release --bin benchmark_runner -- staged_sort --rows 1000 --result-mode persist
+# Candidate branch, same fixture and results directory:
+DATAFUSION_EXECUTION_ENABLE_STAGED_SORT=true cargo run --release --bin benchmark_runner -- staged_sort --rows 1000 --result-mode validate
+```
+
+For a main baseline omit the environment variable on its persist command.
+`--result-mode validate` compares rows in order. Persist again into a fresh
+location whenever fixture options change. The plan assertion checks the
+computed key on the `SortExec` line, not just the presence of a sort. Inspect
+`--debug`/`EXPLAIN` after planner changes: in particular verify that both
+expensive expressions in 01 and 04 remain in the sort, that 07 is TopK, and
+that 08 exercises the intended run/merge path. Plan shape alone does not prove
+which stages execute; compare results and metrics as well.

benchmarks/sql_benchmarks/staged_sort/benchmarks/q01.benchmark

@@ -0,0 +1,9 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: See the NOTICE file distributed with this work for additional information regarding copyright ownership
+# SPDX-FileContributor: Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements
+
+template sql_benchmarks/staged_sort/staged_sort.benchmark.template
+QPAD=01
+NAME=q01_unique_prefix
+THRESHOLD=1073741824
+PLAN=SortExec: expr=[id@0 ASC NULLS LAST, lower(payload@2)

benchmarks/sql_benchmarks/staged_sort/benchmarks/q02.benchmark

@@ -0,0 +1,9 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: See the NOTICE file distributed with this work for additional information regarding copyright ownership
+# SPDX-FileContributor: Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements
+
+template sql_benchmarks/staged_sort/staged_sort.benchmark.template
+QPAD=02
+NAME=q02_partial_ties
+THRESHOLD=1073741824
+PLAN=SortExec: expr=[sparse_key@4 ASC NULLS LAST, md5(payload@2)

benchmarks/sql_benchmarks/staged_sort/benchmarks/q03.benchmark

@@ -0,0 +1,9 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: See the NOTICE file distributed with this work for additional information regarding copyright ownership
+# SPDX-FileContributor: Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements
+
+template sql_benchmarks/staged_sort/staged_sort.benchmark.template
+QPAD=03
+NAME=q03_all_ties
+THRESHOLD=1073741824
+PLAN=SortExec: expr=[tie_key@3 ASC NULLS LAST, md5(payload@2)

benchmarks/sql_benchmarks/staged_sort/benchmarks/q04.benchmark

@@ -0,0 +1,9 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: See the NOTICE file distributed with this work for additional information regarding copyright ownership
+# SPDX-FileContributor: Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements
+
+template sql_benchmarks/staged_sort/staged_sort.benchmark.template
+QPAD=04
+NAME=q04_multiple_stages
+THRESHOLD=1073741824
+PLAN=SortExec: expr=[sparse_key@4 ASC NULLS LAST, lower(payload@2) ASC NULLS LAST, md5(payload@2)

benchmarks/sql_benchmarks/staged_sort/benchmarks/q05.benchmark

@@ -0,0 +1,9 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: See the NOTICE file distributed with this work for additional information regarding copyright ownership
+# SPDX-FileContributor: Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements
+
+template sql_benchmarks/staged_sort/staged_sort.benchmark.template
+QPAD=05
+NAME=q05_null_ties
+THRESHOLD=1073741824
+PLAN=SortExec: expr=[nullable_key@1 DESC, md5(payload@2)

benchmarks/sql_benchmarks/staged_sort/benchmarks/q06.benchmark

@@ -0,0 +1,9 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: See the NOTICE file distributed with this work for additional information regarding copyright ownership
+# SPDX-FileContributor: Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements
+
+template sql_benchmarks/staged_sort/staged_sort.benchmark.template
+QPAD=06
+NAME=q06_cheap_control
+THRESHOLD=1073741824
+PLAN=SortExec: expr=[nullable_key@1 ASC NULLS LAST, id@0

benchmarks/sql_benchmarks/staged_sort/benchmarks/q07.benchmark

@@ -0,0 +1,9 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: See the NOTICE file distributed with this work for additional information regarding copyright ownership
+# SPDX-FileContributor: Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements
+
+template sql_benchmarks/staged_sort/staged_sort.benchmark.template
+QPAD=07
+NAME=q07_topk_control
+THRESHOLD=1073741824
+PLAN=SortExec: TopK(fetch=100), expr=[id@0 % 100 ASC NULLS LAST, md5(payload@2)

benchmarks/sql_benchmarks/staged_sort/benchmarks/q08.benchmark

@@ -0,0 +1,9 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: See the NOTICE file distributed with this work for additional information regarding copyright ownership
+# SPDX-FileContributor: Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements
+
+template sql_benchmarks/staged_sort/staged_sort.benchmark.template
+QPAD=08
+NAME=q08_multiple_runs
+THRESHOLD=0
+PLAN=SortExec: expr=[sparse_key@4 ASC NULLS LAST, md5(payload@2)

benchmarks/sql_benchmarks/staged_sort/load.sql

@@ -0,0 +1,16 @@
+-- SPDX-License-Identifier: Apache-2.0
+-- SPDX-FileCopyrightText: See the NOTICE file distributed with this work for additional information regarding copyright ownership
+-- SPDX-FileContributor: Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements
+
+-- Set this before CTAS as well as before planning the measured query.
+SET datafusion.execution.target_partitions = 1;
+
+-- Unique IDs also make every ORDER BY deterministic when used as the last key.
+CREATE TABLE bench AS
+SELECT value AS id,
+       CASE WHEN value % 10 = 0 THEN NULL ELSE value % 100 END AS nullable_key,
+       concat(CAST(value AS VARCHAR), repeat('x', ${STAGED_SORT_PAYLOAD_WIDTH:-32})) AS payload,
+       0 AS tie_key,
+       CASE WHEN value % 10 = 0 THEN 0 ELSE value END AS sparse_key
+FROM (SELECT ${STAGED_SORT_ROWS:-200000} - 1 - value AS value
+      FROM range(${STAGED_SORT_ROWS:-200000}));

benchmarks/sql_benchmarks/staged_sort/queries/q01.sql

@@ -0,0 +1,6 @@
+-- SPDX-License-Identifier: Apache-2.0
+-- SPDX-FileCopyrightText: See the NOTICE file distributed with this work for additional information regarding copyright ownership
+-- SPDX-FileContributor: Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements
+
+-- Unique cheap prefix: expensive suffix can be skipped.
+SELECT * FROM bench ORDER BY id, lower(payload), md5(payload);

benchmarks/sql_benchmarks/staged_sort/queries/q02.sql

@@ -0,0 +1,6 @@
+-- SPDX-License-Identifier: Apache-2.0
+-- SPDX-FileCopyrightText: See the NOTICE file distributed with this work for additional information regarding copyright ownership
+-- SPDX-FileContributor: Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements
+
+-- Sparse ties: only the zero-key group reaches the expensive suffix.
+SELECT * FROM bench ORDER BY sparse_key, md5(payload), id;

benchmarks/sql_benchmarks/staged_sort/queries/q03.sql

@@ -0,0 +1,6 @@
+-- SPDX-License-Identifier: Apache-2.0
+-- SPDX-FileCopyrightText: See the NOTICE file distributed with this work for additional information regarding copyright ownership
+-- SPDX-FileContributor: Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements
+
+-- All prefix values tie; suffix evaluation cannot be avoided.
+SELECT * FROM bench ORDER BY tie_key, md5(payload), id;

benchmarks/sql_benchmarks/staged_sort/queries/q04.sql

@@ -0,0 +1,6 @@
+-- SPDX-License-Identifier: Apache-2.0
+-- SPDX-FileCopyrightText: See the NOTICE file distributed with this work for additional information regarding copyright ownership
+-- SPDX-FileContributor: Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements
+
+-- Sparse ties reach lower; unique payloads resolve them before md5.
+SELECT * FROM bench ORDER BY sparse_key, lower(payload), md5(payload), id;

benchmarks/sql_benchmarks/staged_sort/queries/q05.sql

@@ -0,0 +1,6 @@
+-- SPDX-License-Identifier: Apache-2.0
+-- SPDX-FileCopyrightText: See the NOTICE file distributed with this work for additional information regarding copyright ownership
+-- SPDX-FileContributor: Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements
+
+-- NULLs form a tied prefix group; descending direction and NULLS FIRST.
+SELECT * FROM bench ORDER BY nullable_key DESC NULLS FIRST, md5(payload), id;

benchmarks/sql_benchmarks/staged_sort/queries/q06.sql

@@ -0,0 +1,6 @@
+-- SPDX-License-Identifier: Apache-2.0
+-- SPDX-FileCopyrightText: See the NOTICE file distributed with this work for additional information regarding copyright ownership
+-- SPDX-FileContributor: Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements
+
+-- All-cheap control, including a deterministic final key.
+SELECT * FROM bench ORDER BY nullable_key, id;

benchmarks/sql_benchmarks/staged_sort/queries/q07.sql

@@ -0,0 +1,6 @@
+-- SPDX-License-Identifier: Apache-2.0
+-- SPDX-FileCopyrightText: See the NOTICE file distributed with this work for additional information regarding copyright ownership
+-- SPDX-FileContributor: Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements
+
+-- TopK uses the eager fetch path, not the full staged sort.
+SELECT * FROM bench ORDER BY id % 100, md5(payload), id LIMIT 100;

benchmarks/sql_benchmarks/staged_sort/queries/q08.sql

@@ -0,0 +1,6 @@
+-- SPDX-License-Identifier: Apache-2.0
+-- SPDX-FileCopyrightText: See the NOTICE file distributed with this work for additional information regarding copyright ownership
+-- SPDX-FileContributor: Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements
+
+-- Same keys as q02, but threshold zero creates multiple sorted runs.
+SELECT * FROM bench ORDER BY sparse_key, md5(payload), id;

benchmarks/sql_benchmarks/staged_sort/staged_sort.benchmark.template

@@ -0,0 +1,53 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: See the NOTICE file distributed with this work for additional information regarding copyright ownership
+# SPDX-FileContributor: Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements
+
+# The load and assertions are outside the timed region.
+load sql_benchmarks/staged_sort/load.sql
+
+name ${NAME}
+group staged_sort
+
+init
+SET datafusion.execution.target_partitions = 1;
+
+init
+SET datafusion.execution.sort_in_place_threshold_bytes = ${THRESHOLD};
+
+assert I
+SELECT count(*) = ${STAGED_SORT_ROWS:-200000} AND count(DISTINCT id) = count(*) AND count(*) >= 100 FROM bench;
+----
+true
+
+assert I
+SELECT count(*) > 1 AND count(*) < (SELECT count(*) FROM bench) FROM bench WHERE nullable_key IS NULL;
+----
+true
+
+assert I
+SELECT count(DISTINCT tie_key) = 1 AND count(DISTINCT id % 100) = 100 FROM bench;
+----
+true
+
+# Nonzero sparse keys equal unique IDs; zero is the only tied group.
+assert I
+SELECT count(*) FILTER (WHERE sparse_key = 0) > 1
+   AND count(*) FILTER (WHERE sparse_key <> 0) > count(*) / 2
+   AND count(*) FILTER (WHERE sparse_key <> 0 AND sparse_key <> id) = 0
+FROM bench;
+----
+true
+
+assert I
+SELECT min(length(payload)) >= ${STAGED_SORT_PAYLOAD_WIDTH:-32} + 1 FROM bench;
+----
+true
+
+expect_plan ${PLAN}
+
+run sql_benchmarks/staged_sort/queries/q${QPAD}.sql
+
+result ${STAGED_SORT_RESULTS_DIR:-sql_benchmarks/staged_sort/results}/q${QPAD}.csv
+
+cleanup
+DROP TABLE bench;

benchmarks/sql_benchmarks/staged_sort/staged_sort.suite

@@ -0,0 +1,22 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: See the NOTICE file distributed with this work for additional information regarding copyright ownership
+# SPDX-FileContributor: Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements
+
+description = "Automatic staged ORDER BY evaluation across tie densities, expensive keys, NULLs, and controls. Data is generated inline."
+
+[[options]]
+name = "rows"
+short = "r"
+env = "STAGED_SORT_ROWS"
+default = "200000"
+help = "Number of fixture rows (at least 100)."
+
+[[options]]
+name = "payload-width"
+env = "STAGED_SORT_PAYLOAD_WIDTH"
+default = "32"
+help = "Number of filler characters appended to each unique numeric payload."
+
+[[examples]]
+command = "cargo run --release --bin benchmark_runner -- staged_sort --query 1"
+description = "Run the unique-prefix case."