diff --git a/datafusion/physical-plan/src/sorts/sort.rs b/datafusion/physical-plan/src/sorts/sort.rs index 96e673e4043bc..c045ccfe3833c 100644 --- a/datafusion/physical-plan/src/sorts/sort.rs +++ b/datafusion/physical-plan/src/sorts/sort.rs @@ -30,10 +30,12 @@ use crate::metrics::{ }; use crate::projection::{ProjectionExec, make_with_child, update_ordering}; use crate::sorts::IncrementalSortIterator; -use crate::sorts::streaming_merge::{SortedSpillFile, StreamingMergeBuilder}; -use crate::spill::get_record_batch_memory_size; +use crate::sorts::streaming_merge::{ + SortedSpillFile, StreamingMergeBuilder, merges_in_row_format, +}; use crate::spill::in_progress_spill_file::InProgressSpillFile; use crate::spill::spill_manager::{GetSlicedSize, SpillManager}; +use crate::spill::{gc_view_arrays, get_record_batch_memory_size, views_share_bytes}; use crate::statistics::{ChildStats, StatisticsArgs, with_per_partition_fetch}; use crate::stream::{ObservedStream, RecordBatchStreamAdapter}; use crate::topk::TopK; @@ -59,12 +61,12 @@ use datafusion_common::{ use datafusion_common_runtime::SpawnedTask; use datafusion_execution::TaskContext; use datafusion_execution::memory_pool::{ - MemoryConsumer, MemoryPool, MemoryReservation, MergeMemoryPool, + MemoryConsumer, MemoryLimit, MemoryPool, MemoryReservation, MergeMemoryPool, }; use datafusion_execution::runtime_env::RuntimeEnv; use datafusion_physical_expr::LexOrdering; use datafusion_physical_expr::PhysicalExpr; -use datafusion_physical_expr::expressions::{DynamicFilterPhysicalExpr, lit}; +use datafusion_physical_expr::expressions::{Column, DynamicFilterPhysicalExpr, lit}; use futures::{Stream, StreamExt, TryStreamExt}; use log::{debug, trace}; @@ -707,7 +709,8 @@ impl ExternalSorter { if self.in_mem_batches.len() == 1 { let batch = self.in_mem_batches.swap_remove(0); let reservation = self.reservation.take(); - let sorted_stream = self.sort_batch_stream(batch, reservation, None)?; + let sorted_stream = + self.sort_batch_stream(batch, reservation, None, self.batch_size)?; return Ok(self.observe_if_output(sorted_stream, is_output_stream)); } @@ -724,7 +727,8 @@ impl ExternalSorter { .try_resize(get_reserved_bytes_for_record_batch(&batch)?) .map_err(Self::err_with_oom_context)?; let reservation = self.reservation.take(); - let sorted_stream = self.sort_batch_stream(batch, reservation, None)?; + let sorted_stream = + self.sort_batch_stream(batch, reservation, None, self.batch_size)?; return Ok(self.observe_if_output(sorted_stream, is_output_stream)); } @@ -786,7 +790,13 @@ impl ExternalSorter { .enumerate() .map(move |(index, (batch, size))| { let reservation = self.reservation.split(size); - let sort = self.sort_run(batch, reservation, shared_input.clone())?; + let batch_size = self.merge_run_batch_size(&batch)?; + let sort = self.sort_run( + batch, + reservation, + shared_input.clone(), + batch_size, + )?; let permits = Arc::clone(&concurrency.permits); Ok::<_, DataFusionError>(async move { let stream = Self::prepare_sort_run(sort, permits).await?; @@ -972,7 +982,8 @@ impl ExternalSorter { .zip(sizes) .map(|(run, size)| { let reservation = self.reservation.split(size); - self.sort_batch_stream(run, reservation, shared_input.clone()) + let batch_size = self.merge_run_batch_size(&run)?; + self.sort_batch_stream(run, reservation, shared_input.clone(), batch_size) }) .collect() } @@ -995,11 +1006,11 @@ impl ExternalSorter { batch: RecordBatch, reservation: MemoryReservation, shared_input: Option>, + batch_size: usize, ) -> Result { let schema = batch.schema(); - let stream = - futures::stream::once(self.sort_run(batch, reservation, shared_input)?) - .try_flatten(); + let sort = self.sort_run(batch, reservation, shared_input, batch_size)?; + let stream = futures::stream::once(sort).try_flatten(); Ok(Box::pin(RecordBatchStreamAdapter::new(schema, stream))) } @@ -1012,6 +1023,7 @@ impl ExternalSorter { batch: RecordBatch, reservation: MemoryReservation, shared_input: Option>, + batch_size: usize, ) -> Result> + Send + use<>> { let reserved_bytes = get_reserved_bytes_for_record_batch(&batch)?; @@ -1024,9 +1036,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; Ok(async move { // Held until `batch` is sorted, when this block ends. The last run to @@ -1037,7 +1050,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 { + sorted_batches = sorted_batches + .iter() + .map(gc_view_arrays) + .collect::>()?; + } // Chunked output can retain shared buffers in every batch and // exceed the input estimate. Borrow only already-reserved spill @@ -1068,6 +1090,67 @@ impl ExternalSorter { }) } + /// 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 { + 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.encoded_keys_may_exceed(run)? + && !views_share_bytes(run) + { + Ok((self.batch_size / 8).max(1024).min(self.batch_size)) + } else { + Ok(self.batch_size) + } + } + + /// Returns whether the row-format keys of `run` may be larger than `run`. A + /// computed key is sized from its type if that has a fixed width, and + /// otherwise counts as larger, since sizing it would evaluate it again. + fn encoded_keys_may_exceed(&self, run: &RecordBatch) -> Result { + // An 8-byte offset per row, then a null byte and the bytes of each key + let mut encoded = run.num_rows() * size_of::(); + for sort_expr in self.expr.iter() { + let key_bytes = match sort_expr.expr.downcast_ref::() { + Some(column) => run + .column(column.index()) + .to_data() + .get_slice_memory_size()?, + None => { + let width = sort_expr + .expr + .data_type(self.schema.as_ref()) + .ok() + .and_then(|data_type| data_type.primitive_width()); + let Some(width) = width else { + return Ok(true); + }; + run.num_rows() * width + } + }; + encoded += run.num_rows() + key_bytes; + } + Ok(encoded > run.get_sliced_size()?) + } + /// If this sort may spill, pre-allocates /// `sort_spill_reservation_bytes` of memory to guarantee memory /// left for the in memory sort/merge. @@ -2601,7 +2684,8 @@ mod tests { let reservation = sorter.reservation.new_empty(); reservation.try_grow(get_reserved_bytes_for_record_batch(&batch)?)?; let before = elapsed_compute.value(); - let sort = sorter.sort_run(batch.clone(), reservation, None)?; + let sort = + sorter.sort_run(batch.clone(), reservation, None, sorter.batch_size)?; assert_eq!(elapsed_compute.value(), before); let stream = if spawned { diff --git a/datafusion/physical-plan/src/sorts/sort/spill_tests.rs b/datafusion/physical-plan/src/sorts/sort/spill_tests.rs index 625f96be682cb..fa132fb91b410 100644 --- a/datafusion/physical-plan/src/sorts/sort/spill_tests.rs +++ b/datafusion/physical-plan/src/sorts/sort/spill_tests.rs @@ -9,10 +9,10 @@ use crate::metrics::ExecutionPlanMetricsSet; use crate::spill::get_record_batch_memory_size; use crate::spill::spill_manager::GetSlicedSize; use arrow::array::{ - ArrayRef, Decimal128Array, DictionaryArray, Int8Array, Int64Array, StringArray, - StringViewArray, + ArrayRef, Decimal128Array, DictionaryArray, FixedSizeBinaryArray, Int8Array, + Int32Array, Int64Array, StringArray, StringViewArray, }; -use arrow::compute::{SortOptions, concat_batches}; +use arrow::compute::{SortOptions, cast, concat_batches}; use arrow::datatypes::{DataType, Field, Int8Type, Schema}; use arrow::record_batch::RecordBatch; use datafusion_common::config::{ExecutionOptions, SpillCompression}; @@ -22,8 +22,8 @@ use datafusion_execution::memory_pool::{ GreedyMemoryPool, MemoryConsumer, MemoryLimit, MemoryPool, MemoryReservation, }; use datafusion_execution::runtime_env::{RuntimeEnv, RuntimeEnvBuilder}; -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}; use futures::TryStreamExt; use std::fmt::{Display, Formatter}; use std::sync::{Arc, Mutex}; @@ -215,6 +215,13 @@ fn sorter( } async fn assert_released(pool: &AdjustablePool, runtime: &RuntimeEnv) { + wait_until_released(pool, runtime).await; + let state = pool.state.lock().unwrap(); + assert!(state.peak <= pool.capacity); + assert_eq!(state.unchecked_over_limit, 0); +} + +async fn wait_until_released(pool: &AdjustablePool, runtime: &RuntimeEnv) { // Buffered readers may finish dropping their state after the owning stream. tokio::time::timeout(Duration::from_secs(5), async { while pool.reserved() != 0 @@ -227,9 +234,6 @@ async fn assert_released(pool: &AdjustablePool, runtime: &RuntimeEnv) { .await .expect("sort reservations and spill files must be released"); assert_eq!(runtime.disk_manager.used_disk_space(), 0); - let state = pool.state.lock().unwrap(); - assert!(state.peak <= pool.capacity); - assert_eq!(state.unchecked_over_limit, 0); } async fn check_spill_preserves_merge_workspace_after_limit_decreases() -> Result<()> { @@ -679,6 +683,228 @@ 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 { + 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. A computed key without a + // fixed width counts as larger than the run, so only the shared views keep + // these runs at `batch_size`. Sorting the buffered batches before a spill + // grows past the pool, on `main` too. + let fields = [&int_fields[..], &[("c", DataType::Utf8View)]].concat(); + let a: Arc = Arc::new(Column::new("a", 0)); + let keys = vec![ + Arc::new(CastExpr::new(a, DataType::Utf8, None)) as Arc, + Arc::new(Column::new("b", 1)), + ]; + check_merge_exprs(&fields, keys, 16, false, |batch| { + let values: StringArray = (0..1024) + .map(|v| Some(format!("{batch:>4}-{v:0>59}"))) + .collect(); + let keys = + Int32Array::from_iter_values((0..rows).map(|i| (i * 7 + batch) % 1024)); + let views = cast( + &DictionaryArray::try_new(keys, Arc::new(values))?, + &DataType::Utf8View, + )?; + Ok([int_keys(batch), vec![views]].concat()) + }) + .await?; + // Batches that do not overlap, so a merged batch takes the strings of + // several split batches of one run, in as many buffers. + let view_keys = [("a", DataType::Utf8View), ("b", DataType::Utf8View)]; + check_merge_keys(&view_keys, 2, 8, true, |batch| { + let keys = |key: &str| -> Result { + let strings: StringArray = (0..rows) + .rev() + .map(|i| Some(format!("{batch:0>4}{key}{i:0>61}"))) + .collect(); + Ok(cast(&strings, &DataType::Utf8View)?) + }; + Ok(vec![keys("a")?, keys("b")?]) + }) + .await?; + check_merge_keys( + &[("k", DataType::FixedSizeBinary(16))], + 1, + 16, + true, + |batch| { + let keys = (0..rows).map(|i| u128::from(((i * 31 + batch) % 1000) as u32)); + let keys = FixedSizeBinaryArray::try_from_iter(keys.map(u128::to_be_bytes))?; + Ok(vec![Arc::new(keys)]) + }, + ) + .await +} + +/// Sorts 64 batches of `columns` by the first `keys` fields with room to buffer +/// `buffered` of them, then checks the spilled output and that all memory is +/// released. Unless `within_pool`, the pool may have been exceeded by an +/// unchecked `grow`. +async fn check_merge_keys( + fields: &[(&str, DataType)], + keys: usize, + buffered: usize, + within_pool: bool, + columns: impl Fn(i32) -> Result>, +) -> Result<()> { + let keys = (0..keys) + .map(|i| Arc::new(Column::new(fields[i].0, i)) as Arc) + .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>, + buffered: usize, + within_pool: bool, + columns: impl Fn(i32) -> Result>, +) -> Result<()> { + let schema = Arc::new(Schema::new( + fields + .iter() + .map(|(name, data_type)| Field::new(*name, data_type.clone(), false)) + .collect::>(), + )); + let batches = (0..64) + .map(|batch| Ok(RecordBatch::try_new(Arc::clone(&schema), columns(batch)?)?)) + .collect::>>()?; + let ordering = LexOrdering::new(keys.into_iter().map(PhysicalSortExpr::new_default)) + .expect("at least one sort key"); + let capacity = + MERGE_BYTES + buffered * get_reserved_bytes_for_record_batch(&batches[0])?; + let pool = AdjustablePool::new(capacity); + let runtime = RuntimeEnvBuilder::new() + .with_memory_pool(Arc::clone(&pool) as Arc) + .build_arc()?; + let mut sorter = ExternalSorter::new( + 0, + Arc::clone(&schema), + ordering.clone(), + batches[0].num_rows(), + MERGE_BYTES, + 0, + SpillCompression::Uncompressed, + &ExecutionPlanMetricsSet::new(), + Arc::clone(&runtime), + )?; + for batch in &batches { + sorter.insert_batch(batch.clone()).await?; + } + assert!(sorter.spilled_before()); + let output: Vec = sorter.sort().await?.try_collect().await?; + let expected = sort_batch(&concat_batches(&schema, &batches)?, &ordering, None)?; + assert_eq!(concat_batches(&schema, &output)?, expected); + drop(sorter); + if within_pool { + assert_released(&pool, &runtime).await; + } else { + wait_until_released(&pool, &runtime).await; + } + Ok(()) +} + +#[tokio::test] +async fn test_merge_keys_larger_than_batches() -> Result<()> { + check_merge_keys_larger_than_batches().await +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn test_merge_keys_larger_than_batches_multithreaded() -> Result<()> { + check_merge_keys_larger_than_batches().await +} + +/// 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; + // Sorts by `first`, then `b` + let check = |first: Arc, + payload: Option, + 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 b: Arc = Arc::new(Column::new("b", 1)); + let ordering = + LexOrdering::new([first, b].into_iter().map(PhysicalSortExpr::new_default)) + .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(()) + }; + let a = || Arc::new(Column::new("a", 0)) as Arc; + // Encoded keys: 26 bytes a row, against 16 bytes a row + check(a(), 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(a(), 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(); + let long: ArrayRef = Arc::new(long); + check(a(), Some(Arc::clone(&long)), batch_size)?; + // A computed key is sized from its type if that has a fixed width, and + // otherwise counts as larger than the run + let cast = |to| Arc::new(CastExpr::new(a(), to, None)) as Arc; + check(cast(DataType::Int64), Some(Arc::clone(&long)), batch_size)?; + check(cast(DataType::Utf8), Some(long), batch_size / 8) +} + #[tokio::test] async fn test_single_batch_spill_preserves_workspace_after_limit_decreases() -> Result<()> { diff --git a/datafusion/physical-plan/src/sorts/streaming_merge.rs b/datafusion/physical-plan/src/sorts/streaming_merge.rs index 2e5e0790accec..1668a55ac923d 100644 --- a/datafusion/physical-plan/src/sorts/streaming_merge.rs +++ b/datafusion/physical-plan/src/sorts/streaming_merge.rs @@ -13,7 +13,7 @@ use crate::sorts::{ }; use crate::{EmptyRecordBatchStream, SendableRecordBatchStream, SpillManager}; use arrow::array::*; -use arrow::datatypes::{DataType, SchemaRef}; +use arrow::datatypes::{DataType, Schema, SchemaRef}; use datafusion_common::human_readable_size; use datafusion_common::{Result, assert_or_internal_err, internal_err}; use datafusion_execution::SpillFile; @@ -273,7 +273,8 @@ impl<'a> StreamingMergeBuilder<'a> { let reservation = reservation.expect("Reservation cannot be empty for streaming merge"); - // Special case single column comparisons with optimized cursor implementations + // Special case single column comparisons with optimized cursor implementations. + // `merges_in_row_format` must match this choice. if expressions.len() == 1 { let sort = expressions[0].clone(); let data_type = sort.expr.data_type(schema.as_ref())?; @@ -308,6 +309,26 @@ impl<'a> StreamingMergeBuilder<'a> { } } +/// Returns whether [`StreamingMergeBuilder::build`] merges on `expressions` in +/// the row format rather than with a cursor for a single column. +pub(crate) fn merges_in_row_format(expressions: &LexOrdering, schema: &Schema) -> bool { + if expressions.len() > 1 { + return true; + } + let Ok(data_type) = expressions[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 + )) +} + #[cfg(test)] mod tests { use crate::{common::collect, stream::RecordBatchStreamAdapter}; diff --git a/datafusion/physical-plan/src/spill/mod.rs b/datafusion/physical-plan/src/spill/mod.rs index f2ec70c6f3544..3bfd0c37eeebb 100644 --- a/datafusion/physical-plan/src/spill/mod.rs +++ b/datafusion/physical-plan/src/spill/mod.rs @@ -833,15 +833,17 @@ const VIEW_SIZE_BYTES: usize = 16; /// /// # When GC is triggered /// -/// View GC is only performed when data buffers exceed a threshold (currently 10KB). -/// This balances memory savings against the CPU overhead of garbage collection. -/// Small view buffers are passed through unchanged since the GC overhead would exceed -/// any memory savings. +/// View GC runs when data buffers above a threshold (currently 10KB) hold more +/// bytes than the views use. This balances memory savings against the CPU +/// overhead of garbage collection. It also runs at any size when several data +/// buffers remain and the views share no bytes, since each buffer read back +/// from the spill is charged for the whole IPC message. /// /// # Performance considerations /// /// - If no arrays need compaction, the original batch is cloned cheaply -/// - View GC is skipped for small buffers to avoid unnecessary CPU overhead +/// - View GC is skipped for a single small data buffer to avoid unnecessary CPU +/// overhead /// - Nested container types are traversed recursively so view arrays inside /// `List`, `Map`, `Union`, `Dictionary`, and other child-bearing arrays are compacted too /// - Traversal uses physical children, not logical reachability. For example, @@ -962,9 +964,11 @@ fn gc_array(array: &ArrayRef, dictionaries: bool) -> Result<(ArrayRef, bool)> { /// batch, and `gc()` copies it once per row /// (). /// -/// Repeated buffer entries are merged first. `gc()` then runs only if it -/// writes less than the merged buffers hold, and runs on the original array -/// so that the views are rewritten at most once. +/// Repeated buffer entries are merged first. `gc()` then runs if it writes +/// less than the merged buffers hold, or if several buffers remain and the +/// views share no bytes, since each buffer read back from the spill is a slice +/// of one IPC message and is charged for all of it. It runs on the original +/// array so that the views are rewritten at most once. fn compact_view_array( array: &GenericByteViewArray, ) -> Option> { @@ -975,7 +979,9 @@ fn compact_view_array( 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) + { return Some(array.gc()); } if unique.len() == buffers.len() { @@ -1034,6 +1040,21 @@ fn dedup_buffer_entries(buffers: &[Buffer]) -> (Vec, Vec) { (remap, unique) } +/// Returns whether a view column of `batch` references more bytes than its +/// data buffers hold, as the views of a dictionary-encoded Parquet column do. +/// Compacting such views would copy the shared bytes once per row. +pub(crate) fn views_share_bytes(batch: &RecordBatch) -> bool { + fn shares_bytes(array: &GenericByteViewArray) -> bool { + let (_, unique) = dedup_buffer_entries(array.data_buffers()); + array.total_buffer_bytes_used() > unique.iter().map(|b| b.len()).sum::() + } + batch.columns().iter().any(|array| match array.data_type() { + DataType::Utf8View => shares_bytes(array.as_string_view()), + DataType::BinaryView => shares_bytes(array.as_binary_view()), + _ => false, + }) +} + fn gc_array_children(array: &ArrayRef, dictionaries: bool) -> Result<(ArrayRef, bool)> { let data = array.to_data(); if data.child_data().is_empty() { @@ -2669,6 +2690,42 @@ mod tests { Ok((dictionary, merged.as_string_view().clone())) } + #[test] + fn test_views_share_bytes() -> Result<()> { + use arrow::array::UInt32Array; + use arrow::compute::take; + + // 8192 views over 1024 values reference 8x the dictionary's bytes + let dictionary = dictionary_array(1024)?; + let indices = UInt32Array::from_iter_values((0..8192).map(|i| i % 1024)); + let batch = + RecordBatch::try_from_iter([("c", take(&dictionary, &indices, None)?)])?; + assert!(views_share_bytes(&batch)); + // 8192 values used once each reference exactly the buffer's bytes + let views: ArrayRef = Arc::new(dictionary_array(8192)?); + let batch = RecordBatch::try_from_iter([("c", views)])?; + assert!(!views_share_bytes(&batch)); + Ok(()) + } + + #[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(()) + } + #[test] fn test_gc_keeps_shared_dictionary_buffer() -> Result<()> { let (dictionary, merged) = shared_dictionary_views(1000, 8, 1024)?;