lance_arrow/
struct.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright The Lance Authors
3
4//! Extension to arrow struct arrays
5
6use arrow_array::{Array, StructArray, cast::AsArray, make_array};
7use arrow_buffer::NullBuffer;
8use arrow_data::{ArrayData, ArrayDataBuilder};
9use arrow_schema::{ArrowError, DataType};
10
11pub trait StructArrayExt {
12    /// Removes the offset / length of the struct array by pushing it into the children
13    ///
14    /// In arrow-rs when slice is called it recursively slices the children.
15    /// In arrow-cpp when slice is called it just sets the offset/length of
16    ///   the struct array and leaves the children as-is
17    ///
18    /// Both are legal approaches (╥﹏╥)
19    ///
20    /// This method helps reduce complexity by folding into the arrow-rs approach
21    fn normalize_slicing(&self) -> Result<Self, ArrowError>
22    where
23        Self: Sized;
24
25    /// Structs are allowed to mask valid items.  For example, a struct array might be:
26    ///
27    /// [ {"items": [1, 2, 3]}, NULL, {"items": NULL}]
28    ///
29    /// However, the underlying items array might be: [[1, 2, 3], [4, 5], NULL]
30    ///
31    /// The [4, 5] list is masked out because the struct array is null.
32    ///
33    /// The struct validity would be [true, false, true] and the list validity would be [true, true, false]
34    ///
35    /// This method pushes nulls down into all children.  In the above example the list validity would become
36    /// [true, false, false].
37    ///
38    /// This method is not recursive.  If a child is a struct array it will not push that child's nulls down.
39    ///
40    /// This method does not remove garbage lists.  It only updates the validity so a future call to
41    /// [crate::list::ListArrayExt::filter_garbage_nulls] will remove the garbage lists (without
42    /// this pushdown it would not)
43    fn pushdown_nulls(&self) -> Result<Self, ArrowError>
44    where
45        Self: Sized;
46}
47
48fn normalized_struct_array_data(data: ArrayData) -> Result<ArrayData, ArrowError> {
49    let parent_offset = data.offset();
50    let parent_len = data.len();
51    let modified_children = data
52        .child_data()
53        .iter()
54        .map(|d| {
55            let d = normalized_struct_array_data(d.clone())?;
56            let offset = d.offset();
57            let len = d.len();
58            if len < parent_len + parent_offset {
59                return Err(ArrowError::InvalidArgumentError(format!(
60                    "Child array {} has length {} which is less than the parent length {} plus the parent offset {}",
61                    d.data_type(),
62                    len,
63                    parent_len,
64                    parent_offset
65                )));
66            }
67            let new_offset = offset + parent_offset;
68            d.into_builder().offset(new_offset)
69                .len(parent_len)
70                .build()
71        })
72        .collect::<Result<Vec<_>, _>>()?;
73    ArrayDataBuilder::new(data.data_type().clone())
74        .len(parent_len)
75        .offset(0)
76        .buffers(data.buffers().to_vec())
77        .child_data(modified_children)
78        .build()
79}
80
81impl StructArrayExt for StructArray {
82    fn normalize_slicing(&self) -> Result<Self, ArrowError>
83    where
84        Self: Sized,
85    {
86        if self.offset() == 0 && self.columns().iter().all(|c| c.len() == self.len()) {
87            return Ok(self.clone());
88        }
89
90        let data = normalized_struct_array_data(self.to_data())?;
91        Ok(Self::from(data))
92    }
93
94    fn pushdown_nulls(&self) -> Result<Self, ArrowError>
95    where
96        Self: Sized,
97    {
98        let Some(validity) = self.nulls() else {
99            return Ok(self.clone());
100        };
101        let data = self.to_data();
102        let children = data
103            .child_data()
104            .iter()
105            .map(|c| {
106                if c.data_type() == &DataType::Null {
107                    // A Null array is already null in every slot, and arrow-rs rejects
108                    // a validity buffer on it.
109                    Ok(c.clone())
110                } else if let Some(child_validity) = c.nulls() {
111                    let new_validity = child_validity.inner() & validity.inner();
112                    c.clone()
113                        .into_builder()
114                        .nulls(Some(NullBuffer::from(new_validity)))
115                        .build()
116                } else {
117                    Ok(c.clone()
118                        .into_builder()
119                        .nulls(Some(validity.clone()))
120                        .build()?)
121                }
122            })
123            .collect::<Result<Vec<_>, _>>()?;
124        let arr = make_array(data.into_builder().child_data(children).build()?);
125        Ok(arr.as_struct().clone())
126    }
127}
128
129#[cfg(test)]
130mod tests {
131    use arrow_array::{
132        Array, Int32Array, NullArray, StructArray, cast::AsArray, make_array, types::Int32Type,
133    };
134    use arrow_buffer::NullBuffer;
135    use arrow_schema::{DataType, Field, Fields};
136    use std::sync::Arc;
137
138    use crate::r#struct::StructArrayExt;
139
140    #[test]
141    fn test_normalize_slicing_no_offset() {
142        let x = Int32Array::from(vec![1, 2, 3]);
143        let y = Int32Array::from(vec![4, 5, 6]);
144        let struct_array = StructArray::new(
145            Fields::from(vec![
146                Field::new("x", DataType::Int32, true),
147                Field::new("y", DataType::Int32, true),
148            ]),
149            vec![Arc::new(x), Arc::new(y)],
150            None,
151        );
152
153        let normalized = struct_array.normalize_slicing().unwrap();
154        assert_eq!(normalized, struct_array);
155    }
156
157    #[test]
158    fn test_arrow_rs_slicing() {
159        let x = Int32Array::from(vec![1, 2, 3, 4]);
160        let y = Int32Array::from(vec![5, 6, 7, 8]);
161        let struct_array = StructArray::new(
162            Fields::from(vec![
163                Field::new("x", DataType::Int32, true),
164                Field::new("y", DataType::Int32, true),
165            ]),
166            vec![Arc::new(x), Arc::new(y)],
167            None,
168        );
169
170        // Slicing with arrow-rs propagates the slicing to the children so there should
171        // be no change needed to the struct array
172        let sliced = struct_array.slice(1, 2);
173        let normalized = sliced.normalize_slicing().unwrap();
174
175        assert_eq!(normalized, sliced);
176    }
177
178    #[test]
179    fn test_arrow_cpp_slicing() {
180        let x = Int32Array::from(vec![1, 2, 3, 4]);
181        let y = Int32Array::from(vec![5, 6, 7, 8]);
182        let struct_array = StructArray::new(
183            Fields::from(vec![
184                Field::new("x", DataType::Int32, true),
185                Field::new("y", DataType::Int32, true),
186            ]),
187            vec![Arc::new(x), Arc::new(y)],
188            None,
189        );
190
191        let data = struct_array.to_data();
192        let sliced = data.into_builder().offset(1).len(2).build().unwrap();
193        let sliced = make_array(sliced);
194        let normalized = sliced.as_struct().clone().normalize_slicing().unwrap();
195
196        assert_eq!(normalized, struct_array.slice(1, 2));
197    }
198
199    #[test]
200    fn test_pushdown_nulls_with_null_typed_child() {
201        // `pa.Table.from_pylist([{"s": {"a": 1, "b": None}}, {"s": None}])` infers a
202        // struct with a Null-typed child. Arrow forbids a validity buffer on a Null
203        // array, so the struct's nulls must not be pushed into it.
204        let struct_array = StructArray::new(
205            Fields::from(vec![
206                Field::new("a", DataType::Int32, true),
207                Field::new("b", DataType::Null, true),
208            ]),
209            vec![
210                Arc::new(Int32Array::from(vec![1, 2, 3])),
211                Arc::new(NullArray::new(3)),
212            ],
213            Some(NullBuffer::from(vec![true, false, true])),
214        );
215
216        let pushed = struct_array.pushdown_nulls().unwrap();
217
218        assert_eq!(pushed.nulls(), struct_array.nulls());
219        assert_eq!(
220            pushed.column(0).as_primitive::<Int32Type>(),
221            &Int32Array::from(vec![Some(1), None, Some(3)])
222        );
223        assert_eq!(pushed.column(1).data_type(), &DataType::Null);
224        assert_eq!(pushed.column(1).logical_null_count(), 3);
225    }
226}
lance_arrow/struct.rs

lance_arrow/
struct.rs
