Skip to content

Commit 8eb93eb

Browse files
committed
perf: split only runs whose encoded keys outgrow them
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.
1 parent 6eef7cb commit 8eb93eb

3 files changed

Lines changed: 119 additions & 31 deletions

File tree

‎datafusion/physical-plan/src/sorts/sort.rs‎

Lines changed: 40 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -66,7 +66,7 @@ use datafusion_execution::memory_pool::{
6666
use datafusion_execution::runtime_env::RuntimeEnv;
6767
use datafusion_physical_expr::LexOrdering;
6868
use datafusion_physical_expr::PhysicalExpr;
69-
use datafusion_physical_expr::expressions::{DynamicFilterPhysicalExpr, lit};
69+
use datafusion_physical_expr::expressions::{Column, DynamicFilterPhysicalExpr, lit};
7070

7171
use futures::{Stream, StreamExt, TryStreamExt};
7272
use log::{debug, trace};
@@ -784,13 +784,13 @@ impl ExternalSorter {
784784
let (sizes, shared_input_size) = self.sort_run_sizes(&runs)?;
785785
let shared_input =
786786
shared_input_size.map(|size| Arc::new(self.reservation.split(size)));
787-
let batch_size = self.merge_run_batch_size();
788787
let jobs =
789788
runs.into_iter()
790789
.zip(sizes)
791790
.enumerate()
792791
.map(move |(index, (batch, size))| {
793792
let reservation = self.reservation.split(size);
793+
let batch_size = self.merge_run_batch_size(&batch)?;
794794
let sort = self.sort_run(
795795
batch,
796796
reservation,
@@ -978,11 +978,11 @@ impl ExternalSorter {
978978
let (sizes, shared_input_size) = self.sort_run_sizes(&runs)?;
979979
let shared_input =
980980
shared_input_size.map(|size| Arc::new(self.reservation.split(size)));
981-
let batch_size = self.merge_run_batch_size();
982981
runs.into_iter()
983982
.zip(sizes)
984983
.map(|(run, size)| {
985984
let reservation = self.reservation.split(size);
985+
let batch_size = self.merge_run_batch_size(&run)?;
986986
self.sort_batch_stream(run, reservation, shared_input.clone(), batch_size)
987987
})
988988
.collect()
@@ -1038,12 +1038,6 @@ impl ExternalSorter {
10381038
let expressions = self.expr.clone();
10391039
let merge_pool = Arc::clone(&self.merge_pool);
10401040
let elapsed_compute = self.metrics.baseline.elapsed_compute().clone();
1041-
// Split batches of views sharing bytes would each be charged for them
1042-
let batch_size = if batch_size < self.batch_size && views_share_bytes(&batch) {
1043-
self.batch_size
1044-
} else {
1045-
batch_size
1046-
};
10471041
// Set when `merge_run_batch_size` split this run
10481042
let split = batch_size < self.batch_size;
10491043

@@ -1096,15 +1090,17 @@ impl ExternalSorter {
10961090
})
10971091
}
10981092

1099-
/// Batch size of the sorted runs that feed a merge. A merge in the row format
1100-
/// holds the keys of every run's current batch at once, and these keys can be
1101-
/// larger than the batches. With a memory limit, runs therefore use an eighth
1102-
/// of `batch_size`, but at least 1024 rows to keep the per-batch merge cost
1103-
/// small. The merge output still uses `batch_size`. Runs keep `batch_size`
1104-
/// when their batches would share buffers other than the data of top-level
1105-
/// view columns, such as dictionary values, since every batch is charged for
1106-
/// them. `sort_run` also keeps it for a run whose views share bytes.
1107-
fn merge_run_batch_size(&self) -> usize {
1093+
/// Batch size of the sorted run `run` when it feeds a merge. A merge in the
1094+
/// row format holds the keys of every run's current batch at once. The sort
1095+
/// reserves about twice the run, so the keys fit next to the sorted copy if
1096+
/// they are at most as large as the run. Otherwise, with a memory limit, the
1097+
/// run uses an eighth of `batch_size`, but at least 1024 rows to keep the
1098+
/// per-batch merge cost small. The merge output still uses `batch_size`.
1099+
///
1100+
/// Runs keep `batch_size` when their batches would share buffers other than
1101+
/// the data of top-level view columns, such as dictionary values, or when the
1102+
/// views of `run` share bytes, since every batch is charged for them.
1103+
fn merge_run_batch_size(&self, run: &RecordBatch) -> Result<usize> {
11081104
let limited = !matches!(
11091105
self.runtime.memory_pool.memory_limit(),
11101106
MemoryLimit::Infinite
@@ -1114,11 +1110,34 @@ impl ExternalSorter {
11141110
sorted_copy_keeps_buffers(data_type)
11151111
&& !matches!(data_type, DataType::Utf8View | DataType::BinaryView)
11161112
});
1117-
if limited && !shares_buffers && merges_in_row_format(&self.expr, &self.schema) {
1118-
(self.batch_size / 8).max(1024).min(self.batch_size)
1113+
if limited
1114+
&& !shares_buffers
1115+
&& merges_in_row_format(&self.expr, &self.schema)
1116+
&& self.encoded_keys_may_exceed(run)?
1117+
&& !views_share_bytes(run)
1118+
{
1119+
Ok((self.batch_size / 8).max(1024).min(self.batch_size))
11191120
} else {
1120-
self.batch_size
1121+
Ok(self.batch_size)
1122+
}
1123+
}
1124+
1125+
/// Returns whether the row-format keys of `run` may be larger than `run`. A
1126+
/// computed key counts as larger, since sizing it would evaluate it again.
1127+
fn encoded_keys_may_exceed(&self, run: &RecordBatch) -> Result<bool> {
1128+
// An 8-byte offset per row, then a null byte and the bytes of each key
1129+
let mut encoded = run.num_rows() * size_of::<usize>();
1130+
for sort_expr in self.expr.iter() {
1131+
let Some(column) = sort_expr.expr.downcast_ref::<Column>() else {
1132+
return Ok(true);
1133+
};
1134+
encoded += run.num_rows()
1135+
+ run
1136+
.column(column.index())
1137+
.to_data()
1138+
.get_slice_memory_size()?;
11211139
}
1140+
Ok(encoded > run.get_sliced_size()?)
11221141
}
11231142

11241143
/// If this sort may spill, pre-allocates

‎datafusion/physical-plan/src/sorts/sort/spill_tests.rs‎

Lines changed: 54 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -685,10 +685,11 @@ async fn test_sort_preparation_keeps_output_reserved() -> Result<()> {
685685

686686
/// Merging sorted runs in the row format holds the keys of every run's current
687687
/// batch. Two `Int32` keys encode to 18 bytes per row against 8 bytes in the
688-
/// batches, and a single `FixedSizeBinary(16)` key also merges in the row
689-
/// format, so merging full-size runs needs more memory than the sort reserves
690-
/// while buffering them. The `Utf8View` and dictionary payloads keep buffers
691-
/// shared after sorting, which must not be charged once per split batch.
688+
/// batches, and a single `FixedSizeBinary(16)` key and two long `Utf8View` keys
689+
/// also encode to more than their runs, so merging full-size runs needs more
690+
/// memory than the sort reserves while buffering them. The payloads make runs
691+
/// larger than their keys or keep buffers shared after sorting, so these runs
692+
/// keep `batch_size` and must not be charged once per batch either.
692693
async fn check_merge_keys_larger_than_batches() -> Result<()> {
693694
let rows = 8192;
694695
let int_keys = |batch: i32| -> Vec<ArrayRef> {
@@ -711,7 +712,9 @@ async fn check_merge_keys_larger_than_batches() -> Result<()> {
711712
DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8));
712713
for payload in [DataType::Utf8View, dictionary] {
713714
let fields = [&int_fields[..], &[("c", payload.clone())]].concat();
714-
check_merge_keys(&fields, 2, 16, true, |batch| {
715+
// Runs of this `Utf8View` payload are not split, as on `main`
716+
let within_pool = payload != DataType::Utf8View;
717+
check_merge_keys(&fields, 2, 16, within_pool, |batch| {
715718
let payload = cast(&strings(batch), &payload)?;
716719
Ok([int_keys(batch), vec![payload]].concat())
717720
})
@@ -830,6 +833,52 @@ async fn test_merge_keys_larger_than_batches_multithreaded() -> Result<()> {
830833
check_merge_keys_larger_than_batches().await
831834
}
832835

836+
/// Only runs whose encoded keys are larger than the run are split.
837+
#[test]
838+
fn test_merge_run_batch_size_follows_key_size() -> Result<()> {
839+
let batch_size = 8192;
840+
let rows = batch_size as i32;
841+
let check = |payload: Option<ArrayRef>, expected: usize| -> Result<()> {
842+
let mut columns: Vec<(&str, ArrayRef)> = vec![
843+
("a", Arc::new(Int64Array::from_iter_values(0..rows as i64))),
844+
("b", Arc::new(Int64Array::from_iter_values(0..rows as i64))),
845+
];
846+
columns.extend(payload.map(|payload| ("c", payload)));
847+
let batch = RecordBatch::try_from_iter(columns)?;
848+
let ordering = LexOrdering::new((0..2).map(|i| {
849+
PhysicalSortExpr::new_default(Arc::new(Column::new(
850+
batch.schema().field(i).name(),
851+
i,
852+
)))
853+
}))
854+
.expect("two sort keys");
855+
let runtime = RuntimeEnvBuilder::new()
856+
.with_memory_limit(1024 * 1024 * 1024, 1.0)
857+
.build_arc()?;
858+
let sorter = ExternalSorter::new(
859+
0,
860+
batch.schema(),
861+
ordering,
862+
batch_size,
863+
0,
864+
0,
865+
SpillCompression::Uncompressed,
866+
&ExecutionPlanMetricsSet::new(),
867+
runtime,
868+
)?;
869+
assert_eq!(sorter.merge_run_batch_size(&batch)?, expected);
870+
Ok(())
871+
};
872+
// Encoded keys: 26 bytes a row, against 16 bytes a row
873+
check(None, batch_size / 8)?;
874+
// Against 16 + 16 bytes a row: inlined views add more than the keys need
875+
let short: StringViewArray = (0..rows).map(|i| Some(format!("{i}"))).collect();
876+
check(Some(Arc::new(short)), batch_size)?;
877+
// Against 16 + 16 + 64 bytes a row
878+
let long: StringViewArray = (0..rows).map(|i| Some(format!("{i:0>64}"))).collect();
879+
check(Some(Arc::new(long)), batch_size)
880+
}
881+
833882
#[tokio::test]
834883
async fn test_single_batch_spill_preserves_workspace_after_limit_decreases() -> Result<()>
835884
{

‎datafusion/physical-plan/src/spill/mod.rs‎

Lines changed: 25 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -833,15 +833,17 @@ const VIEW_SIZE_BYTES: usize = 16;
833833
///
834834
/// # When GC is triggered
835835
///
836-
/// View GC is only performed when data buffers exceed a threshold (currently 10KB).
837-
/// This balances memory savings against the CPU overhead of garbage collection.
838-
/// Small view buffers are passed through unchanged since the GC overhead would exceed
839-
/// any memory savings.
836+
/// View GC runs when data buffers above a threshold (currently 10KB) hold more
837+
/// bytes than the views use. This balances memory savings against the CPU
838+
/// overhead of garbage collection. It also runs at any size when several data
839+
/// buffers remain and the views share no bytes, since each buffer read back
840+
/// from the spill is charged for the whole IPC message.
840841
///
841842
/// # Performance considerations
842843
///
843844
/// - If no arrays need compaction, the original batch is cloned cheaply
844-
/// - View GC is skipped for small buffers to avoid unnecessary CPU overhead
845+
/// - View GC is skipped for a single small data buffer to avoid unnecessary CPU
846+
/// overhead
845847
/// - Nested container types are traversed recursively so view arrays inside
846848
/// `List`, `Map`, `Union`, `Dictionary`, and other child-bearing arrays are compacted too
847849
/// - Traversal uses physical children, not logical reachability. For example,
@@ -2706,6 +2708,24 @@ mod tests {
27062708
Ok(())
27072709
}
27082710

2711+
#[test]
2712+
fn test_gc_keeps_shared_views_in_two_buffers() -> Result<()> {
2713+
use arrow::array::AsArray;
2714+
use arrow::compute::concat;
2715+
2716+
let (_, first) = shared_dictionary_views(1000, 8, 1024)?;
2717+
let (_, second) = shared_dictionary_views(1000, 8, 1024)?;
2718+
let merged = concat(&[&first as &dyn Array, &second])?;
2719+
let merged = merged.as_string_view();
2720+
2721+
// `gc()` would copy 16384 values, 8x what the two dictionaries hold
2722+
let compacted = compact_view_array(merged).expect("buffer entries repeat");
2723+
assert_eq!(compacted.data_buffers().len(), 2);
2724+
assert_eq!(&compacted, merged);
2725+
2726+
Ok(())
2727+
}
2728+
27092729
#[test]
27102730
fn test_gc_keeps_shared_dictionary_buffer() -> Result<()> {
27112731
let (dictionary, merged) = shared_dictionary_views(1000, 8, 1024)?;

0 commit comments

Comments
 (0)