Jefffrey commented on code in PR #11743: URL: https://github.com/apache/arrow-rs/pull/11743#discussion_r4468384074
########## 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, Review Comment: could we just slice `nulls` as well instead of needing this separate field to track the offset? i feel itll also be less confusing that way ########## 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); Review Comment: any reason we dont use `NullBufferBuilder` here instead? ########## 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}"), + } +} Review Comment: ```suggestion pub(crate) fn create_in_progress_dictionary( key_type: &DataType, value_type: &DataType, ) -> Box<dyn InProgressArray> { macro_rules! create { ($key:ty, $value_type:ident) => { Box::new(InProgressDictionaryArray::<$key>::new($value_type)) }; } match value_type { DataType::Utf8 | DataType::LargeUtf8 | DataType::Binary | DataType::LargeBinary => { downcast_integer! { key_type => (create, value_type), _ => unreachable!("illegal dictionary key type {key_type}"), } } _ => Box::new(GenericInProgressArray::new()), } } ``` - utilize `downcast_integer!{}` for boilerplate - remove unused `_batch_size` ########## arrow-select/src/dictionary.rs: ########## @@ -291,6 +291,85 @@ pub(crate) fn merge_dictionary_values<K: ArrowDictionaryKeyType>( }) } +/// Merges referenced entries from multiple dictionary values arrays. +/// +/// `values_masks[i]` corresponds to `values_arrays[i]`, with each set bit identifying a +/// value to retain. Returns the merged values and a mapping from each source value index +/// to its index in the merged array. Mapping entries for unreferenced values are zero and +/// must not be used. +/// +/// Values are gathered without hashing when they fit in `K`, preserving duplicates across +/// sources. Otherwise, values are deduplicated on a best-effort basis to avoid key overflow. +pub(crate) fn merge_dictionary_values_with_masks<K: ArrowDictionaryKeyType>( + values_arrays: &[&dyn Array], + values_masks: &[BooleanBuffer], +) -> Result<MergedDictionaries<K>, ArrowError> { + assert_eq!(values_arrays.len(), values_masks.len()); + let num_values = values_masks.iter().map(BooleanBuffer::count_set_bits).sum(); + + // Dictionary values need not be unique. Avoid hashing when the referenced + // values already fit in the key type, and only deduplicate to avoid overflow. + if num_values == 0 || K::Native::from_usize(num_values - 1).is_some() { + let zero = K::Native::from_usize(0).unwrap(); + let mut indices = Vec::with_capacity(num_values); + let key_mappings = values_arrays + .iter() + .zip(values_masks) + .enumerate() + .map(|(dictionary_idx, (values, mask))| { + let mut mapping = vec![zero; values.len()]; + for value_idx in mask.set_indices() { + mapping[value_idx] = K::Native::from_usize(indices.len()).unwrap(); + indices.push((dictionary_idx, value_idx)); + } + mapping + }) + .collect(); + return Ok(MergedDictionaries { + key_mappings, + values: interleave(values_arrays, &indices)?, + }); + } + + let value_slices = values_arrays + .iter() + .zip(values_masks) + .map(|(values, mask)| get_masked_values(*values, mask)) + .collect::<Vec<_>>(); + // Map from value to new index + let mut interner = Interner::new(num_values); + // Interleave indices for new values array + let mut indices = Vec::with_capacity(num_values); + + // Compute the mapping for each dictionary + let key_mappings = values_arrays + .iter() + .enumerate() + .zip(value_slices) + .map(|((dictionary_idx, value_array), values)| { + let zero = K::Native::from_usize(0).unwrap(); Review Comment: ```suggestion let zero = K::Native::ZERO; ``` -- 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]
