arrow_cast/cast/
list_view.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/// Helper function to cast a list view to a list
21pub(crate) fn cast_list_view_to_list<O: OffsetSizeTrait>(
22    array: &dyn Array,
23    to: &FieldRef,
24    cast_options: &CastOptions,
25) -> Result<ArrayRef, ArrowError> {
26    let list_view = array.as_list_view::<O>();
27    let list_view_offsets = list_view.offsets();
28    let sizes = list_view.sizes();
29    let source_values = list_view.values();
30
31    // Construct the indices and offsets for the new list array by iterating over the list view subarrays
32    let mut indices = Vec::with_capacity(list_view.values().len());
33    let mut offsets = Vec::with_capacity(list_view.len() + 1);
34    // Add the offset for the first subarray
35    offsets.push(O::usize_as(0));
36    for i in 0..list_view.len() {
37        // For each subarray, add the indices of the values to take
38        let offset = list_view_offsets[i].as_usize();
39        let size = sizes[i].as_usize();
40        let end = offset + size;
41        for j in offset..end {
42            indices.push(j as i32);
43        }
44        // Add the offset for the next subarray
45        offsets.push(O::usize_as(indices.len()));
46    }
47
48    // Take the values from the source values using the indices, creating a new array
49    let values = arrow_select::take::take(source_values, &Int32Array::from(indices), None)?;
50
51    // Cast the values to the target data type
52    let values = cast_with_options(&values, to.data_type(), cast_options)?;
53
54    Ok(Arc::new(GenericListArray::<O>::try_new(
55        to.clone(),
56        OffsetBuffer::new(offsets.into()),
57        values,
58        list_view.nulls().cloned(),
59    )?))
60}
61
62pub(crate) fn cast_list_view<I: OffsetSizeTrait, O: OffsetSizeTrait>(
63    array: &dyn Array,
64    to_field: &FieldRef,
65    cast_options: &CastOptions,
66) -> Result<ArrayRef, ArrowError> {
67    let list_view = array.as_list_view::<I>();
68    let (_field, offsets, sizes, values, nulls) = list_view.clone().into_parts();
69
70    // Recursively cast values
71    let values = cast_with_options(&values, to_field.data_type(), cast_options)?;
72
73    let new_offsets: Vec<_> = offsets.iter().map(|x| O::usize_as(x.as_usize())).collect();
74    let new_sizes: Vec<_> = sizes.iter().map(|x| O::usize_as(x.as_usize())).collect();
75    Ok(Arc::new(GenericListViewArray::<O>::try_new(
76        to_field.clone(),
77        new_offsets.into(),
78        new_sizes.into(),
79        values,
80        nulls,
81    )?))
82}
83
84pub(crate) fn cast_list_to_list_view<OffsetSize>(array: &dyn Array) -> Result<ArrayRef, ArrowError>
85where
86    OffsetSize: OffsetSizeTrait,
87{
88    let list = array.as_list::<OffsetSize>();
89    let list_view: GenericListViewArray<OffsetSize> = list.clone().into();
90    Ok(Arc::new(list_view))
91}