arrow_data/transform/
variable_size.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 arrow_buffer::{ArrowNativeType, MutableBuffer};
7use num_integer::Integer;
8use num_traits::{AsPrimitive, CheckedAdd};
9
10use super::{
11    _MutableArrayData, Extend,
12    utils::{get_last_offset, try_extend_offsets},
13};
14
15#[inline]
16fn extend_offset_values<T: ArrowNativeType + AsPrimitive<usize>>(
17    buffer: &mut MutableBuffer,
18    offsets: &[T],
19    values: &[u8],
20    start: usize,
21    len: usize,
22) {
23    let start_values = offsets[start].as_();
24    let end_values: usize = offsets[start + len].as_();
25    let new_values = &values[start_values..end_values];
26    buffer.extend_from_slice(new_values);
27}
28
29pub(super) fn build_extend<T: ArrowNativeType + Integer + CheckedAdd + AsPrimitive<usize>>(
30    array: &ArrayData,
31) -> Extend<'_> {
32    let offsets = array.buffer::<T>(0);
33    let values = array.buffers()[1].as_slice();
34    Box::new(
35        move |mutable: &mut _MutableArrayData, _, start: usize, len: usize| {
36            let offset_buffer = &mut mutable.buffer1;
37            let values_buffer = &mut mutable.buffer2;
38
39            // this is safe due to how offset is built. See details on `get_last_offset`
40            let last_offset = unsafe { get_last_offset(offset_buffer) };
41
42            try_extend_offsets::<T>(offset_buffer, last_offset, &offsets[start..start + len + 1])?;
43            // values
44            extend_offset_values::<T>(values_buffer, offsets, values, start, len);
45            Ok(())
46        },
47    )
48}
49
50pub(super) fn extend_nulls<T: ArrowNativeType>(
51    mutable: &mut _MutableArrayData,
52    len: usize,
53) -> Result<(), arrow_schema::ArrowError> {
54    let offset_buffer = &mut mutable.buffer1;
55
56    // this is safe due to how offset is built. See details on `get_last_offset`
57    let last_offset: T = unsafe { get_last_offset(offset_buffer) };
58
59    (0..len).for_each(|_| offset_buffer.push(last_offset));
60    Ok(())
61}
arrow_data/transform/variable_size.rs

arrow_data/transform/
variable_size.rs
