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

arrow_cast/cast/
list.rs
