Rich-T-kid commented on code in PR #11743: URL: https://github.com/apache/arrow-rs/pull/11743#discussion_r4382908164
########## arrow-select/src/coalesce/dictionary.rs: ########## @@ -0,0 +1,424 @@ +// 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 super::{InProgressArray, generic::GenericInProgressArray}; +use crate::concat::concat; +use crate::filter::{FilterPredicate, FilterSelection}; +use arrow_array::cast::AsArray; +use arrow_array::types::{ + ArrowDictionaryKeyType, BinaryType, ByteArrayType, Int8Type, Int16Type, Int32Type, Int64Type, + LargeBinaryType, LargeUtf8Type, UInt8Type, UInt16Type, UInt32Type, UInt64Type, Utf8Type, +}; +use arrow_array::{Array, ArrayRef, DictionaryArray, GenericByteArray, PrimitiveArray}; +use arrow_buffer::{NullBufferBuilder, ScalarBuffer}; +use arrow_schema::{ArrowError, DataType}; +use std::marker::PhantomData; +use std::mem::size_of; +use std::sync::{Arc, Weak}; + +// Bound values-mapping comparisons without retaining input dictionaries. +const MAX_CACHED_VALUE_ARRAYS: usize = 64; + +struct CachedValues { + /// Source identity used to skip repeated value comparisons. + source: Weak<dyn Array>, + /// Canonical values array shared by equivalent dictionaries. + canonical: Weak<dyn Array>, +} + +struct DirectState<K: ArrowDictionaryKeyType> { + /// Shared dictionary values for the buffered keys. + values: ArrayRef, + /// Keys accumulated without rebuilding dictionary values. + keys: Vec<K::Native>, + /// Validity of `keys`. + nulls: NullBufferBuilder, +} + +fn create<K: ArrowDictionaryKeyType>( + value_type: &DataType, + batch_size: usize, +) -> Box<dyn InProgressArray> { + match value_type { + DataType::Utf8 => Box::new(InProgressDictionaryArray::<K, Utf8Type>::new(batch_size)), + DataType::LargeUtf8 => Box::new(InProgressDictionaryArray::<K, LargeUtf8Type>::new( + batch_size, + )), + DataType::Binary => Box::new(InProgressDictionaryArray::<K, BinaryType>::new(batch_size)), + DataType::LargeBinary => Box::new(InProgressDictionaryArray::<K, LargeBinaryType>::new( + batch_size, + )), + _ => Box::new(GenericInProgressArray::new()), + } +} + +pub(crate) fn create_in_progress_dictionary( + key_type: &DataType, + value_type: &DataType, + batch_size: usize, +) -> Box<dyn InProgressArray> { + match key_type { + DataType::Int8 => create::<Int8Type>(value_type, batch_size), + DataType::Int16 => create::<Int16Type>(value_type, batch_size), + DataType::Int32 => create::<Int32Type>(value_type, batch_size), + DataType::Int64 => create::<Int64Type>(value_type, batch_size), + DataType::UInt8 => create::<UInt8Type>(value_type, batch_size), + DataType::UInt16 => create::<UInt16Type>(value_type, batch_size), + DataType::UInt32 => create::<UInt32Type>(value_type, batch_size), + DataType::UInt64 => create::<UInt64Type>(value_type, batch_size), + _ => unreachable!("illegal dictionary key type {key_type}"), + } +} + +/// Coalesces dictionaries with byte-array values without rebuilding shared mappings. +pub(crate) struct InProgressDictionaryArray<K: ArrowDictionaryKeyType, T: ByteArrayType> { + /// Current source used by [`Self::copy_rows`]. + source: Option<ArrayRef>, + /// Target output row count. + batch_size: usize, + /// Keys buffered while inputs share one dictionary values array. + direct: Option<DirectState<K>>, + /// Array slices retained after the dictionary values mapping changes. + buffered_arrays: Vec<ArrayRef>, + /// Upper bound for the bytes retained by `buffered_arrays`. + buffered_array_bytes: usize, + /// Canonical values arrays for recently seen input dictionaries. + value_cache: Vec<CachedValues>, + /// Retains the byte-array value type in the state type. + value_type: PhantomData<T>, +} + +impl<K: ArrowDictionaryKeyType, T: ByteArrayType> std::fmt::Debug + for InProgressDictionaryArray<K, T> +{ + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("InProgressDictionaryArray").finish() + } +} + +impl<K: ArrowDictionaryKeyType, T: ByteArrayType> InProgressDictionaryArray<K, T> { + pub(crate) fn new(batch_size: usize) -> Self { + Self { + source: None, + batch_size, + direct: None, + buffered_arrays: Vec::new(), + buffered_array_bytes: 0, + value_cache: Vec::new(), + value_type: PhantomData, + } + } + + fn canonical_values(&mut self, values: ArrayRef) -> ArrayRef { + if let Some(canonical) = self.value_cache.iter().find_map(|cached| { + let source = cached.source.upgrade()?; + Arc::ptr_eq(&source, &values) + .then(|| cached.canonical.upgrade()) + .flatten() + }) { + return canonical; + } + + let canonical = self + .value_cache + .iter() + .filter_map(|cached| cached.canonical.upgrade()) + .find(|cached| cached.as_ref() == values.as_ref()) + .map_or_else(|| Arc::clone(&values), |canonical| Arc::clone(&canonical)); + + if self.value_cache.len() < MAX_CACHED_VALUE_ARRAYS { + self.value_cache.push(CachedValues { + source: Arc::downgrade(&values), + canonical: Arc::downgrade(&canonical), + }); + } + canonical + } + + /// Returns the direct state while `values` matches the buffered mapping. + /// A mapping change switches the current output batch to buffered `concat`. + fn direct_for_values( + &mut self, + values: ArrayRef, + ) -> Result<Option<&mut DirectState<K>>, ArrowError> { + if !self.buffered_arrays.is_empty() { + return Ok(None); + } + if self.value_cache.len() >= MAX_CACHED_VALUE_ARRAYS + && self + .direct + .as_ref() + .is_none_or(|direct| direct.keys.is_empty()) + { + self.direct = None; + self.value_cache.clear(); + } + let values = self.canonical_values(values); + if let Some(direct) = self.direct.as_ref() + && !direct.keys.is_empty() + && !Arc::ptr_eq(&direct.values, &values) + { + let array = Arc::new(self.take_direct()?.unwrap()) as ArrayRef; + self.push_buffered(array); + return Ok(None); + } + if !self + .direct + .as_ref() + .is_some_and(|direct| Arc::ptr_eq(&direct.values, &values)) + { + self.direct = Some(DirectState { + values, + keys: Vec::with_capacity(self.batch_size), + nulls: NullBufferBuilder::new(self.batch_size), + }); + } + Ok(self.direct.as_mut()) + } + + fn take_direct(&mut self) -> Result<Option<DictionaryArray<K>>, ArrowError> { + let Some(direct) = self.direct.as_mut() else { + return Ok(None); + }; + if direct.keys.is_empty() { + return Ok(None); + } + let keys = std::mem::take(&mut direct.keys); + let nulls = direct.nulls.finish(); + direct.nulls = NullBufferBuilder::new(self.batch_size); + let keys = PrimitiveArray::<K>::try_new(ScalarBuffer::from(keys), nulls)?; + Ok(Some(DictionaryArray::try_new( + keys, + Arc::clone(&direct.values), + )?)) + } + + fn push_buffered(&mut self, array: ArrayRef) { + self.buffered_array_bytes += array.get_array_memory_size(); + self.buffered_arrays.push(array); + } +} + +impl<K: ArrowDictionaryKeyType, T: ByteArrayType> InProgressArray + for InProgressDictionaryArray<K, T> +{ + fn set_source(&mut self, source: Option<ArrayRef>) { + self.source = source; + } + + fn copy_rows(&mut self, offset: usize, len: usize) -> Result<(), ArrowError> { + let source = Arc::clone(self.source.as_ref().ok_or_else(|| { + ArrowError::InvalidArgumentError( + "Internal Error: InProgressDictionaryArray: source not set".to_string(), + ) + })?); + let dictionary = source.as_dictionary::<K>(); + let values = Arc::clone(dictionary.values()); + let keys = dictionary.keys(); + + let Some(direct) = self.direct_for_values(values)? else { + self.push_buffered(source.slice(offset, len)); + return Ok(()); + }; + if let Some(nulls) = keys.nulls().map(|nulls| nulls.slice(offset, len)) { + direct.nulls.append_buffer(&nulls); + } else { + direct.nulls.append_n_non_nulls(len); + } + direct + .keys + .extend_from_slice(&keys.values()[offset..offset + len]); + Ok(()) + } + + fn copy_rows_by_filter_from( + &mut self, + source: ArrayRef, + filter: &FilterPredicate, + ) -> Result<(), ArrowError> { + let dictionary = source.as_dictionary::<K>(); + dictionary + .downcast_dict::<GenericByteArray<T>>() + .expect("dictionary value type matched at construction"); + let values = Arc::clone(dictionary.values()); + let keys = dictionary.keys(); + if filter.count() == 0 { + return Ok(()); + } + let Some(direct) = self.direct_for_values(values)? else { + self.push_buffered(filter.filter(source.as_ref())?); + return Ok(()); + }; + if let Some(nulls) = filter.filter_nulls(keys.nulls()) { + direct.nulls.append_buffer(&nulls); + } else { + direct.nulls.append_n_non_nulls(filter.count()); + } + direct.keys.reserve(filter.count()); + + match filter.selection() { + FilterSelection::None => {} + FilterSelection::All { len } => { + direct.keys.extend_from_slice(&keys.values()[..len]); + } + FilterSelection::Slices(slices) => { + slices.for_each(|(start, end)| { + direct.keys.extend_from_slice(&keys.values()[start..end]); + }); + } + FilterSelection::Indices(indices) => { + indices.for_each(|index| direct.keys.push(keys.values()[index])); + } + } + Ok(()) + } + + fn finish(&mut self) -> Result<ArrayRef, ArrowError> { + if self.buffered_arrays.is_empty() { + return match self.take_direct()? { + Some(array) => Ok(Arc::new(array)), + None => Ok(Arc::new(DictionaryArray::<K>::try_new( + PrimitiveArray::new_null(0), + Arc::new(GenericByteArray::<T>::new_null(0)), + )?)), + }; + } + if let Some(array) = self.take_direct()? { + self.push_buffered(Arc::new(array)); + } + let arrays = self + .buffered_arrays + .iter() + .map(|array| array.as_ref()) + .collect::<Vec<_>>(); + let output = concat(&arrays)?; Review Comment: a future implementation should try to avoid this concat but I'd rather land this first as a baseline and build on top -- 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]
