From 89554196b7405d303f19584d8ee01522825023fe Mon Sep 17 00:00:00 2001 From: Siew Kam Onn Date: Thu, 5 Nov 2026 18:48:39 +0800 Subject: [PATCH 01/12] Implement staged_sort benchmark with automatic ORDER BY evaluation MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Add staged_sort benchmark with automatic ORDER BY evaluation across ties, NULLs, expensive keys, and controls - Include inline data with knobs: STAGED_SORT_ROWS=200000, STAGED_SORT_PAYLOAD_WIDTH=32 - Update bench.sh to include staged_sort in staged workloads - Document staged_sort as a separate workload with deterministic in-memory table generation - Add staged_sort to README.md with detailed workload descriptions and query variants - Create new README.md section for staged_sort with query variants 01–08 - Define query 01: unique id prefix, optional lower and md5 skips - Define query 02: one sparse tie group (~10% rows), md5 on tied rows - Define query 03: prefix group, all rows reach md5 (fallback) - Define query 04: sparse ties resolve via lower(payload), skips md5 - Define query 05: nullable prefix, descending with NULLS FIRST, NULLs tie - Define query 06: cheap keys only, control with no expensive stage - Define query 07: TopK (LIMIT 100), eager fetch-path control - Define query 08: same keys as 02, multiple sorted runs and merge - Set defaults: STAGED_SORT_ROWS=200000, STAGED_SORT_PAYLOAD_WIDTH=32 - Ensure in-place sort threshold of 1 GiB for queries 01–07, zero for 08 - Use consistent fixture assertions for row count, unique IDs, payload width - Verify plan shape: expensive expressions in 01 and 04 remain, 07 is TopK - Confirm 08 exercises intended run/merge path - Update bench.sh to skip external data for staged_sort - Add debug options via --debug or EXPLAIN ANALYZE for timing insights - Add new benchmark files for q04_multiple_stages, q05_null_ties, q06_cheap_control, q07_topk_control, and q08_multiple_runs - Update template and add QPAD, NAME, and THRESHOLD values for each benchmark - Modify SQL queries in queries/ to reflect new ordering expressions and payloads - Update benchmarks/staged_sort.benchmark.template with load command and group metadata - Ensure consistent data generation in load.sql with deterministic keys and payloads - Maintain test assertions for row counts and NULL handling across all staged sorts - Update staged ORDER BY evaluation logic to handle sparse keys, expensive keys, NULLs, and control conditions - Modify test assertions to validate distinct tie_key and id % 100 conditions - Ensure payload width calculation respects environment variable STAGED_SORT_PAYLOAD_WIDTH - Add inline data generation for fixture rows with at least 100 entries - Adjust suite description to reflect automatic evaluation across tie densities and key types - Set default fixture rows to 200,000 with payload width of 32 - Update command example to reflect unique-prefix case - Confirm cleanup step removes bench table after test run - Maintain test coverage for COUNT, DISTINCT, and filtering logic - Ensure result output path aligns with staged sort directory structure --- benchmarks/bench.sh | 19 +++++ benchmarks/sql_benchmarks/README.md | 1 + .../sql_benchmarks/staged_sort/README.md | 83 +++++++++++++++++++ .../staged_sort/benchmarks/q01.benchmark | 9 ++ .../staged_sort/benchmarks/q02.benchmark | 9 ++ .../staged_sort/benchmarks/q03.benchmark | 9 ++ .../staged_sort/benchmarks/q04.benchmark | 9 ++ .../staged_sort/benchmarks/q05.benchmark | 9 ++ .../staged_sort/benchmarks/q06.benchmark | 9 ++ .../staged_sort/benchmarks/q07.benchmark | 9 ++ .../staged_sort/benchmarks/q08.benchmark | 9 ++ .../sql_benchmarks/staged_sort/load.sql | 16 ++++ .../staged_sort/queries/q01.sql | 6 ++ .../staged_sort/queries/q02.sql | 6 ++ .../staged_sort/queries/q03.sql | 6 ++ .../staged_sort/queries/q04.sql | 6 ++ .../staged_sort/queries/q05.sql | 6 ++ .../staged_sort/queries/q06.sql | 6 ++ .../staged_sort/queries/q07.sql | 6 ++ .../staged_sort/queries/q08.sql | 6 ++ .../staged_sort.benchmark.template | 51 ++++++++++++ .../staged_sort/staged_sort.suite | 22 +++++ 22 files changed, 312 insertions(+) create mode 100644 benchmarks/sql_benchmarks/staged_sort/README.md create mode 100644 benchmarks/sql_benchmarks/staged_sort/benchmarks/q01.benchmark create mode 100644 benchmarks/sql_benchmarks/staged_sort/benchmarks/q02.benchmark create mode 100644 benchmarks/sql_benchmarks/staged_sort/benchmarks/q03.benchmark create mode 100644 benchmarks/sql_benchmarks/staged_sort/benchmarks/q04.benchmark create mode 100644 benchmarks/sql_benchmarks/staged_sort/benchmarks/q05.benchmark create mode 100644 benchmarks/sql_benchmarks/staged_sort/benchmarks/q06.benchmark create mode 100644 benchmarks/sql_benchmarks/staged_sort/benchmarks/q07.benchmark create mode 100644 benchmarks/sql_benchmarks/staged_sort/benchmarks/q08.benchmark create mode 100644 benchmarks/sql_benchmarks/staged_sort/load.sql create mode 100644 benchmarks/sql_benchmarks/staged_sort/queries/q01.sql create mode 100644 benchmarks/sql_benchmarks/staged_sort/queries/q02.sql create mode 100644 benchmarks/sql_benchmarks/staged_sort/queries/q03.sql create mode 100644 benchmarks/sql_benchmarks/staged_sort/queries/q04.sql create mode 100644 benchmarks/sql_benchmarks/staged_sort/queries/q05.sql create mode 100644 benchmarks/sql_benchmarks/staged_sort/queries/q06.sql create mode 100644 benchmarks/sql_benchmarks/staged_sort/queries/q07.sql create mode 100644 benchmarks/sql_benchmarks/staged_sort/queries/q08.sql create mode 100644 benchmarks/sql_benchmarks/staged_sort/staged_sort.benchmark.template create mode 100644 benchmarks/sql_benchmarks/staged_sort/staged_sort.suite diff --git a/benchmarks/bench.sh b/benchmarks/bench.sh index 4d5044857b462..e7c03c747f73e 100755 --- a/benchmarks/bench.sh +++ b/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 diff --git a/benchmarks/sql_benchmarks/README.md b/benchmarks/sql_benchmarks/README.md index 3eb59ce0cc45f..c0c959be39eec 100644 --- a/benchmarks/sql_benchmarks/README.md +++ b/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 | diff --git a/benchmarks/sql_benchmarks/staged_sort/README.md b/benchmarks/sql_benchmarks/staged_sort/README.md new file mode 100644 index 0000000000000..b12fef9dac950 --- /dev/null +++ b/benchmarks/sql_benchmarks/staged_sort/README.md @@ -0,0 +1,83 @@ + + +# 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. diff --git a/benchmarks/sql_benchmarks/staged_sort/benchmarks/q01.benchmark b/benchmarks/sql_benchmarks/staged_sort/benchmarks/q01.benchmark new file mode 100644 index 0000000000000..e325cba8f6c2b --- /dev/null +++ b/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) diff --git a/benchmarks/sql_benchmarks/staged_sort/benchmarks/q02.benchmark b/benchmarks/sql_benchmarks/staged_sort/benchmarks/q02.benchmark new file mode 100644 index 0000000000000..c1fb25640f5b8 --- /dev/null +++ b/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) diff --git a/benchmarks/sql_benchmarks/staged_sort/benchmarks/q03.benchmark b/benchmarks/sql_benchmarks/staged_sort/benchmarks/q03.benchmark new file mode 100644 index 0000000000000..984e1a4042773 --- /dev/null +++ b/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) diff --git a/benchmarks/sql_benchmarks/staged_sort/benchmarks/q04.benchmark b/benchmarks/sql_benchmarks/staged_sort/benchmarks/q04.benchmark new file mode 100644 index 0000000000000..3c9bb8cbb925d --- /dev/null +++ b/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) diff --git a/benchmarks/sql_benchmarks/staged_sort/benchmarks/q05.benchmark b/benchmarks/sql_benchmarks/staged_sort/benchmarks/q05.benchmark new file mode 100644 index 0000000000000..3c3d8f27498cf --- /dev/null +++ b/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) diff --git a/benchmarks/sql_benchmarks/staged_sort/benchmarks/q06.benchmark b/benchmarks/sql_benchmarks/staged_sort/benchmarks/q06.benchmark new file mode 100644 index 0000000000000..e959c43a4ca13 --- /dev/null +++ b/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 diff --git a/benchmarks/sql_benchmarks/staged_sort/benchmarks/q07.benchmark b/benchmarks/sql_benchmarks/staged_sort/benchmarks/q07.benchmark new file mode 100644 index 0000000000000..05d44afeee593 --- /dev/null +++ b/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) diff --git a/benchmarks/sql_benchmarks/staged_sort/benchmarks/q08.benchmark b/benchmarks/sql_benchmarks/staged_sort/benchmarks/q08.benchmark new file mode 100644 index 0000000000000..bb2b62a6c05fe --- /dev/null +++ b/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) diff --git a/benchmarks/sql_benchmarks/staged_sort/load.sql b/benchmarks/sql_benchmarks/staged_sort/load.sql new file mode 100644 index 0000000000000..2bb4e6c8aec50 --- /dev/null +++ b/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})); diff --git a/benchmarks/sql_benchmarks/staged_sort/queries/q01.sql b/benchmarks/sql_benchmarks/staged_sort/queries/q01.sql new file mode 100644 index 0000000000000..a15e806ed3b5e --- /dev/null +++ b/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); diff --git a/benchmarks/sql_benchmarks/staged_sort/queries/q02.sql b/benchmarks/sql_benchmarks/staged_sort/queries/q02.sql new file mode 100644 index 0000000000000..8136ff124f095 --- /dev/null +++ b/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; diff --git a/benchmarks/sql_benchmarks/staged_sort/queries/q03.sql b/benchmarks/sql_benchmarks/staged_sort/queries/q03.sql new file mode 100644 index 0000000000000..8d71185bc7249 --- /dev/null +++ b/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; diff --git a/benchmarks/sql_benchmarks/staged_sort/queries/q04.sql b/benchmarks/sql_benchmarks/staged_sort/queries/q04.sql new file mode 100644 index 0000000000000..b508fd4284b58 --- /dev/null +++ b/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; diff --git a/benchmarks/sql_benchmarks/staged_sort/queries/q05.sql b/benchmarks/sql_benchmarks/staged_sort/queries/q05.sql new file mode 100644 index 0000000000000..100d1b9bf2026 --- /dev/null +++ b/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; diff --git a/benchmarks/sql_benchmarks/staged_sort/queries/q06.sql b/benchmarks/sql_benchmarks/staged_sort/queries/q06.sql new file mode 100644 index 0000000000000..67ea32a266d16 --- /dev/null +++ b/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; diff --git a/benchmarks/sql_benchmarks/staged_sort/queries/q07.sql b/benchmarks/sql_benchmarks/staged_sort/queries/q07.sql new file mode 100644 index 0000000000000..8933873f4dfe3 --- /dev/null +++ b/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; diff --git a/benchmarks/sql_benchmarks/staged_sort/queries/q08.sql b/benchmarks/sql_benchmarks/staged_sort/queries/q08.sql new file mode 100644 index 0000000000000..23d9d5d72f4ba --- /dev/null +++ b/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; diff --git a/benchmarks/sql_benchmarks/staged_sort/staged_sort.benchmark.template b/benchmarks/sql_benchmarks/staged_sort/staged_sort.benchmark.template new file mode 100644 index 0000000000000..66edc5f999bcb --- /dev/null +++ b/benchmarks/sql_benchmarks/staged_sort/staged_sort.benchmark.template @@ -0,0 +1,51 @@ +# 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; +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; diff --git a/benchmarks/sql_benchmarks/staged_sort/staged_sort.suite b/benchmarks/sql_benchmarks/staged_sort/staged_sort.suite new file mode 100644 index 0000000000000..5d3fde5cd2301 --- /dev/null +++ b/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." From 24d41f29a011cbcbc941f0eaf55dfbebaa263fa6 Mon Sep 17 00:00:00 2001 From: Siew Kam Onn Date: Thu, 5 Nov 2026 20:12:30 +0800 Subject: [PATCH 02/12] Update staged sort benchmark template with sort threshold setting - Update the benchmark template to include the sort threshold setting. - Add initialization for sort_in_place_threshold_bytes in the template. --- .../sql_benchmarks/staged_sort/staged_sort.benchmark.template | 2 ++ 1 file changed, 2 insertions(+) diff --git a/benchmarks/sql_benchmarks/staged_sort/staged_sort.benchmark.template b/benchmarks/sql_benchmarks/staged_sort/staged_sort.benchmark.template index 66edc5f999bcb..8265870b9709c 100644 --- a/benchmarks/sql_benchmarks/staged_sort/staged_sort.benchmark.template +++ b/benchmarks/sql_benchmarks/staged_sort/staged_sort.benchmark.template @@ -10,6 +10,8 @@ group staged_sort init SET datafusion.execution.target_partitions = 1; + +init SET datafusion.execution.sort_in_place_threshold_bytes = ${THRESHOLD}; assert I From 64bd3eb2610b35d2ee1ebf94cd4a95a4afc5cade Mon Sep 17 00:00:00 2001 From: yashrb24 Date: Wed, 16 Sep 2026 08:58:00 +0530 Subject: [PATCH 03/12] feat: stage ORDER BY expression evaluation --- datafusion/common/src/config.rs | 12 + datafusion/physical-plan/src/sorts/sort.rs | 205 +++++++++++++++++- datafusion/physical-plan/src/sorts/stream.rs | 12 +- .../test_files/information_schema.slt | 4 + 4 files changed, 221 insertions(+), 12 deletions(-) diff --git a/datafusion/common/src/config.rs b/datafusion/common/src/config.rs index 5106ea51ea6ba..ecf4253095df6 100644 --- a/datafusion/common/src/config.rs +++ b/datafusion/common/src/config.rs @@ -998,6 +998,18 @@ config_namespace! { /// batches and merged. pub sort_in_place_threshold_bytes: usize, default = 1024 * 1024 + /// Enables staged evaluation for eligible final, single-run, in-memory sorts. + /// Later sort keys are evaluated only for rows still tied on earlier keys. + /// Requires a key after the first group that is neither a column nor a literal. + /// Spill preparation, multi-run merges, and TopK sorting continue to use eager + /// key evaluation. This experimental option can suppress expression errors + /// or change volatile expression results by skipping evaluations. + pub enable_staged_sort: bool, default = false + + /// Number of sort keys evaluated together by staged in-memory sorting. + /// Must be greater than zero when `enable_staged_sort` is true. + pub sort_key_group_size: usize, default = 3 + /// Maximum buffer capacity (in bytes) per partition for BufferExec /// inserted during sort pushdown optimization. /// diff --git a/datafusion/physical-plan/src/sorts/sort.rs b/datafusion/physical-plan/src/sorts/sort.rs index 53c2929d4aca9..2776fbf781de9 100644 --- a/datafusion/physical-plan/src/sorts/sort.rs +++ b/datafusion/physical-plan/src/sorts/sort.rs @@ -6,9 +6,11 @@ //! It will do in-memory sorting if it has enough memory budget //! but spills to disk if needed. +use std::cmp::Ordering as CmpOrdering; use std::fmt; use std::fmt::{Debug, Formatter}; use std::future::Future; +use std::ops::Range; use std::sync::{Arc, OnceLock}; use parking_lot::RwLock; @@ -45,9 +47,10 @@ use crate::{ PlanProperties, ReplaceChildrenOptions, SendableRecordBatchStream, Statistics, }; -use arrow::array::{RecordBatch, RecordBatchOptions}; -use arrow::compute::{concat_batches, lexsort_to_indices, take_arrays}; +use arrow::array::{Array, RecordBatch, RecordBatchOptions, UInt32Array}; +use arrow::compute::{SortColumn, concat_batches, lexsort_to_indices, take_arrays}; use arrow::datatypes::{DataType, SchemaRef}; +use arrow_ord::sort::LexicographicalComparator; use datafusion_common::config::SpillCompression; use datafusion_common::tree_node::TreeNodeRecursion; use datafusion_common::utils::{ @@ -65,7 +68,9 @@ use datafusion_execution::memory_pool::{ use datafusion_execution::runtime_env::RuntimeEnv; use datafusion_physical_expr::LexOrdering; use datafusion_physical_expr::PhysicalExpr; -use datafusion_physical_expr::expressions::{DynamicFilterPhysicalExpr, lit}; +use datafusion_physical_expr::expressions::{ + Column, DynamicFilterPhysicalExpr, Literal, lit, +}; use futures::{Stream, StreamExt, TryStreamExt}; use log::{debug, trace}; @@ -260,6 +265,9 @@ struct ExternalSorter { /// the data will be concatenated and sorted in place rather than /// sort/merged. sort_in_place_threshold_bytes: usize, + /// Staged evaluation is restricted to final, single-run in-memory sorts. + enable_staged_sort: bool, + sort_key_group_size: usize, // ======================================================================== // STATE BUFFERS: @@ -373,9 +381,22 @@ impl ExternalSorter { sort_spill_reservation_bytes, sort_in_place_threshold_bytes, sort_concurrency: Arc::new(OnceLock::new()), + enable_staged_sort: false, + sort_key_group_size: 3, }) } + fn with_staged_sort(mut self, enabled: bool, group_size: usize) -> Result { + if enabled && group_size == 0 { + return Err(DataFusionError::Configuration( + "datafusion.execution.sort_key_group_size must be greater than zero when datafusion.execution.enable_staged_sort is true".to_string(), + )); + } + self.enable_staged_sort = enabled; + self.sort_key_group_size = group_size; + Ok(self) + } + /// Appends an unsorted [`RecordBatch`] to `in_mem_batches` /// /// Updates memory usage metrics, and possibly triggers spilling to disk @@ -708,7 +729,7 @@ impl ExternalSorter { if self.in_mem_batches.len() == 1 { let batch = self.in_mem_batches.swap_remove(0); let reservation = self.reservation.take(); - let sorted_stream = self.sort_batch_stream(batch, reservation, None)?; + let sorted_stream = self.sort_batch_stream(batch, reservation, None, is_output_stream && self.use_staged_sort())?; return Ok(self.observe_if_output(sorted_stream, is_output_stream)); } @@ -725,7 +746,7 @@ impl ExternalSorter { .try_resize(get_reserved_bytes_for_record_batch(&batch)?) .map_err(Self::err_with_oom_context)?; let reservation = self.reservation.take(); - let sorted_stream = self.sort_batch_stream(batch, reservation, None)?; + let sorted_stream = self.sort_batch_stream(batch, reservation, None, is_output_stream && self.use_staged_sort())?; return Ok(self.observe_if_output(sorted_stream, is_output_stream)); } @@ -787,7 +808,7 @@ impl ExternalSorter { .enumerate() .map(move |(index, (batch, size))| { let reservation = self.reservation.split(size); - let sort = self.sort_run(batch, reservation, shared_input.clone())?; + let sort = self.sort_run(batch, reservation, shared_input.clone(), false)?; let permits = Arc::clone(&concurrency.permits); Ok::<_, DataFusionError>(async move { let stream = Self::prepare_sort_run(sort, permits).await?; @@ -973,7 +994,7 @@ impl ExternalSorter { .zip(sizes) .map(|(run, size)| { let reservation = self.reservation.split(size); - self.sort_batch_stream(run, reservation, shared_input.clone()) + self.sort_batch_stream(run, reservation, shared_input.clone(), false) }) .collect() } @@ -996,10 +1017,11 @@ impl ExternalSorter { batch: RecordBatch, reservation: MemoryReservation, shared_input: Option>, + staged: bool, ) -> Result { let schema = batch.schema(); let stream = - futures::stream::once(self.sort_run(batch, reservation, shared_input)?) + futures::stream::once(self.sort_run(batch, reservation, shared_input, staged)?) .try_flatten(); Ok(Box::pin(RecordBatchStreamAdapter::new(schema, stream))) } @@ -1013,6 +1035,7 @@ impl ExternalSorter { batch: RecordBatch, reservation: MemoryReservation, shared_input: Option>, + staged: bool, ) -> Result> + Send + use<>> { let reserved_bytes = get_reserved_bytes_for_record_batch(&batch)?; @@ -1024,6 +1047,7 @@ impl ExternalSorter { let expressions = self.expr.clone(); let batch_size = self.batch_size; + let sort_key_group_size = self.sort_key_group_size; let merge_pool = Arc::clone(&self.merge_pool); let elapsed_compute = self.metrics.baseline.elapsed_compute().clone(); @@ -1036,7 +1060,16 @@ impl ExternalSorter { let schema = batch.schema(); // Sort the batch immediately and get all output batches - let sorted_batches = sort_batch_chunked(&batch, &expressions, batch_size)?; + let sorted_batches = if staged { + staged_sort_batch_chunked( + &batch, + &expressions, + batch_size, + sort_key_group_size, + )? + } else { + sort_batch_chunked(&batch, &expressions, batch_size)? + }; // Chunked output can retain shared buffers in every batch and // exceed the input estimate. Borrow only already-reserved spill @@ -1078,6 +1111,17 @@ impl ExternalSorter { }) } + fn use_staged_sort(&self) -> bool { + // The first group is always evaluated for every row. Only computation + // in later groups can be avoided. + self.enable_staged_sort + && self + .expr + .iter() + .skip(self.sort_key_group_size) + .any(|key| !key.expr.is::() && !key.expr.is::()) + } + /// If this sort may spill, pre-allocates /// `sort_spill_reservation_bytes` of memory to guarantee memory /// left for the in memory sort/merge. @@ -1255,6 +1299,143 @@ impl Debug for ExternalSorter { } } +fn staged_sort_batch_chunked( + batch: &RecordBatch, + expressions: &LexOrdering, + batch_size: usize, + group_size: usize, +) -> Result> { + let indices = staged_sort_indices(batch, expressions, group_size)?; + IncrementalSortIterator::new(batch.clone(), expressions.clone(), batch_size) + .with_sorted_indices(indices) + .collect() +} + +/// Compute a complete permutation by evaluating successive key groups only for ties. +fn staged_sort_indices( + batch: &RecordBatch, + expressions: &LexOrdering, + group_size: usize, +) -> Result { + let mut indices = Vec::new(); + // Ranges address positions in `indices`, whose values address the input batch. + let mut unresolved: Vec<_> = std::iter::once(0..batch.num_rows()).collect(); + let stage_count = expressions.len().div_ceil(group_size); + for (stage, keys) in expressions.chunks(group_size).enumerate() { + if unresolved.is_empty() { + break; + } + unresolved = refine_tied_groups( + batch, + keys, + &mut indices, + &unresolved, + stage == 0, + stage + 1 < stage_count, + )?; + } + + Ok(UInt32Array::from(indices)) +} + +/// Evaluate keys once for all unresolved rows, then sort each tied group independently. +fn refine_tied_groups( + batch: &RecordBatch, + keys: &[PhysicalSortExpr], + indices: &mut Vec, + unresolved: &[Range], + first_stage: bool, + find_ties: bool, +) -> Result>> { + let selected = if first_stage { + batch.clone() + } else { + // Gather the original row indices for all tied groups into one buffer. + let mut selection = Vec::with_capacity(unresolved.iter().map(Range::len).sum()); + for range in unresolved { + selection.extend_from_slice(&indices[range.clone()]); + } + arrow::compute::take_record_batch(batch, &UInt32Array::from(selection))? + }; + let sort_columns = keys + .iter() + .map(|expr| expr.evaluate_to_sort_column(&selected)) + .collect::>>()?; + + // Key arrays are already evaluated above. This comparator finds ties that + // need the next stage; the final stage does not need it. + let comparator = find_ties + .then(|| LexicographicalComparator::try_new(&sort_columns)) + .transpose()?; + + let mut next_unresolved = Vec::new(); + if first_stage { + let order = lexsort_to_indices(&sort_columns, None)?; + if let Some(comparator) = &comparator { + append_tied_ranges(comparator, &order, 0, 0, &mut next_unresolved); + } + let (_, values, _) = order.into_parts(); + *indices = values.into(); + return Ok(next_unresolved); + } + + let mut selected_offset = 0; + for range in unresolved { + let columns = sort_columns + .iter() + .map(|column| SortColumn { + values: column.values.slice(selected_offset, range.len()), + options: column.options, + }) + .collect::>(); + let local_order = lexsort_to_indices(&columns, None)?; + let previous = indices[range.clone()].to_vec(); + for (output, local) in indices[range.clone()].iter_mut().zip(local_order.values()) + { + *output = previous[*local as usize]; + } + if let Some(comparator) = &comparator { + append_tied_ranges( + comparator, + &local_order, + selected_offset, + range.start, + &mut next_unresolved, + ); + } + selected_offset += range.len(); + } + Ok(next_unresolved) +} + +/// Find consecutive equal keys in a sorted group; only these rows need later keys. +/// For sorted keys `[a, a, b]` at output offset 5, append the range `5..7`. +/// `selected_offset` locates this group's key values in the gathered batch; +/// `output_offset` locates its positions in the complete output permutation. +fn append_tied_ranges( + comparator: &LexicographicalComparator, + order: &UInt32Array, + selected_offset: usize, + output_offset: usize, + ties: &mut Vec>, +) { + let mut start = 0; + while start < order.len() { + let first = selected_offset + order.value(start) as usize; + let mut end = start + 1; + while end < order.len() + && comparator.compare(first, selected_offset + order.value(end) as usize) + == CmpOrdering::Equal + { + end += 1; + } + if end - start > 1 { + ties.push((output_offset + start)..(output_offset + end)); + } + start = end; + } +} + pub fn sort_batch( batch: &RecordBatch, expressions: &LexOrdering, @@ -1839,6 +2020,10 @@ impl ExecutionPlan for SortExec { context.session_config().spill_compression(), &self.metrics_set, context.runtime_env(), + )? + .with_staged_sort( + execution_options.enable_staged_sort, + execution_options.sort_key_group_size, )?; sorter.sort_concurrency = Arc::clone(&self.sort_concurrency); Ok(Box::pin(RecordBatchStreamAdapter::new( @@ -2557,7 +2742,7 @@ mod tests { let reservation = sorter.reservation.new_empty(); reservation.try_grow(get_reserved_bytes_for_record_batch(&batch)?)?; let before = elapsed_compute.value(); - let sort = sorter.sort_run(batch.clone(), reservation, None)?; + let sort = sorter.sort_run(batch.clone(), reservation, None, false)?; assert_eq!(elapsed_compute.value(), before); let stream = if spawned { diff --git a/datafusion/physical-plan/src/sorts/stream.rs b/datafusion/physical-plan/src/sorts/stream.rs index bdd3cb16dee6e..c5b806020f305 100644 --- a/datafusion/physical-plan/src/sorts/stream.rs +++ b/datafusion/physical-plan/src/sorts/stream.rs @@ -327,8 +327,9 @@ impl PartitionedStream for FieldCursorStream { /// spill when there is not enough memory for an eager sort (which requires ~2x /// peak memory to hold both the unsorted and sorted copies simultaneously). /// -/// On the first call to `next()`, a sorted index array (`UInt32Array`) is -/// computed via `lexsort_to_indices`. Subsequent calls yield chunks of +/// Unless supplied via [`Self::with_sorted_indices`], a sorted index array +/// (`UInt32Array`) is computed via `lexsort_to_indices` on the first call to +/// `next()`. The iterator yields chunks of /// `batch_size` rows by `take`-ing from the original batch using slices of /// this index array. Each `take` copies data for the chunk (not zero-copy), /// but only one chunk is live at a time since the caller consumes it before @@ -359,6 +360,13 @@ impl IncrementalSortIterator { indices: None, } } + + /// Use a precomputed sort permutation instead of evaluating the expressions. + /// Must be called before iteration; indices must contain every input row once. + pub(crate) fn with_sorted_indices(mut self, indices: UInt32Array) -> Self { + self.indices = Some(indices); + self + } } impl Iterator for IncrementalSortIterator { diff --git a/datafusion/sqllogictest/test_files/information_schema.slt b/datafusion/sqllogictest/test_files/information_schema.slt index 4355051ef0254..62135d11835f6 100644 --- a/datafusion/sqllogictest/test_files/information_schema.slt +++ b/datafusion/sqllogictest/test_files/information_schema.slt @@ -213,6 +213,7 @@ datafusion.execution.enable_file_stream_work_stealing true datafusion.execution.enable_migration_aggregate true datafusion.execution.enable_nlj_coordinated_fallback true datafusion.execution.enable_recursive_ctes true +datafusion.execution.enable_staged_sort false datafusion.execution.enforce_batch_size_in_joins false datafusion.execution.hash_join_buffering_capacity 0 datafusion.execution.keep_partition_by_columns false @@ -274,6 +275,7 @@ datafusion.execution.skip_physical_aggregate_schema_check false datafusion.execution.soft_max_bytes_per_output_file 4294967295 datafusion.execution.soft_max_rows_per_output_file 50000000 datafusion.execution.sort_in_place_threshold_bytes 1048576 +datafusion.execution.sort_key_group_size 3 datafusion.execution.sort_pushdown_buffer_capacity 1073741824 datafusion.execution.sort_spill_reservation_bytes 10485760 datafusion.execution.spill_compression uncompressed @@ -380,6 +382,7 @@ datafusion.execution.enable_file_stream_work_stealing true When `true` (the defa datafusion.execution.enable_migration_aggregate true Whether aggregation uses the implementation from the major refactor completed in the 56.0.0 release. When set to `false`, aggregation falls back to the implementation used before 55.0.0. The fallback exists only as a workaround for bugs in the new implementation and will be removed, together with this option, after the 56.0.0 release. See for details. datafusion.execution.enable_nlj_coordinated_fallback true Enables the memory-limited fallback for `NestedLoopJoinExec` join types that emit unmatched left rows in the final output (LEFT, LEFT SEMI, LEFT ANTI, LEFT MARK, FULL) when the right side has multiple partitions. This fallback shares left-side state (the current left chunk, the visited bitmap and the probe-thread counter) across all right-side partitions, which assumes every partition runs in the same process. Distributed engines that execute each output partition as an independent task (e.g. Ballista, datafusion-distributed) give each task its own copy of this state and poll only one partition, so the cross-partition counter never reaches zero and the fallback would stall. Such engines should set this to `false`: the coordinated fallback is then disabled for left-emitting multi-partition joins, which instead fail with a resource-exhaustion error under memory pressure rather than deadlocking. Single-partition and non-left-emitting joins are unaffected and always keep the fallback. datafusion.execution.enable_recursive_ctes true Should DataFusion support recursive CTEs +datafusion.execution.enable_staged_sort false Enables staged evaluation for eligible final, single-run, in-memory sorts. Later sort keys are evaluated only for rows still tied on earlier keys. Requires a key after the first group that is neither a column nor a literal. Spill preparation, multi-run merges, and TopK sorting continue to use eager key evaluation. This experimental option can suppress expression errors or change volatile expression results by skipping evaluations. datafusion.execution.enforce_batch_size_in_joins false Should DataFusion enforce batch size in joins or not. By default, DataFusion will not enforce batch size in joins. Enforcing batch size in joins can reduce memory usage when joining large tables with a highly-selective join filter, but is also slightly slower. Note: this option currently only applies to the symmetric hash join. datafusion.execution.hash_join_buffering_capacity 0 How many bytes to buffer in the probe side of hash joins while the build side is concurrently being built. Without this, hash joins will wait until the full materialization of the build side before polling the probe side. This is useful in scenarios where the query is not completely CPU bounded, allowing to do some early work concurrently and reducing the latency of the query. Note that when hash join buffering is enabled, the probe side will start eagerly polling data, not giving time for the producer side of dynamic filters to produce any meaningful predicate. Queries with dynamic filters might see performance degradation. Disabled by default, set to a number greater than 0 for enabling it. datafusion.execution.keep_partition_by_columns false Should DataFusion keep the columns used for partition_by in the output RecordBatches @@ -441,6 +444,7 @@ datafusion.execution.skip_physical_aggregate_schema_check false When set to true datafusion.execution.soft_max_bytes_per_output_file 4294967295 Target encoded size in bytes of output files when writing multiple. Writers asynchronously report the cumulative encoded size as they process RecordBatches. The final file size may exceed this limit due to batches buffered before the limit is observed, the size of a batch, and file metadata written when the file is finalized. datafusion.execution.soft_max_rows_per_output_file 50000000 Target number of rows in output files when writing multiple. This is a soft max, so it can be exceeded slightly. There also will be one file smaller than the limit if the total number of rows written is not roughly divisible by the soft max datafusion.execution.sort_in_place_threshold_bytes 1048576 When sorting, below what size should data be concatenated and sorted in a single RecordBatch rather than sorted in batches and merged. +datafusion.execution.sort_key_group_size 3 Number of sort keys evaluated together by staged in-memory sorting. Must be greater than zero when `enable_staged_sort` is true. datafusion.execution.sort_pushdown_buffer_capacity 1073741824 Maximum buffer capacity (in bytes) per partition for BufferExec inserted during sort pushdown optimization. When PushdownSort eliminates a SortExec under SortPreservingMergeExec, a BufferExec is inserted to replace SortExec's buffering role. This prevents I/O stalls by allowing the scan to run ahead of the merge. This uses strictly less memory than the SortExec it replaces (which buffers the entire partition). The buffer respects the global memory pool limit. Setting this to a large value is safe — actual memory usage is bounded by partition size and global memory limits. datafusion.execution.sort_spill_reservation_bytes 10485760 Specifies the reserved memory for each spillable sort operation to facilitate an in-memory merge. When a sort operation spills to disk, the in-memory data must be sorted and merged before being written to a file. This setting reserves a specific amount of memory for that in-memory sort/merge process. Note: This setting is irrelevant if the sort operation cannot spill (i.e., if there's no `DiskManager` configured). datafusion.execution.spill_compression uncompressed Sets the compression codec used when spilling data to disk. Since datafusion writes spill files using the Arrow IPC Stream format, only codecs supported by the Arrow IPC Stream Writer are allowed. Valid values are: uncompressed, lz4_frame, zstd. Note: lz4_frame offers faster (de)compression, but typically results in larger spill files. In contrast, zstd achieves higher compression ratios at the cost of slower (de)compression speed. From 35155ce327f1a24895d93c19d6d26c22b68a62e6 Mon Sep 17 00:00:00 2001 From: yashrb24 Date: Wed, 16 Sep 2026 09:06:26 +0530 Subject: [PATCH 04/12] chore: remove staged-sort config docs and snapshot changes --- datafusion/sqllogictest/test_files/information_schema.slt | 4 ---- 1 file changed, 4 deletions(-) diff --git a/datafusion/sqllogictest/test_files/information_schema.slt b/datafusion/sqllogictest/test_files/information_schema.slt index 62135d11835f6..4355051ef0254 100644 --- a/datafusion/sqllogictest/test_files/information_schema.slt +++ b/datafusion/sqllogictest/test_files/information_schema.slt @@ -213,7 +213,6 @@ datafusion.execution.enable_file_stream_work_stealing true datafusion.execution.enable_migration_aggregate true datafusion.execution.enable_nlj_coordinated_fallback true datafusion.execution.enable_recursive_ctes true -datafusion.execution.enable_staged_sort false datafusion.execution.enforce_batch_size_in_joins false datafusion.execution.hash_join_buffering_capacity 0 datafusion.execution.keep_partition_by_columns false @@ -275,7 +274,6 @@ datafusion.execution.skip_physical_aggregate_schema_check false datafusion.execution.soft_max_bytes_per_output_file 4294967295 datafusion.execution.soft_max_rows_per_output_file 50000000 datafusion.execution.sort_in_place_threshold_bytes 1048576 -datafusion.execution.sort_key_group_size 3 datafusion.execution.sort_pushdown_buffer_capacity 1073741824 datafusion.execution.sort_spill_reservation_bytes 10485760 datafusion.execution.spill_compression uncompressed @@ -382,7 +380,6 @@ datafusion.execution.enable_file_stream_work_stealing true When `true` (the defa datafusion.execution.enable_migration_aggregate true Whether aggregation uses the implementation from the major refactor completed in the 56.0.0 release. When set to `false`, aggregation falls back to the implementation used before 55.0.0. The fallback exists only as a workaround for bugs in the new implementation and will be removed, together with this option, after the 56.0.0 release. See for details. datafusion.execution.enable_nlj_coordinated_fallback true Enables the memory-limited fallback for `NestedLoopJoinExec` join types that emit unmatched left rows in the final output (LEFT, LEFT SEMI, LEFT ANTI, LEFT MARK, FULL) when the right side has multiple partitions. This fallback shares left-side state (the current left chunk, the visited bitmap and the probe-thread counter) across all right-side partitions, which assumes every partition runs in the same process. Distributed engines that execute each output partition as an independent task (e.g. Ballista, datafusion-distributed) give each task its own copy of this state and poll only one partition, so the cross-partition counter never reaches zero and the fallback would stall. Such engines should set this to `false`: the coordinated fallback is then disabled for left-emitting multi-partition joins, which instead fail with a resource-exhaustion error under memory pressure rather than deadlocking. Single-partition and non-left-emitting joins are unaffected and always keep the fallback. datafusion.execution.enable_recursive_ctes true Should DataFusion support recursive CTEs -datafusion.execution.enable_staged_sort false Enables staged evaluation for eligible final, single-run, in-memory sorts. Later sort keys are evaluated only for rows still tied on earlier keys. Requires a key after the first group that is neither a column nor a literal. Spill preparation, multi-run merges, and TopK sorting continue to use eager key evaluation. This experimental option can suppress expression errors or change volatile expression results by skipping evaluations. datafusion.execution.enforce_batch_size_in_joins false Should DataFusion enforce batch size in joins or not. By default, DataFusion will not enforce batch size in joins. Enforcing batch size in joins can reduce memory usage when joining large tables with a highly-selective join filter, but is also slightly slower. Note: this option currently only applies to the symmetric hash join. datafusion.execution.hash_join_buffering_capacity 0 How many bytes to buffer in the probe side of hash joins while the build side is concurrently being built. Without this, hash joins will wait until the full materialization of the build side before polling the probe side. This is useful in scenarios where the query is not completely CPU bounded, allowing to do some early work concurrently and reducing the latency of the query. Note that when hash join buffering is enabled, the probe side will start eagerly polling data, not giving time for the producer side of dynamic filters to produce any meaningful predicate. Queries with dynamic filters might see performance degradation. Disabled by default, set to a number greater than 0 for enabling it. datafusion.execution.keep_partition_by_columns false Should DataFusion keep the columns used for partition_by in the output RecordBatches @@ -444,7 +441,6 @@ datafusion.execution.skip_physical_aggregate_schema_check false When set to true datafusion.execution.soft_max_bytes_per_output_file 4294967295 Target encoded size in bytes of output files when writing multiple. Writers asynchronously report the cumulative encoded size as they process RecordBatches. The final file size may exceed this limit due to batches buffered before the limit is observed, the size of a batch, and file metadata written when the file is finalized. datafusion.execution.soft_max_rows_per_output_file 50000000 Target number of rows in output files when writing multiple. This is a soft max, so it can be exceeded slightly. There also will be one file smaller than the limit if the total number of rows written is not roughly divisible by the soft max datafusion.execution.sort_in_place_threshold_bytes 1048576 When sorting, below what size should data be concatenated and sorted in a single RecordBatch rather than sorted in batches and merged. -datafusion.execution.sort_key_group_size 3 Number of sort keys evaluated together by staged in-memory sorting. Must be greater than zero when `enable_staged_sort` is true. datafusion.execution.sort_pushdown_buffer_capacity 1073741824 Maximum buffer capacity (in bytes) per partition for BufferExec inserted during sort pushdown optimization. When PushdownSort eliminates a SortExec under SortPreservingMergeExec, a BufferExec is inserted to replace SortExec's buffering role. This prevents I/O stalls by allowing the scan to run ahead of the merge. This uses strictly less memory than the SortExec it replaces (which buffers the entire partition). The buffer respects the global memory pool limit. Setting this to a large value is safe — actual memory usage is bounded by partition size and global memory limits. datafusion.execution.sort_spill_reservation_bytes 10485760 Specifies the reserved memory for each spillable sort operation to facilitate an in-memory merge. When a sort operation spills to disk, the in-memory data must be sorted and merged before being written to a file. This setting reserves a specific amount of memory for that in-memory sort/merge process. Note: This setting is irrelevant if the sort operation cannot spill (i.e., if there's no `DiskManager` configured). datafusion.execution.spill_compression uncompressed Sets the compression codec used when spilling data to disk. Since datafusion writes spill files using the Arrow IPC Stream format, only codecs supported by the Arrow IPC Stream Writer are allowed. Valid values are: uncompressed, lz4_frame, zstd. Note: lz4_frame offers faster (de)compression, but typically results in larger spill files. In contrast, zstd achieves higher compression ratios at the cost of slower (de)compression speed. From c6de09688876694f1844e35b46fcf2bb9718727f Mon Sep 17 00:00:00 2001 From: yashrb24 Date: Wed, 16 Sep 2026 09:17:14 +0530 Subject: [PATCH 05/12] chore: simplify staged-sort comments --- datafusion/common/src/config.rs | 9 ++------- datafusion/physical-plan/src/sorts/sort.rs | 5 +---- datafusion/physical-plan/src/sorts/stream.rs | 2 -- 3 files changed, 3 insertions(+), 13 deletions(-) diff --git a/datafusion/common/src/config.rs b/datafusion/common/src/config.rs index ecf4253095df6..cb73f978f63cc 100644 --- a/datafusion/common/src/config.rs +++ b/datafusion/common/src/config.rs @@ -998,16 +998,11 @@ config_namespace! { /// batches and merged. pub sort_in_place_threshold_bytes: usize, default = 1024 * 1024 - /// Enables staged evaluation for eligible final, single-run, in-memory sorts. - /// Later sort keys are evaluated only for rows still tied on earlier keys. - /// Requires a key after the first group that is neither a column nor a literal. - /// Spill preparation, multi-run merges, and TopK sorting continue to use eager - /// key evaluation. This experimental option can suppress expression errors - /// or change volatile expression results by skipping evaluations. + /// Enables staged in-memory sorting, evaluating later sort keys only for rows + /// tied on earlier keys. pub enable_staged_sort: bool, default = false /// Number of sort keys evaluated together by staged in-memory sorting. - /// Must be greater than zero when `enable_staged_sort` is true. pub sort_key_group_size: usize, default = 3 /// Maximum buffer capacity (in bytes) per partition for BufferExec diff --git a/datafusion/physical-plan/src/sorts/sort.rs b/datafusion/physical-plan/src/sorts/sort.rs index 2776fbf781de9..ae7879472a6e9 100644 --- a/datafusion/physical-plan/src/sorts/sort.rs +++ b/datafusion/physical-plan/src/sorts/sort.rs @@ -265,7 +265,6 @@ struct ExternalSorter { /// the data will be concatenated and sorted in place rather than /// sort/merged. sort_in_place_threshold_bytes: usize, - /// Staged evaluation is restricted to final, single-run in-memory sorts. enable_staged_sort: bool, sort_key_group_size: usize, @@ -1363,7 +1362,7 @@ fn refine_tied_groups( .collect::>>()?; // Key arrays are already evaluated above. This comparator finds ties that - // need the next stage; the final stage does not need it. + // need the next stage. let comparator = find_ties .then(|| LexicographicalComparator::try_new(&sort_columns)) .transpose()?; @@ -1410,8 +1409,6 @@ fn refine_tied_groups( /// Find consecutive equal keys in a sorted group; only these rows need later keys. /// For sorted keys `[a, a, b]` at output offset 5, append the range `5..7`. -/// `selected_offset` locates this group's key values in the gathered batch; -/// `output_offset` locates its positions in the complete output permutation. fn append_tied_ranges( comparator: &LexicographicalComparator, order: &UInt32Array, diff --git a/datafusion/physical-plan/src/sorts/stream.rs b/datafusion/physical-plan/src/sorts/stream.rs index c5b806020f305..fc19b3307f00a 100644 --- a/datafusion/physical-plan/src/sorts/stream.rs +++ b/datafusion/physical-plan/src/sorts/stream.rs @@ -361,8 +361,6 @@ impl IncrementalSortIterator { } } - /// Use a precomputed sort permutation instead of evaluating the expressions. - /// Must be called before iteration; indices must contain every input row once. pub(crate) fn with_sorted_indices(mut self, indices: UInt32Array) -> Self { self.indices = Some(indices); self From 251f3a4373d1aa933074d2104499d114092752e7 Mon Sep 17 00:00:00 2001 From: yashrb24 Date: Wed, 16 Sep 2026 10:53:13 +0530 Subject: [PATCH 06/12] feat: stage sort keys for in-memory runs --- datafusion/physical-plan/src/sorts/sort.rs | 16 ++++++++++------ 1 file changed, 10 insertions(+), 6 deletions(-) diff --git a/datafusion/physical-plan/src/sorts/sort.rs b/datafusion/physical-plan/src/sorts/sort.rs index ae7879472a6e9..75ea054694469 100644 --- a/datafusion/physical-plan/src/sorts/sort.rs +++ b/datafusion/physical-plan/src/sorts/sort.rs @@ -721,6 +721,8 @@ impl ExternalSorter { let elapsed_compute = self.metrics.baseline.elapsed_compute().clone(); let mut timer = elapsed_compute.timer(); + let staged = is_output_stream && self.use_staged_sort(); + // Sorting can resize reservations and borrow spill workspace. // Prepared output stays reserved until each batch is read from its run stream. // The consumer of the merged stream reserves its final output. @@ -728,7 +730,7 @@ impl ExternalSorter { if self.in_mem_batches.len() == 1 { let batch = self.in_mem_batches.swap_remove(0); let reservation = self.reservation.take(); - let sorted_stream = self.sort_batch_stream(batch, reservation, None, is_output_stream && self.use_staged_sort())?; + let sorted_stream = self.sort_batch_stream(batch, reservation, None, staged)?; return Ok(self.observe_if_output(sorted_stream, is_output_stream)); } @@ -745,7 +747,7 @@ impl ExternalSorter { .try_resize(get_reserved_bytes_for_record_batch(&batch)?) .map_err(Self::err_with_oom_context)?; let reservation = self.reservation.take(); - let sorted_stream = self.sort_batch_stream(batch, reservation, None, is_output_stream && self.use_staged_sort())?; + let sorted_stream = self.sort_batch_stream(batch, reservation, None, staged)?; return Ok(self.observe_if_output(sorted_stream, is_output_stream)); } @@ -766,14 +768,14 @@ impl ExternalSorter { // Time this task's setup and collection work, excluding waits for jobs. timer.stop(); let streams = self - .prepare_sort_runs(runs, &concurrency) + .prepare_sort_runs(runs, &concurrency, staged) .with_elapsed_compute(elapsed_compute.clone()) .await?; timer.restart(); streams } else { // Outside a multithreaded runtime, sort each run when the merge polls it. - self.sort_run_streams(runs)? + self.sort_run_streams(runs, staged)? }; StreamingMergeBuilder::new() @@ -796,6 +798,7 @@ impl ExternalSorter { &mut self, runs: Vec, concurrency: &SortConcurrency, + staged: bool, ) -> Result> { let run_count = runs.len(); let (sizes, shared_input_size) = self.sort_run_sizes(&runs)?; @@ -807,7 +810,7 @@ impl ExternalSorter { .enumerate() .map(move |(index, (batch, size))| { let reservation = self.reservation.split(size); - let sort = self.sort_run(batch, reservation, shared_input.clone(), false)?; + let sort = self.sort_run(batch, reservation, shared_input.clone(), staged)?; let permits = Arc::clone(&concurrency.permits); Ok::<_, DataFusionError>(async move { let stream = Self::prepare_sort_run(sort, permits).await?; @@ -985,6 +988,7 @@ impl ExternalSorter { fn sort_run_streams( &self, runs: Vec, + staged: bool, ) -> Result> { let (sizes, shared_input_size) = self.sort_run_sizes(&runs)?; let shared_input = @@ -993,7 +997,7 @@ impl ExternalSorter { .zip(sizes) .map(|(run, size)| { let reservation = self.reservation.split(size); - self.sort_batch_stream(run, reservation, shared_input.clone(), false) + self.sort_batch_stream(run, reservation, shared_input.clone(), staged) }) .collect() } From bfaf35f791af2acca4b7ef9d66e724a6400e12a6 Mon Sep 17 00:00:00 2001 From: yashrb24 Date: Wed, 16 Sep 2026 11:05:52 +0530 Subject: [PATCH 07/12] fix: synchronize staged-sort configuration entries --- datafusion/sqllogictest/test_files/information_schema.slt | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/datafusion/sqllogictest/test_files/information_schema.slt b/datafusion/sqllogictest/test_files/information_schema.slt index 4355051ef0254..dcde0b98d4004 100644 --- a/datafusion/sqllogictest/test_files/information_schema.slt +++ b/datafusion/sqllogictest/test_files/information_schema.slt @@ -213,6 +213,7 @@ datafusion.execution.enable_file_stream_work_stealing true datafusion.execution.enable_migration_aggregate true datafusion.execution.enable_nlj_coordinated_fallback true datafusion.execution.enable_recursive_ctes true +datafusion.execution.enable_staged_sort false datafusion.execution.enforce_batch_size_in_joins false datafusion.execution.hash_join_buffering_capacity 0 datafusion.execution.keep_partition_by_columns false @@ -274,6 +275,7 @@ datafusion.execution.skip_physical_aggregate_schema_check false datafusion.execution.soft_max_bytes_per_output_file 4294967295 datafusion.execution.soft_max_rows_per_output_file 50000000 datafusion.execution.sort_in_place_threshold_bytes 1048576 +datafusion.execution.sort_key_group_size 3 datafusion.execution.sort_pushdown_buffer_capacity 1073741824 datafusion.execution.sort_spill_reservation_bytes 10485760 datafusion.execution.spill_compression uncompressed @@ -380,6 +382,7 @@ datafusion.execution.enable_file_stream_work_stealing true When `true` (the defa datafusion.execution.enable_migration_aggregate true Whether aggregation uses the implementation from the major refactor completed in the 56.0.0 release. When set to `false`, aggregation falls back to the implementation used before 55.0.0. The fallback exists only as a workaround for bugs in the new implementation and will be removed, together with this option, after the 56.0.0 release. See for details. datafusion.execution.enable_nlj_coordinated_fallback true Enables the memory-limited fallback for `NestedLoopJoinExec` join types that emit unmatched left rows in the final output (LEFT, LEFT SEMI, LEFT ANTI, LEFT MARK, FULL) when the right side has multiple partitions. This fallback shares left-side state (the current left chunk, the visited bitmap and the probe-thread counter) across all right-side partitions, which assumes every partition runs in the same process. Distributed engines that execute each output partition as an independent task (e.g. Ballista, datafusion-distributed) give each task its own copy of this state and poll only one partition, so the cross-partition counter never reaches zero and the fallback would stall. Such engines should set this to `false`: the coordinated fallback is then disabled for left-emitting multi-partition joins, which instead fail with a resource-exhaustion error under memory pressure rather than deadlocking. Single-partition and non-left-emitting joins are unaffected and always keep the fallback. datafusion.execution.enable_recursive_ctes true Should DataFusion support recursive CTEs +datafusion.execution.enable_staged_sort false Enables staged in-memory sorting, evaluating later sort keys only for rows tied on earlier keys. datafusion.execution.enforce_batch_size_in_joins false Should DataFusion enforce batch size in joins or not. By default, DataFusion will not enforce batch size in joins. Enforcing batch size in joins can reduce memory usage when joining large tables with a highly-selective join filter, but is also slightly slower. Note: this option currently only applies to the symmetric hash join. datafusion.execution.hash_join_buffering_capacity 0 How many bytes to buffer in the probe side of hash joins while the build side is concurrently being built. Without this, hash joins will wait until the full materialization of the build side before polling the probe side. This is useful in scenarios where the query is not completely CPU bounded, allowing to do some early work concurrently and reducing the latency of the query. Note that when hash join buffering is enabled, the probe side will start eagerly polling data, not giving time for the producer side of dynamic filters to produce any meaningful predicate. Queries with dynamic filters might see performance degradation. Disabled by default, set to a number greater than 0 for enabling it. datafusion.execution.keep_partition_by_columns false Should DataFusion keep the columns used for partition_by in the output RecordBatches @@ -441,6 +444,7 @@ datafusion.execution.skip_physical_aggregate_schema_check false When set to true datafusion.execution.soft_max_bytes_per_output_file 4294967295 Target encoded size in bytes of output files when writing multiple. Writers asynchronously report the cumulative encoded size as they process RecordBatches. The final file size may exceed this limit due to batches buffered before the limit is observed, the size of a batch, and file metadata written when the file is finalized. datafusion.execution.soft_max_rows_per_output_file 50000000 Target number of rows in output files when writing multiple. This is a soft max, so it can be exceeded slightly. There also will be one file smaller than the limit if the total number of rows written is not roughly divisible by the soft max datafusion.execution.sort_in_place_threshold_bytes 1048576 When sorting, below what size should data be concatenated and sorted in a single RecordBatch rather than sorted in batches and merged. +datafusion.execution.sort_key_group_size 3 Number of sort keys evaluated together by staged in-memory sorting. datafusion.execution.sort_pushdown_buffer_capacity 1073741824 Maximum buffer capacity (in bytes) per partition for BufferExec inserted during sort pushdown optimization. When PushdownSort eliminates a SortExec under SortPreservingMergeExec, a BufferExec is inserted to replace SortExec's buffering role. This prevents I/O stalls by allowing the scan to run ahead of the merge. This uses strictly less memory than the SortExec it replaces (which buffers the entire partition). The buffer respects the global memory pool limit. Setting this to a large value is safe — actual memory usage is bounded by partition size and global memory limits. datafusion.execution.sort_spill_reservation_bytes 10485760 Specifies the reserved memory for each spillable sort operation to facilitate an in-memory merge. When a sort operation spills to disk, the in-memory data must be sorted and merged before being written to a file. This setting reserves a specific amount of memory for that in-memory sort/merge process. Note: This setting is irrelevant if the sort operation cannot spill (i.e., if there's no `DiskManager` configured). datafusion.execution.spill_compression uncompressed Sets the compression codec used when spilling data to disk. Since datafusion writes spill files using the Arrow IPC Stream format, only codecs supported by the Arrow IPC Stream Writer are allowed. Valid values are: uncompressed, lz4_frame, zstd. Note: lz4_frame offers faster (de)compression, but typically results in larger spill files. In contrast, zstd achieves higher compression ratios at the cost of slower (de)compression speed. From b9eab96c07970b20f6c1f3b52843912c2861c1db Mon Sep 17 00:00:00 2001 From: yashrb24 Date: Sat, 3 Oct 2026 20:04:34 +0530 Subject: [PATCH 08/12] chore: default staged sort key groups to five --- datafusion/common/src/config.rs | 2 +- datafusion/physical-plan/src/sorts/sort.rs | 2 +- datafusion/sqllogictest/test_files/information_schema.slt | 4 ++-- 3 files changed, 4 insertions(+), 4 deletions(-) diff --git a/datafusion/common/src/config.rs b/datafusion/common/src/config.rs index cb73f978f63cc..698448b7f17d5 100644 --- a/datafusion/common/src/config.rs +++ b/datafusion/common/src/config.rs @@ -1003,7 +1003,7 @@ config_namespace! { pub enable_staged_sort: bool, default = false /// Number of sort keys evaluated together by staged in-memory sorting. - pub sort_key_group_size: usize, default = 3 + pub sort_key_group_size: usize, default = 5 /// Maximum buffer capacity (in bytes) per partition for BufferExec /// inserted during sort pushdown optimization. diff --git a/datafusion/physical-plan/src/sorts/sort.rs b/datafusion/physical-plan/src/sorts/sort.rs index 75ea054694469..1e9fafd58fa9f 100644 --- a/datafusion/physical-plan/src/sorts/sort.rs +++ b/datafusion/physical-plan/src/sorts/sort.rs @@ -381,7 +381,7 @@ impl ExternalSorter { sort_in_place_threshold_bytes, sort_concurrency: Arc::new(OnceLock::new()), enable_staged_sort: false, - sort_key_group_size: 3, + sort_key_group_size: 5, }) } diff --git a/datafusion/sqllogictest/test_files/information_schema.slt b/datafusion/sqllogictest/test_files/information_schema.slt index dcde0b98d4004..fae5cd7306d91 100644 --- a/datafusion/sqllogictest/test_files/information_schema.slt +++ b/datafusion/sqllogictest/test_files/information_schema.slt @@ -275,7 +275,7 @@ datafusion.execution.skip_physical_aggregate_schema_check false datafusion.execution.soft_max_bytes_per_output_file 4294967295 datafusion.execution.soft_max_rows_per_output_file 50000000 datafusion.execution.sort_in_place_threshold_bytes 1048576 -datafusion.execution.sort_key_group_size 3 +datafusion.execution.sort_key_group_size 5 datafusion.execution.sort_pushdown_buffer_capacity 1073741824 datafusion.execution.sort_spill_reservation_bytes 10485760 datafusion.execution.spill_compression uncompressed @@ -444,7 +444,7 @@ datafusion.execution.skip_physical_aggregate_schema_check false When set to true datafusion.execution.soft_max_bytes_per_output_file 4294967295 Target encoded size in bytes of output files when writing multiple. Writers asynchronously report the cumulative encoded size as they process RecordBatches. The final file size may exceed this limit due to batches buffered before the limit is observed, the size of a batch, and file metadata written when the file is finalized. datafusion.execution.soft_max_rows_per_output_file 50000000 Target number of rows in output files when writing multiple. This is a soft max, so it can be exceeded slightly. There also will be one file smaller than the limit if the total number of rows written is not roughly divisible by the soft max datafusion.execution.sort_in_place_threshold_bytes 1048576 When sorting, below what size should data be concatenated and sorted in a single RecordBatch rather than sorted in batches and merged. -datafusion.execution.sort_key_group_size 3 Number of sort keys evaluated together by staged in-memory sorting. +datafusion.execution.sort_key_group_size 5 Number of sort keys evaluated together by staged in-memory sorting. datafusion.execution.sort_pushdown_buffer_capacity 1073741824 Maximum buffer capacity (in bytes) per partition for BufferExec inserted during sort pushdown optimization. When PushdownSort eliminates a SortExec under SortPreservingMergeExec, a BufferExec is inserted to replace SortExec's buffering role. This prevents I/O stalls by allowing the scan to run ahead of the merge. This uses strictly less memory than the SortExec it replaces (which buffers the entire partition). The buffer respects the global memory pool limit. Setting this to a large value is safe — actual memory usage is bounded by partition size and global memory limits. datafusion.execution.sort_spill_reservation_bytes 10485760 Specifies the reserved memory for each spillable sort operation to facilitate an in-memory merge. When a sort operation spills to disk, the in-memory data must be sorted and merged before being written to a file. This setting reserves a specific amount of memory for that in-memory sort/merge process. Note: This setting is irrelevant if the sort operation cannot spill (i.e., if there's no `DiskManager` configured). datafusion.execution.spill_compression uncompressed Sets the compression codec used when spilling data to disk. Since datafusion writes spill files using the Arrow IPC Stream format, only codecs supported by the Arrow IPC Stream Writer are allowed. Valid values are: uncompressed, lz4_frame, zstd. Note: lz4_frame offers faster (de)compression, but typically results in larger spill files. In contrast, zstd achieves higher compression ratios at the cost of slower (de)compression speed. From 35dfbb0659de196633f57d2871a29661ef49badc Mon Sep 17 00:00:00 2001 From: yashrb24 Date: Tue, 13 Oct 2026 00:34:46 +0530 Subject: [PATCH 09/12] perf: cache staged sort projections and add focused tests --- datafusion/physical-plan/src/sorts/sort.rs | 386 ++++++++++++++++++--- 1 file changed, 344 insertions(+), 42 deletions(-) diff --git a/datafusion/physical-plan/src/sorts/sort.rs b/datafusion/physical-plan/src/sorts/sort.rs index 1e9fafd58fa9f..2311165543829 100644 --- a/datafusion/physical-plan/src/sorts/sort.rs +++ b/datafusion/physical-plan/src/sorts/sort.rs @@ -71,6 +71,8 @@ use datafusion_physical_expr::PhysicalExpr; use datafusion_physical_expr::expressions::{ Column, DynamicFilterPhysicalExpr, Literal, lit, }; +use datafusion_physical_expr::projection::ProjectionExprs; +use datafusion_physical_expr::utils::collect_columns; use futures::{Stream, StreamExt, TryStreamExt}; use log::{debug, trace}; @@ -265,8 +267,7 @@ struct ExternalSorter { /// the data will be concatenated and sorted in place rather than /// sort/merged. sort_in_place_threshold_bytes: usize, - enable_staged_sort: bool, - sort_key_group_size: usize, + sort_stages: Option>, // ======================================================================== // STATE BUFFERS: @@ -380,8 +381,7 @@ impl ExternalSorter { sort_spill_reservation_bytes, sort_in_place_threshold_bytes, sort_concurrency: Arc::new(OnceLock::new()), - enable_staged_sort: false, - sort_key_group_size: 5, + sort_stages: None, }) } @@ -391,8 +391,52 @@ impl ExternalSorter { "datafusion.execution.sort_key_group_size must be greater than zero when datafusion.execution.enable_staged_sort is true".to_string(), )); } - self.enable_staged_sort = enabled; - self.sort_key_group_size = group_size; + self.sort_stages = if enabled + && self + .expr + .iter() + .skip(group_size) + .any(|key| !key.expr.is::() && !key.expr.is::()) + { + let stages = self + .expr + .chunks(group_size) + .enumerate() + .map(|(index, keys)| { + if index == 0 { + return Ok(SortStage { + columns: vec![], + expressions: keys.to_vec(), + }); + } + let mut columns = keys + .iter() + .flat_map(|key| collect_columns(&key.expr)) + .map(|column| column.index()) + .collect::>(); + columns.sort_unstable(); + columns.dedup(); + let projection = + ProjectionExprs::from_indices(&columns, &self.schema); + let expressions = keys + .iter() + .map(|key| { + Ok(PhysicalSortExpr::new( + projection.project_expr(&key.expr)?, + key.options, + )) + }) + .collect::>>()?; + Ok(SortStage { + columns, + expressions, + }) + }) + .collect::>>()?; + Some(stages.into()) + } else { + None + }; Ok(self) } @@ -721,7 +765,7 @@ impl ExternalSorter { let elapsed_compute = self.metrics.baseline.elapsed_compute().clone(); let mut timer = elapsed_compute.timer(); - let staged = is_output_stream && self.use_staged_sort(); + let staged = is_output_stream && self.sort_stages.is_some(); // Sorting can resize reservations and borrow spill workspace. // Prepared output stays reserved until each batch is read from its run stream. @@ -1050,7 +1094,11 @@ impl ExternalSorter { let expressions = self.expr.clone(); let batch_size = self.batch_size; - let sort_key_group_size = self.sort_key_group_size; + let stages = if staged { + self.sort_stages.clone() + } else { + None + }; let merge_pool = Arc::clone(&self.merge_pool); let elapsed_compute = self.metrics.baseline.elapsed_compute().clone(); @@ -1063,13 +1111,8 @@ impl ExternalSorter { let schema = batch.schema(); // Sort the batch immediately and get all output batches - let sorted_batches = if staged { - staged_sort_batch_chunked( - &batch, - &expressions, - batch_size, - sort_key_group_size, - )? + let sorted_batches = if let Some(stages) = stages { + staged_sort_batch_chunked(&batch, &expressions, batch_size, &stages)? } else { sort_batch_chunked(&batch, &expressions, batch_size)? }; @@ -1114,17 +1157,6 @@ impl ExternalSorter { }) } - fn use_staged_sort(&self) -> bool { - // The first group is always evaluated for every row. Only computation - // in later groups can be avoided. - self.enable_staged_sort - && self - .expr - .iter() - .skip(self.sort_key_group_size) - .any(|key| !key.expr.is::() && !key.expr.is::()) - } - /// If this sort may spill, pre-allocates /// `sort_spill_reservation_bytes` of memory to guarantee memory /// left for the in memory sort/merge. @@ -1302,39 +1334,39 @@ impl Debug for ExternalSorter { } } +struct SortStage { + columns: Vec, + expressions: Vec, +} + fn staged_sort_batch_chunked( batch: &RecordBatch, expressions: &LexOrdering, batch_size: usize, - group_size: usize, + stages: &[SortStage], ) -> Result> { - let indices = staged_sort_indices(batch, expressions, group_size)?; + let indices = staged_sort_indices(batch, stages)?; IncrementalSortIterator::new(batch.clone(), expressions.clone(), batch_size) .with_sorted_indices(indices) .collect() } /// Compute a complete permutation by evaluating successive key groups only for ties. -fn staged_sort_indices( - batch: &RecordBatch, - expressions: &LexOrdering, - group_size: usize, -) -> Result { +fn staged_sort_indices(batch: &RecordBatch, stages: &[SortStage]) -> Result { let mut indices = Vec::new(); // Ranges address positions in `indices`, whose values address the input batch. let mut unresolved: Vec<_> = std::iter::once(0..batch.num_rows()).collect(); - let stage_count = expressions.len().div_ceil(group_size); - for (stage, keys) in expressions.chunks(group_size).enumerate() { + for (index, stage) in stages.iter().enumerate() { if unresolved.is_empty() { break; } unresolved = refine_tied_groups( batch, - keys, + stage, &mut indices, &unresolved, - stage == 0, - stage + 1 < stage_count, + index == 0, + index + 1 < stages.len(), )?; } @@ -1344,7 +1376,7 @@ fn staged_sort_indices( /// Evaluate keys once for all unresolved rows, then sort each tied group independently. fn refine_tied_groups( batch: &RecordBatch, - keys: &[PhysicalSortExpr], + stage: &SortStage, indices: &mut Vec, unresolved: &[Range], first_stage: bool, @@ -1353,14 +1385,17 @@ fn refine_tied_groups( let selected = if first_stage { batch.clone() } else { - // Gather the original row indices for all tied groups into one buffer. let mut selection = Vec::with_capacity(unresolved.iter().map(Range::len).sum()); for range in unresolved { selection.extend_from_slice(&indices[range.clone()]); } - arrow::compute::take_record_batch(batch, &UInt32Array::from(selection))? + arrow::compute::take_record_batch( + &batch.project(&stage.columns)?, + &UInt32Array::from(selection), + )? }; - let sort_columns = keys + let sort_columns = stage + .expressions .iter() .map(|expr| expr.evaluate_to_sort_column(&selected)) .collect::>>()?; @@ -4272,6 +4307,273 @@ mod tests { Ok((sorted_batches, metrics)) } + // Exercise SortExec directly so the planner cannot project the sort expressions + // into columns before they reach the sorter. Each key is a separate stage. + async fn staged_sort_test_output( + batch: &RecordBatch, + ordering: &LexOrdering, + enabled: bool, + input_batches: usize, + in_place_threshold: usize, + ) -> Result> { + let rows_per_batch = batch.num_rows() / input_batches; + let batches = (0..input_batches) + .map(|i| batch.slice(i * rows_per_batch, rows_per_batch)) + .collect(); + let input = TestMemoryExec::try_new_exec(&[batches], batch.schema(), None)?; + let sort = Arc::new(SortExec::new(ordering.clone(), input)); + let config = SessionConfig::new() + .with_batch_size(2) + .with_sort_in_place_threshold_bytes(in_place_threshold) + .set_bool("datafusion.execution.enable_staged_sort", enabled) + .set_usize("datafusion.execution.sort_key_group_size", 1); + let result = collect( + Arc::clone(&sort) as Arc, + Arc::new(TaskContext::default().with_session_config(config)), + ) + .await; + assert_eq!(sort.metrics().unwrap().spill_count().unwrap_or(0), 0); + if let Ok(batches) = &result { + assert_eq!( + batches.iter().map(RecordBatch::num_rows).sum::(), + batch.num_rows() + ); + assert!(batches.iter().all(|b| (1..=2).contains(&b.num_rows()))); + } + result + } + + #[tokio::test] + async fn test_staged_sort_matches_eager() -> Result<()> { + use datafusion_physical_expr::expressions::CastExpr; + + // Keep payload columns between the keys to exercise projected column remapping. + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new("suffix", DataType::Utf8, true), + Field::new("payload", DataType::Utf8, false), + Field::new("prefix", DataType::Int32, true), + ])); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int32Array::from_iter_values(0..10)), + Arc::new(StringArray::from(vec![ + Some("3"), + Some("2"), + Some("5"), + None, + Some("3"), + Some("10"), + None, + Some("2"), + Some("0"), + Some("10"), + ])), + Arc::new(StringArray::from_iter_values( + (0..10).map(|i| format!("payload-{i}")), + )), + Arc::new(Int32Array::from(vec![ + Some(2), + Some(1), + None, + Some(1), + Some(2), + Some(1), + None, + Some(1), + Some(3), + Some(1), + ])), + ], + )?; + let ordering = [ + PhysicalSortExpr::new( + col("prefix", &schema)?, + SortOptions { + descending: false, + nulls_first: false, + }, + ), + PhysicalSortExpr::new( + Arc::new(CastExpr::new( + col("suffix", &schema)?, + DataType::Int32, + None, + )), + SortOptions { + descending: true, + nulls_first: true, + }, + ), + PhysicalSortExpr::new_default(col("id", &schema)?), + ] + .into(); + // Prefix ties survive the numeric suffix stage and must be ordered by id. + let expected = arrow::compute::take_record_batch( + &batch, + &UInt32Array::from(vec![3, 5, 9, 1, 7, 0, 4, 8, 6, 2]), + )?; + + // Single batch, concatenated batches, and separate runs followed by a merge. + for (input_batches, threshold) in [(1, 1), (2, usize::MAX), (2, 1)] { + for enabled in [false, true] { + let output = staged_sort_test_output( + &batch, + &ordering, + enabled, + input_batches, + threshold, + ) + .await?; + assert_eq!( + concat_batches(&schema, &output)?, + expected, + "enabled={enabled}, input_batches={input_batches}, threshold={threshold}" + ); + } + } + Ok(()) + } + + #[tokio::test] + async fn test_staged_sort_expression_errors() -> Result<()> { + use datafusion_physical_expr::expressions::CastExpr; + + let schema = Arc::new(Schema::new(vec![ + Field::new("prefix", DataType::Int32, false), + Field::new("invalid", DataType::Utf8, false), + ])); + let ordering = [ + PhysicalSortExpr::new_default(col("prefix", &schema)?), + PhysicalSortExpr::new_default(Arc::new(CastExpr::new( + col("invalid", &schema)?, + DataType::Int32, + None, + ))), + ] + .into(); + for all_tied in [false, true] { + let prefix = if all_tied { + vec![1; 6] + } else { + vec![3, 1, 6, 2, 5, 4] + }; + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int32Array::from(prefix)), + Arc::new(StringArray::from(vec!["invalid"; 6])), + ], + )?; + for (input_batches, threshold) in [(1, 1), (2, usize::MAX), (2, 1)] { + for enabled in [false, true] { + let result = staged_sort_test_output( + &batch, + &ordering, + enabled, + input_batches, + threshold, + ) + .await; + // A merge evaluates all keys again, even if each run skipped the cast. + let merge = input_batches == 2 && threshold == 1; + if !enabled || all_tied || merge { + let error = result.unwrap_err().strip_backtrace(); + assert!( + error.contains("Cannot cast string 'invalid'"), + "{error}" + ); + } else { + let output = concat_batches(&schema, &result?)?; + assert_eq!( + output.column(0).as_ref(), + &Int32Array::from_iter_values(1..=6), + ); + } + } + } + } + Ok(()) + } + + #[tokio::test] + async fn test_staged_sort_volatile_evaluation() -> Result<()> { + use datafusion_expr::{Volatility, create_udf}; + use datafusion_physical_expr::ScalarFunctionExpr; + + let schema = Arc::new(Schema::new(vec![ + Field::new("prefix", DataType::Int32, false), + Field::new("id", DataType::Int32, false), + ])); + // Count rows, not calls: expression evaluation is batched. The middle + // case has four tied rows, with two in each of the split input batches. + for (prefix, tied_rows) in [ + (vec![3, 1, 6, 2, 5, 4], 0), + (vec![2, 1, 1, 3, 4, 4], 4), + (vec![1; 6], 6), + ] { + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int32Array::from(prefix)), + Arc::new(Int32Array::from_iter_values(0..6)), + ], + )?; + for (input_batches, threshold) in [(1, 1), (2, usize::MAX), (2, 1)] { + for enabled in [false, true] { + let rows_evaluated = Arc::new(AtomicUsize::new(0)); + let counter = Arc::clone(&rows_evaluated); + let function = create_udf( + "count_evaluations", + vec![DataType::Int32], + DataType::Int32, + Volatility::Volatile, + Arc::new(move |args| { + counter.fetch_add( + args[0].clone().into_array(1)?.len(), + AtomicOrdering::Relaxed, + ); + Ok(args[0].clone()) + }), + ); + let ordering = [ + PhysicalSortExpr::new_default(col("prefix", &schema)?), + PhysicalSortExpr::new_default(Arc::new( + ScalarFunctionExpr::try_new( + Arc::new(function), + vec![col("id", &schema)?], + &schema, + Arc::new(ConfigOptions::default()), + )?, + )), + ] + .into(); + staged_sort_test_output( + &batch, + &ordering, + enabled, + input_batches, + threshold, + ) + .await?; + let merge_rows = if input_batches == 2 && threshold == 1 { + 6 + } else { + 0 + }; + let sort_rows = if enabled { tied_rows } else { 6 }; + assert_eq!( + rows_evaluated.load(AtomicOrdering::Relaxed), + sort_rows + merge_rows, + "enabled={enabled}, tied_rows={tied_rows}, input_batches={input_batches}, threshold={threshold}" + ); + } + } + } + Ok(()) + } + #[tokio::test] async fn test_sort_batch_chunked_basic() -> Result<()> { let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); From fea4c42a3e0c3a3a2e339f07e46e2f87d569205b Mon Sep 17 00:00:00 2001 From: yashrb24 Date: Tue, 3 Nov 2026 19:55:45 +0530 Subject: [PATCH 10/12] fix: refine automatic staged sort evaluation --- datafusion/common/src/config.rs | 8 +- .../physical-expr/src/expressions/binary.rs | 2 +- .../physical-expr/src/expressions/mod.rs | 2 +- datafusion/physical-plan/src/sorts/sort.rs | 1199 +++++++++++++---- datafusion/physical-plan/src/sorts/stream.rs | 5 +- .../test_files/information_schema.slt | 4 +- docs/source/user-guide/configs.md | 1 + 7 files changed, 914 insertions(+), 307 deletions(-) diff --git a/datafusion/common/src/config.rs b/datafusion/common/src/config.rs index 698448b7f17d5..11c5f1380102d 100644 --- a/datafusion/common/src/config.rs +++ b/datafusion/common/src/config.rs @@ -998,13 +998,11 @@ config_namespace! { /// batches and merged. pub sort_in_place_threshold_bytes: usize, default = 1024 * 1024 - /// Enables staged in-memory sorting, evaluating later sort keys only for rows - /// tied on earlier keys. + /// Enables staged sorting for eligible in-memory sorts. + /// Later keys may be skipped for rows already in order + /// TopK and spill sorting do not use staged sorting. pub enable_staged_sort: bool, default = false - /// Number of sort keys evaluated together by staged in-memory sorting. - pub sort_key_group_size: usize, default = 5 - /// Maximum buffer capacity (in bytes) per partition for BufferExec /// inserted during sort pushdown optimization. /// diff --git a/datafusion/physical-expr/src/expressions/binary.rs b/datafusion/physical-expr/src/expressions/binary.rs index d3cb36e3f6ab9..9f7efd0c5408c 100644 --- a/datafusion/physical-expr/src/expressions/binary.rs +++ b/datafusion/physical-expr/src/expressions/binary.rs @@ -1443,7 +1443,7 @@ fn rows_to_filter_before( /// Whether evaluating `expr` on all rows is cheap and cannot fail. /// Inspired by `reorder_predicates`' cheap forms, with fallible forms excluded. -fn is_cheap_and_infallible(expr: &Arc, schema: &Schema) -> bool { +pub fn is_cheap_and_infallible(expr: &Arc, schema: &Schema) -> bool { use crate::expressions::{ CastExpr, Column, InListExpr, IsNotNullExpr, IsNullExpr, Literal, NotExpr, }; diff --git a/datafusion/physical-expr/src/expressions/mod.rs b/datafusion/physical-expr/src/expressions/mod.rs index 767c0b4ce40f3..d4f933c91de69 100644 --- a/datafusion/physical-expr/src/expressions/mod.rs +++ b/datafusion/physical-expr/src/expressions/mod.rs @@ -28,7 +28,7 @@ pub use crate::PhysicalSortExpr; /// Module with some convenient methods used in expression building pub use crate::aggregate::stats::StatsType; -pub use binary::{BinaryExpr, binary, similar_to}; +pub use binary::{BinaryExpr, binary, is_cheap_and_infallible, similar_to}; pub use case::{CaseExpr, case}; pub use cast::{CastExpr, cast}; pub use column::{Column, col, with_new_schema}; diff --git a/datafusion/physical-plan/src/sorts/sort.rs b/datafusion/physical-plan/src/sorts/sort.rs index 2311165543829..04e59c6a2182b 100644 --- a/datafusion/physical-plan/src/sorts/sort.rs +++ b/datafusion/physical-plan/src/sorts/sort.rs @@ -69,10 +69,11 @@ use datafusion_execution::runtime_env::RuntimeEnv; use datafusion_physical_expr::LexOrdering; use datafusion_physical_expr::PhysicalExpr; use datafusion_physical_expr::expressions::{ - Column, DynamicFilterPhysicalExpr, Literal, lit, + DynamicFilterPhysicalExpr, is_cheap_and_infallible, lit, }; use datafusion_physical_expr::projection::ProjectionExprs; use datafusion_physical_expr::utils::collect_columns; +use datafusion_physical_expr_common::physical_expr::is_volatile; use futures::{Stream, StreamExt, TryStreamExt}; use log::{debug, trace}; @@ -385,58 +386,10 @@ impl ExternalSorter { }) } - fn with_staged_sort(mut self, enabled: bool, group_size: usize) -> Result { - if enabled && group_size == 0 { - return Err(DataFusionError::Configuration( - "datafusion.execution.sort_key_group_size must be greater than zero when datafusion.execution.enable_staged_sort is true".to_string(), - )); + fn with_staged_sort(mut self, enabled: bool) -> Result { + if enabled { + self.sort_stages = create_sort_stages(&self.expr, &self.schema)?; } - self.sort_stages = if enabled - && self - .expr - .iter() - .skip(group_size) - .any(|key| !key.expr.is::() && !key.expr.is::()) - { - let stages = self - .expr - .chunks(group_size) - .enumerate() - .map(|(index, keys)| { - if index == 0 { - return Ok(SortStage { - columns: vec![], - expressions: keys.to_vec(), - }); - } - let mut columns = keys - .iter() - .flat_map(|key| collect_columns(&key.expr)) - .map(|column| column.index()) - .collect::>(); - columns.sort_unstable(); - columns.dedup(); - let projection = - ProjectionExprs::from_indices(&columns, &self.schema); - let expressions = keys - .iter() - .map(|key| { - Ok(PhysicalSortExpr::new( - projection.project_expr(&key.expr)?, - key.options, - )) - }) - .collect::>>()?; - Ok(SortStage { - columns, - expressions, - }) - }) - .collect::>>()?; - Some(stages.into()) - } else { - None - }; Ok(self) } @@ -1238,105 +1191,85 @@ impl ExternalSorter { } } -/// Estimate how much memory is needed to sort a `RecordBatch`. -/// -/// This is used to pre-reserve memory for the sort/merge. The sort/merge process involves -/// creating sorted copies of sorted columns in record batches for speeding up comparison -/// in sorting and merging. The sorted copies are in either row format or array format. -/// Please refer to cursor.rs and stream.rs for more details. No matter what format the -/// sorted copies are, they will use more memory than the original record batch. -/// -/// This can basically be calculated as the sum of the actual space it takes in -/// memory (which would be larger for a sliced batch), and the size of the actual data. -pub(crate) fn get_reserved_bytes_for_record_batch_size( - record_batch_size: usize, - sliced_size: usize, -) -> usize { - // Even 2x may not be enough for some cases, but it's a good enough estimation as a baseline. - // If 2x is not enough, user can set a larger value for `sort_spill_reservation_bytes` - // to compensate for the extra memory needed. - record_batch_size + sliced_size -} - -/// Estimate how much memory is needed to sort a `RecordBatch`. -/// This will just call `get_reserved_bytes_for_record_batch_size` with the -/// memory size of the record batch and its sliced size. -pub(crate) fn get_reserved_bytes_for_record_batch(batch: &RecordBatch) -> Result { - batch.get_sliced_size().map(|sliced_size| { - get_reserved_bytes_for_record_batch_size( - get_record_batch_memory_size(batch), - sliced_size, - ) - }) +struct SortStage { + input_columns: Vec, + expressions: Vec, } -/// Estimate how much memory is needed to sort `batch` when it is buffered with -/// the batches `counter` has already counted. +/// Preserve key order and start a group before each expensive expression. /// -/// Like [`get_reserved_bytes_for_record_batch`], but only counts the buffers of -/// `batch` that `counter` has not seen. Buffered batches can share buffers, for -/// example zero-copy slices of one larger batch, and a shared buffer only needs -/// to be reserved once. Without a `counter`, every buffer of `batch` is counted. -fn get_reserved_bytes_for_next_record_batch( - batch: &RecordBatch, - counter: Option<&mut RecordBatchMemoryCounter>, -) -> Result { - let Some(counter) = counter else { - return get_reserved_bytes_for_record_batch(batch); - }; - let sliced_size = batch.get_sliced_size()?; - Ok(get_reserved_bytes_for_record_batch_size( - counter.count_batch(batch), - sliced_size, - )) -} - -/// Returns whether the sorted copy of an array of `data_type` keeps -/// referencing some of the array's buffers. +/// Examples: /// -/// Sorting takes the selected rows into new buffers, except that `take` keeps -/// the data buffers of view arrays and the values of dictionaries and list -/// views. Every sorted run is charged for those buffers in full, so buffered -/// batches that share them must be charged for them each, too. Otherwise the -/// sorter buffers more batches than the merge of their sorted runs can hold. -fn sorted_copy_keeps_buffers(data_type: &DataType) -> bool { - match data_type { - DataType::Utf8View - | DataType::BinaryView - | DataType::Dictionary(_, _) - | DataType::ListView(_) - | DataType::LargeListView(_) => true, - DataType::List(field) - | DataType::LargeList(field) - | DataType::FixedSizeList(field, _) - | DataType::Map(field, _) => sorted_copy_keeps_buffers(field.data_type()), - DataType::Struct(fields) => fields - .iter() - .any(|field| sorted_copy_keeps_buffers(field.data_type())), - DataType::Union(fields, _) => fields - .iter() - .any(|(_, field)| sorted_copy_keeps_buffers(field.data_type())), - DataType::RunEndEncoded(_, values) => { - sorted_copy_keeps_buffers(values.data_type()) +/// ```text +/// a, lower(name), id => [a] -> [lower(name), id] +/// a, b, lower(name), upper(city), id => [a, b] -> [lower(name)] -> [upper(city), id] +/// lower(name), upper(city), id => [lower(name)] -> [upper(city), id] +/// a, CAST(text AS INT), id => [a] -> [CAST(text AS INT), id] +/// ``` +fn create_sort_stages( + expressions: &LexOrdering, + schema: &SchemaRef, +) -> Result>> { + if expressions.iter().any(|key| is_volatile(&key.expr)) { + return Ok(None); + } + + let mut starts = vec![0]; + for (index, key) in expressions.iter().enumerate().skip(1) { + if !is_cheap_and_infallible(&key.expr, schema) { + starts.push(index); } - _ => false, } -} -impl Debug for ExternalSorter { - fn fmt(&self, f: &mut Formatter) -> fmt::Result { - f.debug_struct("ExternalSorter") - .field("memory_used", &self.used()) - .field("spilled_bytes", &self.spilled_bytes()) - .field("spilled_rows", &self.spilled_rows()) - .field("spill_count", &self.spill_count()) - .finish() + if starts.len() <= 1 { + return Ok(None); } + starts.push(expressions.len()); + + let mut stages = vec![SortStage { + input_columns: vec![], + expressions: expressions[starts[0]..starts[1]].to_vec(), + }]; + for bounds in starts[1..].windows(2) { + stages.push(project_sort_stage( + &expressions[bounds[0]..bounds[1]], + schema, + )?); + } + + Ok(Some(stages.into())) } -struct SortStage { - columns: Vec, - expressions: Vec, +/// Prepare the columns and expressions needed for a later sort stage +fn project_sort_stage( + keys: &[PhysicalSortExpr], + schema: &SchemaRef, +) -> Result { + let mut columns = keys + .iter() + .flat_map(|key| collect_columns(&key.expr)) + .map(|column| column.index()) + .collect::>(); + + // Collect each input column only once + columns.sort_unstable(); + columns.dedup(); + + let projection = ProjectionExprs::from_indices(&columns, schema); + let expressions = keys + .iter() + .map(|key| { + Ok(PhysicalSortExpr::new( + projection.project_expr(&key.expr)?, + key.options, + )) + }) + .collect::>>()?; + + Ok(SortStage { + input_columns: columns, + expressions, + }) } fn staged_sort_batch_chunked( @@ -1345,55 +1278,159 @@ fn staged_sort_batch_chunked( batch_size: usize, stages: &[SortStage], ) -> Result> { - let indices = staged_sort_indices(batch, stages)?; + let indices = staged_sort_indices(batch, stages, expressions)?; IncrementalSortIterator::new(batch.clone(), expressions.clone(), batch_size) .with_sorted_indices(indices) .collect() } -/// Compute a complete permutation by evaluating successive key groups only for ties. -fn staged_sort_indices(batch: &RecordBatch, stages: &[SortStage]) -> Result { - let mut indices = Vec::new(); - // Ranges address positions in `indices`, whose values address the input batch. - let mut unresolved: Vec<_> = std::iter::once(0..batch.num_rows()).collect(); - for (index, stage) in stages.iter().enumerate() { - if unresolved.is_empty() { +/// Return row indices in sort order, evaluating later stages only for ties. +/// +/// For `ORDER BY department, lower(name)`: +/// - Sort all rows by `department`. +/// - Evaluate `lower(name)` only for rows with the same department, then +/// sort those rows within each department. +/// - Return the original row indices in the resulting order. +/// +/// If there are many tied groups, we evaluate the remaining keys together +/// for the rows still tied. +fn staged_sort_indices( + batch: &RecordBatch, + stages: &[SortStage], + expressions: &LexOrdering, +) -> Result { + let (mut indices, mut tied_groups) = sort_first_stage(batch, &stages[0])?; + let mut stage_rows = batch.num_rows(); + let mut key_offset = stages[0].expressions.len(); + + for (index, stage) in stages.iter().enumerate().skip(1) { + if tied_groups.is_empty() { + break; + } + + let tied_rows: usize = tied_groups.iter().map(Range::len).sum(); + let resolved_rows = stage_rows - tied_rows; + + // Compare the number of groups with the reduction from the last stage. + if tied_groups.len() > resolved_rows { + let remaining = + project_sort_stage(&expressions[key_offset..], &batch.schema())?; + sort_remaining_keys(batch, &remaining, &mut indices, &tied_groups)?; break; } - unresolved = refine_tied_groups( + + stage_rows = tied_rows; + tied_groups = refine_tied_groups( batch, stage, &mut indices, - &unresolved, - index == 0, + &tied_groups, index + 1 < stages.len(), )?; + key_offset += stage.expressions.len(); } Ok(UInt32Array::from(indices)) } -/// Evaluate keys once for all unresolved rows, then sort each tied group independently. +/// Sort the first stage and find the groups needing later keys. +fn sort_first_stage( + batch: &RecordBatch, + stage: &SortStage, +) -> Result<(Vec, Vec>)> { + let columns = stage + .expressions + .iter() + .map(|expr| expr.evaluate_to_sort_column(batch)) + .collect::>>()?; + let order = lexsort_to_indices(&columns, None)?; + let comparator = LexicographicalComparator::try_new(&columns)?; + + let mut tied_groups = Vec::new(); + append_tied_ranges(&comparator, &order, 0, 0, &mut tied_groups); + let (_, values, _) = order.into_parts(); + Ok((values.into(), tied_groups)) +} + +/// Gather the columns and rows needed to evaluate the next keys. +fn gather_tied_rows( + batch: &RecordBatch, + stage: &SortStage, + indices: &[u32], + tied_groups: &[Range], +) -> Result<(RecordBatch, UInt32Array)> { + let selected_rows = tied_groups.iter().map(Range::len).sum(); + let mut selection = Vec::with_capacity(selected_rows); + for range in tied_groups { + selection.extend_from_slice(&indices[range.clone()]); + } + let selection = UInt32Array::from(selection); + let projected = batch.project(&stage.input_columns)?; + + let selected = if stage.input_columns.is_empty() { + RecordBatch::try_new_with_options( + projected.schema(), + vec![], + &RecordBatchOptions::new().with_row_count(Some(selected_rows)), + )? + } else { + arrow::compute::take_record_batch(&projected, &selection)? + }; + Ok((selected, selection)) +} + +/// Evaluate all remaining keys for tied rows and sort within each group. +fn sort_remaining_keys( + batch: &RecordBatch, + stage: &SortStage, + indices: &mut [u32], + tied_groups: &[Range], +) -> Result<()> { + let (selected, original_rows) = gather_tied_rows(batch, stage, indices, tied_groups)?; + let columns = stage + .expressions + .iter() + .map(|expr| expr.evaluate_to_sort_column(&selected)) + .collect::>>()?; + let comparator = LexicographicalComparator::try_new(&columns)?; + + // Comparator indices address the compact batch; map them back afterward. + let mut compact_order: Vec = (0..original_rows.len() as u32).collect(); + let mut offset = 0; + for range in tied_groups { + let end = offset + range.len(); + compact_order[offset..end] + .sort_unstable_by(|a, b| comparator.compare(*a as usize, *b as usize)); + for (output, compact_row) in indices[range.clone()] + .iter_mut() + .zip(&compact_order[offset..end]) + { + *output = original_rows.value(*compact_row as usize); + } + offset = end; + } + Ok(()) +} + +/// Sort each tied group using the next stage's keys. +/// +/// For example, given ties on a given key: +/// +/// ```text +/// Before: A: [Zoe, Ada] B: [Mia, Bob] C: [Lea] +/// After: A: [Ada, Zoe] B: [Bob, Mia] C: [Lea] +/// ``` +/// +/// This updates `indices` in place. fn refine_tied_groups( batch: &RecordBatch, stage: &SortStage, - indices: &mut Vec, + indices: &mut [u32], unresolved: &[Range], - first_stage: bool, find_ties: bool, ) -> Result>> { - let selected = if first_stage { - batch.clone() - } else { - let mut selection = Vec::with_capacity(unresolved.iter().map(Range::len).sum()); - for range in unresolved { - selection.extend_from_slice(&indices[range.clone()]); - } - arrow::compute::take_record_batch( - &batch.project(&stage.columns)?, - &UInt32Array::from(selection), - )? - }; + let (selected, _) = gather_tied_rows(batch, stage, indices, unresolved)?; + let sort_columns = stage .expressions .iter() @@ -1407,16 +1444,6 @@ fn refine_tied_groups( .transpose()?; let mut next_unresolved = Vec::new(); - if first_stage { - let order = lexsort_to_indices(&sort_columns, None)?; - if let Some(comparator) = &comparator { - append_tied_ranges(comparator, &order, 0, 0, &mut next_unresolved); - } - let (_, values, _) = order.into_parts(); - *indices = values.into(); - return Ok(next_unresolved); - } - let mut selected_offset = 0; for range in unresolved { let columns = sort_columns @@ -1426,12 +1453,14 @@ fn refine_tied_groups( options: column.options, }) .collect::>(); + let local_order = lexsort_to_indices(&columns, None)?; let previous = indices[range.clone()].to_vec(); for (output, local) in indices[range.clone()].iter_mut().zip(local_order.values()) { *output = previous[*local as usize]; } + if let Some(comparator) = &comparator { append_tied_ranges( comparator, @@ -1443,11 +1472,14 @@ fn refine_tied_groups( } selected_offset += range.len(); } + Ok(next_unresolved) } -/// Find consecutive equal keys in a sorted group; only these rows need later keys. -/// For sorted keys `[a, a, b]` at output offset 5, append the range `5..7`. +/// Find groups of equal keys that need to be sorted by the next stage. +/// +/// For sorted keys `[A, A, B, C, C]`, append the ranges `0..2` and `3..5`. +/// A single row, like `B`, is already in order and needs no further sorting. fn append_tied_ranges( comparator: &LexicographicalComparator, order: &UInt32Array, @@ -1472,6 +1504,102 @@ fn append_tied_ranges( } } +/// Estimate how much memory is needed to sort a `RecordBatch`. +/// +/// This is used to pre-reserve memory for the sort/merge. The sort/merge process involves +/// creating sorted copies of sorted columns in record batches for speeding up comparison +/// in sorting and merging. The sorted copies are in either row format or array format. +/// Please refer to cursor.rs and stream.rs for more details. No matter what format the +/// sorted copies are, they will use more memory than the original record batch. +/// +/// This can basically be calculated as the sum of the actual space it takes in +/// memory (which would be larger for a sliced batch), and the size of the actual data. +pub(crate) fn get_reserved_bytes_for_record_batch_size( + record_batch_size: usize, + sliced_size: usize, +) -> usize { + // Even 2x may not be enough for some cases, but it's a good enough estimation as a baseline. + // If 2x is not enough, user can set a larger value for `sort_spill_reservation_bytes` + // to compensate for the extra memory needed. + record_batch_size + sliced_size +} + +/// Estimate how much memory is needed to sort a `RecordBatch`. +/// This will just call `get_reserved_bytes_for_record_batch_size` with the +/// memory size of the record batch and its sliced size. +pub(crate) fn get_reserved_bytes_for_record_batch(batch: &RecordBatch) -> Result { + batch.get_sliced_size().map(|sliced_size| { + get_reserved_bytes_for_record_batch_size( + get_record_batch_memory_size(batch), + sliced_size, + ) + }) +} + +/// Estimate how much memory is needed to sort `batch` when it is buffered with +/// the batches `counter` has already counted. +/// +/// Like [`get_reserved_bytes_for_record_batch`], but only counts the buffers of +/// `batch` that `counter` has not seen. Buffered batches can share buffers, for +/// example zero-copy slices of one larger batch, and a shared buffer only needs +/// to be reserved once. Without a `counter`, every buffer of `batch` is counted. +fn get_reserved_bytes_for_next_record_batch( + batch: &RecordBatch, + counter: Option<&mut RecordBatchMemoryCounter>, +) -> Result { + let Some(counter) = counter else { + return get_reserved_bytes_for_record_batch(batch); + }; + let sliced_size = batch.get_sliced_size()?; + Ok(get_reserved_bytes_for_record_batch_size( + counter.count_batch(batch), + sliced_size, + )) +} + +/// Returns whether the sorted copy of an array of `data_type` keeps +/// referencing some of the array's buffers. +/// +/// Sorting takes the selected rows into new buffers, except that `take` keeps +/// the data buffers of view arrays and the values of dictionaries and list +/// views. Every sorted run is charged for those buffers in full, so buffered +/// batches that share them must be charged for them each, too. Otherwise the +/// sorter buffers more batches than the merge of their sorted runs can hold. +fn sorted_copy_keeps_buffers(data_type: &DataType) -> bool { + match data_type { + DataType::Utf8View + | DataType::BinaryView + | DataType::Dictionary(_, _) + | DataType::ListView(_) + | DataType::LargeListView(_) => true, + DataType::List(field) + | DataType::LargeList(field) + | DataType::FixedSizeList(field, _) + | DataType::Map(field, _) => sorted_copy_keeps_buffers(field.data_type()), + DataType::Struct(fields) => fields + .iter() + .any(|field| sorted_copy_keeps_buffers(field.data_type())), + DataType::Union(fields, _) => fields + .iter() + .any(|(_, field)| sorted_copy_keeps_buffers(field.data_type())), + DataType::RunEndEncoded(_, values) => { + sorted_copy_keeps_buffers(values.data_type()) + } + _ => false, + } +} + +impl Debug for ExternalSorter { + fn fmt(&self, f: &mut Formatter) -> fmt::Result { + f.debug_struct("ExternalSorter") + .field("memory_used", &self.used()) + .field("spilled_bytes", &self.spilled_bytes()) + .field("spilled_rows", &self.spilled_rows()) + .field("spill_count", &self.spill_count()) + .finish() + } +} + pub fn sort_batch( batch: &RecordBatch, expressions: &LexOrdering, @@ -2057,10 +2185,7 @@ impl ExecutionPlan for SortExec { &self.metrics_set, context.runtime_env(), )? - .with_staged_sort( - execution_options.enable_staged_sort, - execution_options.sort_key_group_size, - )?; + .with_staged_sort(execution_options.enable_staged_sort)?; sorter.sort_concurrency = Arc::clone(&self.sort_concurrency); Ok(Box::pin(RecordBatchStreamAdapter::new( self.schema(), @@ -2710,7 +2835,6 @@ mod tests { use std::collections::HashMap; use std::path::Path; use std::pin::Pin; - use std::sync::atomic::{AtomicUsize, Ordering as AtomicOrdering}; use std::task::{Context, Poll}; use super::*; @@ -2748,12 +2872,15 @@ mod tests { use datafusion_execution::{ AsyncSpillWriter, SpillFile, SpillWriter, TempFileFactory, }; + use datafusion_expr::Operator; + use datafusion_physical_expr::expressions::{BinaryExpr, CastExpr}; use datafusion_physical_expr::expressions::{Column, Literal}; use datafusion_physical_expr::{DynamicFilterTracking, EquivalenceProperties}; use datafusion_physical_expr_common::metrics::MetricValue; use futures::{FutureExt, Stream, TryStreamExt}; use insta::assert_snapshot; + use std::sync::atomic::{AtomicUsize, Ordering as AtomicOrdering}; use tokio::sync::Notify; #[tokio::test] @@ -2911,6 +3038,52 @@ mod tests { abort_count: Arc, } + #[test] + fn test_staged_sort_automatic_groups() -> Result<()> { + let schema = Arc::new(Schema::new(vec![ + Field::new("i", DataType::Int32, false), + Field::new("s", DataType::Utf8, false), + ])); + let cheap = col("i", &schema)?; + let expensive: Arc = + Arc::new(CastExpr::new(col("s", &schema)?, DataType::Int32, None)); + let group_sizes = |classes: &[bool]| -> Result> { + let keys = + LexOrdering::new(classes.iter().enumerate().map(|(index, cheap_key)| { + let expr = if *cheap_key { + Arc::new(BinaryExpr::new( + Arc::clone(&cheap), + Operator::Plus, + lit(index as i32), + )) as Arc + } else { + Arc::new(BinaryExpr::new( + Arc::clone(&expensive), + Operator::Plus, + lit(index as i32), + )) + }; + PhysicalSortExpr::new_default(expr) + })) + .unwrap(); + Ok(create_sort_stages(&keys, &schema)? + .map(|plan| plan.iter().map(|stage| stage.expressions.len()).collect()) + .unwrap_or_default()) + }; + assert_eq!( + group_sizes(&[true, true, false, false, true, false])?, + vec![2, 1, 2, 1] + ); + assert_eq!(group_sizes(&[false, false, true, false])?, vec![1, 2, 1]); + assert_eq!(group_sizes(&[true, false])?, vec![1, 1]); + assert_eq!(group_sizes(&[false, false])?, vec![1, 1]); + assert_eq!(group_sizes(&[false; 5])?, vec![1, 1, 1, 1, 1]); + assert!(group_sizes(&[true; 5])?.is_empty()); + assert!(group_sizes(&[false])?.is_empty()); + + Ok(()) + } + impl TempFileFactory for PendingWriteTempFileFactory { fn create_temp_file(&self, _description: &str) -> Result> { Ok(Arc::new(PendingWriteSpillFile { @@ -4307,8 +4480,41 @@ mod tests { Ok((sorted_batches, metrics)) } + fn counting_identity( + name: &str, + input: Arc, + schema: &SchemaRef, + volatility: datafusion_expr::Volatility, + ) -> Result<(Arc, Arc)> { + use datafusion_expr::create_udf; + use datafusion_physical_expr::ScalarFunctionExpr; + + let rows = Arc::new(AtomicUsize::new(0)); + let counter = Arc::clone(&rows); + let udf = create_udf( + name, + vec![DataType::Int32], + DataType::Int32, + volatility, + Arc::new(move |args| { + counter.fetch_add( + args[0].clone().into_array(1)?.len(), + AtomicOrdering::Relaxed, + ); + Ok(args[0].clone()) + }), + ); + let expression = Arc::new(ScalarFunctionExpr::try_new( + Arc::new(udf), + vec![input], + schema, + Arc::new(ConfigOptions::default()), + )?); + Ok((expression, rows)) + } + // Exercise SortExec directly so the planner cannot project the sort expressions - // into columns before they reach the sorter. Each key is a separate stage. + // into columns before they reach the sorter. async fn staged_sort_test_output( batch: &RecordBatch, ordering: &LexOrdering, @@ -4325,8 +4531,7 @@ mod tests { let config = SessionConfig::new() .with_batch_size(2) .with_sort_in_place_threshold_bytes(in_place_threshold) - .set_bool("datafusion.execution.enable_staged_sort", enabled) - .set_usize("datafusion.execution.sort_key_group_size", 1); + .set_bool("datafusion.execution.enable_staged_sort", enabled); let result = collect( Arc::clone(&sort) as Arc, Arc::new(TaskContext::default().with_session_config(config)), @@ -4343,21 +4548,340 @@ mod tests { result } + async fn check_staged_multi_run_path(concurrent: bool) -> Result<()> { + use datafusion_expr::Volatility; + + let batch = record_batch!( + ("p", Int32, [8, 6, 4, 2, 7, 5, 3, 1]), + ("v", Int32, [8, 6, 4, 2, 7, 5, 3, 1]) + )?; + let schema = batch.schema(); + for enabled in [false, true] { + let (suffix, rows) = counting_identity( + "count_suffix", + col("v", &schema)?, + &schema, + Volatility::Immutable, + )?; + + let ordering = [ + PhysicalSortExpr::new_default(col("p", &schema)?), + PhysicalSortExpr::new_default(suffix), + ] + .into(); + let input = TestMemoryExec::try_new_exec( + &[vec![batch.slice(0, 4), batch.slice(4, 4)]], + Arc::clone(&schema), + None, + )?; + let sort = Arc::new(SortExec::new(ordering, input)); + let config = SessionConfig::new() + .with_batch_size(2) + .with_sort_in_place_threshold_bytes(1) + .set_bool("datafusion.execution.enable_staged_sort", enabled); + let output = collect( + Arc::clone(&sort) as Arc, + Arc::new(TaskContext::default().with_session_config(config)), + ) + .await?; + assert_eq!(sort.sort_concurrency.get().is_some(), concurrent); + assert_eq!( + rows.load(AtomicOrdering::Relaxed), + if enabled { 8 } else { 16 } + ); + let expected = arrow::compute::take_record_batch( + &batch, + &UInt32Array::from(vec![7, 3, 6, 2, 5, 1, 4, 0]), + )?; + assert_eq!(concat_batches(&schema, &output)?, expected); + } + Ok(()) + } + + #[tokio::test] + async fn test_staged_sort_multi_run_inline() -> Result<()> { + check_staged_multi_run_path(false).await + } + + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn test_staged_sort_multi_run_concurrent() -> Result<()> { + check_staged_multi_run_path(true).await + } + + #[tokio::test] + async fn test_staged_sort_zero_arg_udf() -> Result<()> { + use datafusion_expr::{ + ColumnarValue, ScalarFunctionArgs, ScalarUDF, ScalarUDFImpl, Signature, + Volatility, + }; + use datafusion_physical_expr::ScalarFunctionExpr; + static ROWS: AtomicUsize = AtomicUsize::new(0); + #[derive(Debug, PartialEq, Eq, Hash)] + struct ZeroArg { + signature: Signature, + } + impl ScalarUDFImpl for ZeroArg { + fn name(&self) -> &str { + "staged_zero_arg" + } + fn signature(&self) -> &Signature { + &self.signature + } + fn return_type(&self, _: &[DataType]) -> Result { + Ok(DataType::Int32) + } + fn invoke_with_args( + &self, + args: ScalarFunctionArgs, + ) -> Result { + assert!(args.args.is_empty()); + ROWS.fetch_add(args.number_rows, AtomicOrdering::Relaxed); + Ok(ColumnarValue::Array(Arc::new(Int32Array::from( + vec![7; args.number_rows], + )))) + } + } + let schema = Arc::new(Schema::new(vec![Field::new("p", DataType::Int32, false)])); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(Int32Array::from(vec![5, 1, 3, 2, 4, 1, 6, 7]))], + )?; + let udf = Arc::new(ScalarUDF::from(ZeroArg { + signature: Signature::exact(vec![], Volatility::Immutable), + })); + let expression = Arc::new(ScalarFunctionExpr::try_new( + udf, + vec![], + &schema, + Arc::new(ConfigOptions::default()), + )?); + let ordering = [ + PhysicalSortExpr::new_default(col("p", &schema)?), + PhysicalSortExpr::new_default(expression), + ] + .into(); + for (input_batches, threshold) in [(1, 1), (2, usize::MAX)] { + ROWS.store(0, AtomicOrdering::Relaxed); + let result = staged_sort_test_output( + &batch, + &ordering, + true, + input_batches, + threshold, + ) + .await?; + assert_eq!(ROWS.load(AtomicOrdering::Relaxed), 2); + let actual = concat_batches(&schema, &result)?; + assert_eq!( + actual.column(0).as_ref(), + &Int32Array::from(vec![1, 1, 2, 3, 4, 5, 6, 7]) + ); + } + Ok(()) + } + + #[tokio::test] + async fn test_staged_sort_fallback_reuses_prefix() -> Result<()> { + use datafusion_expr::Volatility; + + let schema = Arc::new(Schema::new(vec![ + Field::new("p", DataType::Int32, false), + Field::new("s", DataType::Utf8, false), + ])); + let (prefix, rows) = counting_identity( + "count_prefix", + col("p", &schema)?, + &schema, + Volatility::Immutable, + )?; + let suffix = Arc::new(CastExpr::new(col("s", &schema)?, DataType::Int32, None)); + let ordering = [ + PhysicalSortExpr::new_default(prefix), + PhysicalSortExpr::new_default(suffix), + ] + .into(); + for (prefix, suffix, expected_indices) in [ + ( + vec![1; 8], + vec!["8", "7", "6", "5", "4", "3", "2", "1"], + vec![7, 6, 5, 4, 3, 2, 1, 0], + ), + ( + vec![0, 0, 1, 1, 2, 2, 3, 3], + vec!["8", "7", "6", "5", "4", "3", "2", "1"], + vec![1, 0, 3, 2, 5, 4, 7, 6], + ), + ] { + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int32Array::from(prefix)), + Arc::new(StringArray::from(suffix)), + ], + )?; + let expected = arrow::compute::take_record_batch( + &batch, + &UInt32Array::from(expected_indices), + )?; + rows.store(0, AtomicOrdering::Relaxed); + let result = staged_sort_test_output(&batch, &ordering, true, 1, 1).await?; + assert_eq!(rows.load(AtomicOrdering::Relaxed), 8); + assert_eq!(concat_batches(&schema, &result)?, expected); + } + Ok(()) + } + + #[tokio::test] + async fn test_staged_sort_fallback_ignores_resolved_rows() -> Result<()> { + let batch = record_batch!( + ("prefix", Int32, [0, 0, 1, 1, 2, 3, 3, 4, 4]), + ( + "suffix", + Utf8, + ["8", "7", "6", "5", "invalid", "4", "3", "2", "1"] + ) + )?; + let schema = batch.schema(); + let suffix = Arc::new(CastExpr::new( + col("suffix", &schema)?, + DataType::Int32, + None, + )); + let ordering = [ + PhysicalSortExpr::new_default(col("prefix", &schema)?), + PhysicalSortExpr::new_default(suffix), + ] + .into(); + let expected = arrow::compute::take_record_batch( + &batch, + &UInt32Array::from(vec![1, 0, 3, 2, 4, 6, 5, 8, 7]), + )?; + + let output = staged_sort_test_output(&batch, &ordering, true, 1, 1).await?; + assert_eq!(concat_batches(&schema, &output)?, expected); + Ok(()) + } + + #[tokio::test] + async fn test_staged_sort_fallback_after_second_stage() -> Result<()> { + use datafusion_expr::Volatility; + + let batch = record_batch!( + ("prefix", Int32, [0, 0, 0, 0, 1, 1, 1, 1, 2, 3, 4, 5]), + ("second", Int32, [0, 0, 1, 1, 0, 0, 1, 1, 0, 0, 0, 0]), + ("third", Int32, [2, 1, 2, 1, 2, 1, 2, 1, 0, 0, 0, 0]), + ("fourth", Int32, [9, 8, 7, 6, 5, 4, 3, 2, 0, 0, 0, 0]) + )?; + let schema = batch.schema(); + let (second, second_rows) = counting_identity( + "fallback_second", + col("second", &schema)?, + &schema, + Volatility::Immutable, + )?; + let (third, third_rows) = counting_identity( + "fallback_third", + col("third", &schema)?, + &schema, + Volatility::Immutable, + )?; + let (fourth, fourth_rows) = counting_identity( + "fallback_fourth", + col("fourth", &schema)?, + &schema, + Volatility::Immutable, + )?; + let ordering = [ + PhysicalSortExpr::new_default(col("prefix", &schema)?), + PhysicalSortExpr::new_default(second), + PhysicalSortExpr::new_default(third), + PhysicalSortExpr::new_default(fourth), + ] + .into(); + let expected = arrow::compute::take_record_batch( + &batch, + &UInt32Array::from(vec![1, 0, 3, 2, 5, 4, 7, 6, 8, 9, 10, 11]), + )?; + + let output = staged_sort_test_output(&batch, &ordering, true, 1, 1).await?; + assert_eq!(concat_batches(&schema, &output)?, expected); + assert_eq!(second_rows.load(AtomicOrdering::Relaxed), 8); + assert_eq!(third_rows.load(AtomicOrdering::Relaxed), 8); + assert_eq!(fourth_rows.load(AtomicOrdering::Relaxed), 8); + Ok(()) + } + + #[tokio::test] + async fn test_staged_sort_successive_deferred_groups() -> Result<()> { + use datafusion_expr::Volatility; + + let batch = record_batch!( + ("prefix", Int32, [0, 0, 0, 1, 1, 1, 2, 3, 4, 5, 6, 7]), + ("second", Int32, [1, 1, 2, 1, 1, 2, 0, 0, 0, 0, 0, 0]), + ("third", Int32, [2, 1, 0, 2, 1, 0, 0, 0, 0, 0, 0, 0]) + )?; + let schema = batch.schema(); + let (second, second_rows) = counting_identity( + "second_stage", + col("second", &schema)?, + &schema, + Volatility::Immutable, + )?; + let (third, third_rows) = counting_identity( + "third_stage", + col("third", &schema)?, + &schema, + Volatility::Immutable, + )?; + let ordering = [ + PhysicalSortExpr::new_default(col("prefix", &schema)?), + PhysicalSortExpr::new_default(second), + PhysicalSortExpr::new_default(third), + ] + .into(); + let expected = arrow::compute::take_record_batch( + &batch, + &UInt32Array::from(vec![1, 0, 2, 4, 3, 5, 6, 7, 8, 9, 10, 11]), + )?; + for (input_batches, threshold) in [(1, 1), (2, usize::MAX)] { + second_rows.store(0, AtomicOrdering::Relaxed); + third_rows.store(0, AtomicOrdering::Relaxed); + let output = staged_sort_test_output( + &batch, + &ordering, + true, + input_batches, + threshold, + ) + .await?; + assert_eq!(concat_batches(&schema, &output)?, expected); + assert_eq!(second_rows.load(AtomicOrdering::Relaxed), 6); + assert_eq!(third_rows.load(AtomicOrdering::Relaxed), 4); + } + Ok(()) + } + #[tokio::test] async fn test_staged_sort_matches_eager() -> Result<()> { - use datafusion_physical_expr::expressions::CastExpr; + use datafusion_expr::Volatility; - // Keep payload columns between the keys to exercise projected column remapping. + // Deferred keys reference source columns 1 and 3, which are remapped + // to columns 0 and 1 after projecting away the payload columns. let schema = Arc::new(Schema::new(vec![ + Field::new("payload", DataType::Utf8, false), Field::new("id", DataType::Int32, false), + Field::new("other_payload", DataType::Int32, false), Field::new("suffix", DataType::Utf8, true), - Field::new("payload", DataType::Utf8, false), Field::new("prefix", DataType::Int32, true), ])); let batch = RecordBatch::try_new( Arc::clone(&schema), vec![ - Arc::new(Int32Array::from_iter_values(0..10)), + Arc::new(StringArray::from_iter_values( + (0..16).map(|i| format!("payload-{i}")), + )), + Arc::new(Int32Array::from_iter_values(0..16)), + Arc::new(Int32Array::from_iter_values(100..116)), Arc::new(StringArray::from(vec![ Some("3"), Some("2"), @@ -4369,10 +4893,13 @@ mod tests { Some("2"), Some("0"), Some("10"), + Some("1"), + Some("1"), + Some("1"), + Some("1"), + Some("1"), + Some("1"), ])), - Arc::new(StringArray::from_iter_values( - (0..10).map(|i| format!("payload-{i}")), - )), Arc::new(Int32Array::from(vec![ Some(2), Some(1), @@ -4384,9 +4911,22 @@ mod tests { Some(1), Some(3), Some(1), + Some(4), + Some(5), + Some(6), + Some(7), + Some(8), + Some(9), ])), ], )?; + let cast = Arc::new(CastExpr::new( + col("suffix", &schema)?, + DataType::Int32, + None, + )); + let (suffix, suffix_rows) = + counting_identity("count_cast_rows", cast, &schema, Volatility::Immutable)?; let ordering = [ PhysicalSortExpr::new( col("prefix", &schema)?, @@ -4396,11 +4936,7 @@ mod tests { }, ), PhysicalSortExpr::new( - Arc::new(CastExpr::new( - col("suffix", &schema)?, - DataType::Int32, - None, - )), + suffix, SortOptions { descending: true, nulls_first: true, @@ -4409,15 +4945,20 @@ mod tests { PhysicalSortExpr::new_default(col("id", &schema)?), ] .into(); - // Prefix ties survive the numeric suffix stage and must be ordered by id. + + // Nine rows remain tied in three groups; seven rows are already resolved. + // The prefix and numeric suffix still tie for some rows, so id breaks ties. let expected = arrow::compute::take_record_batch( &batch, - &UInt32Array::from(vec![3, 5, 9, 1, 7, 0, 4, 8, 6, 2]), + &UInt32Array::from(vec![ + 3, 5, 9, 1, 7, 0, 4, 8, 10, 11, 12, 13, 14, 15, 6, 2, + ]), )?; // Single batch, concatenated batches, and separate runs followed by a merge. for (input_batches, threshold) in [(1, 1), (2, usize::MAX), (2, 1)] { for enabled in [false, true] { + suffix_rows.store(0, AtomicOrdering::Relaxed); let output = staged_sort_test_output( &batch, &ordering, @@ -4426,11 +4967,15 @@ mod tests { threshold, ) .await?; + assert_eq!( concat_batches(&schema, &output)?, expected, "enabled={enabled}, input_batches={input_batches}, threshold={threshold}" ); + if enabled && (input_batches == 1 || threshold == usize::MAX) { + assert_eq!(suffix_rows.load(AtomicOrdering::Relaxed), 9); + } } } Ok(()) @@ -4453,45 +4998,98 @@ mod tests { ))), ] .into(); - for all_tied in [false, true] { - let prefix = if all_tied { - vec![1; 6] - } else { - vec![3, 1, 6, 2, 5, 4] - }; + struct Case<'a> { + name: &'a str, + prefix: &'a [i32; 6], + enabled: bool, + input_batches: usize, + threshold: usize, + expect_error: bool, + } + + let distinct = [3, 1, 6, 2, 5, 4]; + let all_tied = [1; 6]; + let cases = [ + Case { + name: "disabled evaluates eagerly", + prefix: &distinct, + enabled: false, + input_batches: 1, + threshold: 1, + expect_error: true, + }, + Case { + name: "distinct rows skip suffix", + prefix: &distinct, + enabled: true, + input_batches: 1, + threshold: 1, + expect_error: false, + }, + Case { + name: "ties evaluate suffix", + prefix: &all_tied, + enabled: true, + input_batches: 1, + threshold: 1, + expect_error: true, + }, + Case { + name: "concatenated rows skip suffix", + prefix: &distinct, + enabled: true, + input_batches: 2, + threshold: usize::MAX, + expect_error: false, + }, + Case { + name: "merge reevaluates suffix", + prefix: &distinct, + enabled: true, + input_batches: 2, + threshold: 1, + expect_error: true, + }, + ]; + + for Case { + name, + prefix, + enabled, + input_batches, + threshold, + expect_error, + } in cases + { let batch = RecordBatch::try_new( Arc::clone(&schema), vec![ - Arc::new(Int32Array::from(prefix)), + Arc::new(Int32Array::from(prefix.to_vec())), Arc::new(StringArray::from(vec!["invalid"; 6])), ], )?; - for (input_batches, threshold) in [(1, 1), (2, usize::MAX), (2, 1)] { - for enabled in [false, true] { - let result = staged_sort_test_output( - &batch, - &ordering, - enabled, - input_batches, - threshold, - ) - .await; - // A merge evaluates all keys again, even if each run skipped the cast. - let merge = input_batches == 2 && threshold == 1; - if !enabled || all_tied || merge { - let error = result.unwrap_err().strip_backtrace(); - assert!( - error.contains("Cannot cast string 'invalid'"), - "{error}" - ); - } else { - let output = concat_batches(&schema, &result?)?; - assert_eq!( - output.column(0).as_ref(), - &Int32Array::from_iter_values(1..=6), - ); - } - } + let result = staged_sort_test_output( + &batch, + &ordering, + enabled, + input_batches, + threshold, + ) + .await; + + if expect_error { + let error = result.unwrap_err().strip_backtrace(); + assert!( + error.contains("Cannot cast string 'invalid'"), + "{name}: {error}" + ); + } else { + let output = concat_batches(&schema, &result?)?; + assert_eq!( + output.column(0).as_ref(), + &Int32Array::from_iter_values(1..=6), + "{name}" + ); } } Ok(()) @@ -4499,77 +5097,90 @@ mod tests { #[tokio::test] async fn test_staged_sort_volatile_evaluation() -> Result<()> { - use datafusion_expr::{Volatility, create_udf}; - use datafusion_physical_expr::ScalarFunctionExpr; + use datafusion_expr::Volatility; let schema = Arc::new(Schema::new(vec![ Field::new("prefix", DataType::Int32, false), Field::new("id", DataType::Int32, false), ])); - // Count rows, not calls: expression evaluation is batched. The middle - // case has four tied rows, with two in each of the split input batches. - for (prefix, tied_rows) in [ - (vec![3, 1, 6, 2, 5, 4], 0), - (vec![2, 1, 1, 3, 4, 4], 4), - (vec![1; 6], 6), - ] { + struct Case<'a> { + name: &'a str, + prefix: &'a [i32; 6], + enabled: bool, + input_batches: usize, + threshold: usize, + expected_rows: usize, + } + + let distinct = [3, 1, 6, 2, 5, 4]; + let partly_tied = [2, 1, 1, 3, 4, 4]; + let cases = [ + Case { + name: "distinct enabled", + prefix: &distinct, + enabled: true, + input_batches: 1, + threshold: 1, + expected_rows: 6, + }, + Case { + name: "partly tied disabled", + prefix: &partly_tied, + enabled: false, + input_batches: 1, + threshold: 1, + expected_rows: 6, + }, + Case { + name: "partly tied enabled", + prefix: &partly_tied, + enabled: true, + input_batches: 1, + threshold: 1, + expected_rows: 6, + }, + Case { + name: "merge reevaluates", + prefix: &partly_tied, + enabled: true, + input_batches: 2, + threshold: 1, + expected_rows: 12, + }, + ]; + + for Case { + name, + prefix, + enabled, + input_batches, + threshold, + expected_rows, + } in cases + { let batch = RecordBatch::try_new( Arc::clone(&schema), vec![ - Arc::new(Int32Array::from(prefix)), + Arc::new(Int32Array::from(prefix.to_vec())), Arc::new(Int32Array::from_iter_values(0..6)), ], )?; - for (input_batches, threshold) in [(1, 1), (2, usize::MAX), (2, 1)] { - for enabled in [false, true] { - let rows_evaluated = Arc::new(AtomicUsize::new(0)); - let counter = Arc::clone(&rows_evaluated); - let function = create_udf( - "count_evaluations", - vec![DataType::Int32], - DataType::Int32, - Volatility::Volatile, - Arc::new(move |args| { - counter.fetch_add( - args[0].clone().into_array(1)?.len(), - AtomicOrdering::Relaxed, - ); - Ok(args[0].clone()) - }), - ); - let ordering = [ - PhysicalSortExpr::new_default(col("prefix", &schema)?), - PhysicalSortExpr::new_default(Arc::new( - ScalarFunctionExpr::try_new( - Arc::new(function), - vec![col("id", &schema)?], - &schema, - Arc::new(ConfigOptions::default()), - )?, - )), - ] - .into(); - staged_sort_test_output( - &batch, - &ordering, - enabled, - input_batches, - threshold, - ) - .await?; - let merge_rows = if input_batches == 2 && threshold == 1 { - 6 - } else { - 0 - }; - let sort_rows = if enabled { tied_rows } else { 6 }; - assert_eq!( - rows_evaluated.load(AtomicOrdering::Relaxed), - sort_rows + merge_rows, - "enabled={enabled}, tied_rows={tied_rows}, input_batches={input_batches}, threshold={threshold}" - ); - } - } + let (volatile, rows) = counting_identity( + "count_evaluations", + col("id", &schema)?, + &schema, + Volatility::Volatile, + )?; + let nested = Arc::new(BinaryExpr::new(volatile, Operator::Plus, lit(0i32))); + let ordering = [ + PhysicalSortExpr::new_default(col("prefix", &schema)?), + PhysicalSortExpr::new_default(nested), + ] + .into(); + + staged_sort_test_output(&batch, &ordering, enabled, input_batches, threshold) + .await?; + assert_eq!(rows.load(AtomicOrdering::Relaxed), expected_rows, "{name}"); } Ok(()) } diff --git a/datafusion/physical-plan/src/sorts/stream.rs b/datafusion/physical-plan/src/sorts/stream.rs index fc19b3307f00a..c0f92d785357d 100644 --- a/datafusion/physical-plan/src/sorts/stream.rs +++ b/datafusion/physical-plan/src/sorts/stream.rs @@ -327,9 +327,8 @@ impl PartitionedStream for FieldCursorStream { /// spill when there is not enough memory for an eager sort (which requires ~2x /// peak memory to hold both the unsorted and sorted copies simultaneously). /// -/// Unless supplied via [`Self::with_sorted_indices`], a sorted index array -/// (`UInt32Array`) is computed via `lexsort_to_indices` on the first call to -/// `next()`. The iterator yields chunks of +/// On the first call to `next()`, a sorted index array (`UInt32Array`) is +/// computed via `lexsort_to_indices`. Subsequent calls yield chunks of /// `batch_size` rows by `take`-ing from the original batch using slices of /// this index array. Each `take` copies data for the chunk (not zero-copy), /// but only one chunk is live at a time since the caller consumes it before diff --git a/datafusion/sqllogictest/test_files/information_schema.slt b/datafusion/sqllogictest/test_files/information_schema.slt index fae5cd7306d91..2d742902621d3 100644 --- a/datafusion/sqllogictest/test_files/information_schema.slt +++ b/datafusion/sqllogictest/test_files/information_schema.slt @@ -275,7 +275,6 @@ datafusion.execution.skip_physical_aggregate_schema_check false datafusion.execution.soft_max_bytes_per_output_file 4294967295 datafusion.execution.soft_max_rows_per_output_file 50000000 datafusion.execution.sort_in_place_threshold_bytes 1048576 -datafusion.execution.sort_key_group_size 5 datafusion.execution.sort_pushdown_buffer_capacity 1073741824 datafusion.execution.sort_spill_reservation_bytes 10485760 datafusion.execution.spill_compression uncompressed @@ -382,7 +381,7 @@ datafusion.execution.enable_file_stream_work_stealing true When `true` (the defa datafusion.execution.enable_migration_aggregate true Whether aggregation uses the implementation from the major refactor completed in the 56.0.0 release. When set to `false`, aggregation falls back to the implementation used before 55.0.0. The fallback exists only as a workaround for bugs in the new implementation and will be removed, together with this option, after the 56.0.0 release. See for details. datafusion.execution.enable_nlj_coordinated_fallback true Enables the memory-limited fallback for `NestedLoopJoinExec` join types that emit unmatched left rows in the final output (LEFT, LEFT SEMI, LEFT ANTI, LEFT MARK, FULL) when the right side has multiple partitions. This fallback shares left-side state (the current left chunk, the visited bitmap and the probe-thread counter) across all right-side partitions, which assumes every partition runs in the same process. Distributed engines that execute each output partition as an independent task (e.g. Ballista, datafusion-distributed) give each task its own copy of this state and poll only one partition, so the cross-partition counter never reaches zero and the fallback would stall. Such engines should set this to `false`: the coordinated fallback is then disabled for left-emitting multi-partition joins, which instead fail with a resource-exhaustion error under memory pressure rather than deadlocking. Single-partition and non-left-emitting joins are unaffected and always keep the fallback. datafusion.execution.enable_recursive_ctes true Should DataFusion support recursive CTEs -datafusion.execution.enable_staged_sort false Enables staged in-memory sorting, evaluating later sort keys only for rows tied on earlier keys. +datafusion.execution.enable_staged_sort false Enables staged sorting for eligible in-memory sorts. Later keys may be skipped for rows already in order TopK and spill sorting do not use staged sorting. datafusion.execution.enforce_batch_size_in_joins false Should DataFusion enforce batch size in joins or not. By default, DataFusion will not enforce batch size in joins. Enforcing batch size in joins can reduce memory usage when joining large tables with a highly-selective join filter, but is also slightly slower. Note: this option currently only applies to the symmetric hash join. datafusion.execution.hash_join_buffering_capacity 0 How many bytes to buffer in the probe side of hash joins while the build side is concurrently being built. Without this, hash joins will wait until the full materialization of the build side before polling the probe side. This is useful in scenarios where the query is not completely CPU bounded, allowing to do some early work concurrently and reducing the latency of the query. Note that when hash join buffering is enabled, the probe side will start eagerly polling data, not giving time for the producer side of dynamic filters to produce any meaningful predicate. Queries with dynamic filters might see performance degradation. Disabled by default, set to a number greater than 0 for enabling it. datafusion.execution.keep_partition_by_columns false Should DataFusion keep the columns used for partition_by in the output RecordBatches @@ -444,7 +443,6 @@ datafusion.execution.skip_physical_aggregate_schema_check false When set to true datafusion.execution.soft_max_bytes_per_output_file 4294967295 Target encoded size in bytes of output files when writing multiple. Writers asynchronously report the cumulative encoded size as they process RecordBatches. The final file size may exceed this limit due to batches buffered before the limit is observed, the size of a batch, and file metadata written when the file is finalized. datafusion.execution.soft_max_rows_per_output_file 50000000 Target number of rows in output files when writing multiple. This is a soft max, so it can be exceeded slightly. There also will be one file smaller than the limit if the total number of rows written is not roughly divisible by the soft max datafusion.execution.sort_in_place_threshold_bytes 1048576 When sorting, below what size should data be concatenated and sorted in a single RecordBatch rather than sorted in batches and merged. -datafusion.execution.sort_key_group_size 5 Number of sort keys evaluated together by staged in-memory sorting. datafusion.execution.sort_pushdown_buffer_capacity 1073741824 Maximum buffer capacity (in bytes) per partition for BufferExec inserted during sort pushdown optimization. When PushdownSort eliminates a SortExec under SortPreservingMergeExec, a BufferExec is inserted to replace SortExec's buffering role. This prevents I/O stalls by allowing the scan to run ahead of the merge. This uses strictly less memory than the SortExec it replaces (which buffers the entire partition). The buffer respects the global memory pool limit. Setting this to a large value is safe — actual memory usage is bounded by partition size and global memory limits. datafusion.execution.sort_spill_reservation_bytes 10485760 Specifies the reserved memory for each spillable sort operation to facilitate an in-memory merge. When a sort operation spills to disk, the in-memory data must be sorted and merged before being written to a file. This setting reserves a specific amount of memory for that in-memory sort/merge process. Note: This setting is irrelevant if the sort operation cannot spill (i.e., if there's no `DiskManager` configured). datafusion.execution.spill_compression uncompressed Sets the compression codec used when spilling data to disk. Since datafusion writes spill files using the Arrow IPC Stream format, only codecs supported by the Arrow IPC Stream Writer are allowed. Valid values are: uncompressed, lz4_frame, zstd. Note: lz4_frame offers faster (de)compression, but typically results in larger spill files. In contrast, zstd achieves higher compression ratios at the cost of slower (de)compression speed. diff --git a/docs/source/user-guide/configs.md b/docs/source/user-guide/configs.md index 445c63a05c93f..89479dbce2e21 100644 --- a/docs/source/user-guide/configs.md +++ b/docs/source/user-guide/configs.md @@ -116,6 +116,7 @@ The following configuration settings are available: | datafusion.execution.spill_compression | uncompressed | Sets the compression codec used when spilling data to disk. Since datafusion writes spill files using the Arrow IPC Stream format, only codecs supported by the Arrow IPC Stream Writer are allowed. Valid values are: uncompressed, lz4_frame, zstd. Note: lz4_frame offers faster (de)compression, but typically results in larger spill files. In contrast, zstd achieves higher compression ratios at the cost of slower (de)compression speed. | | datafusion.execution.sort_spill_reservation_bytes | 10485760 | Specifies the reserved memory for each spillable sort operation to facilitate an in-memory merge. When a sort operation spills to disk, the in-memory data must be sorted and merged before being written to a file. This setting reserves a specific amount of memory for that in-memory sort/merge process. Note: This setting is irrelevant if the sort operation cannot spill (i.e., if there's no `DiskManager` configured). | | datafusion.execution.sort_in_place_threshold_bytes | 1048576 | When sorting, below what size should data be concatenated and sorted in a single RecordBatch rather than sorted in batches and merged. | +| datafusion.execution.enable_staged_sort | false | Enables staged sorting for eligible in-memory sorts. Later keys may be skipped for rows already in order TopK and spill sorting do not use staged sorting. | | datafusion.execution.sort_pushdown_buffer_capacity | 1073741824 | Maximum buffer capacity (in bytes) per partition for BufferExec inserted during sort pushdown optimization. When PushdownSort eliminates a SortExec under SortPreservingMergeExec, a BufferExec is inserted to replace SortExec's buffering role. This prevents I/O stalls by allowing the scan to run ahead of the merge. This uses strictly less memory than the SortExec it replaces (which buffers the entire partition). The buffer respects the global memory pool limit. Setting this to a large value is safe — actual memory usage is bounded by partition size and global memory limits. | | datafusion.execution.max_spill_file_size_bytes | 134217728 | Maximum size in bytes for individual spill files before rotating to a new file. When operators spill data to disk (e.g., RepartitionExec), they write multiple batches to the same file until this size limit is reached, then rotate to a new file. This reduces syscall overhead compared to one-file-per-batch while preventing files from growing too large. A larger value reduces file creation overhead but may hold more disk space. A smaller value creates more files but allows finer-grained space reclamation as files can be deleted once fully consumed. Now only `RepartitionExec` supports this spill file rotation feature, other spilling operators may create spill files larger than the limit. Default: 128 MB | | datafusion.execution.enable_nlj_coordinated_fallback | true | Enables the memory-limited fallback for `NestedLoopJoinExec` join types that emit unmatched left rows in the final output (LEFT, LEFT SEMI, LEFT ANTI, LEFT MARK, FULL) when the right side has multiple partitions. This fallback shares left-side state (the current left chunk, the visited bitmap and the probe-thread counter) across all right-side partitions, which assumes every partition runs in the same process. Distributed engines that execute each output partition as an independent task (e.g. Ballista, datafusion-distributed) give each task its own copy of this state and poll only one partition, so the cross-partition counter never reaches zero and the fallback would stall. Such engines should set this to `false`: the coordinated fallback is then disabled for left-emitting multi-partition joins, which instead fail with a resource-exhaustion error under memory pressure rather than deadlocking. Single-partition and non-left-emitting joins are unaffected and always keep the fallback. | From db36fdfc484a6c83ce3efd03a6f819c0d3aa3e71 Mon Sep 17 00:00:00 2001 From: yashrb24 Date: Tue, 3 Nov 2026 22:30:15 +0530 Subject: [PATCH 11/12] test: trim staged sort coverage --- datafusion/physical-plan/src/sorts/sort.rs | 495 ++++----------------- 1 file changed, 77 insertions(+), 418 deletions(-) diff --git a/datafusion/physical-plan/src/sorts/sort.rs b/datafusion/physical-plan/src/sorts/sort.rs index 04e59c6a2182b..d0cc3d4807e63 100644 --- a/datafusion/physical-plan/src/sorts/sort.rs +++ b/datafusion/physical-plan/src/sorts/sort.rs @@ -3038,52 +3038,6 @@ mod tests { abort_count: Arc, } - #[test] - fn test_staged_sort_automatic_groups() -> Result<()> { - let schema = Arc::new(Schema::new(vec![ - Field::new("i", DataType::Int32, false), - Field::new("s", DataType::Utf8, false), - ])); - let cheap = col("i", &schema)?; - let expensive: Arc = - Arc::new(CastExpr::new(col("s", &schema)?, DataType::Int32, None)); - let group_sizes = |classes: &[bool]| -> Result> { - let keys = - LexOrdering::new(classes.iter().enumerate().map(|(index, cheap_key)| { - let expr = if *cheap_key { - Arc::new(BinaryExpr::new( - Arc::clone(&cheap), - Operator::Plus, - lit(index as i32), - )) as Arc - } else { - Arc::new(BinaryExpr::new( - Arc::clone(&expensive), - Operator::Plus, - lit(index as i32), - )) - }; - PhysicalSortExpr::new_default(expr) - })) - .unwrap(); - Ok(create_sort_stages(&keys, &schema)? - .map(|plan| plan.iter().map(|stage| stage.expressions.len()).collect()) - .unwrap_or_default()) - }; - assert_eq!( - group_sizes(&[true, true, false, false, true, false])?, - vec![2, 1, 2, 1] - ); - assert_eq!(group_sizes(&[false, false, true, false])?, vec![1, 2, 1]); - assert_eq!(group_sizes(&[true, false])?, vec![1, 1]); - assert_eq!(group_sizes(&[false, false])?, vec![1, 1]); - assert_eq!(group_sizes(&[false; 5])?, vec![1, 1, 1, 1, 1]); - assert!(group_sizes(&[true; 5])?.is_empty()); - assert!(group_sizes(&[false])?.is_empty()); - - Ok(()) - } - impl TempFileFactory for PendingWriteTempFileFactory { fn create_temp_file(&self, _description: &str) -> Result> { Ok(Arc::new(PendingWriteSpillFile { @@ -4548,66 +4502,6 @@ mod tests { result } - async fn check_staged_multi_run_path(concurrent: bool) -> Result<()> { - use datafusion_expr::Volatility; - - let batch = record_batch!( - ("p", Int32, [8, 6, 4, 2, 7, 5, 3, 1]), - ("v", Int32, [8, 6, 4, 2, 7, 5, 3, 1]) - )?; - let schema = batch.schema(); - for enabled in [false, true] { - let (suffix, rows) = counting_identity( - "count_suffix", - col("v", &schema)?, - &schema, - Volatility::Immutable, - )?; - - let ordering = [ - PhysicalSortExpr::new_default(col("p", &schema)?), - PhysicalSortExpr::new_default(suffix), - ] - .into(); - let input = TestMemoryExec::try_new_exec( - &[vec![batch.slice(0, 4), batch.slice(4, 4)]], - Arc::clone(&schema), - None, - )?; - let sort = Arc::new(SortExec::new(ordering, input)); - let config = SessionConfig::new() - .with_batch_size(2) - .with_sort_in_place_threshold_bytes(1) - .set_bool("datafusion.execution.enable_staged_sort", enabled); - let output = collect( - Arc::clone(&sort) as Arc, - Arc::new(TaskContext::default().with_session_config(config)), - ) - .await?; - assert_eq!(sort.sort_concurrency.get().is_some(), concurrent); - assert_eq!( - rows.load(AtomicOrdering::Relaxed), - if enabled { 8 } else { 16 } - ); - let expected = arrow::compute::take_record_batch( - &batch, - &UInt32Array::from(vec![7, 3, 6, 2, 5, 1, 4, 0]), - )?; - assert_eq!(concat_batches(&schema, &output)?, expected); - } - Ok(()) - } - - #[tokio::test] - async fn test_staged_sort_multi_run_inline() -> Result<()> { - check_staged_multi_run_path(false).await - } - - #[tokio::test(flavor = "multi_thread", worker_threads = 2)] - async fn test_staged_sort_multi_run_concurrent() -> Result<()> { - check_staged_multi_run_path(true).await - } - #[tokio::test] async fn test_staged_sort_zero_arg_udf() -> Result<()> { use datafusion_expr::{ @@ -4707,9 +4601,9 @@ mod tests { vec![7, 6, 5, 4, 3, 2, 1, 0], ), ( - vec![0, 0, 1, 1, 2, 2, 3, 3], - vec!["8", "7", "6", "5", "4", "3", "2", "1"], - vec![1, 0, 3, 2, 5, 4, 7, 6], + vec![0, 0, 1, 1, 2, 3, 3, 4, 4], + vec!["8", "7", "6", "5", "invalid", "4", "3", "2", "1"], + vec![1, 0, 3, 2, 4, 6, 5, 8, 7], ), ] { let batch = RecordBatch::try_new( @@ -4725,43 +4619,12 @@ mod tests { )?; rows.store(0, AtomicOrdering::Relaxed); let result = staged_sort_test_output(&batch, &ordering, true, 1, 1).await?; - assert_eq!(rows.load(AtomicOrdering::Relaxed), 8); + assert_eq!(rows.load(AtomicOrdering::Relaxed), batch.num_rows()); assert_eq!(concat_batches(&schema, &result)?, expected); } Ok(()) } - #[tokio::test] - async fn test_staged_sort_fallback_ignores_resolved_rows() -> Result<()> { - let batch = record_batch!( - ("prefix", Int32, [0, 0, 1, 1, 2, 3, 3, 4, 4]), - ( - "suffix", - Utf8, - ["8", "7", "6", "5", "invalid", "4", "3", "2", "1"] - ) - )?; - let schema = batch.schema(); - let suffix = Arc::new(CastExpr::new( - col("suffix", &schema)?, - DataType::Int32, - None, - )); - let ordering = [ - PhysicalSortExpr::new_default(col("prefix", &schema)?), - PhysicalSortExpr::new_default(suffix), - ] - .into(); - let expected = arrow::compute::take_record_batch( - &batch, - &UInt32Array::from(vec![1, 0, 3, 2, 4, 6, 5, 8, 7]), - )?; - - let output = staged_sort_test_output(&batch, &ordering, true, 1, 1).await?; - assert_eq!(concat_batches(&schema, &output)?, expected); - Ok(()) - } - #[tokio::test] async fn test_staged_sort_fallback_after_second_stage() -> Result<()> { use datafusion_expr::Volatility; @@ -4811,98 +4674,30 @@ mod tests { Ok(()) } - #[tokio::test] - async fn test_staged_sort_successive_deferred_groups() -> Result<()> { - use datafusion_expr::Volatility; - - let batch = record_batch!( - ("prefix", Int32, [0, 0, 0, 1, 1, 1, 2, 3, 4, 5, 6, 7]), - ("second", Int32, [1, 1, 2, 1, 1, 2, 0, 0, 0, 0, 0, 0]), - ("third", Int32, [2, 1, 0, 2, 1, 0, 0, 0, 0, 0, 0, 0]) - )?; - let schema = batch.schema(); - let (second, second_rows) = counting_identity( - "second_stage", - col("second", &schema)?, - &schema, - Volatility::Immutable, - )?; - let (third, third_rows) = counting_identity( - "third_stage", - col("third", &schema)?, - &schema, - Volatility::Immutable, - )?; - let ordering = [ - PhysicalSortExpr::new_default(col("prefix", &schema)?), - PhysicalSortExpr::new_default(second), - PhysicalSortExpr::new_default(third), - ] - .into(); - let expected = arrow::compute::take_record_batch( - &batch, - &UInt32Array::from(vec![1, 0, 2, 4, 3, 5, 6, 7, 8, 9, 10, 11]), - )?; - for (input_batches, threshold) in [(1, 1), (2, usize::MAX)] { - second_rows.store(0, AtomicOrdering::Relaxed); - third_rows.store(0, AtomicOrdering::Relaxed); - let output = staged_sort_test_output( - &batch, - &ordering, - true, - input_batches, - threshold, - ) - .await?; - assert_eq!(concat_batches(&schema, &output)?, expected); - assert_eq!(second_rows.load(AtomicOrdering::Relaxed), 6); - assert_eq!(third_rows.load(AtomicOrdering::Relaxed), 4); - } - Ok(()) - } - #[tokio::test] async fn test_staged_sort_matches_eager() -> Result<()> { use datafusion_expr::Volatility; - // Deferred keys reference source columns 1 and 3, which are remapped - // to columns 0 and 1 after projecting away the payload columns. - let schema = Arc::new(Schema::new(vec![ - Field::new("payload", DataType::Utf8, false), - Field::new("id", DataType::Int32, false), - Field::new("other_payload", DataType::Int32, false), - Field::new("suffix", DataType::Utf8, true), - Field::new("prefix", DataType::Int32, true), - ])); - let batch = RecordBatch::try_new( - Arc::clone(&schema), - vec![ - Arc::new(StringArray::from_iter_values( - (0..16).map(|i| format!("payload-{i}")), - )), - Arc::new(Int32Array::from_iter_values(0..16)), - Arc::new(Int32Array::from_iter_values(100..116)), - Arc::new(StringArray::from(vec![ + let batch = record_batch!( + ("id", Int32, [7, 6, 5, 4, 3, 2, 1, 0]), + ( + "suffix", + Utf8, + [ Some("3"), Some("2"), - Some("5"), + Some("0"), None, Some("3"), Some("10"), - None, - Some("2"), - Some("0"), - Some("10"), - Some("1"), Some("1"), - Some("1"), - Some("1"), - Some("1"), - Some("1"), - ])), - Arc::new(Int32Array::from(vec![ - Some(2), - Some(1), + Some("2") + ] + ), + ( + "prefix", + Int32, + [ None, Some(1), Some(2), @@ -4910,16 +4705,11 @@ mod tests { None, Some(1), Some(3), - Some(1), - Some(4), - Some(5), - Some(6), - Some(7), - Some(8), - Some(9), - ])), - ], + Some(1) + ] + ) )?; + let schema = batch.schema(); let cast = Arc::new(CastExpr::new( col("suffix", &schema)?, DataType::Int32, @@ -4927,6 +4717,12 @@ mod tests { )); let (suffix, suffix_rows) = counting_identity("count_cast_rows", cast, &schema, Volatility::Immutable)?; + let (id, id_rows) = counting_identity( + "count_id_rows", + col("id", &schema)?, + &schema, + Volatility::Immutable, + )?; let ordering = [ PhysicalSortExpr::new( col("prefix", &schema)?, @@ -4942,39 +4738,29 @@ mod tests { nulls_first: true, }, ), - PhysicalSortExpr::new_default(col("id", &schema)?), + PhysicalSortExpr::new_default(id), ] .into(); - - // Nine rows remain tied in three groups; seven rows are already resolved. - // The prefix and numeric suffix still tie for some rows, so id breaks ties. let expected = arrow::compute::take_record_batch( &batch, - &UInt32Array::from(vec![ - 3, 5, 9, 1, 7, 0, 4, 8, 10, 11, 12, 13, 14, 15, 6, 2, - ]), + &UInt32Array::from(vec![3, 5, 7, 1, 2, 6, 4, 0]), )?; - // Single batch, concatenated batches, and separate runs followed by a merge. - for (input_batches, threshold) in [(1, 1), (2, usize::MAX), (2, 1)] { + for input_batches in [1, 2] { for enabled in [false, true] { suffix_rows.store(0, AtomicOrdering::Relaxed); - let output = staged_sort_test_output( - &batch, - &ordering, - enabled, - input_batches, - threshold, - ) - .await?; - + id_rows.store(0, AtomicOrdering::Relaxed); + let output = + staged_sort_test_output(&batch, &ordering, enabled, input_batches, 1) + .await?; assert_eq!( concat_batches(&schema, &output)?, expected, - "enabled={enabled}, input_batches={input_batches}, threshold={threshold}" + "enabled={enabled}, input_batches={input_batches}" ); - if enabled && (input_batches == 1 || threshold == usize::MAX) { - assert_eq!(suffix_rows.load(AtomicOrdering::Relaxed), 9); + if enabled && input_batches == 1 { + assert_eq!(suffix_rows.load(AtomicOrdering::Relaxed), 6); + assert_eq!(id_rows.load(AtomicOrdering::Relaxed), 4); } } } @@ -4983,12 +4769,11 @@ mod tests { #[tokio::test] async fn test_staged_sort_expression_errors() -> Result<()> { - use datafusion_physical_expr::expressions::CastExpr; - - let schema = Arc::new(Schema::new(vec![ - Field::new("prefix", DataType::Int32, false), - Field::new("invalid", DataType::Utf8, false), - ])); + let batch = record_batch!( + ("prefix", Int32, [3, 1, 6, 2, 5, 4]), + ("invalid", Utf8, vec!["invalid"; 6]) + )?; + let schema = batch.schema(); let ordering = [ PhysicalSortExpr::new_default(col("prefix", &schema)?), PhysicalSortExpr::new_default(Arc::new(CastExpr::new( @@ -4998,99 +4783,34 @@ mod tests { ))), ] .into(); - struct Case<'a> { - name: &'a str, - prefix: &'a [i32; 6], - enabled: bool, - input_batches: usize, - threshold: usize, - expect_error: bool, - } - let distinct = [3, 1, 6, 2, 5, 4]; - let all_tied = [1; 6]; - let cases = [ - Case { - name: "disabled evaluates eagerly", - prefix: &distinct, - enabled: false, - input_batches: 1, - threshold: 1, - expect_error: true, - }, - Case { - name: "distinct rows skip suffix", - prefix: &distinct, - enabled: true, - input_batches: 1, - threshold: 1, - expect_error: false, - }, - Case { - name: "ties evaluate suffix", - prefix: &all_tied, - enabled: true, - input_batches: 1, - threshold: 1, - expect_error: true, - }, - Case { - name: "concatenated rows skip suffix", - prefix: &distinct, - enabled: true, - input_batches: 2, - threshold: usize::MAX, - expect_error: false, - }, - Case { - name: "merge reevaluates suffix", - prefix: &distinct, - enabled: true, - input_batches: 2, - threshold: 1, - expect_error: true, - }, - ]; + let eager = staged_sort_test_output(&batch, &ordering, false, 1, 1).await; + let skipped = staged_sort_test_output(&batch, &ordering, true, 1, 1).await?; + assert_eq!( + concat_batches(&schema, &skipped)?.column(0).as_ref(), + &Int32Array::from_iter_values(1..=6) + ); - for Case { - name, - prefix, - enabled, - input_batches, - threshold, - expect_error, - } in cases - { - let batch = RecordBatch::try_new( - Arc::clone(&schema), - vec![ - Arc::new(Int32Array::from(prefix.to_vec())), - Arc::new(StringArray::from(vec!["invalid"; 6])), - ], - )?; - let result = staged_sort_test_output( - &batch, - &ordering, - enabled, - input_batches, - threshold, - ) - .await; + let all_tied = RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int32Array::from(vec![1; 6])), + Arc::clone(batch.column(1)), + ], + )?; + let tied = staged_sort_test_output(&all_tied, &ordering, true, 1, 1).await; + let merged = staged_sort_test_output(&batch, &ordering, true, 2, 1).await; - if expect_error { - let error = result.unwrap_err().strip_backtrace(); - assert!( - error.contains("Cannot cast string 'invalid'"), - "{name}: {error}" - ); - } else { - let output = concat_batches(&schema, &result?)?; - assert_eq!( - output.column(0).as_ref(), - &Int32Array::from_iter_values(1..=6), - "{name}" - ); - } + for (name, result) in [ + ("disabled evaluates eagerly", eager), + ("ties evaluate suffix", tied), + ("merge reevaluates suffix", merged), + ] { + let error = result.unwrap_err().strip_backtrace(); + assert!( + error.contains("Cannot cast string 'invalid'"), + "{name}: {error}" + ); } Ok(()) } @@ -5099,72 +4819,12 @@ mod tests { async fn test_staged_sort_volatile_evaluation() -> Result<()> { use datafusion_expr::Volatility; - let schema = Arc::new(Schema::new(vec![ - Field::new("prefix", DataType::Int32, false), - Field::new("id", DataType::Int32, false), - ])); - struct Case<'a> { - name: &'a str, - prefix: &'a [i32; 6], - enabled: bool, - input_batches: usize, - threshold: usize, - expected_rows: usize, - } - - let distinct = [3, 1, 6, 2, 5, 4]; - let partly_tied = [2, 1, 1, 3, 4, 4]; - let cases = [ - Case { - name: "distinct enabled", - prefix: &distinct, - enabled: true, - input_batches: 1, - threshold: 1, - expected_rows: 6, - }, - Case { - name: "partly tied disabled", - prefix: &partly_tied, - enabled: false, - input_batches: 1, - threshold: 1, - expected_rows: 6, - }, - Case { - name: "partly tied enabled", - prefix: &partly_tied, - enabled: true, - input_batches: 1, - threshold: 1, - expected_rows: 6, - }, - Case { - name: "merge reevaluates", - prefix: &partly_tied, - enabled: true, - input_batches: 2, - threshold: 1, - expected_rows: 12, - }, - ]; - - for Case { - name, - prefix, - enabled, - input_batches, - threshold, - expected_rows, - } in cases - { - let batch = RecordBatch::try_new( - Arc::clone(&schema), - vec![ - Arc::new(Int32Array::from(prefix.to_vec())), - Arc::new(Int32Array::from_iter_values(0..6)), - ], - )?; + let batch = record_batch!( + ("prefix", Int32, [2, 1, 1, 3, 4, 4]), + ("id", Int32, [0, 1, 2, 3, 4, 5]) + )?; + let schema = batch.schema(); + for enabled in [false, true] { let (volatile, rows) = counting_identity( "count_evaluations", col("id", &schema)?, @@ -5178,9 +4838,8 @@ mod tests { ] .into(); - staged_sort_test_output(&batch, &ordering, enabled, input_batches, threshold) - .await?; - assert_eq!(rows.load(AtomicOrdering::Relaxed), expected_rows, "{name}"); + staged_sort_test_output(&batch, &ordering, enabled, 1, 1).await?; + assert_eq!(rows.load(AtomicOrdering::Relaxed), 6, "enabled={enabled}"); } Ok(()) } From 99fce229527c40cb313f152932093e5b2b48171e Mon Sep 17 00:00:00 2001 From: Siew Kam Onn Date: Thu, 5 Nov 2026 18:32:11 +0800 Subject: [PATCH 12/12] rebased onto main --- datafusion/physical-plan/src/sorts/sort.rs | 19 +++++++++++++------ 1 file changed, 13 insertions(+), 6 deletions(-) diff --git a/datafusion/physical-plan/src/sorts/sort.rs b/datafusion/physical-plan/src/sorts/sort.rs index d0cc3d4807e63..58780e6745ca4 100644 --- a/datafusion/physical-plan/src/sorts/sort.rs +++ b/datafusion/physical-plan/src/sorts/sort.rs @@ -727,7 +727,8 @@ impl ExternalSorter { if self.in_mem_batches.len() == 1 { let batch = self.in_mem_batches.swap_remove(0); let reservation = self.reservation.take(); - let sorted_stream = self.sort_batch_stream(batch, reservation, None, staged)?; + let sorted_stream = + self.sort_batch_stream(batch, reservation, None, staged)?; return Ok(self.observe_if_output(sorted_stream, is_output_stream)); } @@ -744,7 +745,8 @@ impl ExternalSorter { .try_resize(get_reserved_bytes_for_record_batch(&batch)?) .map_err(Self::err_with_oom_context)?; let reservation = self.reservation.take(); - let sorted_stream = self.sort_batch_stream(batch, reservation, None, staged)?; + let sorted_stream = + self.sort_batch_stream(batch, reservation, None, staged)?; return Ok(self.observe_if_output(sorted_stream, is_output_stream)); } @@ -807,7 +809,8 @@ impl ExternalSorter { .enumerate() .map(move |(index, (batch, size))| { let reservation = self.reservation.split(size); - let sort = self.sort_run(batch, reservation, shared_input.clone(), staged)?; + let sort = + self.sort_run(batch, reservation, shared_input.clone(), staged)?; let permits = Arc::clone(&concurrency.permits); Ok::<_, DataFusionError>(async move { let stream = Self::prepare_sort_run(sort, permits).await?; @@ -1020,9 +1023,13 @@ impl ExternalSorter { staged: bool, ) -> Result { let schema = batch.schema(); - let stream = - futures::stream::once(self.sort_run(batch, reservation, shared_input, staged)?) - .try_flatten(); + let stream = futures::stream::once(self.sort_run( + batch, + reservation, + shared_input, + staged, + )?) + .try_flatten(); Ok(Box::pin(RecordBatchStreamAdapter::new(schema, stream))) }