Skip to main content

arrow_data/transform/
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 super::{
6    _MutableArrayData, Extend,
7    utils::{get_last_offset, try_extend_offsets},
8};
9use crate::ArrayData;
10use arrow_buffer::ArrowNativeType;
11use arrow_schema::ArrowError;
12use num_integer::Integer;
13use num_traits::CheckedAdd;
14
15pub(super) fn build_extend<T: ArrowNativeType + Integer + CheckedAdd>(
16    array: &ArrayData,
17) -> Extend<'_> {
18    let offsets = array.buffer::<T>(0);
19    Box::new(
20        move |mutable: &mut _MutableArrayData, index: usize, start: usize, len: usize| {
21            let offset_buffer = &mut mutable.buffer1;
22
23            // this is safe due to how offset is built. See details on `get_last_offset`
24            let last_offset: T = unsafe { get_last_offset(offset_buffer) };
25
26            // offsets
27            try_extend_offsets::<T>(offset_buffer, last_offset, &offsets[start..start + len + 1])?;
28
29            mutable.child_data[0].try_extend(
30                index,
31                offsets[start].as_usize(),
32                offsets[start + len].as_usize(),
33            )
34        },
35    )
36}
37
38pub(super) fn extend_nulls<T: ArrowNativeType>(
39    mutable: &mut _MutableArrayData,
40    len: usize,
41) -> Result<(), ArrowError> {
42    let offset_buffer = &mut mutable.buffer1;
43
44    // this is safe due to how offset is built. See details on `get_last_offset`
45    let last_offset: T = unsafe { get_last_offset(offset_buffer) };
46
47    (0..len).for_each(|_| offset_buffer.push(last_offset));
48    Ok(())
49}