adriangb opened a new pull request, #25644: URL: https://github.com/apache/datafusion/pull/25644
## Which issue does this PR close? - None. This is follow-up work to https://github.com/apache/datafusion/pull/23985, which added per-query `pool_peak_bytes` to the `dfbench` results JSON. ## Rationale for this change The suites that `bench.sh` runs through the Criterion SQL harness (`cargo bench --bench sql`: `spill_views`, `wide_schema`, `predicate_eval`, and others) do not report a per-query memory pool peak. The harness prints one number to stdout and writes no results JSON. For example, this benchmark bot run of `spill_views` with `DATAFUSION_RUNTIME_MEMORY_LIMIT: "512M"` has no pool peak section: https://github.com/apache/datafusion/pull/23565#issuecomment-5783697679 A second problem: a suite that sets its own limit with `SET datafusion.runtime.memory_limit` records nothing, even if the output were written. `spill_views` does this in its `init` script. The `SET` makes `SessionContext` build a new `RuntimeEnv` with a new pool, and that drops the `PeakRecordingPool` that `CommonOpt::runtime_env_builder` installed. ## What changes are included in this PR? Two commits: 1. **Keep the recorder across a SQL `SET`.** After a benchmark's `load` and `init` steps, `prepare_benchmark` puts a `PeakRecordingPool` in front of the session's pool again if the pool has a finite limit and no recorder. The new `RuntimeEnv` is installed the same way `SET datafusion.runtime.*` does it (`SessionStateBuilder::from(state).with_runtime_env(..)`), so tables and config are kept. The simple runner (`benchmark_runner -o`) also goes through `prepare_benchmark`, so it gets this fix too. 2. **Write a results JSON from the Criterion harness.** When `BENCH_RESULTS_FILE` is set, the harness writes the same JSON format as `dfbench -o`. Each case is the Criterion id (for example, `spill_views/q01_sort_string_1_distinct_repeated`, the same as in `critcmp`). `pool_peak_bytes` covers every execution Criterion makes of the query, including warm-up. The `load`, `init` and `assert` steps are not included. `iterations` is empty, because Criterion keeps the timings. `bench.sh` passes `results/<name>/criterion/<file>.json` at each SQL-harness call site. The subdirectory keeps the file out of `bench.sh compare`, which reads `results/<name>/*.json` as timing results. This means no second timing table. The benchmark bot reads the new file in https://github.com/adriangb/datafusion-benchmarking/pull/38. The bot takes `bench.sh` from `main`, so it shows these tables only after this PR is merged. Until then, and for any side that is older than this PR, the bot shows a note that the pool peaks are not available. Not changed here: `SET datafusion.runtime.memory_limit` always builds a `GreedyMemoryPool` and drops any wrapper, so `--mem-pool-type` has no effect on suites that set their own limit. The fix in this PR stays in the benchmarks crate. A core change that keeps a pool wrapper across `SET` is possible, but it is out of scope for this PR. ## What is the testing strategy for this PR? New unit tests in `benchmarks/src/sql_benchmark_runner.rs`: - `sql_memory_limit_keeps_the_peak_recorder`: a SQL `SET` replaces the harness pool, and the recorder is put back and records the next query. - `no_memory_limit_gets_no_recorder`: without any limit, nothing changes. - `criterion_harness_writes_pool_peak_per_case`: an end-to-end Criterion run of a suite whose `init` sets the limit writes a JSON with a non-zero peak and empty `iterations`. If the rewrap is disabled, the first and third tests fail. `./bench.sh run spill_views`, the command that the bot runs (`SQL_CARGO_COMMAND="cargo bench --bench sql -- --save-baseline ..."`), writes `results/<name>/criterion/spill_views.json`. The peaks are the same with `DATAFUSION_RUNTIME_MEMORY_LIMIT=512M` and with no env limit, because the suite's own limits apply (40M for q01 to q03, 96M for q04 and q05): | Query | `pool_peak_bytes` | | --- | --- | | `spill_views/q01_sort_string_1_distinct_repeated` | 42614784 (40.6 MiB) | | `spill_views/q02_sort_string_1000_distinct_repeated` | 41709504 (39.8 MiB) | | `spill_views/q03_sort_binary_1000_distinct_repeated` | 41709504 (39.8 MiB) | | `spill_views/q04_sort_string_all_distinct_distinct` | 107544576 (102.6 MiB) | | `spill_views/q05_group_by_string_all_distinct_distinct` | 100633968 (96.0 MiB) | The q01 and q04 peaks are a little above their limits. The pool allows this: infallible `grow` calls are not limited, and the recorder counts what the pool grants. `./bench.sh compare` with two such result sets reads no file. On `main`, the output is also empty for a Criterion-only run. `sort_tpch` (dfbench, queries 1 and 2, `DATAFUSION_RUNTIME_MEMORY_LIMIT=512M`) is unchanged. It writes `results/<name>/sort_tpch1.json` with 5 iterations and `pool_peak_bytes` for each query, and it writes no `criterion/` directory. ## Are there any user-facing changes? Only for benchmark tooling. There is a new optional `BENCH_RESULTS_FILE` env var for `cargo bench --bench sql` (documented in `benchmarks/sql_benchmarks/README.md`). `bench.sh run` now writes `results/<name>/criterion/*.json` for SQL-harness suites. A SQL-set memory limit now reports a pool peak. There are no public API changes outside the benchmarks crate. 🤖 Generated with [Claude Code](https://claude.com/claude-code) -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
