comphead commented on code in PR #26953:
URL: https://github.com/apache/datafusion/pull/26953#discussion_r4540821459
##########
datafusion/physical-plan/src/sorts/sort/spill_tests.rs:
##########
@@ -679,6 +683,202 @@ async fn test_sort_preparation_keeps_output_reserved() ->
Result<()> {
Ok(())
}
+/// Merging sorted runs in the row format holds the keys of every run's current
+/// batch. Two `Int32` keys encode to 18 bytes per row against 8 bytes in the
+/// batches, and a single `FixedSizeBinary(16)` key and two long `Utf8View`
keys
+/// also encode to more than their runs, so merging full-size runs needs more
+/// memory than the sort reserves while buffering them. The payloads make runs
+/// larger than their keys or keep buffers shared after sorting, so these runs
+/// keep `batch_size` and must not be charged once per batch either.
+async fn check_merge_keys_larger_than_batches() -> Result<()> {
+ let rows = 8192;
+ let int_keys = |batch: i32| -> Vec<ArrayRef> {
+ vec![
+ Arc::new(Int32Array::from_iter_values(
+ (0..rows).map(|i| (i * 31 + batch) % 1000),
+ )),
+ Arc::new(Int32Array::from_iter_values((0..rows).rev())),
+ ]
+ };
+ let int_fields = [("a", DataType::Int32), ("b", DataType::Int32)];
+ check_merge_keys(&int_fields, 2, 16, true, |batch|
Ok(int_keys(batch))).await?;
+ // Longer than 12 bytes, so views keep them in data buffers.
+ let strings = |batch: i32| -> StringArray {
+ (0..rows)
+ .map(|i| Some(format!("{batch}-{i:0>64}")))
+ .collect()
+ };
+ let dictionary =
+ DataType::Dictionary(Box::new(DataType::Int32),
Box::new(DataType::Utf8));
+ for payload in [DataType::Utf8View, dictionary] {
+ let fields = [&int_fields[..], &[("c", payload.clone())]].concat();
+ // Runs of this `Utf8View` payload are not split, as on `main`
+ let within_pool = payload != DataType::Utf8View;
+ check_merge_keys(&fields, 2, 16, within_pool, |batch| {
+ let payload = cast(&strings(batch), &payload)?;
+ Ok([int_keys(batch), vec![payload]].concat())
+ })
+ .await?;
+ }
+ // Views into 1024 dictionary values share their bytes, as a
+ // dictionary-encoded Parquet column produces. Sorting the buffered batches
+ // before a spill grows past the pool, on `main` too.
+ let fields = [&int_fields[..], &[("c", DataType::Utf8View)]].concat();
+ check_merge_keys(&fields, 2, 16, false, |batch| {
Review Comment:
One more change in 0e71291f0. Computed keys counted as larger, so a computed
integer key with a `Utf8View` payload still split, and `ORDER BY a + 1, b` on
your `unique.parquet` at `-m 16G` took 11% longer than `main` with 35% more
peak RSS. Keys with a fixed-width type are now sized from it, which brings that
sort back to `main`. A computed-key version of the issue's reproducer, `ORDER
BY a + 1, b` over `range(30000000)` at 512M, still fails on `main` and passes
here.
The shared-views case now sorts by `CAST(a AS Utf8), b`, which still counts
as larger and still fails without the gate, and
`test_merge_run_batch_size_follows_key_size` now checks both kinds of computed
key.
--
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]
