jayshrivastava commented on code in PR #23187: URL: https://github.com/apache/datafusion/pull/23187#discussion_r3547391162
########## datafusion/physical-plan/src/aggregates/group_values/multi_group_by/dictionary.rs: ########## @@ -0,0 +1,341 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use crate::aggregates::group_values::multi_group_by::GroupColumn; +use arrow::array::{ + Array, ArrayRef, AsArray, BooleanBufferBuilder, DictionaryArray, PrimitiveArray, +}; +use arrow::compute::concat; +use arrow::datatypes::{ArrowDictionaryKeyType, ArrowNativeType, DataType, Field}; +use datafusion_common::Result; +use std::marker::PhantomData; +use std::sync::Arc; + +pub struct DictionaryGroupValuesColumn<K: ArrowDictionaryKeyType + Send + Sync> { + inner: Box<dyn GroupColumn>, + null_array: ArrayRef, + _phantom: PhantomData<K>, +} + +impl<K: ArrowDictionaryKeyType + Send + Sync> DictionaryGroupValuesColumn<K> { + pub fn new(inner: Box<dyn GroupColumn>, field: &Field) -> Self { + let null_array = arrow::array::new_null_array(field.data_type(), 1); + Self { + inner, + null_array, + _phantom: PhantomData, + } + } + + fn into_dict(values: ArrayRef) -> ArrayRef { + let num_values = values.len(); + assert!( + Self::valid_bounds::<K>(num_values), + "Dictionary key type {:?} cannot hold {} values", + K::DATA_TYPE, + num_values + ); + let keys: PrimitiveArray<K> = (0..num_values) + .map(|i| (!values.is_null(i)).then(|| K::Native::usize_as(i))) + .collect(); + Arc::new(DictionaryArray::<K>::new(keys, values)) + } + + fn valid_bounds<T: ArrowDictionaryKeyType>(num_values: usize) -> bool { + let max: usize = match T::DATA_TYPE { + DataType::Int8 => i8::MAX as usize, + DataType::Int16 => i16::MAX as usize, + DataType::Int32 => i32::MAX as usize, + DataType::Int64 => i64::MAX as usize, + DataType::UInt8 => u8::MAX as usize, + DataType::UInt16 => u16::MAX as usize, + DataType::UInt32 => u32::MAX as usize, + DataType::UInt64 => usize::MAX, + _ => return false, + }; + num_values == 0 || num_values - 1 <= max + } +} + +impl<K: ArrowDictionaryKeyType + Send + Sync> GroupColumn + for DictionaryGroupValuesColumn<K> +{ + fn equal_to(&self, lhs_row: usize, array: &ArrayRef, rhs_row: usize) -> bool { + let dict = array.as_dictionary::<K>(); + match dict.key(rhs_row) { + None => self.inner.equal_to(lhs_row, &self.null_array, 0), + Some(key) => self.inner.equal_to(lhs_row, dict.values(), key), + } + } + + fn append_val(&mut self, array: &ArrayRef, row: usize) -> Result<()> { + let dict = array.as_dictionary::<K>(); + match dict.key(row) { + None => self.inner.append_val(&self.null_array, 0), + Some(key) => self.inner.append_val(dict.values(), key), + } + } + + fn vectorized_equal_to( + &self, + lhs_rows: &[usize], + array: &ArrayRef, + rhs_rows: &[usize], + equal_to_results: &mut BooleanBufferBuilder, + ) { + let dict = array.as_dictionary::<K>(); + let dict_values = dict.values(); + let dict_keys = dict.keys(); + + if dict_keys.null_count() == 0 { + // Fast path: no nulls in the key array, resolve all indices up front and + // delegate to the inner column's vectorized comparison. + let value_indices: Vec<usize> = rhs_rows + .iter() + .map(|&row_index| dict_keys.value(row_index).as_usize()) + .collect(); + self.inner.vectorized_equal_to( + lhs_rows, + dict_values, + &value_indices, + equal_to_results, + ); + } else { + // Null keys must be routed to a null sentinel. + let combined = concat(&[dict_values.as_ref(), self.null_array.as_ref()]) + .expect("concat of dict values and null sentinel should not fail"); + let null_sentinel_index = combined.len() - 1; + let value_indices: Vec<usize> = rhs_rows + .iter() + .map(|&row_index| dict.key(row_index).unwrap_or(null_sentinel_index)) + .collect(); + self.inner.vectorized_equal_to( + lhs_rows, + &combined, + &value_indices, + equal_to_results, + ); + } + } + + fn vectorized_append(&mut self, array: &ArrayRef, rows: &[usize]) -> Result<()> { + let dict = array.as_dictionary::<K>(); + let dict_keys = dict.keys(); + + if dict_keys.null_count() == 0 { + let value_indices: Vec<usize> = rows + .iter() + .map(|&row_index| dict_keys.value(row_index).as_usize()) + .collect(); + self.inner.vectorized_append(dict.values(), &value_indices) + } else { + // `get_combined` appends the null sentinel as the final element, so any + // null dictionary key maps to that last index. + let combined_with_null_sentinel = + concat(&[dict.values(), self.null_array.as_ref()]) Review Comment: The concat isn't great tbh bc it reallocates the whole array, but I don't see how we can avoid it. Since the benchmarks are good, it's probably fine. It's unfortunate that keys might be null. Ideally the dict arrives "normalized" where all keys are valid and may point to a null value. How likely is it that there's null keys? Would it help to see how often `.unwrap_or(null_sentinel_index)` happens? If it happens 0 times, we can avoid `concat`? -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
