arrow_select/coalesce/
dictionary.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: See the NOTICE file distributed with this work for additional information regarding copyright ownership
3// SPDX-FileContributor: Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements
4
5use super::{InProgressArray, generic::GenericInProgressArray};
6use crate::dictionary::{KeyMapping, child_values_ptr_eq, merge_dictionary_values_with_masks};
7use ahash::RandomState;
8use arrow_array::cast::AsArray;
9use arrow_array::types::ArrowDictionaryKeyType;
10use arrow_array::{
11    Array, ArrayRef, DictionaryArray, PrimitiveArray, downcast_integer, new_empty_array,
12};
13use arrow_buffer::{
14    ArrowNativeType, BooleanBufferBuilder, NullBuffer, NullBufferBuilder, ScalarBuffer,
15};
16use arrow_schema::{ArrowError, DataType};
17use std::collections::HashMap;
18use std::mem::size_of;
19use std::sync::Arc;
20
21struct KeyChunk<K: ArrowDictionaryKeyType> {
22    /// Zero-copy slice of source-local keys for the retained rows.
23    keys: ScalarBuffer<K::Native>,
24    /// Shared validity buffer from the source keys, if present.
25    nulls: Option<NullBuffer>,
26    /// Index of the values array in `InProgressDictionaryArray::sources`.
27    source: usize,
28}
29
30pub(crate) fn create_in_progress_dictionary(
31    key_type: &DataType,
32    value_type: &DataType,
33) -> Box<dyn InProgressArray> {
34    macro_rules! create {
35        ($key:ty, $value_type:ident) => {
36            Box::new(InProgressDictionaryArray::<$key>::new($value_type))
37        };
38    }
39
40    match value_type {
41        DataType::Utf8 | DataType::LargeUtf8 | DataType::Binary | DataType::LargeBinary => {
42            downcast_integer! {
43                key_type => (create, value_type),
44                _ => unreachable!("illegal dictionary key type {key_type}"),
45            }
46        }
47        _ => Box::new(GenericInProgressArray::new()),
48    }
49}
50
51struct InProgressDictionaryArray<K: ArrowDictionaryKeyType> {
52    /// Dictionary currently selected by [`InProgressArray::set_source`].
53    source: Option<ArrayRef>,
54    /// Number of output rows retained across `chunks`.
55    len: usize,
56    /// Typed empty values array used when finishing without any rows.
57    empty_values: ArrayRef,
58    /// Distinct dictionary values arrays retained by buffer identity.
59    sources: Vec<ArrayRef>,
60    /// Maps each retained values `Arc` address to its index in `sources`.
61    source_ids: HashMap<usize, usize, RandomState>,
62    /// Zero-copy key slices in output order, tagged with their values source.
63    chunks: Vec<KeyChunk<K>>,
64}
65
66impl<K: ArrowDictionaryKeyType> std::fmt::Debug for InProgressDictionaryArray<K> {
67    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
68        f.debug_struct("InProgressDictionaryArray").finish()
69    }
70}
71
72impl<K: ArrowDictionaryKeyType> InProgressDictionaryArray<K> {
73    fn new(value_type: &DataType) -> Self {
74        Self {
75            source: None,
76            len: 0,
77            empty_values: new_empty_array(value_type),
78            sources: Vec::new(),
79            source_ids: HashMap::with_hasher(RandomState::with_seeds(0, 0, 0, 0)),
80            chunks: Vec::new(),
81        }
82    }
83
84    /// Returns a stable index for `values`, reusing a source when either its `Arc`
85    /// pointer or its backing buffer pointers match.
86    fn source_id(&mut self, values: ArrayRef) -> usize {
87        if let Some(last) = self.chunks.last()
88            && child_values_ptr_eq(&self.sources[last.source], &values)
89        {
90            return last.source;
91        }
92        let values_id = Arc::as_ptr(&values).cast::<()>() as usize;
93        if let Some(&source) = self.source_ids.get(&values_id) {
94            return source;
95        }
96        if let Some(source) = self
97            .sources
98            .iter()
99            .position(|existing| child_values_ptr_eq(existing, &values))
100        {
101            return source;
102        }
103        let source = self.sources.len();
104        self.sources.push(values);
105        self.source_ids.insert(values_id, source);
106        source
107    }
108
109    /// Compacts referenced values and remaps source-local keys while preserving chunk order.
110    /// Equal referenced values are only deduplicated when their count would overflow
111    /// `K`'s key space.
112    fn merge_sources(
113        sources: &[ArrayRef],
114        chunks: &[KeyChunk<K>],
115        len: usize,
116    ) -> Result<(Vec<K::Native>, ArrayRef), ArrowError> {
117        let mut referenced = sources
118            .iter()
119            .map(|source| {
120                let mut mask = BooleanBufferBuilder::new(source.len());
121                mask.advance(source.len());
122                mask
123            })
124            .collect::<Vec<_>>();
125
126        for chunk in chunks {
127            let mask = &mut referenced[chunk.source];
128            match &chunk.nulls {
129                Some(nulls) => nulls
130                    .inner()
131                    .set_indices()
132                    .for_each(|offset| mask.set_bit(chunk.keys[offset].as_usize(), true)),
133                None => chunk
134                    .keys
135                    .iter()
136                    .for_each(|key| mask.set_bit(key.as_usize(), true)),
137            }
138        }
139
140        let referenced = referenced
141            .into_iter()
142            .map(BooleanBufferBuilder::build)
143            .collect::<Vec<_>>();
144        let source_refs = sources
145            .iter()
146            .map(|source| source.as_ref() as &dyn Array)
147            .collect::<Vec<_>>();
148        let merged = merge_dictionary_values_with_masks::<K>(&source_refs, &referenced)?;
149
150        let mut keys = Vec::with_capacity(len);
151        for chunk in chunks {
152            let KeyMapping::Dense(mapping) = &merged.key_mappings[chunk.source] else {
153                unreachable!("merge_dictionary_values_with_masks returns dense mappings");
154            };
155            match &chunk.nulls {
156                Some(nulls) => keys.extend(chunk.keys.iter().enumerate().map(|(offset, key)| {
157                    if nulls.is_valid(offset) {
158                        mapping[key.as_usize()]
159                    } else {
160                        K::Native::default()
161                    }
162                })),
163                None => keys.extend(chunk.keys.iter().map(|key| mapping[key.as_usize()])),
164            }
165        }
166        Ok((keys, merged.values))
167    }
168
169    /// Concatenates chunk validity in output order, omitting it when all keys are valid.
170    fn build_nulls(chunks: &[KeyChunk<K>], len: usize) -> Option<NullBuffer> {
171        let mut nulls = NullBufferBuilder::new(len);
172        for chunk in chunks {
173            match &chunk.nulls {
174                Some(chunk_nulls) => nulls.append_buffer(chunk_nulls),
175                None => nulls.append_n_non_nulls(chunk.keys.len()),
176            }
177        }
178        nulls.finish()
179    }
180}
181
182impl<K: ArrowDictionaryKeyType> InProgressArray for InProgressDictionaryArray<K> {
183    fn set_source(&mut self, source: Option<ArrayRef>) {
184        self.source = source;
185    }
186
187    fn copy_rows(&mut self, offset: usize, len: usize) -> Result<(), ArrowError> {
188        if len == 0 {
189            return Ok(());
190        }
191        let source = Arc::clone(self.source.as_ref().ok_or_else(|| {
192            ArrowError::InvalidArgumentError(
193                "Internal Error: InProgressDictionaryArray: source not set".to_string(),
194            )
195        })?);
196        let dictionary = source.as_dictionary::<K>();
197        let keys = dictionary.keys();
198        let source = self.source_id(Arc::clone(dictionary.values()));
199        self.chunks.push(KeyChunk {
200            keys: keys.values().slice(offset, len),
201            nulls: keys.nulls().map(|nulls| nulls.slice(offset, len)),
202            source,
203        });
204        self.len += len;
205        Ok(())
206    }
207
208    fn finish(&mut self) -> Result<ArrayRef, ArrowError> {
209        let sources = std::mem::take(&mut self.sources);
210        let chunks = std::mem::take(&mut self.chunks);
211        let len = std::mem::take(&mut self.len);
212        self.source_ids.clear();
213
214        let (keys, values) = match sources.len() {
215            0 => (Vec::new(), Arc::clone(&self.empty_values)),
216            1 => (
217                chunks
218                    .iter()
219                    .flat_map(|chunk| chunk.keys.iter().copied())
220                    .collect(),
221                Arc::clone(&sources[0]),
222            ),
223            _ => Self::merge_sources(&sources, &chunks, len)?,
224        };
225        let nulls = Self::build_nulls(&chunks, len);
226        let keys = PrimitiveArray::<K>::try_new(ScalarBuffer::from(keys), nulls)?;
227        Ok(Arc::new(DictionaryArray::try_new(keys, values)?))
228    }
229
230    /// Estimates memory retained by the source, value arrays, key slices, and bookkeeping.
231    fn size(&self) -> usize {
232        self.source
233            .as_ref()
234            .map_or(0, |source| source.get_array_memory_size())
235            + self
236                .sources
237                .iter()
238                .map(|source| source.get_array_memory_size())
239                .sum::<usize>()
240            + self.sources.capacity() * size_of::<ArrayRef>()
241            + self.source_ids.capacity() * size_of::<(usize, usize)>()
242            + self.chunks.capacity() * size_of::<KeyChunk<K>>()
243            + self
244                .chunks
245                .iter()
246                .map(|chunk| {
247                    chunk.keys.inner().capacity()
248                        + chunk
249                            .nulls
250                            .as_ref()
251                            .map_or(0, |nulls| nulls.inner().inner().capacity())
252                })
253                .sum::<usize>()
254    }
255}
256
257#[cfg(test)]
258mod tests {
259    use super::*;
260    use crate::coalesce::BatchCoalescer;
261    use arrow_array::types::{Int8Type, Int32Type};
262    use arrow_array::{Int8Array, Int32Array, RecordBatch, StringArray};
263
264    fn batch(values: &ArrayRef, keys: Vec<Option<i32>>) -> RecordBatch {
265        let dictionary = Arc::new(DictionaryArray::<Int32Type>::new(
266            Int32Array::from(keys),
267            Arc::clone(values),
268        )) as ArrayRef;
269        RecordBatch::try_from_iter([("dict", dictionary)]).unwrap()
270    }
271
272    #[test]
273    fn coalesces_repeated_and_distinct_values() {
274        let first_values = Arc::new(StringArray::from(vec!["a", "b"])) as ArrayRef;
275        let second_values = Arc::new(StringArray::from(vec!["b", "c"])) as ArrayRef;
276        let batches = [
277            batch(&first_values, vec![Some(0), None]),
278            batch(&second_values, vec![Some(0), Some(1)]),
279            batch(&first_values, vec![Some(1), Some(0)]),
280        ];
281        let mut coalescer = BatchCoalescer::new(batches[0].schema(), 6);
282        for batch in batches {
283            coalescer.push_batch(batch).unwrap();
284        }
285
286        let output = coalescer.next_completed_batch().unwrap();
287        let output = output.column(0).as_dictionary::<Int32Type>();
288        let values = output.values().as_string::<i32>();
289        let decoded = output
290            .keys_iter()
291            .map(|key| key.map(|key| values.value(key)))
292            .collect::<Vec<_>>();
293        assert_eq!(
294            decoded,
295            vec![Some("a"), None, Some("b"), Some("c"), Some("b"), Some("a")]
296        );
297        assert_eq!(values.len(), 4);
298    }
299
300    #[test]
301    fn reuses_one_values_array() {
302        let values = Arc::new(StringArray::from(vec!["a", "b"])) as ArrayRef;
303        let mut coalescer = BatchCoalescer::new(batch(&values, vec![Some(0)]).schema(), 2);
304        coalescer.push_batch(batch(&values, vec![Some(0)])).unwrap();
305        coalescer.push_batch(batch(&values, vec![Some(1)])).unwrap();
306
307        let output = coalescer.next_completed_batch().unwrap();
308        let output = output.column(0).as_dictionary::<Int32Type>();
309        assert!(Arc::ptr_eq(output.values(), &values));
310    }
311
312    #[test]
313    fn reuses_values_with_shared_buffers() {
314        let values = StringArray::from_iter_values((0..128).map(|i| format!("value-{i}")));
315        let first_values = Arc::new(values.clone()) as ArrayRef;
316        let second_values = Arc::new(values) as ArrayRef;
317        let make_batch = |values| {
318            let dictionary =
319                DictionaryArray::<Int8Type>::new(Int8Array::from_iter_values(0..=127), values);
320            RecordBatch::try_from_iter([("dict", Arc::new(dictionary) as ArrayRef)]).unwrap()
321        };
322
323        let first = make_batch(Arc::clone(&first_values));
324        let mut coalescer = BatchCoalescer::new(first.schema(), 256);
325        coalescer.push_batch(first).unwrap();
326        coalescer.push_batch(make_batch(second_values)).unwrap();
327
328        let output = coalescer.next_completed_batch().unwrap();
329        let output = output.column(0).as_dictionary::<Int8Type>();
330        assert!(Arc::ptr_eq(output.values(), &first_values));
331        assert_eq!(
332            output.keys().values(),
333            &(0..=127).chain(0..=127).collect::<Vec<_>>()
334        );
335    }
336
337    #[test]
338    fn remaps_sparse_values_across_sources() {
339        let first = Arc::new(DictionaryArray::<Int32Type>::new(
340            Int32Array::from_iter_values([1, 2, 0, 2, 1]),
341            Arc::new(StringArray::from(vec!["a", "unused-a", "b"])),
342        )) as ArrayRef;
343        let second = Arc::new(DictionaryArray::<Int32Type>::new(
344            Int32Array::from_iter_values([2, 1, 0, 1, 2]),
345            Arc::new(StringArray::from(vec!["b", "c", "unused-b"])),
346        )) as ArrayRef;
347
348        let mut in_progress = InProgressDictionaryArray::<Int32Type>::new(&DataType::Utf8);
349        in_progress.set_source(Some(first));
350        in_progress.copy_rows(1, 3).unwrap();
351        in_progress.set_source(Some(second));
352        in_progress.copy_rows(1, 3).unwrap();
353
354        let output = in_progress.finish().unwrap();
355        let output = output.as_dictionary::<Int32Type>();
356        assert_eq!(output.keys().values(), &[1, 0, 1, 3, 2, 3]);
357        assert_eq!(
358            output
359                .values()
360                .as_string::<i32>()
361                .iter()
362                .collect::<Vec<_>>(),
363            vec![Some("a"), Some("b"), Some("b"), Some("c")]
364        );
365    }
366
367    #[test]
368    fn remaps_keys_when_deduplication_is_required() {
369        let make_values = || {
370            Arc::new(StringArray::from_iter_values(
371                (0..70).map(|idx| format!("value-{idx}")),
372            )) as ArrayRef
373        };
374        let first_values = make_values();
375        let second_values = make_values();
376        let make_dictionary = |values| {
377            Arc::new(DictionaryArray::<Int8Type>::new(
378                Int8Array::from_iter_values(0..70),
379                values,
380            )) as ArrayRef
381        };
382
383        let mut in_progress = InProgressDictionaryArray::<Int8Type>::new(&DataType::Utf8);
384        in_progress.set_source(Some(make_dictionary(first_values)));
385        in_progress.copy_rows(0, 70).unwrap();
386        in_progress.set_source(Some(make_dictionary(second_values)));
387        in_progress.copy_rows(0, 70).unwrap();
388
389        let output = in_progress.finish().unwrap();
390        let output = output.as_dictionary::<Int8Type>();
391        let values = output.values().as_string::<i32>();
392        for (idx, key) in output.keys().values().iter().enumerate() {
393            assert_eq!(values.value(*key as usize), format!("value-{}", idx % 70));
394        }
395        assert!(values.len() <= i8::MAX as usize + 1);
396    }
397}
arrow_select/coalesce/dictionary.rs

arrow_select/coalesce/
dictionary.rs
