arrow_data/transform/
union.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
9fn nullable_union_field(fields: &arrow_schema::UnionFields) -> Result<(usize, i8), ArrowError> {
10    fields
11        .iter()
12        .enumerate()
13        .find_map(|(index, (type_id, field))| field.is_nullable().then_some((index, type_id)))
14        .ok_or_else(|| {
15            ArrowError::InvalidArgumentError(if fields.is_empty() {
16                "Cannot extend a union with no fields with nulls".to_owned()
17            } else {
18                "Cannot extend a union with no nullable fields with nulls".to_owned()
19            })
20        })
21}
22
23pub(super) fn build_extend_sparse(array: &ArrayData) -> Extend<'_> {
24    // `buffer` already applies the parent offset to the type ids
25    let type_ids = array.buffer::<i8>(0);
26    // but a sparse union's children are not sliced by the parent offset
27    let offset = array.offset();
28
29    Box::new(
30        move |mutable: &mut _MutableArrayData, index: usize, start: usize, len: usize| {
31            // extends type_ids
32            mutable
33                .buffer1
34                .extend_from_slice(&type_ids[start..start + len]);
35
36            let child_start = start.checked_add(offset).ok_or_else(|| {
37                ArrowError::InvalidArgumentError(format!(
38                    "sparse union offset {offset} with start {start} overflows usize"
39                ))
40            })?;
41            let child_end = child_start.checked_add(len).ok_or_else(|| {
42                ArrowError::InvalidArgumentError(format!(
43                    "sparse union start {child_start} with length {len} overflows usize"
44                ))
45            })?;
46            for child in &mut mutable.child_data {
47                child.try_extend(index, child_start, child_end)?;
48            }
49            Ok(())
50        },
51    )
52}
53
54pub(super) fn build_extend_dense(array: &ArrayData) -> Extend<'_> {
55    let type_ids = array.buffer::<i8>(0);
56    let offsets = array.buffer::<i32>(1);
57    let arrow_schema::DataType::Union(src_fields, _) = array.data_type() else {
58        unreachable!();
59    };
60
61    Box::new(
62        move |mutable: &mut _MutableArrayData, index: usize, start: usize, len: usize| {
63            // extends type_ids
64            mutable
65                .buffer1
66                .extend_from_slice(&type_ids[start..start + len]);
67
68            for i in start..start + len {
69                let type_id = type_ids[i];
70                let child_index = src_fields
71                    .iter()
72                    .position(|(r, _)| r == type_id)
73                    .expect("invalid union type ID");
74                let src_offset = offsets[i] as usize;
75                let child_data = &mut mutable.child_data[child_index];
76                let dst_offset = child_data.len();
77
78                // Extend offsets
79                mutable.buffer2.push(dst_offset as i32);
80                mutable.child_data[child_index].try_extend(index, src_offset, src_offset + 1)?;
81            }
82            Ok(())
83        },
84    )
85}
86
87pub(super) fn extend_nulls_dense(
88    mutable: &mut _MutableArrayData,
89    len: usize,
90) -> Result<(), ArrowError> {
91    let DataType::Union(fields, _) = &mutable.data_type else {
92        unreachable!()
93    };
94    let (child_index, null_type_id) = nullable_union_field(fields)?;
95
96    // Extend type_ids buffer
97    mutable
98        .buffer1
99        .try_extend_from_slice(&vec![null_type_id; len])
100        .map_err(|e| ArrowError::MemoryError(e.to_string()))?;
101
102    // Dense: point into a nullable child and extend its values with nulls.
103    let child_offset = mutable.child_data[child_index].len();
104    let (start, end) = (child_offset as i32, (child_offset + len) as i32);
105    mutable.buffer2.extend(start..end);
106    mutable.child_data[child_index].try_extend_nulls(len)?;
107    Ok(())
108}
109
110pub(super) fn extend_nulls_sparse(
111    mutable: &mut _MutableArrayData,
112    len: usize,
113) -> Result<(), ArrowError> {
114    let DataType::Union(fields, _) = &mutable.data_type else {
115        unreachable!()
116    };
117    let (_, null_type_id) = nullable_union_field(fields)?;
118
119    // Extend type_ids buffer
120    mutable
121        .buffer1
122        .try_extend_from_slice(&vec![null_type_id; len])
123        .map_err(|e| ArrowError::MemoryError(e.to_string()))?;
124
125    // Sparse: extend nulls in ALL children
126    for child in &mut mutable.child_data {
127        child.try_extend_nulls(len)?;
128    }
129    Ok(())
130}
arrow_data/transform/union.rs

arrow_data/transform/
union.rs
