Skip to main content

arrow_data/transform/
structure.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::{_MutableArrayData, Extend};
6use crate::ArrayData;
7use arrow_schema::{ArrowError, DataType};
8
9pub(super) fn build_extend(array: &ArrayData) -> Extend<'_> {
10    let offset = array.offset();
11    Box::new(
12        move |mutable: &mut _MutableArrayData, index: usize, start: usize, len: usize| {
13            let start = start.checked_add(offset).ok_or_else(|| {
14                ArrowError::InvalidArgumentError(format!(
15                    "struct offset {offset} with start {start} overflows usize"
16                ))
17            })?;
18            let end = start.checked_add(len).ok_or_else(|| {
19                ArrowError::InvalidArgumentError(format!(
20                    "struct start {start} with length {len} overflows usize"
21                ))
22            })?;
23            for (col_idx, child) in mutable.child_data.iter_mut().enumerate() {
24                child
25                    .try_extend(index, start, end)
26                    .map_err(|e| wrap_column_error(e, col_idx, &mutable.data_type))?
27            }
28            Ok(())
29        },
30    )
31}
32
33pub(super) fn extend_nulls(mutable: &mut _MutableArrayData, len: usize) -> Result<(), ArrowError> {
34    for (col_idx, child) in mutable.child_data.iter_mut().enumerate() {
35        child
36            .try_extend_nulls(len)
37            .map_err(|e| wrap_column_error(e, col_idx, &mutable.data_type))?;
38    }
39    Ok(())
40}
41
42fn wrap_column_error(e: ArrowError, col_idx: usize, data_type: &DataType) -> ArrowError {
43    let name_ctx = if let DataType::Struct(fields) = data_type {
44        fields
45            .get(col_idx)
46            .map(|f| format!(" (\"{}\")", f.name()))
47            .unwrap_or_default()
48    } else {
49        String::new()
50    };
51    ArrowError::InvalidArgumentError(format!("struct column {col_idx}{name_ctx} failed: {e}"))
52}