jayzhan211 commented on code in PR #26953:
URL: https://github.com/apache/datafusion/pull/26953#discussion_r4492816067
##########
datafusion/physical-plan/src/sorts/sort.rs:
##########
@@ -1023,9 +1035,10 @@ impl ExternalSorter {
}
let expressions = self.expr.clone();
- let batch_size = self.batch_size;
let merge_pool = Arc::clone(&self.merge_pool);
let elapsed_compute = self.metrics.baseline.elapsed_compute().clone();
+ // Set when `merge_run_batch_size` split this run
+ let split = batch_size < self.batch_size;
Review Comment:
Sorry, my `gc` suggestion misses views that share bytes, which is what
dictionary-encoded Parquet strings produce. A chunk that would write at least
as much as the buffer holds skips `gc`, so all 8 chunks keep the buffer and
each is charged for it. This fails at `-m 256M`, `512M` and `1G` with `Not
enough memory to continue external sort`, and passes on `main`:
```sql
COPY (SELECT value % 1000003 AS a, value AS b,
concat(repeat('x', 240), (value % 1000)::varchar) AS c
FROM range(5000000)) TO 'shared.parquet' STORED AS PARQUET;
SET datafusion.execution.target_partitions = 1;
SELECT * FROM 'shared.parquet' ORDER BY a, b;
```
Could these runs keep `batch_size`? The helper below goes next to
`gc_view_arrays` in `spill/mod.rs`, and `sort.rs` needs to import it.
```suggestion
// Split batches of views sharing bytes would each be charged for
them.
let batch_size = if views_share_bytes(&batch) {
self.batch_size
} else {
batch_size
};
// Set when `merge_run_batch_size` split this run
let split = batch_size < self.batch_size;
```
```rs
pub(crate) fn views_share_bytes(batch: &RecordBatch) -> bool {
use arrow::array::AsArray;
fn shares<T: ByteViewType>(array: &GenericByteViewArray<T>) -> bool {
let (_, unique) = dedup_buffer_entries(array.data_buffers());
array.total_buffer_bytes_used() > unique.iter().map(|b|
b.len()).sum::<usize>()
}
batch.columns().iter().any(|array| match array.data_type() {
DataType::Utf8View => shares(array.as_string_view()),
DataType::BinaryView => shares(array.as_binary_view()),
_ => false,
})
}
```
Nested views already keep `batch_size` through `sorted_copy_keeps_buffers`,
so only these two top-level types need the check. Could you add a test with a
payload of `cast(&dict, &DataType::Utf8View)`, using 8192 keys over 1024 values
of 64 bytes? It should succeed with correct output. Skip the `assert_released`
peak check there, because the `grow` in `try_sort_and_spill_in_mem_batches`
goes over the pool on `main` too. Please also unit test the boundary:
`views_share_bytes` is `true` for that array and `false` for 8192 values each
used once.
##########
datafusion/physical-plan/src/sorts/sort.rs:
##########
@@ -1036,7 +1049,16 @@ 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 mut sorted_batches =
+ sort_batch_chunked(&batch, &expressions, batch_size)?;
+ // `take` keeps the data buffers of view arrays, so splitting would
+ // charge every batch for all of them.
+ if split && sorted_batches.len() > 1 {
Review Comment:
The per-chunk `gc` also causes a second failure, on clustered input. An
output batch taken from one run holds 8 chunk buffers with no unused bytes, so
`gc` skips it. After the spill is read back, each buffer is a slice of one IPC
message and is charged for the whole message, so a batch's
`max_record_batch_memory` reached 22.1 MB against 3.0 MB on `main`. This fails
at `-m 32M` and `64M`, and passes on `main`:
```sql
COPY (SELECT value / 1000 AS a, value AS b,
concat(repeat('x', 60), value::varchar) AS c,
concat(repeat('y', 60), value::varchar) AS d,
concat(repeat('z', 60), value::varchar) AS e
FROM range(5000000)) TO 'clustered3.parquet' STORED AS PARQUET;
SET datafusion.execution.target_partitions = 1;
SELECT * FROM 'clustered3.parquet' ORDER BY a, b;
```
Could spilling compact these arrays to one buffer? With this and the
`views_share_bytes` change, all of my probes match or beat `main`, and all
`sorts::` and `spill::` tests pass, including yours. The issue's reproducer
still passes too.
```diff
--- a/datafusion/physical-plan/src/spill/mod.rs
+++ b/datafusion/physical-plan/src/spill/mod.rs
@@ fn compact_view_array
- if should_gc_view_array(array, unique_size) {
+ // Several buffers read back as slices of one IPC message, each charged
for all of it
+ if should_gc_view_array(array, unique_size)
+ || (unique.len() > 1 && array.total_buffer_bytes_used() <=
unique_size)
+ {
return Some(array.gc());
}
```
The `<=` keeps `gc` away from views that share bytes, where it would copy
once per row (#23564). Could you add a test with two `Utf8View` keys of unique
66-byte strings, prefixed by batch number so the runs don't overlap, in a pool
of 8 buffered batches? It should succeed; today it fails in
`ExternalSorterMerge`.
--
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]
