arrow_cast/cast/
dictionary.rs

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

arrow_cast/cast/
dictionary.rs
