Rich-T-kid commented on code in PR #11743: URL: https://github.com/apache/arrow-rs/pull/11743#discussion_r4403488583
########## arrow-select/src/coalesce/dictionary.rs: ########## @@ -0,0 +1,425 @@ +// 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; + +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, + /// 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_type: PhantomData, + } + } + + /// Returns the direct state while `values` is the same allocation as the + /// buffered mapping. Keys cannot be reused across mappings, so a different + /// `Arc` 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 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(mut direct) = self.direct.take() else { + return Ok(None); + }; + if direct.keys.is_empty() { Review Comment: yeah (https://github.com/apache/arrow-rs/pull/11743/changes/396ab234ef5732669d871423849ad32e9cd83eab) -- 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]
