Repository navigation
Conversation
|
This PR would be good addition to #26931 |
jayzhan211
left a comment
There was a problem hiding this comment.
@comphead thanks for the fix! I found two things, both inline: a case where the split now runs out of memory where main succeeds, and single-key sorts that the gate misses.
| /// 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`. | ||
| fn merge_run_batch_size(&self) -> usize { |
There was a problem hiding this comment.
With a non-inlined Utf8View column, each of the 8 chunks from sort_batch_chunked keeps the input's full view data buffer, and sort_run adds up get_record_batch_memory_size per chunk. With logging added to sort_run, one run had input=1228800 reserved=2457600 total_sorted=8454144, so it asks for about 8x the payload. This test passes on main and fails here with Failed to allocate additional 5.7 MB for ExternalSorter[0]:
let schema = Arc::new(Schema::new(vec![
Field::new("a", DataType::Int32, false),
Field::new("b", DataType::Int32, false),
Field::new("c", DataType::Utf8View, false),
]));
// 64 batches of 8192 rows; c = format!("{i:0>64}"); sort by (a, b)
let capacity = MERGE_BYTES + 16 * get_reserved_bytes_for_record_batch(&batches[0])?;
let runtime = RuntimeEnvBuilder::new()
.with_memory_pool(Arc::new(GreedyMemoryPool::new(capacity)))
.build_arc()?;Running crate::spill::gc_view_arrays on the chunks when sorted_batches.len() > 1 fixed it locally. Three exact-accounting tests in spill_tests.rs then need new expected numbers. Could you add this case as a test with the fix?
There was a problem hiding this comment.
Thanks, confirmed and fixed in 571f7f9. Split runs now go through gc_view_arrays, and a Utf8View payload is a case of test_merge_keys_larger_than_batches. I limited the compaction to runs split for the merge. Runs chunked only because they exceed batch_size keep their current charge, and the three workspace-loan tests stay unchanged, since they create the loan with those unsplit chunks.
A high-cardinality dictionary payload failed the same way, and gc_view_arrays cannot compact shared dictionary values. So schemas whose split batches would still share buffers keep batch_size, and that payload is a test case too.
| self.runtime.memory_pool.memory_limit(), | ||
| MemoryLimit::Infinite | ||
| ); | ||
| if limited && self.expr.len() > 1 { |
There was a problem hiding this comment.
expr.len() > 1 doesn't match when the merge uses the row format: StreamingMergeBuilder::build also uses RowCursorStream for a single key that isn't primitive or a plain string or binary type. A single FixedSizeBinary(16) key (UUIDs) still fails the same way. Could the gate follow the merge's choice, ideally as a helper next to build so the two stay in sync, with a test for a single FixedSizeBinary key? With this, that test passes and the existing sort tests still pass:
fn merges_in_row_format(expr: &LexOrdering, schema: &Schema) -> bool {
if expr.len() > 1 {
return true;
}
let Ok(data_type) = expr[0].expr.data_type(schema) else {
return true;
};
!(data_type.is_primitive()
|| matches!(
data_type,
DataType::Utf8
| DataType::Utf8View
| DataType::LargeUtf8
| DataType::Binary
| DataType::LargeBinary
))
}There was a problem hiding this comment.
Good catch, thanks. merges_in_row_format now sits next to StreamingMergeBuilder::build, with a note in build that the two must match, and merge_run_batch_size uses it. A single FixedSizeBinary(16) key is a case of test_merge_keys_larger_than_batches. It fails on main and on the previous head, and passes now.
|
@jayzhan211 PTAL |
jayzhan211
left a comment
There was a problem hiding this comment.
@comphead thanks for the update! Both earlier asks look good, and the issue's reproducer passes. My gc_view_arrays suggestion had two gaps, though: two real Parquet sorts now run out of memory where main succeeds. Both are inline with a fix.
| // Set when `merge_run_batch_size` split this run | ||
| let split = batch_size < self.batch_size; |
There was a problem hiding this 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:
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.
| // Set when `merge_run_batch_size` split this run | |
| let split = batch_size < self.batch_size; | |
| // 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; |
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.
There was a problem hiding this comment.
Thanks, applied in 0938167. views_share_bytes sits next to gc_view_arrays, and sort_run only checks it when the run would be split, so sorts without a memory limit don't scan the views. The dictionary-backed Utf8View payload is a case of test_merge_keys_larger_than_batches, without the pool checks as you suggested, and without the change it fails with Not enough memory to continue external sort. The boundary test is test_views_share_bytes in spill/mod.rs. Your query now passes at 256M, 512M and 1G with the same spill counts as main.
| 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 { |
There was a problem hiding this 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:
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.
--- 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.
There was a problem hiding this comment.
Applied as suggested in 0938167, and the doc comment of compact_view_array now covers it. The two clustered Utf8View keys with room for 8 batches fail without it and pass now. Your query passes at 32M and 64M with the same 507 and 113 spills as main. Since this also changes spilling for other operators, I checked a GROUP BY on a Utf8View column at 64M and 128M, which spills the same as on main. A BinaryView payload, Utf8View keys with overlapping runs, shared and unique views in one table, nine Int64 keys and a single FixedSizeBinary(16) key also match or beat main.
jayzhan211
left a comment
There was a problem hiding this comment.
@comphead thanks for the update! All of the earlier failures are fixed, and I found no case that fails where main succeeds. I do have one performance ask inline, which I should have raised earlier, and a small test request.
One optional doc fix: in spill/mod.rs around lines 762 and 770, the gc_view_arrays docs still say small buffers skip gc. Arrays with several buffers now compact at any 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 { |
There was a problem hiding this 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:
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;--- 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.
/// 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)
}There was a problem hiding this comment.
Thanks, this is a good catch. Applied in 8eb93eb, with one change: merge_run_batch_size checks the key size before views_share_bytes, so runs that keep batch_size skip the scan of their views. With your repro at -m 16G, the previous version was 14% slower than main with 35% more peak RSS, and this one is within 4% on both. The issue's reproducer and the memory-limited probes from the earlier rounds still pass with the same spill counts. Your boundary test is in as test_merge_run_batch_size_follows_key_size.
For variable-length keys the estimate counts their bytes without the row format's padding, so string keys that only slightly outgrow a run are not split, as on main. I noted that under the known limitations.
| 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) |
There was a problem hiding this 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 <=.
#[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(())
}There was a problem hiding this comment.
Added as test_gc_keeps_shared_views_in_two_buffers in 8eb93eb. It fails when the condition is only unique.len() > 1. Thanks!
… limit A multi-column merge holds the row-format keys of every sorted run's current batch. When the sort keys are most of each row, those keys are larger than the batches, and the merge needs more memory than the sort reserved while buffering them, so the sort fails with ResourcesExhausted instead of spilling. With a memory limit, sort the runs of a multi-column merge into batches of an eighth of `batch_size`, but at least 1024 rows. The merge output keeps `batch_size`, and sorts without a memory limit are unchanged. Closes apache#26940
- Split runs whenever the merge uses the row format, including a single key without a specialized cursor, such as `FixedSizeBinary`. - Compact the view arrays of split runs, so each batch is charged for its own strings rather than for the run's whole data buffers. - Keep `batch_size` when split batches would still share buffers, such as dictionary values.
- Keep `batch_size` for a run whose views share bytes, as dictionary-encoded Parquet strings do. Compacting such views copies the shared bytes once per row, so every split batch would keep and be charged for the whole buffer. - Compact view arrays with several data buffers before spilling unless their views share bytes. Read back, each buffer is a slice of one IPC message and is charged for all of it.
The sort reserves about twice each run, so a merge in the row format fits when the keys are at most as large as the run. Splitting smaller keys only added work: with integer keys and a `Utf8View` payload, the compaction of split batches copied every string. `merge_run_batch_size` now checks each run, from the bytes of its key columns, and also covers views that share bytes. Also test that spill compaction leaves views sharing bytes across several buffers alone, and update the `gc_view_arrays` docs.
0938167 to
8eb93eb
Compare
|
@jayzhan211 thank you for the careful reviews, your probes made this PR much better. I rebased onto |
jayzhan211
left a comment
There was a problem hiding this comment.
@comphead thanks for applying all of the last round's changes. One gap is left: a test for the views_share_bytes gate, details inline.
| // 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| { |
There was a problem hiding this comment.
Since the key-size check went in, no test reaches views_share_bytes: removing && !views_share_bytes(run) from merge_run_batch_size still passes every sorts:: and spill:: test. It's still needed for computed keys, which always count as larger than the run. With round 2's shared.parquet, this passes on main and here at -m 256M, 512M and 1G, and fails without the gate with Not enough memory to continue external sort:
SET datafusion.execution.target_partitions = 1;
SELECT * FROM 'shared.parquet' ORDER BY a + 1, b;Could you add a computed-key case? The diff below passes here and fails without the gate.
Test diff for spill_tests.rs
-use datafusion_physical_expr::expressions::Column;
-use datafusion_physical_expr::{LexOrdering, PhysicalSortExpr};
+use datafusion_physical_expr::expressions::{CastExpr, Column};
+use datafusion_physical_expr::{LexOrdering, PhysicalExpr, PhysicalSortExpr};
@@ async fn check_merge_keys_larger_than_batches
let fields = [&int_fields[..], &[("c", DataType::Utf8View)]].concat();
- check_merge_keys(&fields, 2, 16, false, |batch| {
+ let shared_views = |batch: i32| -> Result<Vec<ArrayRef>> {
let values: StringArray = (0..1024)
@@
Ok([int_keys(batch), vec![views]].concat())
- })
- .await?;
+ };
+ check_merge_keys(&fields, 2, 16, false, shared_views).await?;
+ // A computed key counts as larger than the run, so only the shared views
+ // keep this run at `batch_size`.
+ let a: Arc<dyn PhysicalExpr> = Arc::new(Column::new("a", 0));
+ let keys = vec![
+ Arc::new(CastExpr::new(a, DataType::Int64, None)) as Arc<dyn PhysicalExpr>,
+ Arc::new(Column::new("b", 1)),
+ ];
+ check_merge_exprs(&fields, keys, 16, false, shared_views).await?;
@@ async fn check_merge_keys(
columns: impl Fn(i32) -> Result<Vec<ArrayRef>>,
+) -> Result<()> {
+ let keys = (0..keys)
+ .map(|i| Arc::new(Column::new(fields[i].0, i)) as Arc<dyn PhysicalExpr>)
+ .collect();
+ check_merge_exprs(fields, keys, buffered, within_pool, columns).await
+}
+
+/// Like [`check_merge_keys`], sorting by `keys`.
+async fn check_merge_exprs(
+ fields: &[(&str, DataType)],
+ keys: Vec<Arc<dyn PhysicalExpr>>,
+ buffered: usize,
+ within_pool: bool,
+ columns: impl Fn(i32) -> Result<Vec<ArrayRef>>,
) -> Result<()> {
@@
- let ordering =
- LexOrdering::new((0..keys).map(|i| {
- PhysicalSortExpr::new_default(Arc::new(Column::new(fields[i].0, i)))
- }))
+ let ordering = LexOrdering::new(keys.into_iter().map(PhysicalSortExpr::new_default))
.expect("at least one sort key");There was a problem hiding this 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, b29d5b0 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.
There was a problem hiding this comment.
One more change in 0e71291. 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.
Since runs whose keys fit keep `batch_size`, the shared-views case sorted by column keys never reached `views_share_bytes`. Sort it by a computed key, which counts as larger than the run, so the case fails without the gate.
A computed key counted as larger than the run, so its runs were always split under a memory limit. With a computed integer key and a `Utf8View` payload, the compaction of split batches copied every string. Size such keys from their type when it has a fixed width. The shared-views case now sorts by a string cast, which still counts as larger.
Which issue does this PR close?
Rationale for this change
A sort under a memory limit can fail with
ResourcesExhaustedforExternalSorterMergeinstead of spilling when its merge uses the row format, as with several sort keys. The merge holds the row-format keys of every sorted run at once. When the keys are most of each row, that needs more than the 2x the sort reserves for each buffered batch. The issue has the details and a reproducer.What changes are included in this PR?
With a memory limit,
ExternalSortersorts a run that feeds a merge in the row format into batches of an eighth ofbatch_size, but at least 1024 rows, when the run's encoded keys may be larger than the run, so the merge holds the keys of one such batch per run. Smaller keys fit next to the sorted copy, since the sort reserves about twice the run. That covers several keys and a single key without a specialized cursor, such asFixedSizeBinary. A computed key is sized from its type if that has a fixed width, and otherwise counts as larger. Split runs compact their view arrays, so each batch is charged only for its own strings. Runs whose views share bytes, as dictionary-encoded Parquet strings do, keepbatch_size, because compacting them would copy the shared bytes once per row. Spilling also compacts view arrays with several data buffers, since each buffer read back from the spill is a slice of one IPC message and is charged for all of it. The merge output keepsbatch_size, and sorts without a memory limit are unchanged. This adds abatch_sizeparameter tosort_run.What is the testing strategy for this PR?
test_merge_keys_larger_than_batches(current-thread and multi-thread) sorts in a pool that holds 16 buffered batches, or 8 for the clustered keys. TwoInt32keys and a singleFixedSizeBinary(16)key fail onmainwithResourcesExhaustedand pass here. TwoInt32keys with aUtf8View, a dictionary or a dictionary-backedUtf8Viewpayload check that splitting does not charge each batch for the shared buffers, and two clusteredUtf8Viewkeys check the spilled buffers.test_merge_run_batch_size_follows_key_sizechecks that only runs whose keys outgrow them are split. The reproducer from the issue and two Parquet sorts from review, withdatafusion-cli:main-m 512M-m 1G-m 1G-m 256M-m 32MWith
dfbench,sort-tpchQ5 from #16979 (SF1, 1 partition,--memory-limit 100M) fails onmainand passes here. With the same settings indatafusion-cli, it also fails onmainat 50M and 200M and passes here with 6 and 2 spills.The criterion
sortbenchmark has no memory limit and is unchanged. With the split forced on atbatch_size = 8192, its multi-column in-memory cases took 2 to 12% longer, which is the cost under a memory limit. No-spill Parquet sorts at-m 16Gwith integer or computed integer keys and aUtf8Viewpayload are within 4% ofmainin time and peak RSS.Are there any user-facing changes?
No. Known limitations: runs of at most 1024 rows are not split. That covers sorts with a
batch_sizeof 1024 or less, and small input batches, which #26931 would coalesce into larger runs. Sorts with dictionary, list view or nested view columns, and runs whose views share bytes, are not split, since each batch would be charged for the shared buffers. The key size counts the bytes of variable-length keys without the row format's padding, so string keys that only slightly outgrow their run are not split. Runs sorted by a computed key without a fixed width, such as a string expression, are split under a memory limit even when their keys fit. Each run's existing one-batch read-ahead now holds one of the smaller batches with its reservation already released.