Kristofferjoha opened a new pull request, #26592: URL: https://github.com/apache/datafusion/pull/26592
## Which issue does this PR close? Closes #26478. ## Rationale for this change ASOF joins currently collect the complete right input into a non-spillable reservation. Queries fail when that input exceeds the memory pool, even when upstream sorts can spill. Each left partition also scans the collected right input independently. ## What changes are included in this PR? - Add `AsOfJoinMode::Partitioned`, requiring both inputs to be co-partitioned on equality keys and sorted within each partition. Each left partition streams its corresponding right partition. - Account for retained streaming buffers in the memory pool, including candidate matches and pending output rows, without counting shared buffers twice. - Preserve the mode through projection, child replacement, and protobuf/JSON serialization. Older serialized plans default to broadcast. - Add the opt-in `datafusion.optimizer.repartition_asof_joins` option, defaulting to `false`. Joins without equality keys continue to use broadcast execution. The join itself does not spill: upstream sorts handle spilling, and the join still needs memory for its active working set. Skewed keys can reduce parallelism. Existing ASOF SQL benchmarks compared with local `main` (`6c1727fcd`), using four partitions and 12 measured samples per case after warm-up: | Case | Main median | Partitioned median | | --- | ---: | ---: | | Q02: balanced equality keys | 183.55 ms | 108.14 ms | | Q03: wide payload | 63.86 ms | 46.93 ms | | Q06: skewed equality keys | 175.70 ms | 335.76 ms | | Q11: large right input | 2687.45 ms | 1045.61 ms | These timings used no memory limit. Broadcast timings on this branch showed no substantial regression. The skew result supports keeping partitioned execution opt-in. ## What is the testing strategy for this PR? - SQL logic tests cover all four match operators, composite and computed equality keys, NULLs, empty inputs, ties, and small batches. - SQL integration tests check strategy selection, input distribution and ordering, plan display, and query results. - A memory-limit integration test verifies that partitioned execution succeeds with spilled sort output under a budget that rejects broadcast collection. - Unit tests cover partition selection, mode preservation, buffer accounting, bounded streaming, and reservation cleanup on completion, cancellation, and error. - Protobuf and JSON tests cover mode round trips, compatibility with missing mode fields, and invalid serialized modes. Local validation: the required extended workspace suite passed (12,814 Rust tests and all 538 SQL test files). After consolidating redundant unit tests, all 21 ASOF unit tests passed. Formatting, all-targets/all-features Clippy with warnings denied, generated documentation, HTML/API documentation, and dependency structure checks passed. The full lint script's security audit reports RUSTSEC-2026-0382 (`rustls-webpki`) and RUSTSEC-2026-0383 (`webpki-root-certs`). Both findings were reproduced using the unchanged lockfile on local `main`. ## Are there any user-facing changes? Yes. Enable partitioned execution for ASOF joins with equality keys using: ```sql SET datafusion.optimizer.repartition_asof_joins = true; ``` The default execution strategy remains broadcast. The SQL guide and generated configuration reference document the option. This contribution used AI assistance. Keeping this PR in draft while I finish reviewing and walking through the implementation before requesting maintainer review. -- 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]
