Skip to main content

arrow_data/transform/
list_view.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::ArrayData;
6use crate::transform::_MutableArrayData;
7use arrow_buffer::ArrowNativeType;
8use arrow_schema::ArrowError;
9use num_integer::Integer;
10use num_traits::CheckedAdd;
11
12pub(super) fn build_extend<T: ArrowNativeType + Integer + CheckedAdd>(
13    array: &ArrayData,
14) -> crate::transform::Extend<'_> {
15    let offsets = array.buffer::<T>(0);
16    let sizes = array.buffer::<T>(1);
17    Box::new(
18        move |mutable: &mut _MutableArrayData, index: usize, start: usize, len: usize| {
19            let mut new_offset = T::usize_as(mutable.child_data[0].len());
20
21            for i in start..start + len {
22                mutable.buffer1.push(new_offset);
23                mutable.buffer2.push(sizes[i]);
24                new_offset = new_offset.checked_add(&sizes[i]).ok_or_else(|| {
25                    ArrowError::InvalidArgumentError(
26                        "offset overflow: data exceeds the capacity of the offset type. \
27                         Try splitting into smaller batches or using a larger type \
28                         (e.g. LargeListView instead of ListView)"
29                            .to_string(),
30                    )
31                })?;
32
33                let size = sizes[i].as_usize();
34                if size > 0 {
35                    let child_start = offsets[i].as_usize();
36                    mutable.child_data[0].try_extend(index, child_start, child_start + size)?;
37                }
38            }
39            Ok(())
40        },
41    )
42}
43
44pub(super) fn extend_nulls<T: ArrowNativeType>(
45    mutable: &mut _MutableArrayData,
46    len: usize,
47) -> Result<(), ArrowError> {
48    let offset_buffer = &mut mutable.buffer1;
49    let sizes_buffer = &mut mutable.buffer2;
50
51    // We push 0 as a placeholder for NULL values in both the offsets and sizes
52    (0..len).for_each(|_| offset_buffer.push(T::default()));
53    (0..len).for_each(|_| sizes_buffer.push(T::default()));
54    Ok(())
55}