jayzhan211 commented on code in PR #26953:
URL: https://github.com/apache/datafusion/pull/26953#discussion_r4499604259
##########
datafusion/physical-plan/src/spill/mod.rs:
##########
@@ -836,7 +838,9 @@ fn compact_view_array<T: ByteViewType>(
let (remap, unique) = dedup_buffer_entries(buffers);
let unique_size: usize = unique.iter().map(|b| b.len()).sum();
- if should_gc_view_array(array, unique_size) {
+ if should_gc_view_array(array, unique_size)
+ || (unique.len() > 1 && array.total_buffer_bytes_used() <= unique_size)
Review Comment:
Nothing tests the `<=` here: changing this to `unique.len() > 1` still
passes every `sorts::` and `spill::` test. It's what keeps `gc` from copying
shared bytes once per row when views span several distinct buffers. Could you
add this next to `test_gc_keeps_shared_dictionary_buffer`? It passes today and
fails without the `<=`.
```rs
#[test]
fn test_gc_keeps_shared_views_in_two_buffers() -> Result<()> {
use arrow::array::AsArray;
use arrow::compute::concat;
let (_, first) = shared_dictionary_views(1000, 8, 1024)?;
let (_, second) = shared_dictionary_views(1000, 8, 1024)?;
let merged = concat(&[&first as &dyn Array, &second])?;
let merged = merged.as_string_view();
// `gc()` would copy 16384 values, 8x what the two dictionaries hold
let compacted = compact_view_array(merged).expect("buffer entries
repeat");
assert_eq!(compacted.data_buffers().len(), 2);
assert_eq!(&compacted, merged);
Ok(())
}
```
##########
datafusion/physical-plan/src/sorts/sort.rs:
##########
@@ -1078,6 +1106,31 @@ impl ExternalSorter {
})
}
+ /// Batch size of the sorted runs that feed a merge. A merge in the row
format
+ /// holds the keys of every run's current batch at once, and these keys
can be
+ /// larger than the batches. With a memory limit, runs therefore use an
eighth
+ /// of `batch_size`, but at least 1024 rows to keep the per-batch merge
cost
+ /// small. The merge output still uses `batch_size`. Runs keep `batch_size`
+ /// when their batches would share buffers other than the data of top-level
+ /// view columns, such as dictionary values, since every batch is charged
for
+ /// them. `sort_run` also keeps it for a run whose views share bytes.
+ fn merge_run_batch_size(&self) -> usize {
Review Comment:
Sorry, I should have raised this last round: splitting also costs runs that
don't need it. With integer keys and a `Utf8View` payload, the per-batch `gc`
copies every string, so Parquet sorts at `-m 16G` with no spill took 9–17%
longer than `main`, and peak RSS rose 13–38%. Their encoded keys are 26 bytes a
row against about 100 for the run, so they already fit in the 2x reservation.
Could the split apply only when the keys outgrow the run? With the diff
below those sorts are within 3% of `main`, and the issue's reproducer still
passes. The trade-off is that the `Utf8View` payload in your test goes back to
`main`'s pool overshoot, hence `within_pool`. Repro:
```sql
COPY (SELECT value % 1000003 AS a, value AS b,
concat(repeat('u', 60), value::varchar) AS c
FROM range(5000000)) TO 'unique.parquet' STORED AS PARQUET;
SET datafusion.execution.target_partitions = 1;
EXPLAIN ANALYZE SELECT * FROM 'unique.parquet' ORDER BY a, b;
```
```diff
--- a/datafusion/physical-plan/src/sorts/sort.rs
+++ b/datafusion/physical-plan/src/sorts/sort.rs
@@ imports
use datafusion_physical_expr::PhysicalExpr;
+use datafusion_physical_expr::expressions::Column;
use datafusion_physical_expr::expressions::{DynamicFilterPhysicalExpr, lit};
@@ fn prepare_sort_runs
- let batch_size = self.merge_run_batch_size();
let jobs =
runs.into_iter()
.zip(sizes)
.enumerate()
.map(move |(index, (batch, size))| {
let reservation = self.reservation.split(size);
+ let batch_size = self.merge_run_batch_size(&batch)?;
let sort = self.sort_run(
@@ fn sort_run_streams
- let batch_size = self.merge_run_batch_size();
runs.into_iter()
.zip(sizes)
.map(|(run, size)| {
let reservation = self.reservation.split(size);
+ let batch_size = self.merge_run_batch_size(&run)?;
self.sort_batch_stream(run, reservation,
shared_input.clone(), batch_size)
@@ fn sort_run
- // Split batches of views sharing bytes would each be charged for
them
- let batch_size = if batch_size < self.batch_size &&
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;
@@ fn merge_run_batch_size
- /// Batch size of the sorted runs that feed a merge. A merge in the row
format
- /// holds the keys of every run's current batch at once, and these keys
can be
- /// larger than the batches. With a memory limit, runs therefore use an
eighth
- /// of `batch_size`, but at least 1024 rows to keep the per-batch merge
cost
- /// small. The merge output still uses `batch_size`. Runs keep
`batch_size`
- /// when their batches would share buffers other than the data of
top-level
- /// view columns, such as dictionary values, since every batch is
charged for
- /// them. `sort_run` also keeps it for a run whose views share bytes.
- fn merge_run_batch_size(&self) -> usize {
+ /// Batch size of the sorted run `run` when it feeds a merge. A merge
in the
+ /// row format holds the keys of every run's current batch at once. The
sort
+ /// reserves about twice the run, so the keys fit next to the sorted
copy if
+ /// they are at most as large as the run. Otherwise, with a memory
limit, the
+ /// run uses an eighth of `batch_size`, but at least 1024 rows to keep
the
+ /// per-batch merge cost small. The merge output still uses
`batch_size`.
+ ///
+ /// Runs keep `batch_size` when their batches would share buffers other
than
+ /// the data of top-level view columns, such as dictionary values, or
when the
+ /// views of `run` share bytes, since every batch is charged for them.
+ fn merge_run_batch_size(&self, run: &RecordBatch) -> Result<usize> {
let limited = !matches!(
self.runtime.memory_pool.memory_limit(),
MemoryLimit::Infinite
);
let shares_buffers = self.schema.fields().iter().any(|field| {
let data_type = field.data_type();
sorted_copy_keeps_buffers(data_type)
&& !matches!(data_type, DataType::Utf8View |
DataType::BinaryView)
});
- if limited && !shares_buffers && merges_in_row_format(&self.expr,
&self.schema) {
- (self.batch_size / 8).max(1024).min(self.batch_size)
+ if limited
+ && !shares_buffers
+ && merges_in_row_format(&self.expr, &self.schema)
+ && !views_share_bytes(run)
+ && self.encoded_keys_may_exceed(run)?
+ {
+ Ok((self.batch_size / 8).max(1024).min(self.batch_size))
} else {
- self.batch_size
+ Ok(self.batch_size)
+ }
+ }
+
+ /// Returns whether the row-format keys of `run` may be larger than
`run`. A
+ /// computed key counts as larger, since sizing it would evaluate it
again.
+ fn encoded_keys_may_exceed(&self, run: &RecordBatch) -> Result<bool> {
+ // An 8-byte offset per row, then a null byte and the bytes of each
key
+ let mut encoded = run.num_rows() * size_of::<usize>();
+ for sort_expr in self.expr.iter() {
+ let Some(column) = sort_expr.expr.downcast_ref::<Column>() else
{
+ return Ok(true);
+ };
+ encoded += run.num_rows()
+ + run
+ .column(column.index())
+ .to_data()
+ .get_slice_memory_size()?;
}
+ Ok(encoded > run.get_sliced_size()?)
}
--- a/datafusion/physical-plan/src/sorts/sort/spill_tests.rs
+++ b/datafusion/physical-plan/src/sorts/sort/spill_tests.rs
@@ fn check_merge_keys_larger_than_batches
- check_merge_keys(&fields, 2, 16, true, |batch| {
+ // 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| {
```
Please also add this boundary test. It separates the rule from "any view
payload": a short inlined payload must not split either.
```rs
/// Only runs whose encoded keys are larger than the run are split.
#[test]
fn test_merge_run_batch_size_follows_key_size() -> Result<()> {
let batch_size = 8192;
let rows = batch_size as i32;
let check = |payload: Option<ArrayRef>, expected: usize| -> Result<()> {
let mut columns: Vec<(&str, ArrayRef)> = vec![
("a", Arc::new(Int64Array::from_iter_values(0..rows as i64))),
("b", Arc::new(Int64Array::from_iter_values(0..rows as i64))),
];
columns.extend(payload.map(|payload| ("c", payload)));
let batch = RecordBatch::try_from_iter(columns)?;
let ordering = LexOrdering::new((0..2).map(|i| {
PhysicalSortExpr::new_default(Arc::new(Column::new(
batch.schema().field(i).name(),
i,
)))
}))
.expect("two sort keys");
let runtime = RuntimeEnvBuilder::new()
.with_memory_limit(1024 * 1024 * 1024, 1.0)
.build_arc()?;
let sorter = ExternalSorter::new(
0,
batch.schema(),
ordering,
batch_size,
0,
0,
SpillCompression::Uncompressed,
&ExecutionPlanMetricsSet::new(),
runtime,
)?;
assert_eq!(sorter.merge_run_batch_size(&batch)?, expected);
Ok(())
};
// Encoded keys: 26 bytes a row, against 16 bytes a row
check(None, batch_size / 8)?;
// Against 16 + 16 bytes a row: inlined views add more than the keys need
let short: StringViewArray = (0..rows).map(|i|
Some(format!("{i}"))).collect();
check(Some(Arc::new(short)), batch_size)?;
// Against 16 + 16 + 64 bytes a row
let long: StringViewArray = (0..rows).map(|i|
Some(format!("{i:0>64}"))).collect();
check(Some(Arc::new(long)), batch_size)
}
```
--
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]
