comphead commented on code in PR #26953:
URL: https://github.com/apache/datafusion/pull/26953#discussion_r4540429372
##########
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:
Thanks, you're right on both points. With the gate removed every `sorts::`
and `spill::` test still passed, and your `ORDER BY a + 1, b` on
`shared.parquet` fails at 256M, 512M and 1G while `main` and this PR pass with
the same spill counts.
Rather than add a second shared-views case, b29d5b02a switches the existing
one to `CAST(a AS Int64), b` through your `check_merge_exprs`. With column keys
that case no longer reached the gate, and the `Utf8View` payload case already
covers keys that fit. It now fails without the gate with `Not enough memory to
continue external sort`.
--
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]
