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..8265870b9709c --- /dev/null +++ b/benchmarks/sql_benchmarks/staged_sort/staged_sort.benchmark.template @@ -0,0 +1,53 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: See the NOTICE file distributed with this work for additional information regarding copyright ownership +# SPDX-FileContributor: Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements + +# The load and assertions are outside the timed region. +load sql_benchmarks/staged_sort/load.sql + +name ${NAME} +group staged_sort + +init +SET datafusion.execution.target_partitions = 1; + +init +SET datafusion.execution.sort_in_place_threshold_bytes = ${THRESHOLD}; + +assert I +SELECT count(*) = ${STAGED_SORT_ROWS:-200000} AND count(DISTINCT id) = count(*) AND count(*) >= 100 FROM bench; +---- +true + +assert I +SELECT count(*) > 1 AND count(*) < (SELECT count(*) FROM bench) FROM bench WHERE nullable_key IS NULL; +---- +true + +assert I +SELECT count(DISTINCT tie_key) = 1 AND count(DISTINCT id % 100) = 100 FROM bench; +---- +true + +# Nonzero sparse keys equal unique IDs; zero is the only tied group. +assert I +SELECT count(*) FILTER (WHERE sparse_key = 0) > 1 + AND count(*) FILTER (WHERE sparse_key <> 0) > count(*) / 2 + AND count(*) FILTER (WHERE sparse_key <> 0 AND sparse_key <> id) = 0 +FROM bench; +---- +true + +assert I +SELECT min(length(payload)) >= ${STAGED_SORT_PAYLOAD_WIDTH:-32} + 1 FROM bench; +---- +true + +expect_plan ${PLAN} + +run sql_benchmarks/staged_sort/queries/q${QPAD}.sql + +result ${STAGED_SORT_RESULTS_DIR:-sql_benchmarks/staged_sort/results}/q${QPAD}.csv + +cleanup +DROP TABLE bench; 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." diff --git a/datafusion/common/src/config.rs b/datafusion/common/src/config.rs index 5106ea51ea6ba..11c5f1380102d 100644 --- a/datafusion/common/src/config.rs +++ b/datafusion/common/src/config.rs @@ -998,6 +998,11 @@ config_namespace! { /// batches and merged. pub sort_in_place_threshold_bytes: usize, default = 1024 * 1024 + /// 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 + /// 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 53c2929d4aca9..58780e6745ca4 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,12 @@ 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::{ + 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}; @@ -260,6 +268,7 @@ struct ExternalSorter { /// the data will be concatenated and sorted in place rather than /// sort/merged. sort_in_place_threshold_bytes: usize, + sort_stages: Option>, // ======================================================================== // STATE BUFFERS: @@ -373,9 +382,17 @@ impl ExternalSorter { sort_spill_reservation_bytes, sort_in_place_threshold_bytes, sort_concurrency: Arc::new(OnceLock::new()), + sort_stages: None, }) } + fn with_staged_sort(mut self, enabled: bool) -> Result { + if enabled { + self.sort_stages = create_sort_stages(&self.expr, &self.schema)?; + } + Ok(self) + } + /// Appends an unsorted [`RecordBatch`] to `in_mem_batches` /// /// Updates memory usage metrics, and possibly triggers spilling to disk @@ -701,6 +718,8 @@ impl ExternalSorter { let elapsed_compute = self.metrics.baseline.elapsed_compute().clone(); let mut timer = elapsed_compute.timer(); + 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. // The consumer of the merged stream reserves its final output. @@ -708,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)?; + let sorted_stream = + self.sort_batch_stream(batch, reservation, None, staged)?; return Ok(self.observe_if_output(sorted_stream, is_output_stream)); } @@ -725,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)?; + let sorted_stream = + self.sort_batch_stream(batch, reservation, None, staged)?; return Ok(self.observe_if_output(sorted_stream, is_output_stream)); } @@ -746,14 +767,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() @@ -776,6 +797,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)?; @@ -787,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())?; + 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?; @@ -965,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 = @@ -973,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()) + self.sort_batch_stream(run, reservation, shared_input.clone(), staged) }) .collect() } @@ -996,11 +1020,16 @@ 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)?) - .try_flatten(); + let stream = futures::stream::once(self.sort_run( + batch, + reservation, + shared_input, + staged, + )?) + .try_flatten(); Ok(Box::pin(RecordBatchStreamAdapter::new(schema, stream))) } @@ -1013,6 +1042,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 +1054,11 @@ impl ExternalSorter { let expressions = self.expr.clone(); let batch_size = self.batch_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(); @@ -1036,7 +1071,11 @@ 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 let Some(stages) = stages { + staged_sort_batch_chunked(&batch, &expressions, batch_size, &stages)? + } 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 @@ -1159,6 +1198,319 @@ impl ExternalSorter { } } +struct SortStage { + input_columns: Vec, + expressions: Vec, +} + +/// Preserve key order and start a group before each expensive expression. +/// +/// Examples: +/// +/// ```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); + } + } + + 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())) +} + +/// 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( + batch: &RecordBatch, + expressions: &LexOrdering, + batch_size: usize, + stages: &[SortStage], +) -> Result> { + let indices = staged_sort_indices(batch, stages, expressions)?; + IncrementalSortIterator::new(batch.clone(), expressions.clone(), batch_size) + .with_sorted_indices(indices) + .collect() +} + +/// 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; + } + + stage_rows = tied_rows; + tied_groups = refine_tied_groups( + batch, + stage, + &mut indices, + &tied_groups, + index + 1 < stages.len(), + )?; + key_offset += stage.expressions.len(); + } + + Ok(UInt32Array::from(indices)) +} + +/// 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 [u32], + unresolved: &[Range], + find_ties: bool, +) -> Result>> { + let (selected, _) = gather_tied_rows(batch, stage, indices, unresolved)?; + + let sort_columns = stage + .expressions + .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. + let comparator = find_ties + .then(|| LexicographicalComparator::try_new(&sort_columns)) + .transpose()?; + + let mut next_unresolved = Vec::new(); + 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 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, + 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; + } +} + /// 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 @@ -1839,7 +2191,8 @@ impl ExecutionPlan for SortExec { context.session_config().spill_compression(), &self.metrics_set, context.runtime_env(), - )?; + )? + .with_staged_sort(execution_options.enable_staged_sort)?; sorter.sort_concurrency = Arc::clone(&self.sort_concurrency); Ok(Box::pin(RecordBatchStreamAdapter::new( self.schema(), @@ -2489,7 +2842,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::*; @@ -2527,12 +2879,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] @@ -2557,7 +2912,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 { @@ -4086,6 +4441,416 @@ 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. + 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); + 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_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, 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( + 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), batch.num_rows()); + assert_eq!(concat_batches(&schema, &result)?, 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_matches_eager() -> Result<()> { + use datafusion_expr::Volatility; + + let batch = record_batch!( + ("id", Int32, [7, 6, 5, 4, 3, 2, 1, 0]), + ( + "suffix", + Utf8, + [ + Some("3"), + Some("2"), + Some("0"), + None, + Some("3"), + Some("10"), + Some("1"), + Some("2") + ] + ), + ( + "prefix", + Int32, + [ + None, + Some(1), + Some(2), + Some(1), + None, + Some(1), + Some(3), + Some(1) + ] + ) + )?; + let schema = batch.schema(); + 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 (id, id_rows) = counting_identity( + "count_id_rows", + col("id", &schema)?, + &schema, + Volatility::Immutable, + )?; + let ordering = [ + PhysicalSortExpr::new( + col("prefix", &schema)?, + SortOptions { + descending: false, + nulls_first: false, + }, + ), + PhysicalSortExpr::new( + suffix, + SortOptions { + descending: true, + nulls_first: true, + }, + ), + PhysicalSortExpr::new_default(id), + ] + .into(); + let expected = arrow::compute::take_record_batch( + &batch, + &UInt32Array::from(vec![3, 5, 7, 1, 2, 6, 4, 0]), + )?; + + for input_batches in [1, 2] { + for enabled in [false, true] { + suffix_rows.store(0, AtomicOrdering::Relaxed); + 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}" + ); + if enabled && input_batches == 1 { + assert_eq!(suffix_rows.load(AtomicOrdering::Relaxed), 6); + assert_eq!(id_rows.load(AtomicOrdering::Relaxed), 4); + } + } + } + Ok(()) + } + + #[tokio::test] + async fn test_staged_sort_expression_errors() -> Result<()> { + 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( + col("invalid", &schema)?, + DataType::Int32, + None, + ))), + ] + .into(); + + 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) + ); + + 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; + + 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(()) + } + + #[tokio::test] + async fn test_staged_sort_volatile_evaluation() -> Result<()> { + use datafusion_expr::Volatility; + + 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)?, + &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, 1, 1).await?; + assert_eq!(rows.load(AtomicOrdering::Relaxed), 6, "enabled={enabled}"); + } + Ok(()) + } + #[tokio::test] async fn test_sort_batch_chunked_basic() -> Result<()> { let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); diff --git a/datafusion/physical-plan/src/sorts/stream.rs b/datafusion/physical-plan/src/sorts/stream.rs index bdd3cb16dee6e..c0f92d785357d 100644 --- a/datafusion/physical-plan/src/sorts/stream.rs +++ b/datafusion/physical-plan/src/sorts/stream.rs @@ -359,6 +359,11 @@ impl IncrementalSortIterator { indices: None, } } + + 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..2d742902621d3 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 @@ -380,6 +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 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 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. |