diff --git a/datafusion/physical-expr-common/src/binary_map.rs b/datafusion/physical-expr-common/src/binary_map.rs index 0bd7f97fb76b3..bfe688afae2aa 100644 --- a/datafusion/physical-expr-common/src/binary_map.rs +++ b/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(&mut self, values: &ArrayRef, observe_payload_fn: OP) + where + OP: FnMut(Option), + { + match self.output_type { + OutputType::Binary => { + assert!(matches!( + values.data_type(), + DataType::Binary | DataType::LargeBinary + )); + self.find_inner::>(values, observe_payload_fn) + } + OutputType::Utf8 => { + assert!(matches!( + values.data_type(), + DataType::Utf8 | DataType::LargeUtf8 + )); + self.find_inner::>(values, observe_payload_fn) + } + _ => unreachable!("View types should use `ArrowBytesViewMap`"), + }; + } + + /// Generic version of [`Self::find`] that handles `ByteArrayType` + fn find_inner(&mut self, values: &ArrayRef, mut observe_payload_fn: OP) + where + OP: FnMut(Option), + 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::(); + 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. diff --git a/datafusion/physical-expr-common/src/binary_view_map.rs b/datafusion/physical-expr-common/src/binary_view_map.rs index 2467b6e7bda09..5b224636303b4 100644 --- a/datafusion/physical-expr-common/src/binary_view_map.rs +++ b/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(&mut self, values: &ArrayRef, observe_payload_fn: OP) + where + OP: FnMut(Option), + { + match self.output_type { + OutputType::BinaryView => { + assert!(matches!(values.data_type(), DataType::BinaryView)); + self.find_inner::(values, observe_payload_fn) + } + OutputType::Utf8View => { + assert!(matches!(values.data_type(), DataType::Utf8View)); + self.find_inner::(values, observe_payload_fn) + } + _ => unreachable!("Utf8/Binary should use `ArrowBytesSet`"), + }; + } + + /// Generic version of [`Self::find`] that handles `ByteViewType` + fn find_inner(&mut self, values: &ArrayRef, mut observe_payload_fn: OP) + where + OP: FnMut(Option), + 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::(); + 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]; diff --git a/datafusion/physical-plan/src/aggregates/group_values/mod.rs b/datafusion/physical-plan/src/aggregates/group_values/mod.rs index 016223170bd26..989ed7dc902f3 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/mod.rs +++ b/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) -> 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>) -> 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, + batches: Vec>, + missing: Vec, + ) { + let mut gv: Box = + 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::>()); + } + 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 { + 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>| -> 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 = 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 = 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) -> GroupValuesRows + let list = |v: Vec>>>| -> ArrayRef { + Arc::new(ListArray::from_iter_primitive::(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 = 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 = 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 = 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( diff --git a/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/mod.rs b/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/mod.rs index 554ab4971abde..af069798e6017 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/mod.rs +++ b/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/mod.rs @@ -1164,6 +1164,95 @@ impl GroupValues for GroupValuesColumn { } } + fn find(&mut self, cols: &[ArrayRef], groups: &mut Vec>) -> Result<()> { + let n_rows = cols[0].len(); + let mut batch_hashes = mem::take(&mut self.hashes_buffer); + batch_hashes.clear(); + batch_hashes.resize(n_rows, 0); + create_hashes(cols, &self.random_state, &mut batch_hashes)?; + + groups.clear(); + groups.resize(n_rows, None); + + // Streaming mode interns with `scalarized_intern`, which stores one + // inlined map entry per group even when hashes collide. Compare values + // while probing so a collision doesn't hide the matching entry. + if STREAMING { + for (row, &target_hash) in batch_hashes.iter().enumerate() { + groups[row] = self + .map + .find(target_hash, |(exist_hash, group_idx_view)| { + target_hash == *exist_hash + && self.group_values.iter().enumerate().all( + |(i, group_val)| { + group_val.equal_to( + group_idx_view.value() as usize, + &cols[i], + row, + ) + }, + ) + }) + .map(|(_, group_idx_view)| group_idx_view.value() as usize); + } + self.hashes_buffer = batch_hashes; + return Ok(()); + } + + // Vectorized mode keeps one map entry per hash, with colliding groups + // in a list. Collect every (row, candidate group) pair whose hash + // matches, then compare them column by column in one vectorized pass. + let buffers = &mut self.vectorized_operation_buffers; + let mut rows = mem::take(&mut buffers.equal_to_row_indices); + let mut candidates = mem::take(&mut buffers.equal_to_group_indices); + rows.clear(); + candidates.clear(); + for (row, &target_hash) in batch_hashes.iter().enumerate() { + let Some((_, view)) = self + .map + .find(target_hash, |(exist_hash, _)| target_hash == *exist_hash) + else { + continue; + }; + if view.is_non_inlined() { + let list = &self.group_index_lists[view.value() as usize]; + candidates.extend_from_slice(list); + rows.extend(std::iter::repeat_n(row, list.len())); + } else { + rows.push(row); + candidates.push(view.value() as usize); + } + } + + if !rows.is_empty() { + let mut results = mem::replace( + &mut self.vectorized_operation_buffers.equal_to_results, + BooleanBufferBuilder::new(0), + ); + results.truncate(0); + results.append_n(rows.len(), true); + for (col_idx, group_col) in self.group_values.iter().enumerate() { + group_col.vectorized_equal_to( + &candidates, + &cols[col_idx], + &rows, + &mut results, + ); + } + for (idx, (&row, &group)) in rows.iter().zip(&candidates).enumerate() { + if results.get_bit(idx) { + groups[row] = Some(group); + } + } + self.vectorized_operation_buffers.equal_to_results = results; + } + + self.vectorized_operation_buffers.equal_to_row_indices = rows; + self.vectorized_operation_buffers.equal_to_group_indices = candidates; + self.hashes_buffer = batch_hashes; + Ok(()) + } + fn size(&self) -> usize { let group_values_size: usize = self.group_values.iter().map(|v| v.size()).sum(); group_values_size + self.map_size + self.hashes_buffer.allocated_size() @@ -2054,6 +2143,54 @@ mod tests { assert_eq!(group_values.len(), data_set.expected_batch.num_rows()); } + #[test] + fn test_find_matches_intern_without_inserting() { + fn check(data_set: &VectorizedTestDataSet) { + let mut group_values = + GroupValuesColumn::::try_new(data_set.schema()).unwrap(); + let mut interned = vec![]; + for batch in &data_set.test_batches { + group_values.intern(batch, &mut interned).unwrap(); + } + let num_groups = group_values.len(); + + // Every stored row is found at the id `intern` assigned. + let mut found = vec![]; + for batch in &data_set.test_batches { + group_values.intern(batch, &mut interned).unwrap(); + group_values.find(batch, &mut found).unwrap(); + let expected = interned.iter().copied().map(Some).collect::>(); + assert_eq!(found, expected); + } + + // Missing rows are `None` and are not inserted. + let missing: Vec = data_set + .schema() + .fields() + .iter() + .map(|field| match field.data_type() { + DataType::Int64 => { + Arc::new(Int64Array::from(vec![Some(i64::MAX)])) as ArrayRef + } + DataType::Utf8 => { + Arc::new(StringArray::from(vec![Some("missing")])) as ArrayRef + } + DataType::Utf8View => { + Arc::new(StringViewArray::from(vec![Some("missing")])) as ArrayRef + } + other => panic!("unexpected test type {other}"), + }) + .collect(); + group_values.find(&missing, &mut found).unwrap(); + assert_eq!(found, vec![None]); + assert_eq!(group_values.len(), num_groups); + } + + let data_set = VectorizedTestDataSet::new(); + check::(&data_set); + check::(&data_set); + } + #[test] fn test_emit_first_n_for_vectorized_group_values() { let data_set = VectorizedTestDataSet::new(); diff --git a/datafusion/physical-plan/src/aggregates/group_values/row.rs b/datafusion/physical-plan/src/aggregates/group_values/row.rs index 01e9f3eaa71ee..35584d1d1fd4d 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/row.rs +++ b/datafusion/physical-plan/src/aggregates/group_values/row.rs @@ -184,6 +184,38 @@ impl GroupValues for GroupValuesRows { Ok(()) } + fn find(&mut self, cols: &[ArrayRef], groups: &mut Vec>) -> Result<()> { + // Same normalization as `intern`, so lookups agree with stored keys. + let normalized_cols: Vec = + cols.iter().map(normalize_float_zero).collect(); + let cols = normalized_cols.as_slice(); + + let group_rows = &mut self.rows_buffer; + group_rows.clear(); + self.row_converter.append(group_rows, cols)?; + let n_rows = group_rows.num_rows(); + + let batch_hashes = &mut self.hashes_buffer; + batch_hashes.clear(); + batch_hashes.resize(n_rows, 0); + create_hashes(cols, &self.random_state, batch_hashes)?; + + groups.clear(); + let Some(group_values) = self.group_values.as_ref() else { + groups.resize(n_rows, None); + return Ok(()); + }; + + for (row, &target_hash) in batch_hashes.iter().enumerate() { + let entry = self.map.find(target_hash, |(exist_hash, group_idx)| { + target_hash == *exist_hash + && group_rows.row(row) == group_values.row(*group_idx) + }); + groups.push(entry.map(|&(_, group_idx)| group_idx)); + } + Ok(()) + } + fn size(&self) -> usize { let group_values_size = self.group_values.as_ref().map(|v| v.size()).unwrap_or(0); self.row_converter.size() diff --git a/datafusion/physical-plan/src/aggregates/group_values/single_group_by/boolean.rs b/datafusion/physical-plan/src/aggregates/group_values/single_group_by/boolean.rs index 76fb0529488f0..6156ecd558c35 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/single_group_by/boolean.rs +++ b/datafusion/physical-plan/src/aggregates/group_values/single_group_by/boolean.rs @@ -83,6 +83,17 @@ impl GroupValues for GroupValuesBoolean { Ok(()) } + fn find(&mut self, cols: &[ArrayRef], groups: &mut Vec>) -> Result<()> { + let array = cols[0].as_boolean(); + groups.clear(); + groups.extend(array.iter().map(|value| match value { + Some(false) => self.false_group, + Some(true) => self.true_group, + None => self.null_group, + })); + Ok(()) + } + fn size(&self) -> usize { size_of::() } diff --git a/datafusion/physical-plan/src/aggregates/group_values/single_group_by/bytes.rs b/datafusion/physical-plan/src/aggregates/group_values/single_group_by/bytes.rs index 34ec36be31d2e..71d951fc6b7ff 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/single_group_by/bytes.rs +++ b/datafusion/physical-plan/src/aggregates/group_values/single_group_by/bytes.rs @@ -72,6 +72,19 @@ impl GroupValues for GroupValuesBytes { Ok(()) } + fn find(&mut self, cols: &[ArrayRef], groups: &mut Vec>) -> Result<()> { + assert_eq!(cols.len(), 1); + let arr = &cols[0]; + + groups.clear(); + self.map.find(arr, |group_idx| { + groups.push(group_idx); + }); + + assert_eq!(groups.len(), arr.len()); + Ok(()) + } + fn size(&self) -> usize { self.map.size() + size_of::() } 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 997a7ce166a71..66faa4674242e 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 @@ -74,6 +74,23 @@ impl GroupValues for GroupValuesBytesView { Ok(()) } + fn find( + &mut self, + cols: &[ArrayRef], + groups: &mut Vec>, + ) -> datafusion_common::Result<()> { + assert_eq!(cols.len(), 1); + let arr = &cols[0]; + + groups.clear(); + self.map.find(arr, |group_idx| { + groups.push(group_idx); + }); + + assert_eq!(groups.len(), arr.len()); + Ok(()) + } + fn size(&self) -> usize { self.map.size() + size_of::() } diff --git a/datafusion/physical-plan/src/aggregates/group_values/single_group_by/primitive.rs b/datafusion/physical-plan/src/aggregates/group_values/single_group_by/primitive.rs index 21b62457e3831..64915335bd13c 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/single_group_by/primitive.rs +++ b/datafusion/physical-plan/src/aggregates/group_values/single_group_by/primitive.rs @@ -177,6 +177,28 @@ where Ok(()) } + fn find(&mut self, cols: &[ArrayRef], groups: &mut Vec>) -> Result<()> { + assert_eq!(cols.len(), 1); + groups.clear(); + + for v in cols[0].as_primitive::() { + let group_id = match v { + None => self.null_group, + Some(key) => { + let key = key.canonicalize(); + let hash = key.hash(&self.random_state); + self.map + .find(hash, |&(g, h)| unsafe { + hash == h && self.values.get_unchecked(g).is_eq(key) + }) + .map(|&(g, _)| g) + } + }; + groups.push(group_id) + } + Ok(()) + } + fn size(&self) -> usize { self.map.capacity() * size_of::<(usize, u64)>() + self.values.allocated_size() } @@ -322,4 +344,26 @@ mod tests { Ok(()) } + + #[test] + fn test_find_does_not_insert() -> Result<()> { + let mut group_values = GroupValuesPrimitive::::new(DataType::Int32); + let stored: ArrayRef = Arc::new(Int32Array::from(vec![Some(7), None, Some(3)])); + let mut groups = vec![]; + group_values.intern(&[stored], &mut groups)?; + + let probe: ArrayRef = + Arc::new(Int32Array::from(vec![Some(3), Some(5), None, Some(7)])); + let mut found = vec![]; + group_values.find(&[probe], &mut found)?; + assert_eq!(found, vec![Some(2), None, Some(1), Some(0)]); + assert_eq!(group_values.len(), 3); + + let mut no_nulls = GroupValuesPrimitive::::new(DataType::Int32); + let probe: ArrayRef = Arc::new(Int32Array::from(vec![None])); + no_nulls.find(&[probe], &mut found)?; + assert_eq!(found, vec![None]); + assert!(no_nulls.is_empty()); + Ok(()) + } }