This is an automated email from the ASF dual-hosted git repository. Rachelint pushed a commit to branch mark-near-unique-to-group-values in repository https://gitbox.apache.org/repos/asf/datafusion.git
commit ae74b2607ec5863774f03b1e22cf73720efbfa43 Author: kamille <[email protected]> AuthorDate: Wed Jul 8 18:20:47 2026 +0800 Optimize near-unique string view group values --- .../physical-expr-common/src/binary_view_map.rs | 92 ++++++++++++++++++++-- .../src/aggregates/aggregate_hash_table/common.rs | 8 +- .../aggregate_hash_table/common_ordered.rs | 11 ++- .../aggregates/aggregate_hash_table/final_table.rs | 18 ++++- .../aggregate_hash_table/ordered_final_table.rs | 14 +++- .../aggregate_hash_table/partial_table.rs | 1 + .../src/aggregates/group_values/mod.rs | 49 ++++++++++++ .../group_values/single_group_by/bytes_view.rs | 4 + .../src/aggregates/grouped_hash_stream.rs | 29 ++++++- 9 files changed, 210 insertions(+), 16 deletions(-) diff --git a/datafusion/physical-expr-common/src/binary_view_map.rs b/datafusion/physical-expr-common/src/binary_view_map.rs index 9d4b556393..8e7670e722 100644 --- a/datafusion/physical-expr-common/src/binary_view_map.rs +++ b/datafusion/physical-expr-common/src/binary_view_map.rs @@ -145,6 +145,9 @@ where /// NOTE null_index is the logical index in the final array, not the index /// in the buffer null: Option<(V, usize)>, + + /// Prefer retaining input buffers instead of compacting copied values. + is_near_unique: bool, } /// The size, in number of entries, of the initial hash table @@ -166,6 +169,7 @@ where random_state: RandomState::default(), hashes_buffer: vec![], null: None, + is_near_unique: false, } } @@ -173,10 +177,21 @@ where /// the same output type pub fn take(&mut self) -> Self { let mut new_self = Self::new(self.output_type); + new_self.is_near_unique = self.is_near_unique; std::mem::swap(self, &mut new_self); new_self } + /// Marks this map as storing near-unique input values. + /// + /// In this mode, newly interned non-inline `StringView` values retain the + /// input array buffers instead of copying bytes into compacted buffers. + pub fn mark_near_unique(&mut self) { + if self.output_type == OutputType::Utf8View { + self.is_near_unique = true; + } + } + /// Inserts each value from `values` into the map, invoking `payload_fn` for /// each value if *not* already present, deferring the allocation of the /// payload until it is needed. @@ -270,6 +285,7 @@ where // Ensure lengths are equivalent assert_eq!(values.len(), self.hashes_buffer.len()); + let mut shared_buffer_index_offset = None; for i in 0..values.len() { let view_u128 = input_views[i]; let hash = self.hashes_buffer[i]; @@ -357,7 +373,22 @@ where } else { let value: &[u8] = values.value(i).as_ref(); let payload = make_payload_fn(Some(value)); - let new_view = self.append_value(value); + let new_view = if self.is_near_unique { + let buffer_index_offset = match shared_buffer_index_offset { + Some(buffer_index_offset) => buffer_index_offset, + None => { + self.flush_in_progress(); + let buffer_index_offset = self.completed.len() as u32; + self.completed + .extend(values.data_buffers().iter().cloned()); + shared_buffer_index_offset = Some(buffer_index_offset); + buffer_index_offset + } + }; + self.append_shared_view(view_u128, buffer_index_offset) + } else { + self.append_value(value) + }; (new_view, payload) }; @@ -383,10 +414,7 @@ where /// they were first seen. pub fn into_state(mut self) -> ArrayRef { // Flush any remaining in-progress buffer - if !self.in_progress.is_empty() { - let flushed = std::mem::take(&mut self.in_progress); - self.completed.push(Buffer::from_vec(flushed)); - } + self.flush_in_progress(); // Build null buffer if we have any nulls let null_buffer = self.nulls.finish(); @@ -426,6 +454,23 @@ where view } + fn flush_in_progress(&mut self) { + if !self.in_progress.is_empty() { + let flushed = std::mem::take(&mut self.in_progress); + self.completed.push(Buffer::from_vec(flushed)); + } + } + + /// Append a non-inline view that points into shared input buffers. + fn append_shared_view(&mut self, view: u128, buffer_index_offset: u32) -> u128 { + let mut byte_view = ByteView::from(view); + byte_view.buffer_index += buffer_index_offset; + let new_view = byte_view.as_u128(); + self.views.push(new_view); + self.nulls.append_non_null(); + new_view + } + /// Append a value to our buffers and return the view pointing to it fn append_value(&mut self, value: &[u8]) -> u128 { let len = value.len(); @@ -660,6 +705,43 @@ mod tests { assert_eq!(&set.into_state(), &expected); } + // Covers near-unique insertion retaining the input data buffer. + // Example: two long unique strings share the original `StringViewArray` buffer. + #[test] + fn test_bytes_view_map_mark_near_unique_retains_input_buffer() { + let input = StringViewArray::from(vec![ + Some("long string value 1"), + Some("long string value 2"), + ]); + let input_buffer_ptr = input.data_buffers()[0].as_slice().as_ptr(); + let values: ArrayRef = Arc::new(input); + + let mut map = ArrowBytesViewMap::<usize>::new(OutputType::Utf8View); + map.mark_near_unique(); + let mut next_index = 0; + map.insert_if_new( + &values, + |_| { + let index = next_index; + next_index += 1; + index + }, + |_| {}, + ); + + let output = map.into_state(); + let output = output.as_string_view(); + assert_eq!( + output.iter().collect::<Vec<_>>(), + vec![Some("long string value 1"), Some("long string value 2"),] + ); + assert_eq!(output.data_buffers().len(), 1); + assert_eq!( + output.data_buffers()[0].as_slice().as_ptr(), + input_buffer_ptr + ); + } + // inserting strings into the set does not increase reported memory #[test] fn test_string_set_memory_usage() { diff --git a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common.rs b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common.rs index f29f3e7ff8..8f2899c92a 100644 --- a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common.rs +++ b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common.rs @@ -27,7 +27,9 @@ use datafusion_expr::{EmitTo, GroupsAccumulator}; use datafusion_physical_expr::aggregate::AggregateFunctionExpr; use crate::PhysicalExpr; -use crate::aggregates::group_values::{GroupByMetrics, GroupValues, new_group_values}; +use crate::aggregates::group_values::{ + GroupByMetrics, GroupValues, NearUniqueGroupValuesProbe, new_group_values, +}; use crate::aggregates::grouped_hash_stream::create_group_accumulator; use crate::aggregates::order::GroupOrdering; use crate::aggregates::{ @@ -135,6 +137,7 @@ impl<AggrMode> AggregateHashTable<AggrMode> { state: AggregateHashTableState::Building(AggregateHashTableBuffer { group_by: Arc::clone(&agg.group_by), group_values, + near_unique_probe: None, batch_group_indices: Default::default(), accumulators, }), @@ -281,6 +284,9 @@ pub(super) struct AggregateHashTableBuffer { /// Interned group keys. Accumulator state is stored separately by group index. pub(super) group_values: Box<dyn GroupValues>, + /// Samples final aggregate input to detect near-unique group keys. + pub(super) near_unique_probe: Option<NearUniqueGroupValuesProbe>, + /// Group index for each row in the current input batch. /// /// Each value indexes into `group_values`, and the same index is used by every diff --git a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common_ordered.rs b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common_ordered.rs index c83303c51d..356701f4cc 100644 --- a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common_ordered.rs +++ b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common_ordered.rs @@ -30,7 +30,9 @@ use datafusion_expr::EmitTo; use crate::InputOrderMode; use crate::PhysicalExpr; -use crate::aggregates::group_values::{GroupByMetrics, GroupValues, new_group_values}; +use crate::aggregates::group_values::{ + GroupByMetrics, GroupValues, NearUniqueGroupValuesProbe, new_group_values, +}; use crate::aggregates::grouped_hash_stream::create_group_accumulator; use crate::aggregates::order::GroupOrdering; use crate::aggregates::{ @@ -111,6 +113,9 @@ pub(super) struct OrderedAggregateTableBuffer { /// Interned group keys, in the same group-id order used by accumulators. pub(super) group_values: Box<dyn GroupValues>, + /// Samples final aggregate input to detect near-unique group keys. + pub(super) near_unique_probe: Option<NearUniqueGroupValuesProbe>, + /// Scratch group id vector for the current input batch. pub(super) group_indices: Vec<usize>, @@ -174,6 +179,7 @@ impl<AggrMode> OrderedAggregateTable<AggrMode> { group_by: Arc::clone(&agg.group_by), group_ordering, group_values, + near_unique_probe: None, group_indices: vec![], accumulators, }, @@ -269,6 +275,9 @@ impl<AggrMode> OrderedAggregateTable<AggrMode> { self.buffer .group_values .intern(group_values, &mut self.buffer.group_indices)?; + if is_final && let Some(probe) = self.buffer.near_unique_probe.as_mut() { + probe.observe(self.buffer.group_values.as_mut(), group_values[0].len()); + } let total_num_groups = self.buffer.group_values.len(); if total_num_groups > starting_num_groups { self.buffer.group_ordering.new_groups( diff --git a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/final_table.rs b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/final_table.rs index 568b866b10..e8473ded5f 100644 --- a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/final_table.rs +++ b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/final_table.rs @@ -17,12 +17,13 @@ use std::sync::Arc; -use arrow::datatypes::SchemaRef; +use arrow::datatypes::{DataType, SchemaRef}; use arrow::record_batch::RecordBatch; use datafusion_common::{Result, internal_err}; use datafusion_expr::EmitTo; use crate::aggregates::AggregateExec; +use crate::aggregates::group_values::NearUniqueGroupValuesProbe; use super::common::{ AggregateHashTable, AggregateHashTableBuffer, AggregateHashTableState, FinalMarker, @@ -43,13 +44,21 @@ impl AggregateHashTable<FinalMarker> { output_schema: SchemaRef, batch_size: usize, ) -> Result<Self> { - Self::new_with_filters( + let mut table = Self::new_with_filters( agg, partition, output_schema, batch_size, vec![None; agg.aggr_expr.len()], - ) + )?; + let group_schema = agg.group_by.group_schema(&agg.input().schema())?; + if group_schema.fields().len() == 1 + && matches!(group_schema.field(0).data_type(), DataType::Utf8View) + { + table.state.building_mut().near_unique_probe = + Some(NearUniqueGroupValuesProbe::new()); + } + Ok(table) } /// Emits the next batch of aggregated group keys and final aggregate values. @@ -131,6 +140,9 @@ impl AggregateHashTable<FinalMarker> { state .group_values .intern(group_values, &mut state.batch_group_indices)?; + if let Some(probe) = state.near_unique_probe.as_mut() { + probe.observe(state.group_values.as_mut(), group_values[0].len()); + } let group_indices = &state.batch_group_indices; let total_num_groups = state.group_values.len(); diff --git a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/ordered_final_table.rs b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/ordered_final_table.rs index b7e3fd38ed..dede4db73b 100644 --- a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/ordered_final_table.rs +++ b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/ordered_final_table.rs @@ -19,12 +19,13 @@ //! //! See comments in [`super::ordered_partial_table`] for details. -use arrow::datatypes::SchemaRef; +use arrow::datatypes::{DataType, SchemaRef}; use arrow::record_batch::RecordBatch; use datafusion_common::Result; use crate::InputOrderMode; use crate::aggregates::aggregate_hash_table::FinalMarker; +use crate::aggregates::group_values::NearUniqueGroupValuesProbe; use crate::aggregates::{AggregateExec, AggregateMode}; use super::common_ordered::OrderedAggregateTable; @@ -47,7 +48,7 @@ impl OrderedAggregateTable<FinalMarker> { batch_size: usize, input_order_mode: &InputOrderMode, ) -> Result<Self> { - Self::new_for_mode( + let mut table = Self::new_for_mode( agg, partition, input_schema, @@ -56,7 +57,14 @@ impl OrderedAggregateTable<FinalMarker> { input_order_mode, &AggregateMode::Final, vec![None; agg.aggr_expr.len()], - ) + )?; + let group_schema = agg.group_by.group_schema(input_schema)?; + if group_schema.fields().len() == 1 + && matches!(group_schema.field(0).data_type(), DataType::Utf8View) + { + table.buffer.near_unique_probe = Some(NearUniqueGroupValuesProbe::new()); + } + Ok(table) } /// Merges one partial-state input batch and updates ordering information for diff --git a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/partial_table.rs b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/partial_table.rs index f11eef8c14..0ce07bd636 100644 --- a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/partial_table.rs +++ b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/partial_table.rs @@ -154,6 +154,7 @@ impl AggregateHashTable<PartialMarker> { state: AggregateHashTableState::Building(AggregateHashTableBuffer { group_by: Arc::clone(&state.group_by), group_values, + near_unique_probe: None, batch_group_indices: Default::default(), accumulators, }), diff --git a/datafusion/physical-plan/src/aggregates/group_values/mod.rs b/datafusion/physical-plan/src/aggregates/group_values/mod.rs index ee253e5d7a..b10d4aa1bb 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/mod.rs +++ b/datafusion/physical-plan/src/aggregates/group_values/mod.rs @@ -51,6 +51,52 @@ mod null_builder; pub(crate) use metrics::GroupByMetrics; +const NEAR_UNIQUE_SAMPLE_ROWS: usize = 100_000; +const NEAR_UNIQUE_RATIO_THRESHOLD: f64 = 0.8; + +/// Samples input row counts to decide whether group values are near-unique. +pub(in crate::aggregates) struct NearUniqueGroupValuesProbe { + /// Number of rows observed in the sample window. + num_sampled_rows: usize, + /// Whether the sampling decision has already been made. + is_complete: bool, +} + +impl NearUniqueGroupValuesProbe { + /// Creates a new probe for the first input rows. + pub(in crate::aggregates) fn new() -> Self { + Self { + num_sampled_rows: 0, + is_complete: false, + } + } + + /// Observes one interned batch and marks near-unique if needed. + pub(in crate::aggregates) fn observe( + &mut self, + group_values: &mut dyn GroupValues, + num_rows: usize, + ) { + if self.is_complete { + return; + } + + let num_remaining_rows = NEAR_UNIQUE_SAMPLE_ROWS - self.num_sampled_rows; + let num_rows_to_sample = num_remaining_rows.min(num_rows); + self.num_sampled_rows += num_rows_to_sample; + + if self.num_sampled_rows < NEAR_UNIQUE_SAMPLE_ROWS { + return; + } + + let near_unique_ratio = group_values.len() as f64 / self.num_sampled_rows as f64; + if near_unique_ratio > NEAR_UNIQUE_RATIO_THRESHOLD { + group_values.mark_near_unique(); + } + self.is_complete = true; + } +} + /// Stores the group values during hash aggregation. /// /// # Background @@ -111,6 +157,9 @@ pub trait GroupValues: Send { /// Emits the group values fn emit(&mut self, emit_to: EmitTo) -> Result<Vec<ArrayRef>>; + /// Marks this [`GroupValues`] as likely to receive mostly unique values. + fn mark_near_unique(&mut self) {} + /// Clear the contents and shrink the capacity to the size of the batch (free up memory usage) fn clear_shrink(&mut self, num_rows: usize); } diff --git a/datafusion/physical-plan/src/aggregates/group_values/single_group_by/bytes_view.rs b/datafusion/physical-plan/src/aggregates/group_values/single_group_by/bytes_view.rs index 7a56f7c52c..7ba2a8ac60 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/single_group_by/bytes_view.rs +++ b/datafusion/physical-plan/src/aggregates/group_values/single_group_by/bytes_view.rs @@ -122,6 +122,10 @@ impl GroupValues for GroupValuesBytesView { Ok(vec![group_values]) } + fn mark_near_unique(&mut self) { + self.map.mark_near_unique(); + } + fn clear_shrink(&mut self, _num_rows: usize) { // in theory we could potentially avoid this reallocation and clear the // contents of the maps, but for now we just reset the map from the beginning diff --git a/datafusion/physical-plan/src/aggregates/grouped_hash_stream.rs b/datafusion/physical-plan/src/aggregates/grouped_hash_stream.rs index 0d00e5c4d0..e176a4b18a 100644 --- a/datafusion/physical-plan/src/aggregates/grouped_hash_stream.rs +++ b/datafusion/physical-plan/src/aggregates/grouped_hash_stream.rs @@ -24,7 +24,9 @@ use std::vec; use super::order::GroupOrdering; use super::skip_partial::SkipAggregationProbe; use super::{AggregateExec, format_human_display}; -use crate::aggregates::group_values::{GroupByMetrics, GroupValues, new_group_values}; +use crate::aggregates::group_values::{ + GroupByMetrics, GroupValues, NearUniqueGroupValuesProbe, new_group_values, +}; use crate::aggregates::order::GroupOrderingFull; use crate::aggregates::{ AggregateInputMode, AggregateMode, AggregateOutputMode, PhysicalGroupBy, @@ -39,7 +41,7 @@ use crate::{PhysicalExpr, aggregates, metrics}; use crate::{RecordBatchStream, SendableRecordBatchStream}; use arrow::array::*; -use arrow::datatypes::SchemaRef; +use arrow::datatypes::{DataType, SchemaRef}; use datafusion_common::{ DataFusionError, Result, assert_eq_or_internal_err, assert_or_internal_err, internal_err, resources_datafusion_err, @@ -330,6 +332,9 @@ pub(crate) struct GroupedHashAggregateStream { /// An interning store of group keys group_values: Box<dyn GroupValues>, + /// Samples final aggregate input to detect near-unique group keys. + near_unique_probe: Option<NearUniqueGroupValuesProbe>, + /// scratch space for the current input [`RecordBatch`] being /// processed. Reused across batches here to avoid reallocations current_group_indices: Vec<usize>, @@ -511,7 +516,17 @@ impl GroupedHashAggregateStream { _ => OutOfMemoryMode::ReportError, }; - let group_values = new_group_values(group_schema, &group_ordering)?; + let group_values = new_group_values(Arc::clone(&group_schema), &group_ordering)?; + let near_unique_probe = if matches!( + agg.mode, + AggregateMode::Final | AggregateMode::FinalPartitioned + ) && group_schema.fields().len() == 1 + && matches!(group_schema.field(0).data_type(), DataType::Utf8View) + { + Some(NearUniqueGroupValuesProbe::new()) + } else { + None + }; let reservation = MemoryConsumer::new(name) // We interpret 'can spill' as 'can handle memory back pressure'. // This value needs to be set to true for the default memory pool implementations @@ -600,6 +615,7 @@ impl GroupedHashAggregateStream { reservation, oom_mode, group_values, + near_unique_probe, current_group_indices: Default::default(), exec_state, baseline_metrics, @@ -885,6 +901,13 @@ impl GroupedHashAggregateStream { let starting_num_groups = self.group_values.len(); self.group_values .intern(group_values, &mut self.current_group_indices)?; + if matches!( + self.mode, + AggregateMode::Final | AggregateMode::FinalPartitioned + ) && let Some(probe) = self.near_unique_probe.as_mut() + { + probe.observe(self.group_values.as_mut(), group_values[0].len()); + } let group_indices = &self.current_group_indices; // Update ordering information if necessary --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
