feat: GroupValues::find to look up groups without inserting

datafusion/physical-expr-common/src/binary_map.rs

@@ -477,6 +477,83 @@ where
         }
     }
 
+    /// Looks up each value in `values` without inserting anything.
+    ///
+    /// `observe_payload_fn` is called once per row, in order, with
+    /// `Some(payload)` if the value is already stored and `None` otherwise. A
+    /// null row maps to the payload of the null entry if one exists, matching
+    /// [`Self::insert_if_new`].
+    pub fn find<OP>(&mut self, values: &ArrayRef, observe_payload_fn: OP)
+    where
+        OP: FnMut(Option<V>),
+    {
+        match self.output_type {
+            OutputType::Binary => {
+                assert!(matches!(
+                    values.data_type(),
+                    DataType::Binary | DataType::LargeBinary
+                ));
+                self.find_inner::<OP, GenericBinaryType<O>>(values, observe_payload_fn)
+            }
+            OutputType::Utf8 => {
+                assert!(matches!(
+                    values.data_type(),
+                    DataType::Utf8 | DataType::LargeUtf8
+                ));
+                self.find_inner::<OP, GenericStringType<O>>(values, observe_payload_fn)
+            }
+            _ => unreachable!("View types should use `ArrowBytesViewMap`"),
+        };
+    }
+
+    /// Generic version of [`Self::find`] that handles `ByteArrayType`
+    fn find_inner<OP, B>(&mut self, values: &ArrayRef, mut observe_payload_fn: OP)
+    where
+        OP: FnMut(Option<V>),
+        B: ByteArrayType,
+    {
+        let batch_hashes = &mut self.hashes_buffer;
+        batch_hashes.clear();
+        batch_hashes.resize(values.len(), 0);
+        create_hashes([values], &self.random_state, batch_hashes)
+            // hash is supported for all types and create_hashes only
+            // returns errors for unsupported types
+            .unwrap();
+
+        let values = values.as_bytes::<B>();
+        assert_eq!(values.len(), batch_hashes.len());
+
+        for (value, &hash) in values.iter().zip(batch_hashes.iter()) {
+            let Some(value) = value else {
+                observe_payload_fn(self.null.as_ref().map(|&(payload, _)| payload));
+                continue;
+            };
+
+            let value: &[u8] = value.as_ref();
+            let value_len = O::usize_as(value.len());
+
+            let entry = if value.len() <= SHORT_VALUE_LEN {
+                let inline = value.iter().fold(0usize, |acc, &x| (acc << 8) | x as usize);
+                self.map.find(hash, |header| {
+                    header.hash == hash
+                        && header.len == value_len
+                        && inline == header.offset_or_inline
+                })
+            } else {
+                self.map.find(hash, |header| {
+                    if header.hash != hash {
+                        return false;
+                    }
+                    // SAFETY: buffer is only appended to, and we correctly inserted values and offsets
+                    let existing_value =
+                        unsafe { self.buffer.get_unchecked(header.range()) };
+                    value == existing_value
+                })
+            };
+            observe_payload_fn(entry.map(|entry| entry.payload));
+        }
+    }
+
     /// Converts this set into a `StringArray`, `LargeStringArray`,
     /// `BinaryArray`, or `LargeBinaryArray` containing each distinct value
     /// that was inserted. This is done without copying the values.

datafusion/physical-expr-common/src/binary_view_map.rs

@@ -379,6 +379,90 @@ where
         }
     }
 
+    /// Looks up each value in `values` without inserting anything.
+    ///
+    /// `observe_payload_fn` is called once per row, in order, with
+    /// `Some(payload)` if the value is already stored and `None` otherwise. A
+    /// null row maps to the payload of the null entry if one exists, matching
+    /// [`Self::insert_if_new`].
+    pub fn find<OP>(&mut self, values: &ArrayRef, observe_payload_fn: OP)
+    where
+        OP: FnMut(Option<V>),
+    {
+        match self.output_type {
+            OutputType::BinaryView => {
+                assert!(matches!(values.data_type(), DataType::BinaryView));
+                self.find_inner::<OP, BinaryViewType>(values, observe_payload_fn)
+            }
+            OutputType::Utf8View => {
+                assert!(matches!(values.data_type(), DataType::Utf8View));
+                self.find_inner::<OP, StringViewType>(values, observe_payload_fn)
+            }
+            _ => unreachable!("Utf8/Binary should use `ArrowBytesSet`"),
+        };
+    }
+
+    /// Generic version of [`Self::find`] that handles `ByteViewType`
+    fn find_inner<OP, B>(&mut self, values: &ArrayRef, mut observe_payload_fn: OP)
+    where
+        OP: FnMut(Option<V>),
+        B: ByteViewType,
+    {
+        let batch_hashes = &mut self.hashes_buffer;
+        batch_hashes.clear();
+        batch_hashes.resize(values.len(), 0);
+        create_hashes([values], &self.random_state, batch_hashes)
+            // hash is supported for all types and create_hashes only
+            // returns errors for unsupported types
+            .unwrap();
+
+        let values = values.as_byte_view::<B>();
+        let input_views = values.views();
+        assert_eq!(values.len(), self.hashes_buffer.len());
+
+        let completed = &self.completed;
+        let in_progress = &self.in_progress;
+        for i in 0..values.len() {
+            if values.is_null(i) {
+                observe_payload_fn(self.null.as_ref().map(|&(payload, _)| payload));
+                continue;
+            }
+
+            let view_u128 = input_views[i];
+            let hash = self.hashes_buffer[i];
+            let len = view_u128 as u32;
+
+            let payload = self
+                .map
+                .find(hash, |header| {
+                    if header.hash != hash {
+                        return false;
+                    }
+                    if len <= 12 {
+                        return header.view == view_u128;
+                    }
+                    let stored_prefix = (header.view >> 32) as u32;
+                    let input_prefix = (view_u128 >> 32) as u32;
+                    if stored_prefix != input_prefix {
+                        return false;
+                    }
+                    let byte_view = ByteView::from(header.view);
+                    let stored_len = byte_view.length as usize;
+                    let buffer_index = byte_view.buffer_index as usize;
+                    let offset = byte_view.offset as usize;
+                    let stored_value = if buffer_index < completed.len() {
+                        &completed[buffer_index].as_slice()[offset..offset + stored_len]
+                    } else {
+                        &in_progress[offset..offset + stored_len]
+                    };
+                    let input_value: &[u8] = values.value(i).as_ref();
+                    stored_value == input_value
+                })
+                .map(|entry| entry.payload);
+            observe_payload_fn(payload);
+        }
+    }
+
     /// Returns the bytes for the key at `index`, irrespective of nullness.
     fn key(&self, index: usize) -> &[u8] {
         let view = &self.views[index];

datafusion/physical-plan/src/aggregates/group_values/mod.rs

@@ -99,6 +99,14 @@ pub trait GroupValues: Send {
     /// assigned.
     fn intern(&mut self, cols: &[ArrayRef], groups: &mut Vec<usize>) -> Result<()>;
 
+    /// Looks up the group id for each input row of `cols` without assigning
+    /// new group ids.
+    ///
+    /// When the function returns, `groups` has been cleared and contains one
+    /// entry per input row: `Some(group_id)` if the value is already stored
+    /// and `None` otherwise. Stored values and group ids are unchanged.
+    fn find(&mut self, cols: &[ArrayRef], groups: &mut Vec<Option<usize>>) -> Result<()>;
+
     /// Returns the number of bytes of memory used by this [`GroupValues`].
     ///
     /// May be expensive; check the implementation before calling on hot paths.
@@ -237,13 +245,203 @@ pub fn new_group_values(
 mod tests {
     use std::sync::Arc;
 
-    use arrow::array::{ArrayRef, AsArray, Int32Array, StringArray, StringViewArray};
+    use arrow::array::{
+        ArrayRef, AsArray, BinaryArray, BinaryViewArray, BooleanArray, Int32Array,
+        Int64Array, LargeStringArray, ListArray, StringArray, StringViewArray,
+    };
     use arrow::datatypes::{DataType, Field, Int32Type, Schema};
     use datafusion_expr::{EmitTo, GroupSelection};
 
-    use super::new_group_values;
+    use super::{GroupValues, new_group_values};
     use crate::aggregates::order::GroupOrdering;
 
+    /// Interns `batches`, then checks that `find` returns the interned id for
+    /// each stored row, `None` for `missing`, and never inserts.
+    fn assert_find_matches_intern(
+        schema: Arc<Schema>,
+        batches: Vec<Vec<ArrayRef>>,
+        missing: Vec<ArrayRef>,
+    ) {
+        let mut gv: Box<dyn GroupValues> =
+            new_group_values(schema, &GroupOrdering::None).unwrap();
+        let mut interned = vec![];
+        let mut expected = vec![];
+        for batch in &batches {
+            gv.intern(batch, &mut interned).unwrap();
+            expected.push(interned.iter().copied().map(Some).collect::<Vec<_>>());
+        }
+        let num_groups = gv.len();
+
+        // `groups` is cleared, so start with stale contents.
+        let mut found = vec![Some(usize::MAX); 3];
+        for (batch, expected) in batches.iter().zip(&expected) {
+            gv.find(batch, &mut found).unwrap();
+            assert_eq!(&found, expected);
+        }
+        assert_eq!(gv.len(), num_groups);
+
+        let n = missing[0].len();
+        gv.find(&missing, &mut found).unwrap();
+        assert_eq!(found, vec![None; n]);
+        assert_eq!(gv.len(), num_groups);
+    }
+
+    fn single(data_type: DataType) -> Arc<Schema> {
+        Arc::new(Schema::new(vec![Field::new("g", data_type, true)]))
+    }
+
+    #[test]
+    fn find_matches_intern_single_column() {
+        let long = "a long value that is not inline";
+        let ints = || -> ArrayRef {
+            Arc::new(Int32Array::from(vec![Some(1), None, Some(2), Some(1)]))
+        };
+        assert_find_matches_intern(
+            single(DataType::Int32),
+            vec![vec![ints()]],
+            vec![Arc::new(Int32Array::from(vec![Some(99), Some(100)]))],
+        );
+
+        let strs = |v: Vec<Option<&str>>| -> ArrayRef { Arc::new(StringArray::from(v)) };
+        assert_find_matches_intern(
+            single(DataType::Utf8),
+            vec![
+                vec![strs(vec![Some("a"), None, Some(long), Some("a")])],
+                vec![strs(vec![Some("b"), Some(long)])],
+            ],
+            vec![strs(vec![
+                Some("zz"),
+                Some("a different long value, not inline"),
+            ])],
+        );
+        // no null stored: a null probe is missing
+        assert_find_matches_intern(
+            single(DataType::Utf8),
+            vec![vec![strs(vec![Some("a")])]],
+            vec![strs(vec![None, Some("b")])],
+        );
+
+        assert_find_matches_intern(
+            single(DataType::LargeUtf8),
+            vec![vec![Arc::new(LargeStringArray::from(vec![
+                Some("a"),
+                None,
+                Some(long),
+            ]))]],
+            vec![Arc::new(LargeStringArray::from(vec![
+                Some("q"),
+                Some("rr"),
+            ]))],
+        );
+
+        assert_find_matches_intern(
+            single(DataType::Binary),
+            vec![vec![Arc::new(BinaryArray::from_opt_vec(vec![
+                Some(b"a".as_ref()),
+                None,
+                Some(long.as_bytes()),
+            ]))]],
+            vec![Arc::new(BinaryArray::from_opt_vec(vec![Some(
+                b"q".as_ref(),
+            )]))],
+        );
+
+        assert_find_matches_intern(
+            single(DataType::Utf8View),
+            vec![vec![Arc::new(StringViewArray::from(vec![
+                Some("a"),
+                None,
+                Some(long),
+                Some("a"),
+            ]))]],
+            vec![Arc::new(StringViewArray::from(vec![
+                Some("q"),
+                Some("a different long value, not inline"),
+            ]))],
+        );
+
+        assert_find_matches_intern(
+            single(DataType::BinaryView),
+            vec![vec![Arc::new(BinaryViewArray::from(vec![
+                Some(b"a".as_ref()),
+                None,
+                Some(long.as_bytes()),
+            ]))]],
+            vec![Arc::new(BinaryViewArray::from(vec![Some(b"q".as_ref())]))],
+        );
+
+        assert_find_matches_intern(
+            single(DataType::Boolean),
+            vec![vec![Arc::new(BooleanArray::from(vec![
+                Some(true),
+                None,
+                Some(true),
+            ]))]],
+            vec![Arc::new(BooleanArray::from(vec![Some(false)]))],
+        );
+        assert_find_matches_intern(
+            single(DataType::Boolean),
+            vec![vec![Arc::new(BooleanArray::from(vec![Some(false)]))]],
+            vec![Arc::new(BooleanArray::from(vec![Some(true), None]))],
+        );
+    }
+
+    #[test]
+    fn find_matches_intern_multi_column() {
+        // (Int64, Utf8) -> GroupValuesColumn
+        let schema = Arc::new(Schema::new(vec![
+            Field::new("a", DataType::Int64, true),
+            Field::new("b", DataType::Utf8, true),
+        ]));
+        let batch: Vec<ArrayRef> = vec![
+            Arc::new(Int64Array::from(vec![Some(1), Some(1), None, Some(2)])),
+            Arc::new(StringArray::from(vec![
+                Some("x"),
+                Some("y"),
+                Some("x"),
+                None,
+            ])),
+        ];
+        let missing: Vec<ArrayRef> = vec![
+            Arc::new(Int64Array::from(vec![Some(1), Some(3)])),
+            Arc::new(StringArray::from(vec![Some("z"), Some("x")])),
+        ];
+        assert_find_matches_intern(schema, vec![batch], missing);
+
+        // (List<Int32>, Int32) -> GroupValuesRows
+        let list = |v: Vec<Option<Vec<Option<i32>>>>| -> ArrayRef {
+            Arc::new(ListArray::from_iter_primitive::<Int32Type, _, _>(v))
+        };
+        let schema = Arc::new(Schema::new(vec![
+            Field::new(
+                "a",
+                DataType::List(Arc::new(Field::new_list_field(DataType::Int32, true))),
+                true,
+            ),
+            Field::new("b", DataType::Int32, true),
+        ]));
+        let batch: Vec<ArrayRef> = vec![
+            list(vec![Some(vec![Some(1)]), None, Some(vec![Some(1)])]),
+            Arc::new(Int32Array::from(vec![Some(1), Some(1), Some(2)])),
+        ];
+        let missing: Vec<ArrayRef> = vec![
+            list(vec![Some(vec![Some(1)]), Some(vec![Some(7)])]),
+            Arc::new(Int32Array::from(vec![Some(3), Some(1)])),
+        ];
+        assert_find_matches_intern(Arc::clone(&schema), vec![batch], missing);
+
+        // find on an empty Rows group values
+        let mut gv = new_group_values(schema, &GroupOrdering::None).unwrap();
+        let probe: Vec<ArrayRef> = vec![
+            list(vec![Some(vec![Some(1)])]),
+            Arc::new(Int32Array::from(vec![Some(1)])),
+        ];
+        let mut found = vec![];
+        gv.find(&probe, &mut found).unwrap();
+        assert_eq!(found, vec![None]);
+        assert!(gv.is_empty());
+    }
+
     #[test]
     fn preserving_values_keep_group_indices_valid() {
         let schema = Arc::new(Schema::new(vec![Field::new(