arrow_cast/cast/
list.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::*;
6use arrow_buffer::ScalarBuffer;
7
8/// Converts a non-list array to a list array where every element is a single element
9/// list. `NULL`s in the original array become `[NULL]` (i.e. output list array
10/// contains no nulls since it wraps all input nulls in a single element list).
11///
12/// For example: `Int32([1, NULL, 2]) -> List<Int32>([[1], [NULL], [2]])`
13pub(crate) fn cast_values_to_list<O: OffsetSizeTrait>(
14    array: &dyn Array,
15    to: &FieldRef,
16    cast_options: &CastOptions,
17) -> Result<ArrayRef, ArrowError> {
18    if array.len() > O::MAX_OFFSET {
19        return Err(ArrowError::ComputeError(format!(
20            "Offset overflow when casting from {} to {}",
21            array.data_type(),
22            to.data_type()
23        )));
24    }
25    let values = cast_with_options(array, to.data_type(), cast_options)?;
26    let offsets = OffsetBuffer::from_repeated_length(1, values.len());
27    let list = GenericListArray::<O>::try_new(to.clone(), offsets, values, None)?;
28    Ok(Arc::new(list))
29}
30
31/// Same as [`cast_values_to_list`] but output list view array.
32pub(crate) fn cast_values_to_list_view<O: OffsetSizeTrait>(
33    array: &dyn Array,
34    to: &FieldRef,
35    cast_options: &CastOptions,
36) -> Result<ArrayRef, ArrowError> {
37    if array.len() > O::MAX_OFFSET {
38        return Err(ArrowError::ComputeError(format!(
39            "Offset overflow when casting from {} to {}",
40            array.data_type(),
41            to.data_type()
42        )));
43    }
44    let values = cast_with_options(array, to.data_type(), cast_options)?;
45    let offsets = (0..values.len())
46        .map(|index| O::usize_as(index))
47        .collect::<Vec<O>>();
48    let list = GenericListViewArray::<O>::try_new(
49        to.clone(),
50        offsets.into(),
51        vec![O::one(); values.len()].into(),
52        values,
53        None,
54    )?;
55    Ok(Arc::new(list))
56}
57
58/// Cast fixed size list array to inner values type, essentially flattening the
59/// lists.
60///
61/// For example: `FixedSizeList<Int32, 2>([[1, 2], [3, 4]]) -> Int32([1, 2, 3, 4])`
62pub(crate) fn cast_single_element_fixed_size_list_to_values(
63    array: &dyn Array,
64    to: &DataType,
65    cast_options: &CastOptions,
66) -> Result<ArrayRef, ArrowError> {
67    let values = array.as_fixed_size_list().values();
68    cast_with_options(values, to, cast_options)
69}
70
71fn cast_fixed_size_list_to_list_inner<OffsetSize: OffsetSizeTrait, const IS_LIST_VIEW: bool>(
72    array: &dyn Array,
73    to: &FieldRef,
74    cast_options: &CastOptions,
75) -> Result<ArrayRef, ArrowError> {
76    let array = array.as_fixed_size_list();
77    let DataType::FixedSizeList(inner_field, size) = array.data_type() else {
78        unreachable!()
79    };
80    let array = if to != inner_field {
81        // To transform inner type, can first cast to FSL with new inner type.
82        let fsl_to = DataType::FixedSizeList(to.clone(), *size);
83        let array = cast_with_options(array, &fsl_to, cast_options)?;
84        array.as_fixed_size_list().clone()
85    } else {
86        array.clone()
87    };
88    if IS_LIST_VIEW {
89        let list: GenericListViewArray<OffsetSize> = array.into();
90        Ok(Arc::new(list))
91    } else {
92        let list: GenericListArray<OffsetSize> = array.into();
93        Ok(Arc::new(list))
94    }
95}
96
97/// Cast fixed size list arrays to list arrays, maintaining the lengths of the inner
98/// lists.
99///
100/// For example: `FixedSizeList<Int32, 2>([[1, 2], [3, 4]]) -> List<Int32>([[1, 2], [3, 4]])`
101pub(crate) fn cast_fixed_size_list_to_list<OffsetSize: OffsetSizeTrait>(
102    array: &dyn Array,
103    to: &FieldRef,
104    cast_options: &CastOptions,
105) -> Result<ArrayRef, ArrowError> {
106    cast_fixed_size_list_to_list_inner::<OffsetSize, false>(array, to, cast_options)
107}
108
109/// Same as [`cast_fixed_size_list_to_list`] but output list view array.
110pub(crate) fn cast_fixed_size_list_to_list_view<OffsetSize: OffsetSizeTrait>(
111    array: &dyn Array,
112    to: &FieldRef,
113    cast_options: &CastOptions,
114) -> Result<ArrayRef, ArrowError> {
115    cast_fixed_size_list_to_list_inner::<OffsetSize, true>(array, to, cast_options)
116}
117
118/// Cast list to fixed size list array. If any inner list size does not match the
119/// size of the output fixed size list array, depending on `cast_options` we either
120/// output `NULL` for that element (safe) or raise an error.
121pub(crate) fn cast_list_to_fixed_size_list<OffsetSize>(
122    array: &dyn Array,
123    field: &FieldRef,
124    size: i32,
125    cast_options: &CastOptions,
126) -> Result<ArrayRef, ArrowError>
127where
128    OffsetSize: OffsetSizeTrait,
129{
130    let array = array.as_list::<OffsetSize>();
131
132    let size_usize = usize::try_from(size).map_err(|_| {
133        ArrowError::CastError(format!(
134            "Cannot cast to FixedSizeList({size}): size must be non-negative"
135        ))
136    })?;
137    let cap = array
138        .len()
139        .checked_mul(size_usize)
140        .ok_or_else(|| ArrowError::MemoryError("fixed-size list capacity overflow".into()))?;
141
142    let mut null_builder = NullBufferBuilder::new(array.len());
143    if let Some(nulls) = array.nulls().filter(|b| b.null_count() > 0) {
144        null_builder.append_buffer(nulls);
145    } else {
146        null_builder.append_n_non_nulls(array.len());
147    }
148
149    // Whether the resulting array may contain null lists
150    let nullable = cast_options.safe || array.null_count() != 0;
151    // Nulls in FixedSizeListArray take up space and so we must pad the values
152    let values = array.values().to_data();
153    let mut mutable = MutableArrayData::try_new(vec![&values], nullable, cap)?;
154    let first_pos = array.offsets()[0].as_usize();
155    // The end position in values of the last incorrectly-sized list slice,
156    // or None if no padding has been needed (including for empty slices).
157    let mut last_pos = None;
158
159    for (idx, w) in array.offsets().windows(2).enumerate() {
160        let start_pos = w[0].as_usize();
161        let end_pos = w[1].as_usize();
162        let len = end_pos - start_pos;
163
164        if len != size_usize {
165            if cast_options.safe || array.is_null(idx) {
166                let copy_start = last_pos.unwrap_or(first_pos);
167                if copy_start != start_pos {
168                    // Extend with valid slices
169                    mutable
170                        .try_extend(0, copy_start, start_pos)
171                        .map_err(|e| ArrowError::CastError(e.to_string()))?;
172                }
173                // Pad this slice with nulls
174                mutable
175                    .try_extend_nulls(size_usize)
176                    .map_err(|e| ArrowError::CastError(e.to_string()))?;
177                null_builder.set_bit(idx, false);
178                // Set last_pos to the end of this slice's values
179                last_pos = Some(end_pos)
180            } else {
181                return Err(ArrowError::CastError(format!(
182                    "Cannot cast to FixedSizeList({size}): value at index {idx} has length {len}",
183                )));
184            }
185        }
186    }
187
188    let values = match last_pos {
189        None => array.values().slice(first_pos, cap), // All slices were the correct length
190        Some(last_pos) => {
191            if mutable.len() != cap {
192                // Remaining slices were all correct length
193                let remaining = cap - mutable.len();
194                mutable
195                    .try_extend(0, last_pos, last_pos + remaining)
196                    .map_err(|e| ArrowError::CastError(e.to_string()))?;
197            }
198            make_array(mutable.freeze())
199        }
200    };
201
202    // Cast the inner values if necessary
203    let values = cast_with_options(values.as_ref(), field.data_type(), cast_options)?;
204
205    let nulls = null_builder.build();
206    // Degenerate case where we may lose length information if there isn't a null
207    // buffer to infer length from
208    let array = if size == 0 && nulls.is_none() {
209        FixedSizeListArray::try_new_with_length(field.clone(), size, values, nulls, array.len())?
210    } else {
211        FixedSizeListArray::try_new(field.clone(), size, values, nulls)?
212    };
213    Ok(Arc::new(array))
214}
215
216/// Same as [`cast_list_to_fixed_size_list`] but for list view arrays.
217pub(crate) fn cast_list_view_to_fixed_size_list<O: OffsetSizeTrait>(
218    array: &dyn Array,
219    field: &FieldRef,
220    size: i32,
221    cast_options: &CastOptions,
222) -> Result<ArrayRef, ArrowError> {
223    let array = array.as_list_view::<O>();
224
225    let mut null_builder = NullBufferBuilder::new(array.len());
226    if let Some(nulls) = array.nulls().filter(|b| b.null_count() > 0) {
227        null_builder.append_buffer(nulls);
228    } else {
229        null_builder.append_n_non_nulls(array.len());
230    }
231
232    let nullable = cast_options.safe || array.null_count() != 0;
233    let values = array.values().to_data();
234    let size_usize = usize::try_from(size).map_err(|_| {
235        ArrowError::CastError(format!(
236            "Cannot cast to FixedSizeList({size}): size must be non-negative"
237        ))
238    })?;
239    let cap = array
240        .len()
241        .checked_mul(size_usize)
242        .ok_or_else(|| ArrowError::MemoryError("fixed-size list capacity overflow".into()))?;
243    let mut mutable = MutableArrayData::try_new(vec![&values], nullable, cap)?;
244
245    for idx in 0..array.len() {
246        let offset = array.value_offset(idx).as_usize();
247        let len = array.value_size(idx).as_usize();
248
249        if len != size_usize {
250            // Nulls in FixedSizeListArray take up space and so we must pad the values
251            if cast_options.safe || array.is_null(idx) {
252                mutable
253                    .try_extend_nulls(size_usize)
254                    .map_err(|e| ArrowError::CastError(e.to_string()))?;
255                null_builder.set_bit(idx, false);
256            } else {
257                return Err(ArrowError::CastError(format!(
258                    "Cannot cast to FixedSizeList({size}): value at index {idx} has length {len}",
259                )));
260            }
261        } else {
262            mutable
263                .try_extend(0, offset, offset + len)
264                .map_err(|e| ArrowError::CastError(e.to_string()))?;
265        }
266    }
267
268    let values = make_array(mutable.freeze());
269    let values = cast_with_options(values.as_ref(), field.data_type(), cast_options)?;
270
271    let nulls = null_builder.build();
272    // Degenerate case where we may lose length information if there isn't a null
273    // buffer to infer length from
274    let array = if size == 0 && nulls.is_none() {
275        FixedSizeListArray::try_new_with_length(field.clone(), size, values, nulls, array.len())?
276    } else {
277        FixedSizeListArray::try_new(field.clone(), size, values, nulls)?
278    };
279    Ok(Arc::new(array))
280}
281
282/// Casting between list arrays of same offset size; we cast only the inner type.
283pub(crate) fn cast_list_values<O: OffsetSizeTrait>(
284    array: &dyn Array,
285    to: &FieldRef,
286    cast_options: &CastOptions,
287) -> Result<ArrayRef, ArrowError> {
288    let list = array.as_list::<O>();
289    let (offsets, values) = list_visible_values(list);
290    let values = cast_with_options(&values, to.data_type(), cast_options)?;
291    Ok(Arc::new(GenericListArray::<O>::try_new(
292        to.clone(),
293        offsets,
294        values,
295        list.nulls().cloned(),
296    )?))
297}
298
299/// Returns the part of `list.values()` that the rows use, with the offsets
300/// adjusted to start at 0. For a sliced list, this drops the values outside
301/// the slice.
302fn list_visible_values<O: OffsetSizeTrait>(
303    list: &GenericListArray<O>,
304) -> (OffsetBuffer<O>, ArrayRef) {
305    let offsets = list.offsets();
306    let first = offsets.first();
307    let values = list
308        .values()
309        .slice(first.as_usize(), (offsets.last() - first).as_usize());
310    (offsets.clone().subtract(first), values)
311}
312
313/// Returns the values of `list_view` from the lowest offset to the highest end of
314/// its rows, and the offsets of the rows into those values. As rows can be in any
315/// order, these values can include some that no row uses.
316fn list_view_visible_values<O: OffsetSizeTrait>(
317    list_view: &GenericListViewArray<O>,
318) -> (ScalarBuffer<O>, ArrayRef) {
319    let offsets = list_view.offsets();
320    let sizes = list_view.sizes();
321    let values = list_view.values();
322
323    let mut start = usize::MAX;
324    let mut end = 0;
325    for (offset, size) in offsets.iter().zip(sizes.iter()) {
326        if size.as_usize() > 0 {
327            start = start.min(offset.as_usize());
328            end = end.max(offset.as_usize() + size.as_usize());
329        }
330    }
331    if end == 0 {
332        // All rows are empty
333        start = 0;
334    }
335
336    if start == 0 && end == values.len() {
337        return (offsets.clone(), Arc::clone(values));
338    }
339
340    // The offsets of empty rows can be outside the new range, so set them to 0
341    let offsets = offsets
342        .iter()
343        .zip(sizes.iter())
344        .map(|(offset, size)| {
345            if size.as_usize() > 0 {
346                O::usize_as(offset.as_usize() - start)
347            } else {
348                O::usize_as(0)
349            }
350        })
351        .collect();
352    (offsets, values.slice(start, end - start))
353}
354
355/// Casting between list view arrays of same offset size; we cast only the inner type.
356pub(crate) fn cast_list_view_values<O: OffsetSizeTrait>(
357    array: &dyn Array,
358    to: &FieldRef,
359    cast_options: &CastOptions,
360) -> Result<ArrayRef, ArrowError> {
361    let list = array.as_list_view::<O>();
362    let (offsets, values) = list_view_visible_values(list);
363    let values = cast_with_options(&values, to.data_type(), cast_options)?;
364    Ok(Arc::new(GenericListViewArray::<O>::try_new(
365        to.clone(),
366        offsets,
367        list.sizes().clone(),
368        values,
369        list.nulls().cloned(),
370    )?))
371}
372
373/// Casting between list arrays of different offset size (e.g. List -> LargeList)
374pub(crate) fn cast_list<I: OffsetSizeTrait, O: OffsetSizeTrait>(
375    array: &dyn Array,
376    field: &FieldRef,
377    cast_options: &CastOptions,
378) -> Result<ArrayRef, ArrowError> {
379    let list = array.as_list::<I>();
380    let (offsets, values) = list_visible_values(list);
381    let nulls = list.nulls().cloned();
382
383    if offsets.last().as_usize() > O::MAX_OFFSET {
384        return Err(ArrowError::ComputeError(format!(
385            "Offset overflow when casting from {} to {}",
386            array.data_type(),
387            field.data_type()
388        )));
389    }
390
391    // Recursively cast values
392    let values = cast_with_options(&values, field.data_type(), cast_options)?;
393    let offsets: Vec<_> = offsets.iter().map(|x| O::usize_as(x.as_usize())).collect();
394
395    // Safety: valid offsets and checked for overflow
396    let offsets = unsafe { OffsetBuffer::new_unchecked(offsets.into()) };
397
398    Ok(Arc::new(GenericListArray::<O>::try_new(
399        field.clone(),
400        offsets,
401        values,
402        nulls,
403    )?))
404}
405
406/// Casting list view arrays to list.
407pub(crate) fn cast_list_view_to_list<I, O>(
408    array: &dyn Array,
409    to: &FieldRef,
410    cast_options: &CastOptions,
411) -> Result<ArrayRef, ArrowError>
412where
413    I: OffsetSizeTrait,
414    // We need ArrowPrimitiveType here to be able to create indices array for the
415    // take kernel.
416    O: ArrowPrimitiveType,
417    O::Native: OffsetSizeTrait,
418{
419    let list_view = array.as_list_view::<I>();
420    let list_view_offsets = list_view.offsets();
421    let sizes = list_view.sizes();
422
423    let mut take_indices: Vec<O::Native> = Vec::with_capacity(list_view.values().len());
424    let mut offsets: Vec<O::Native> = Vec::with_capacity(list_view.len() + 1);
425    use num_traits::Zero;
426    offsets.push(O::Native::zero());
427
428    for i in 0..list_view.len() {
429        if list_view.is_null(i) {
430            offsets.push(O::Native::usize_as(take_indices.len()));
431            continue;
432        }
433
434        let offset = list_view_offsets[i].as_usize();
435        let size = sizes[i].as_usize();
436
437        for value_index in offset..offset + size {
438            take_indices.push(O::Native::usize_as(value_index));
439        }
440
441        // Must guard all cases since ListView<i32> can overflow List<i32>
442        // e.g. if offsets of [0, 0, 0] and sizes [i32::MAX, i32::MAX, i32::MAX]
443        if take_indices.len() > O::Native::MAX_OFFSET {
444            return Err(ArrowError::ComputeError(format!(
445                "Offset overflow when casting from {} to {}",
446                array.data_type(),
447                to.data_type()
448            )));
449        }
450        offsets.push(O::Native::usize_as(take_indices.len()));
451    }
452
453    // Form a contiguous values array
454    let take_indices = PrimitiveArray::<O>::from_iter_values(take_indices);
455    let values = arrow_select::take::take(list_view.values(), &take_indices, None)?;
456    let values = cast_with_options(&values, to.data_type(), cast_options)?;
457
458    Ok(Arc::new(GenericListArray::<O::Native>::try_new(
459        to.clone(),
460        OffsetBuffer::new(offsets.into()),
461        values,
462        list_view.nulls().cloned(),
463    )?))
464}
465
466/// Casting between list view arrays of different offset size (e.g. ListView -> LargeListView)
467pub(crate) fn cast_list_view<I: OffsetSizeTrait, O: OffsetSizeTrait>(
468    array: &dyn Array,
469    to_field: &FieldRef,
470    cast_options: &CastOptions,
471) -> Result<ArrayRef, ArrowError> {
472    let list_view = array.as_list_view::<I>();
473    let (offsets, values) = list_view_visible_values(list_view);
474
475    // Recursively cast values
476    let values = cast_with_options(&values, to_field.data_type(), cast_options)?;
477
478    let offsets = offsets
479        .iter()
480        .map(|offset| {
481            let offset = offset.as_usize();
482            if offset > O::MAX_OFFSET {
483                return Err(ArrowError::ComputeError(format!(
484                    "Offset overflow when casting from {} to {}",
485                    array.data_type(),
486                    to_field.data_type()
487                )));
488            }
489            Ok(O::usize_as(offset))
490        })
491        .collect::<Result<Vec<O>, _>>()?;
492    let sizes = list_view
493        .sizes()
494        .iter()
495        .map(|size| {
496            let size = size.as_usize();
497            if size > O::MAX_OFFSET {
498                return Err(ArrowError::ComputeError(format!(
499                    "Offset overflow when casting from {} to {}",
500                    array.data_type(),
501                    to_field.data_type()
502                )));
503            }
504            Ok(O::usize_as(size))
505        })
506        .collect::<Result<Vec<O>, _>>()?;
507    Ok(Arc::new(GenericListViewArray::<O>::try_new(
508        to_field.clone(),
509        offsets.into(),
510        sizes.into(),
511        values,
512        list_view.nulls().cloned(),
513    )?))
514}
515
516/// Casting list arrays to list view.
517pub(crate) fn cast_list_to_list_view<I: OffsetSizeTrait, O: OffsetSizeTrait>(
518    array: &dyn Array,
519    to_field: &FieldRef,
520    cast_options: &CastOptions,
521) -> Result<ArrayRef, ArrowError> {
522    let list = array.as_list::<I>();
523    let (offsets, values) = list_visible_values(list);
524    let nulls = list.nulls().cloned();
525
526    let len = offsets.len() - 1;
527    let mut sizes = Vec::with_capacity(len);
528    let mut view_offsets = Vec::with_capacity(len);
529    for (i, offset) in offsets.iter().enumerate().take(len) {
530        let offset = offset.as_usize();
531        let size = offsets[i + 1].as_usize() - offset;
532
533        if offset > O::MAX_OFFSET || size > O::MAX_OFFSET {
534            return Err(ArrowError::ComputeError(format!(
535                "Offset overflow when casting from {} to {}",
536                array.data_type(),
537                to_field.data_type()
538            )));
539        }
540
541        view_offsets.push(O::usize_as(offset));
542        sizes.push(O::usize_as(size));
543    }
544    let values = cast_with_options(&values, to_field.data_type(), cast_options)?;
545    let array = GenericListViewArray::<O>::try_new(
546        to_field.clone(),
547        view_offsets.into(),
548        sizes.into(),
549        values,
550        nulls,
551    )?;
552    Ok(Arc::new(array))
553}
arrow_cast/cast/list.rs

arrow_cast/cast/
list.rs
