Rich-T-kid commented on code in PR #11743: URL: https://github.com/apache/arrow-rs/pull/11743#discussion_r4450519239
########## arrow-select/src/coalesce/dictionary.rs: ########## @@ -0,0 +1,389 @@ +// 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::dictionary::merge_dictionary_values_with_masks; +use ahash::RandomState; +use arrow_array::cast::AsArray; +use arrow_array::types::{ + ArrowDictionaryKeyType, Int8Type, Int16Type, Int32Type, Int64Type, UInt8Type, UInt16Type, + UInt32Type, UInt64Type, +}; +use arrow_array::{Array, ArrayRef, DictionaryArray, PrimitiveArray, new_empty_array}; +use arrow_buffer::{ArrowNativeType, BooleanBufferBuilder, NullBuffer, ScalarBuffer}; +use arrow_schema::{ArrowError, DataType}; +use std::collections::HashMap; +use std::mem::size_of; +use std::sync::Arc; + +struct KeyChunk<K: ArrowDictionaryKeyType> { + /// Zero-copy slice of source-local keys for the retained rows. + keys: ScalarBuffer<K::Native>, + /// Shared validity buffer from the source keys, if present. + nulls: Option<NullBuffer>, + /// Position in `nulls` corresponding to `keys[0]`. + null_offset: usize, + /// Index of the values array in `InProgressDictionaryArray::sources`. + source: usize, +} + +fn create<K: ArrowDictionaryKeyType>(value_type: &DataType) -> Box<dyn InProgressArray> { + match value_type { + DataType::Utf8 | DataType::LargeUtf8 | DataType::Binary | DataType::LargeBinary => { + Box::new(InProgressDictionaryArray::<K>::new(value_type)) + } + _ => 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), + DataType::Int16 => create::<Int16Type>(value_type), + DataType::Int32 => create::<Int32Type>(value_type), + DataType::Int64 => create::<Int64Type>(value_type), + DataType::UInt8 => create::<UInt8Type>(value_type), + DataType::UInt16 => create::<UInt16Type>(value_type), + DataType::UInt32 => create::<UInt32Type>(value_type), + DataType::UInt64 => create::<UInt64Type>(value_type), + _ => unreachable!("illegal dictionary key type {key_type}"), + } +} + +struct InProgressDictionaryArray<K: ArrowDictionaryKeyType> { + /// Dictionary currently selected by [`InProgressArray::set_source`]. + source: Option<ArrayRef>, + /// Number of output rows retained across `chunks`. + len: usize, + /// Typed empty values array used when finishing without any rows. + empty_values: ArrayRef, + /// Distinct dictionary values arrays retained by `Arc` identity. + sources: Vec<ArrayRef>, + /// Maps each retained values `Arc` address to its index in `sources`. + source_ids: HashMap<usize, usize, RandomState>, + /// Zero-copy key slices in output order, tagged with their values source. + chunks: Vec<KeyChunk<K>>, +} + +impl<K: ArrowDictionaryKeyType> std::fmt::Debug for InProgressDictionaryArray<K> { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("InProgressDictionaryArray").finish() + } +} + +impl<K: ArrowDictionaryKeyType> InProgressDictionaryArray<K> { + fn new(value_type: &DataType) -> Self { + Self { + source: None, + len: 0, + empty_values: new_empty_array(value_type), + sources: Vec::new(), + source_ids: HashMap::with_hasher(RandomState::with_seeds(0, 0, 0, 0)), + chunks: Vec::new(), + } + } + + /// Returns a stable index for `values`, reusing it when the same `Arc` appears again. + fn source_id(&mut self, values: ArrayRef) -> usize { + if let Some(last) = self.chunks.last() + && Arc::ptr_eq(&self.sources[last.source], &values) + { + return last.source; + } + let values_id = Arc::as_ptr(&values).cast::<()>() as usize; + if let Some(&source) = self.source_ids.get(&values_id) { + return source; + } + let source = self.sources.len(); + self.sources.push(values); + self.source_ids.insert(values_id, source); + source + } + + /// Compacts referenced values and remaps source-local keys while preserving chunk order. + /// Equal values are only deduplicated when the combined values would overflow `K`. + fn merge_sources( + sources: &[ArrayRef], + chunks: &[KeyChunk<K>], + len: usize, + ) -> Result<(Vec<K::Native>, ArrayRef), ArrowError> { + let mut referenced = sources + .iter() + .map(|source| { + let mut mask = BooleanBufferBuilder::new(source.len()); + mask.advance(source.len()); + mask + }) + .collect::<Vec<_>>(); + + for chunk in chunks { + let mask = &mut referenced[chunk.source]; + match &chunk.nulls { + Some(nulls) => nulls + .inner() + .slice(chunk.null_offset, chunk.keys.len()) + .set_indices() + .for_each(|offset| mask.set_bit(chunk.keys[offset].as_usize(), true)), + None => chunk + .keys + .iter() + .for_each(|key| mask.set_bit(key.as_usize(), true)), + } + } + + let referenced = referenced + .into_iter() + .map(BooleanBufferBuilder::build) + .collect::<Vec<_>>(); + let source_refs = sources + .iter() + .map(|source| source.as_ref() as &dyn Array) + .collect::<Vec<_>>(); + let merged = merge_dictionary_values_with_masks::<K>(&source_refs, &referenced)?; + + let mut keys = Vec::with_capacity(len); + for chunk in chunks { + let mapping = &merged.key_mappings[chunk.source]; + match &chunk.nulls { + Some(nulls) => keys.extend(chunk.keys.iter().enumerate().map(|(offset, key)| { + if nulls.is_valid(chunk.null_offset + offset) { + mapping[key.as_usize()] + } else { + K::Native::default() + } + })), + None => keys.extend(chunk.keys.iter().map(|key| mapping[key.as_usize()])), + } + } + Ok((keys, merged.values)) + } + + /// Concatenates chunk validity in output order, omitting it when all keys are valid. + fn build_nulls(chunks: &[KeyChunk<K>], len: usize) -> Option<NullBuffer> { + chunks.iter().any(|chunk| chunk.nulls.is_some()).then(|| { + let mut nulls = BooleanBufferBuilder::new(len); + for chunk in chunks { + match &chunk.nulls { + Some(chunk_nulls) => nulls.append_buffer( + &chunk_nulls + .inner() + .slice(chunk.null_offset, chunk.keys.len()), + ), + None => nulls.append_n(chunk.keys.len(), true), + } + } + NullBuffer::new(nulls.finish()) + }) + } +} + +impl<K: ArrowDictionaryKeyType> InProgressArray for InProgressDictionaryArray<K> { Review Comment: note that there no longer any specialized trait functions implemented like `copy_rows_by_filter_from`, simply because id like to handle it and `copy_rows_by_indices` in one PR after https://github.com/apache/arrow-rs/pull/11174 lands. -- 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]
