POC Optimize near-unique string view group values

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,17 +169,29 @@ where
             random_state: RandomState::default(),
             hashes_buffer: vec![],
             null: None,
+            is_near_unique: false,
         }
     }
 
     /// Return the contents of this map and replace it with a new empty map with
     /// 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() {

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

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(

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();
 

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

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,
             }),

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);
 }

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