arrow_cast/cast/
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 crate::cast::*;
6
7/// Attempts to cast an `ArrayDictionary` with index type K into
8/// `to_type` for supported types.
9///
10/// K is the key type
11pub(crate) fn dictionary_cast<K: ArrowDictionaryKeyType>(
12    array: &dyn Array,
13    to_type: &DataType,
14    cast_options: &CastOptions,
15) -> Result<ArrayRef, ArrowError> {
16    use DataType::*;
17
18    /// Whether the dictionary is sparse; gates short of the measured 0.6x crossover for margin.
19    #[inline]
20    fn is_sparse<K: ArrowDictionaryKeyType>(array: &DictionaryArray<K>) -> bool {
21        array.keys().len() < array.values().len() / 2
22    }
23
24    #[inline]
25    fn values_buffer_fits_in_view<T: ByteArrayType>(values: &GenericByteArray<T>) -> bool {
26        values.values().len() < i32::MAX as usize
27    }
28
29    let array = array.as_dictionary::<K>();
30    let from_child_type = array.values().data_type();
31    match (from_child_type, to_type) {
32        (_, Dictionary(to_index_type, to_value_type)) => {
33            dictionary_to_dictionary_cast(array, to_index_type, to_value_type, cast_options)
34        }
35        // `unpack_dictionary` operates per dictionary value before using take kernel to form
36        // final view. `view_from_dict_values` builds output view directly per row (key index).
37        // Based on benchmarking, `view_from_dict_values` is more efficient for sparse dictionaries
38        // (more dictionary values than there are rows/keys), whilst `unpack_dictionary` is
39        // more efficient for dense dictionaries, where a sparse dictionary is when rows reach
40        // roughly 0.6x the dictionary size.
41        //
42        // Therefore delegate to the faster method based on the density of the input dictionary.
43        (Utf8, Utf8View) if is_sparse(array) => {
44            view_from_dict_values::<K, Utf8Type, StringViewType>(
45                array.keys(),
46                array.values().as_string::<i32>(),
47            )
48        }
49        (Binary, BinaryView) if is_sparse(array) => {
50            view_from_dict_values::<K, BinaryType, BinaryViewType>(
51                array.keys(),
52                array.values().as_binary::<i32>(),
53            )
54        }
55        // `view_from_dict_values` directly appends the values buffer as a block using
56        // `GenericByteViewBuilder::append_block` which asserts length of the buffer; we must
57        // ensure this assertion holds to use it for large variants which may exceed the max allowable length.
58        // If we exceed the length we can simply fallback to `unpack_dictionary` which still builds it
59        // correctly.
60        (LargeUtf8, Utf8View)
61            if is_sparse(array)
62                && values_buffer_fits_in_view(array.values().as_string::<i64>()) =>
63        {
64            view_from_dict_values::<K, LargeUtf8Type, StringViewType>(
65                array.keys(),
66                array.values().as_string::<i64>(),
67            )
68        }
69        (LargeBinary, BinaryView)
70            if is_sparse(array)
71                && values_buffer_fits_in_view(array.values().as_binary::<i64>()) =>
72        {
73            view_from_dict_values::<K, LargeBinaryType, BinaryViewType>(
74                array.keys(),
75                array.values().as_binary::<i64>(),
76            )
77        }
78        // Cross casts to a binary view need no validation: valid UTF-8 is valid binary.
79        (Utf8, BinaryView) if is_sparse(array) => {
80            view_from_dict_values::<K, Utf8Type, BinaryViewType>(
81                array.keys(),
82                array.values().as_string::<i32>(),
83            )
84        }
85        (LargeUtf8, BinaryView)
86            if is_sparse(array)
87                && values_buffer_fits_in_view(array.values().as_string::<i64>()) =>
88        {
89            view_from_dict_values::<K, LargeUtf8Type, BinaryViewType>(
90                array.keys(),
91                array.values().as_string::<i64>(),
92            )
93        }
94        // Cross casts to a string view require UTF-8 validation of the dictionary values.
95        (Binary, Utf8View) if is_sparse(array) => binary_dict_to_string_view::<K, i32>(
96            array.keys(),
97            array.values().as_binary::<i32>(),
98            cast_options,
99        ),
100        (LargeBinary, Utf8View)
101            if is_sparse(array)
102                && values_buffer_fits_in_view(array.values().as_binary::<i64>()) =>
103        {
104            binary_dict_to_string_view::<K, i64>(
105                array.keys(),
106                array.values().as_binary::<i64>(),
107                cast_options,
108            )
109        }
110        _ => unpack_dictionary(array, to_type, cast_options),
111    }
112}
113
114fn dictionary_to_dictionary_cast<K: ArrowDictionaryKeyType>(
115    array: &DictionaryArray<K>,
116    to_index_type: &DataType,
117    to_value_type: &DataType,
118    cast_options: &CastOptions,
119) -> Result<ArrayRef, ArrowError> {
120    use DataType::*;
121
122    // Value type unchanged (and not nested dict): cast keys only, reuse values Arc.
123    if array.values().data_type() == to_value_type
124        && !matches!(array.values().data_type(), Dictionary(_, _))
125    {
126        return cast_keys_only(array, to_index_type, cast_options);
127    }
128
129    // Fast path for a nested dictionary source (`Dictionary<K, Dictionary<K2, V>>`).
130    // Both layers index into the same inner values, so the two index layers can
131    // be composed into one rather than materializing the values: `take` gathers
132    // the inner keys through the outer keys and reuses the inner values buffer
133    // untouched, so no value data is rewritten. The flattened single-level
134    // dictionary is then cast to the requested index/value types.
135    if matches!(array.values().data_type(), Dictionary(_, _)) {
136        let flattened = take(array.values().as_ref(), array.keys(), None)?;
137        return cast_with_options(
138            &flattened,
139            &Dictionary(
140                Box::new(to_index_type.clone()),
141                Box::new(to_value_type.clone()),
142            ),
143            cast_options,
144        );
145    }
146
147    let keys_array: ArrayRef = Arc::new(PrimitiveArray::<K>::from(array.keys().to_data()));
148    let values_array = array.values();
149    let cast_keys = cast_with_options(&keys_array, to_index_type, cast_options)?;
150    let cast_values = cast_with_options(values_array, to_value_type, cast_options)?;
151
152    // Failure to cast keys (because they don't fit in the
153    // target type) results in NULL values;
154    if cast_keys.null_count() > keys_array.null_count() {
155        return Err(ArrowError::ComputeError(format!(
156            "Could not convert {} dictionary indexes from {:?} to {:?}",
157            cast_keys.null_count() - keys_array.null_count(),
158            keys_array.data_type(),
159            to_index_type
160        )));
161    }
162
163    let data = cast_keys.into_data();
164    let builder = data
165        .into_builder()
166        .data_type(Dictionary(
167            Box::new(to_index_type.clone()),
168            Box::new(to_value_type.clone()),
169        ))
170        .child_data(vec![cast_values.into_data()]);
171
172    // Safety
173    // Cast keys are still valid
174    let data = unsafe { builder.build_unchecked() };
175
176    // create the appropriate array type
177    let new_array: ArrayRef = match to_index_type {
178        Int8 => Arc::new(DictionaryArray::<Int8Type>::from(data)),
179        Int16 => Arc::new(DictionaryArray::<Int16Type>::from(data)),
180        Int32 => Arc::new(DictionaryArray::<Int32Type>::from(data)),
181        Int64 => Arc::new(DictionaryArray::<Int64Type>::from(data)),
182        UInt8 => Arc::new(DictionaryArray::<UInt8Type>::from(data)),
183        UInt16 => Arc::new(DictionaryArray::<UInt16Type>::from(data)),
184        UInt32 => Arc::new(DictionaryArray::<UInt32Type>::from(data)),
185        UInt64 => Arc::new(DictionaryArray::<UInt64Type>::from(data)),
186        _ => {
187            return Err(ArrowError::InvalidArgumentError(format!(
188                "Unsupported type {to_index_type} for dictionary index"
189            )));
190        }
191    };
192
193    Ok(new_array)
194}
195
196fn finish_keys_only<K2: ArrowDictionaryKeyType>(
197    cast_keys: ArrayRef,
198    values: ArrayRef,
199) -> Result<ArrayRef, ArrowError> {
200    let new_keys = cast_keys
201        .as_any()
202        .downcast_ref::<PrimitiveArray<K2>>()
203        .ok_or_else(|| {
204            ArrowError::ComputeError(
205                "Internal Error: cast keys did not downcast to target key type".to_string(),
206            )
207        })?
208        .clone();
209    // SAFETY: caller rejects overflowing keys, so every non-null key stays in bounds of `values`.
210    Ok(Arc::new(unsafe {
211        DictionaryArray::<K2>::new_unchecked(new_keys, values)
212    }))
213}
214
215fn cast_keys_only<K: ArrowDictionaryKeyType>(
216    array: &DictionaryArray<K>,
217    to_index_type: &DataType,
218    cast_options: &CastOptions,
219) -> Result<ArrayRef, ArrowError> {
220    let source_keys = array.keys();
221    let cast_keys = cast_with_options(&source_keys, to_index_type, cast_options)?;
222
223    if cast_keys.null_count() > source_keys.null_count() {
224        return Err(ArrowError::ComputeError(format!(
225            "Could not convert {} dictionary indexes from {:?} to {:?}",
226            cast_keys.null_count() - source_keys.null_count(),
227            source_keys.data_type(),
228            to_index_type
229        )));
230    }
231
232    let values = array.values().clone();
233    match to_index_type {
234        DataType::Int8 => finish_keys_only::<Int8Type>(cast_keys, values),
235        DataType::Int16 => finish_keys_only::<Int16Type>(cast_keys, values),
236        DataType::Int32 => finish_keys_only::<Int32Type>(cast_keys, values),
237        DataType::Int64 => finish_keys_only::<Int64Type>(cast_keys, values),
238        DataType::UInt8 => finish_keys_only::<UInt8Type>(cast_keys, values),
239        DataType::UInt16 => finish_keys_only::<UInt16Type>(cast_keys, values),
240        DataType::UInt32 => finish_keys_only::<UInt32Type>(cast_keys, values),
241        DataType::UInt64 => finish_keys_only::<UInt64Type>(cast_keys, values),
242        _ => Err(ArrowError::InvalidArgumentError(format!(
243            "Unsupported type {to_index_type} for dictionary index"
244        ))),
245    }
246}
247
248/// Cast `Dict<K, Binary>` or `Dict<K, LargeBinary>` to `Utf8View`, validating UTF-8 for each
249/// dictionary value.
250///
251/// Fast path when all values are valid UTF-8: reuses the values buffer without copying.
252/// When some values are invalid and `cast_options.safe` is true, rows pointing to those
253/// values become null. When `cast_options.safe` is false, returns an error immediately.
254fn binary_dict_to_string_view<K: ArrowDictionaryKeyType, O: OffsetSizeTrait>(
255    keys: &PrimitiveArray<K>,
256    values: &GenericByteArray<GenericBinaryType<O>>,
257    cast_options: &CastOptions,
258) -> Result<ArrayRef, ArrowError> {
259    match GenericStringArray::<O>::try_from_binary(values.clone()) {
260        Ok(_) => {
261            // All dictionary values are valid UTF-8: reuse the buffer zero-copy.
262            view_from_dict_values::<K, GenericBinaryType<O>, StringViewType>(keys, values)
263        }
264        Err(e) => {
265            if !cast_options.safe {
266                return Err(e);
267            }
268            // safe=true: validate each dictionary value individually so we can nullify
269            // only the rows whose key points to a null or invalid UTF-8 value.
270            let valid: Vec<bool> = (0..values.len())
271                .map(|i| !values.is_null(i) && std::str::from_utf8(values.value(i)).is_ok())
272                .collect();
273
274            let value_buffer = values.values();
275            let value_offsets = values.value_offsets();
276            let mut builder = StringViewBuilder::with_capacity(keys.len());
277            builder.append_block(value_buffer.clone());
278
279            for key in keys {
280                match key {
281                    Some(v) => {
282                        let idx = v.to_usize().ok_or_else(|| {
283                            ArrowError::ComputeError("Invalid dictionary index".to_string())
284                        })?;
285                        let is_valid = *valid.get(idx).ok_or_else(|| {
286                            ArrowError::InvalidArgumentError(format!(
287                                "Dictionary key {idx} out of bounds for dictionary values of length {}",
288                                valid.len()
289                            ))
290                        })?;
291                        if is_valid {
292                            // Safety:
293                            // (1) `idx` and `idx + 1` are in bounds, checked above
294                            // (2) offsets are monotonically increasing, so end >= offset
295                            // (3) the slice [offset..end] is within the buffer
296                            // (4) the bytes are valid UTF-8, checked above
297                            unsafe {
298                                let offset = value_offsets.get_unchecked(idx).as_usize();
299                                let end = value_offsets.get_unchecked(idx + 1).as_usize();
300                                let length = end - offset;
301                                builder.append_view_unchecked(0, offset as u32, length as u32);
302                            }
303                        } else {
304                            builder.append_null();
305                        }
306                    }
307                    None => builder.append_null(),
308                }
309            }
310            Ok(Arc::new(builder.finish()))
311        }
312    }
313}
314
315fn view_from_dict_values<K: ArrowDictionaryKeyType, V: ByteArrayType, T: ByteViewType>(
316    keys: &PrimitiveArray<K>,
317    values: &GenericByteArray<V>,
318) -> Result<ArrayRef, ArrowError> {
319    let value_buffer = values.values();
320    let value_offsets = values.value_offsets();
321    // A null *value* must produce a null row, not the empty slice its offsets happen to span.
322    let values_have_nulls = values.null_count() != 0;
323    let mut builder = GenericByteViewBuilder::<T>::with_capacity(keys.len());
324    builder.append_block(value_buffer.clone());
325    for i in keys {
326        match i {
327            Some(v) => {
328                let idx = v.to_usize().ok_or_else(|| {
329                    ArrowError::ComputeError("Invalid dictionary index".to_string())
330                })?;
331
332                if values_have_nulls && values.is_null(idx) {
333                    builder.append_null();
334                    continue;
335                }
336
337                // Safety
338                // (1) The index is within bounds as they are offsets
339                // (2) The append_view is safe
340                unsafe {
341                    let offset = value_offsets.get_unchecked(idx).as_usize();
342                    let end = value_offsets.get_unchecked(idx + 1).as_usize();
343                    let length = end - offset;
344                    builder.append_view_unchecked(0, offset as u32, length as u32)
345                }
346            }
347            None => {
348                builder.append_null();
349            }
350        }
351    }
352    Ok(Arc::new(builder.finish()))
353}
354
355// Unpack a dictionary into a flattened array of type to_type
356pub(crate) fn unpack_dictionary<K: ArrowDictionaryKeyType>(
357    array: &DictionaryArray<K>,
358    to_type: &DataType,
359    cast_options: &CastOptions,
360) -> Result<ArrayRef, ArrowError> {
361    let cast_dict_values = cast_with_options(array.values(), to_type, cast_options)?;
362    take(cast_dict_values.as_ref(), array.keys(), None)
363}
364
365/// Pack a data type into a dictionary array passing the values through a primitive array
366pub(crate) fn pack_array_to_dictionary_via_primitive<K: ArrowDictionaryKeyType>(
367    array: &dyn Array,
368    primitive_type: DataType,
369    dict_value_type: &DataType,
370    cast_options: &CastOptions,
371) -> Result<ArrayRef, ArrowError> {
372    let primitive = cast_with_options(array, &primitive_type, cast_options)?;
373    let dict = cast_with_options(
374        primitive.as_ref(),
375        &DataType::Dictionary(Box::new(K::DATA_TYPE), Box::new(primitive_type)),
376        cast_options,
377    )?;
378    cast_with_options(
379        dict.as_ref(),
380        &DataType::Dictionary(Box::new(K::DATA_TYPE), Box::new(dict_value_type.clone())),
381        cast_options,
382    )
383}
384
385/// Attempts to encode an array into an `ArrayDictionary` with index
386/// type K and value (dictionary) type value_type
387///
388/// K is the key type
389pub(crate) fn cast_to_dictionary<K: ArrowDictionaryKeyType>(
390    array: &dyn Array,
391    dict_value_type: &DataType,
392    cast_options: &CastOptions,
393) -> Result<ArrayRef, ArrowError> {
394    use DataType::*;
395
396    match *dict_value_type {
397        Int8 => pack_numeric_to_dictionary::<K, Int8Type>(array, dict_value_type, cast_options),
398        Int16 => pack_numeric_to_dictionary::<K, Int16Type>(array, dict_value_type, cast_options),
399        Int32 => pack_numeric_to_dictionary::<K, Int32Type>(array, dict_value_type, cast_options),
400        Int64 => pack_numeric_to_dictionary::<K, Int64Type>(array, dict_value_type, cast_options),
401        UInt8 => pack_numeric_to_dictionary::<K, UInt8Type>(array, dict_value_type, cast_options),
402        UInt16 => pack_numeric_to_dictionary::<K, UInt16Type>(array, dict_value_type, cast_options),
403        UInt32 => pack_numeric_to_dictionary::<K, UInt32Type>(array, dict_value_type, cast_options),
404        UInt64 => pack_numeric_to_dictionary::<K, UInt64Type>(array, dict_value_type, cast_options),
405        Decimal32(p, s) => pack_decimal_to_dictionary::<K, Decimal32Type>(
406            array,
407            dict_value_type,
408            p,
409            s,
410            cast_options,
411        ),
412        Decimal64(p, s) => pack_decimal_to_dictionary::<K, Decimal64Type>(
413            array,
414            dict_value_type,
415            p,
416            s,
417            cast_options,
418        ),
419        Decimal128(p, s) => pack_decimal_to_dictionary::<K, Decimal128Type>(
420            array,
421            dict_value_type,
422            p,
423            s,
424            cast_options,
425        ),
426        Decimal256(p, s) => pack_decimal_to_dictionary::<K, Decimal256Type>(
427            array,
428            dict_value_type,
429            p,
430            s,
431            cast_options,
432        ),
433        Float16 => {
434            pack_numeric_to_dictionary::<K, Float16Type>(array, dict_value_type, cast_options)
435        }
436        Float32 => {
437            pack_numeric_to_dictionary::<K, Float32Type>(array, dict_value_type, cast_options)
438        }
439        Float64 => {
440            pack_numeric_to_dictionary::<K, Float64Type>(array, dict_value_type, cast_options)
441        }
442        Date32 => pack_array_to_dictionary_via_primitive::<K>(
443            array,
444            DataType::Int32,
445            dict_value_type,
446            cast_options,
447        ),
448        Date64 => pack_array_to_dictionary_via_primitive::<K>(
449            array,
450            DataType::Int64,
451            dict_value_type,
452            cast_options,
453        ),
454        Time32(_) => pack_array_to_dictionary_via_primitive::<K>(
455            array,
456            DataType::Int32,
457            dict_value_type,
458            cast_options,
459        ),
460        Time64(_) => pack_array_to_dictionary_via_primitive::<K>(
461            array,
462            DataType::Int64,
463            dict_value_type,
464            cast_options,
465        ),
466        Timestamp(_, _) => pack_array_to_dictionary_via_primitive::<K>(
467            array,
468            DataType::Int64,
469            dict_value_type,
470            cast_options,
471        ),
472        Utf8 => {
473            // If the input is a view type, we can avoid casting (thus copying) the data
474            if array.data_type() == &DataType::Utf8View {
475                return string_view_to_dictionary::<K, i32>(array);
476            }
477            pack_byte_to_dictionary::<K, GenericStringType<i32>>(array, cast_options)
478        }
479        LargeUtf8 => {
480            // If the input is a view type, we can avoid casting (thus copying) the data
481            if array.data_type() == &DataType::Utf8View {
482                return string_view_to_dictionary::<K, i64>(array);
483            }
484            pack_byte_to_dictionary::<K, GenericStringType<i64>>(array, cast_options)
485        }
486        Utf8View => {
487            let base_value_type = match array.data_type() {
488                DataType::LargeUtf8 | DataType::Utf8View => DataType::LargeUtf8,
489                _ => DataType::Utf8,
490            };
491
492            let dict_base = cast_to_dictionary::<K>(array, &base_value_type, cast_options)?;
493            dictionary_cast::<K>(
494                dict_base.as_ref(),
495                &DataType::Dictionary(Box::new(K::DATA_TYPE), Box::new(DataType::Utf8View)),
496                cast_options,
497            )
498        }
499        Binary => {
500            // If the input is a view type, we can avoid casting (thus copying) the data
501            if array.data_type() == &DataType::BinaryView {
502                return binary_view_to_dictionary::<K, i32>(array);
503            }
504            pack_byte_to_dictionary::<K, GenericBinaryType<i32>>(array, cast_options)
505        }
506        LargeBinary => {
507            // If the input is a view type, we can avoid casting (thus copying) the data
508            if array.data_type() == &DataType::BinaryView {
509                return binary_view_to_dictionary::<K, i64>(array);
510            }
511            pack_byte_to_dictionary::<K, GenericBinaryType<i64>>(array, cast_options)
512        }
513        BinaryView => {
514            let base_value_type = match array.data_type() {
515                DataType::LargeBinary | DataType::BinaryView => DataType::LargeBinary,
516                _ => DataType::Binary,
517            };
518
519            let dict_base = cast_to_dictionary::<K>(array, &base_value_type, cast_options)?;
520            dictionary_cast::<K>(
521                dict_base.as_ref(),
522                &DataType::Dictionary(Box::new(K::DATA_TYPE), Box::new(DataType::BinaryView)),
523                cast_options,
524            )
525        }
526        FixedSizeBinary(byte_size) => {
527            pack_byte_to_fixed_size_dictionary::<K>(array, cast_options, byte_size)
528        }
529        Struct(_) => pack_struct_to_dictionary::<K>(array, dict_value_type, cast_options),
530        _ => Err(ArrowError::InvalidArgumentError(format!(
531            "Unsupported output type for dictionary packing: {dict_value_type}"
532        ))),
533    }
534}
535
536/// Wrap a struct-valued array as a `DictionaryArray<K, Struct>` with identity
537/// keys `[0, 1, ..., len-1]`. Unlike the primitive / byte packers above, no
538/// deduplication is performed, since struct values have no general hash/equality
539/// builder in arrow-rs.
540///
541/// Each child field of the source is recursively cast to the matching field of
542/// `dict_value_type` via `cast_with_options` before keys are emitted. If any
543/// child cast fails, the whole pack fails, the same contract as the primitive
544/// packers above.
545fn pack_struct_to_dictionary<K: ArrowDictionaryKeyType>(
546    array: &dyn Array,
547    dict_value_type: &DataType,
548    cast_options: &CastOptions,
549) -> Result<ArrayRef, ArrowError> {
550    let cast_values = cast_with_options(array, dict_value_type, cast_options)?;
551    let len = cast_values.len();
552
553    // Identity keys `[0, 1, ..., len-1]`, with null entries wherever the
554    // source row is null so the dictionary's logical null mask matches.
555    let mut builder = PrimitiveBuilder::<K>::with_capacity(len);
556    for i in 0..len {
557        if cast_values.is_null(i) {
558            builder.append_null();
559        } else {
560            let key = K::Native::from_usize(i).ok_or_else(|| {
561                ArrowError::CastError(format!(
562                    "Cannot fit {len} dictionary keys in {:?}",
563                    K::DATA_TYPE,
564                ))
565            })?;
566            builder.append_value(key);
567        }
568    }
569    let keys = builder.finish();
570
571    Ok(Arc::new(DictionaryArray::<K>::try_new(keys, cast_values)?))
572}
573
574// Packs the data from the primitive array of type <V> to a
575// DictionaryArray with keys of type K and values of value_type V
576pub(crate) fn pack_numeric_to_dictionary<K, V>(
577    array: &dyn Array,
578    dict_value_type: &DataType,
579    cast_options: &CastOptions,
580) -> Result<ArrayRef, ArrowError>
581where
582    K: ArrowDictionaryKeyType,
583    V: ArrowPrimitiveType,
584{
585    // attempt to cast the source array values to the target value type (the dictionary values type)
586    let cast_values = cast_with_options(array, dict_value_type, cast_options)?;
587    let values = cast_values.as_primitive::<V>();
588
589    let mut b = PrimitiveDictionaryBuilder::<K, V>::with_capacity(values.len(), values.len());
590
591    // copy each element one at a time
592    for i in 0..values.len() {
593        if values.is_null(i) {
594            b.append_null();
595        } else {
596            b.append(values.value(i))?;
597        }
598    }
599    Ok(Arc::new(b.finish()))
600}
601
602pub(crate) fn pack_decimal_to_dictionary<K, D>(
603    array: &dyn Array,
604    dict_value_type: &DataType,
605    precision: u8,
606    scale: i8,
607    cast_options: &CastOptions,
608) -> Result<ArrayRef, ArrowError>
609where
610    K: ArrowDictionaryKeyType,
611    D: DecimalType + ArrowPrimitiveType,
612{
613    let dict = pack_numeric_to_dictionary::<K, D>(array, dict_value_type, cast_options)?;
614    let dict = dict.as_dictionary::<K>();
615    let value = dict
616        .values()
617        .as_primitive::<D>()
618        .clone()
619        .with_precision_and_scale(precision, scale)?;
620    Ok(Arc::new(dict.with_values(Arc::new(value))))
621}
622
623pub(crate) fn string_view_to_dictionary<K, O: OffsetSizeTrait>(
624    array: &dyn Array,
625) -> Result<ArrayRef, ArrowError>
626where
627    K: ArrowDictionaryKeyType,
628{
629    let mut b = GenericByteDictionaryBuilder::<K, GenericStringType<O>>::with_capacity(
630        array.len(),
631        1024,
632        1024,
633    );
634    let string_view = array.as_string_view();
635    for v in string_view {
636        match v {
637            Some(v) => {
638                b.append(v)?;
639            }
640            None => {
641                b.append_null();
642            }
643        }
644    }
645
646    Ok(Arc::new(b.finish()))
647}
648
649pub(crate) fn binary_view_to_dictionary<K, O: OffsetSizeTrait>(
650    array: &dyn Array,
651) -> Result<ArrayRef, ArrowError>
652where
653    K: ArrowDictionaryKeyType,
654{
655    let mut b = GenericByteDictionaryBuilder::<K, GenericBinaryType<O>>::with_capacity(
656        array.len(),
657        1024,
658        1024,
659    );
660    let binary_view = array.as_binary_view();
661    for v in binary_view {
662        match v {
663            Some(v) => {
664                b.append(v)?;
665            }
666            None => {
667                b.append_null();
668            }
669        }
670    }
671
672    Ok(Arc::new(b.finish()))
673}
674
675// Packs the data as a GenericByteDictionaryBuilder, if possible, with the
676// key types of K
677pub(crate) fn pack_byte_to_dictionary<K, T>(
678    array: &dyn Array,
679    cast_options: &CastOptions,
680) -> Result<ArrayRef, ArrowError>
681where
682    K: ArrowDictionaryKeyType,
683    T: ByteArrayType,
684{
685    let cast_values = cast_with_options(array, &T::DATA_TYPE, cast_options)?;
686    let values = cast_values.as_bytes::<T>();
687    let mut b = GenericByteDictionaryBuilder::<K, T>::with_capacity(values.len(), 1024, 1024);
688    b.append_array(values)?;
689    Ok(Arc::new(b.finish()))
690}
691
692// Packs the data as a GenericByteDictionaryBuilder, if possible, with the
693// key types of K
694pub(crate) fn pack_byte_to_fixed_size_dictionary<K>(
695    array: &dyn Array,
696    cast_options: &CastOptions,
697    byte_width: i32,
698) -> Result<ArrayRef, ArrowError>
699where
700    K: ArrowDictionaryKeyType,
701{
702    let cast_values =
703        cast_with_options(array, &DataType::FixedSizeBinary(byte_width), cast_options)?;
704    let values = cast_values.as_fixed_size_binary();
705    let mut b = FixedSizeBinaryDictionaryBuilder::<K>::with_capacity(1024, 1024, byte_width);
706
707    // copy each element one at a time
708    for i in 0..values.len() {
709        if values.is_null(i) {
710            b.append_null();
711        } else {
712            b.append(values.value(i))?;
713        }
714    }
715    Ok(Arc::new(b.finish()))
716}
717
718#[cfg(test)]
719mod tests {
720    use super::*;
721    use DataType::*;
722
723    // Too few keys for the dictionary, so `is_sparse` holds and views are built per row.
724    // Covers an inlined value, a buffer backed one, a null key and a null dictionary value.
725    fn sparse_keys(values: &ArrayRef) -> Int32Array {
726        let keys = Int32Array::from(vec![Some(0), Some(2), None, Some(3)]);
727        assert!(
728            keys.len() < values.len() / 2,
729            "keys must reach the direct path"
730        );
731        keys
732    }
733
734    // One key per value, so `is_sparse` fails and `unpack_dictionary` runs.
735    fn dense_keys(values: &ArrayRef) -> Int32Array {
736        let keys = Int32Array::from(vec![
737            Some(0),
738            Some(2),
739            None,
740            Some(3),
741            Some(1),
742            Some(0),
743            Some(4),
744            Some(2),
745            Some(3),
746            Some(5),
747        ]);
748        assert!(
749            keys.len() >= values.len() / 2,
750            "keys must reach unpack_dictionary"
751        );
752        keys
753    }
754
755    fn make_dict(keys: &Int32Array, values: &ArrayRef) -> DictionaryArray<Int32Type> {
756        DictionaryArray::try_new(keys.clone(), values.clone()).unwrap()
757    }
758
759    #[test]
760    fn test_dict_to_view_matches_take_then_cast() {
761        let long = "a value over twelve bytes";
762        let utf8: ArrayRef = Arc::new(StringArray::from(vec![
763            Some("aa"),
764            Some("bb"),
765            Some(long),
766            None,
767            Some("ee"),
768            Some("ff"),
769            Some("gg"),
770            Some("hh"),
771            Some("ii"),
772            Some("jj"),
773        ]));
774
775        for from in [Utf8, LargeUtf8, Binary, LargeBinary] {
776            let values = cast(&utf8, &from).unwrap();
777            for to in [Utf8View, BinaryView] {
778                for keys in [sparse_keys(&values), dense_keys(&values)] {
779                    let expected = cast(&take(&values, &keys, None).unwrap(), &to).unwrap();
780                    let casted = cast(&make_dict(&keys, &values), &to).unwrap();
781                    assert_eq!(casted.as_ref(), expected.as_ref(), "{from:?} -> {to:?}");
782                }
783            }
784        }
785    }
786
787    #[test]
788    fn test_dict_binary_to_utf8view_invalid_utf8() {
789        let bytes: Vec<&[u8]> = vec![
790            b"aa",
791            b"bb",
792            &[0xFF, 0xFE],
793            b"dd",
794            b"ee",
795            b"ff",
796            b"gg",
797            b"hh",
798            b"ii",
799            b"jj",
800        ];
801        let strict = CastOptions {
802            safe: false,
803            ..Default::default()
804        };
805        let safe = CastOptions {
806            safe: true,
807            ..Default::default()
808        };
809
810        for values in [
811            Arc::new(BinaryArray::from_vec(bytes.clone())) as ArrayRef,
812            Arc::new(LargeBinaryArray::from_vec(bytes.clone())) as ArrayRef,
813        ] {
814            // index 2 holds the invalid value, so keys pointing at it are nullified when safe
815            for (keys, expected) in [
816                (
817                    sparse_keys(&values),
818                    vec![Some("aa"), None, None, Some("dd")],
819                ),
820                (
821                    dense_keys(&values),
822                    vec![
823                        Some("aa"),
824                        None,
825                        None,
826                        Some("dd"),
827                        Some("bb"),
828                        Some("aa"),
829                        Some("ee"),
830                        None,
831                        Some("dd"),
832                        Some("ff"),
833                    ],
834                ),
835            ] {
836                let dict = make_dict(&keys, &values);
837                assert!(cast_with_options(&dict, &Utf8View, &strict).is_err());
838
839                let casted = cast_with_options(&dict, &Utf8View, &safe).unwrap();
840                assert_eq!(casted.as_string_view(), &StringViewArray::from(expected));
841            }
842        }
843    }
844}
arrow_cast/cast/dictionary.rs

arrow_cast/cast/
dictionary.rs
