Rich-T-kid commented on code in PR #23187: URL: https://github.com/apache/datafusion/pull/23187#discussion_r3553086238
########## 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. I agree, for a 8192 size batch of keys with just a single null value it will cause the entire 8192 batch to be copied to a new array just to append 1 null at the end. In`vectorized_append()` we can avoid copying value arrays blindly by doing a `Arc::ptr_eq()` check. this is a cheap comparison that can save an `O(target_batch_size) + 1` copy for every call. as for `vectorized_equal_to()` theres not much we can do here since it takes a `&self` meaning we wouldn't be able to re-cache the values array. > 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 what I tried avoiding by adding the null branch case. Fundemantally we have to account for the possibility of keys being null. [arrows format doc:](https://github.com/apache/arrow/blob/main/docs/source/format/Columnar.rst) > The memory layout for a dictionary-encoded array is the same as that of a primitive integer layout. the primitive array section states: > A primitive value array represents an array of values each having the same physical slot width typically measured in bytes, though the spec also provides for bit-packed types (e.g. boolean values encoded in bits). > Internally, the array contains a contiguous memory buffer whose total size is at least as large as the slot width multiplied by the array length. For bit-packed types, the size is rounded up to the nearest byte. > The associated validity bitmap is contiguously allocated (as described above) but does not need to be adjacent in memory to the values buffer. > Example Layout: Int32 Array > For example a primitive array of int32s: > [1, null, 2, 4, 8] Unlike a primitive, where nulls live in a single validity bitmap, a dictionary array's logical nulls can come from two independent places: the keys array's validity bitmap, or a null entry in the values (dictionary) array that a key points to. The field's nullable flag doesn't tell us which it's only an advisory "may contain nulls," same as for any type. So to know the actual null situation we have to inspect both the keys and the values arrays. -- 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]
