This is an automated email from the ASF dual-hosted git repository. github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-11433-1e6aef1c07dc1e36758374ab28bdafb4fd93bead in repository https://gitbox.apache.org/repos/asf/arrow-rs.git
commit fa96419f5cff79d98b7ae77c9fabb62dfd173346 Author: memosmorgan <[email protected]> AuthorDate: Fri Oct 9 20:31:26 2026 +0000 Add `ArrayData::sliced` to avoid allocations when slicing owned ArrayData (#11433) # Which issue does this PR close? - Closes #9140 Replaces #11432. # Rationale for this change Owned `ArrayData` can be sliced without allocating new buffer or child vectors. AI assistance: OpenAI Codex generated the implementation and tests, ran the local checks, and reviewed the diff. # What changes are included in this PR? Add `ArrayData::sliced`, which consumes the array and reuses its allocations, including nested struct children. Use it when converting owned data to `StructArray`. # Are these changes tested? Regression tests cover equivalence with borrowed slicing, offsets, nested structs, lists, nulls, empty/end slices, and bounds. They also check that vector allocations are retained. A local release-mode allocation counter measured the following for 1,024-element inputs sliced at offset 1 for 512 elements, excluding input construction: | Data | `slice` allocations | `sliced` allocations | | --- | ---: | ---: | | Int32 | 1 | 0 | | Nested Struct | 3 | 0 | | List<Int32> | 3 | 0 | # Are there any significant user-facing changes? No breaking changes. Borrowed `slice` retains its existing behavior. Co-authored-by: Andrew Lamb <[email protected]> --- arrow-array/src/array/struct_array.rs | 2 +- arrow-data/src/data.rs | 200 ++++++++++++++++++++++++++++++---- 2 files changed, 181 insertions(+), 21 deletions(-) diff --git a/arrow-array/src/array/struct_array.rs b/arrow-array/src/array/struct_array.rs index 592126318a..266d1abdea 100644 --- a/arrow-array/src/array/struct_array.rs +++ b/arrow-array/src/array/struct_array.rs @@ -435,7 +435,7 @@ impl From<ArrayData> for StructArray { .into_iter() .map(|cd| { if parent_offset != 0 || parent_len != cd.len() { - make_array(cd.slice(parent_offset, parent_len)) + make_array(cd.sliced(parent_offset, parent_len)) } else { make_array(cd) } diff --git a/arrow-data/src/data.rs b/arrow-data/src/data.rs index a8a6970125..82dfefcac0 100644 --- a/arrow-data/src/data.rs +++ b/arrow-data/src/data.rs @@ -688,6 +688,22 @@ impl ArrayData { /// /// Panics if `offset + length` overflows or is greater than `self.len()`. pub fn slice(&self, offset: usize, length: usize) -> ArrayData { + self.clone().sliced(offset, length) + } + + /// Creates a zero-copy slice of this array, reusing its buffer and child allocations. + /// + /// This is equivalent to [`Self::slice`], but consumes `self`. + /// + /// # Panics + /// + /// Panics if `offset + length` overflows or is greater than `self.len()`. + pub fn sliced(mut self, offset: usize, length: usize) -> Self { + self.slice_in_place(offset, length); + self + } + + fn slice_in_place(&mut self, offset: usize, length: usize) { let end = offset .checked_add(length) .expect("offset + length overflow"); @@ -699,29 +715,17 @@ impl ArrayData { // the slice to both would count it twice, so the cumulative offset // goes to the children and this array's offset is reset to 0. let child_offset = self.offset + offset; - ArrayData { - data_type: self.data_type().clone(), - len: length, - offset: 0, - buffers: self.buffers.clone(), - child_data: self - .child_data() - .iter() - .map(|data| data.slice(child_offset, length)) - .collect(), - // `nulls` belongs to this array rather than to the children, so - // it is sliced by `offset` alone. - nulls: self.nulls.as_ref().map(|x| x.slice(offset, length)), + for child in &mut self.child_data { + child.slice_in_place(child_offset, length); } + self.offset = 0; } else { - let mut new_data = self.clone(); - - new_data.len = length; - new_data.offset = offset + self.offset; - new_data.nulls = self.nulls.as_ref().map(|x| x.slice(offset, length)); - - new_data + self.offset += offset; } + self.len = length; + // `nulls` belongs to this array rather than to the children, so + // it is sliced by `offset` alone. + self.nulls = self.nulls.as_ref().map(|x| x.slice(offset, length)); } /// Returns the `buffer` as a slice of type `T` starting at self.offset @@ -2803,6 +2807,162 @@ mod tests { sliced.slice(1, usize::MAX); } + #[test] + fn test_sliced() { + for nulls in [ + None, + Some(NullBuffer::from(vec![true, false, true, false, true])), + ] { + let data = ArrayData::builder(DataType::Int32) + .len(5) + .offset(1) + .add_buffer(make_i32_buffer(6)) + .nulls(nulls) + .build() + .unwrap(); + + for offset in 0..=data.len() { + for length in 0..=data.len() - offset { + let owned = data.clone(); + let buffers = owned.buffers().as_ptr(); + let sliced = owned.sliced(offset, length); + assert_eq!(sliced, data.slice(offset, length)); + assert_eq!(sliced.len(), length); + assert_eq!(sliced.offset(), 1 + offset); + assert_eq!(sliced.buffers().as_ptr(), buffers); + assert_eq!(sliced.buffers()[0].as_ptr(), data.buffers()[0].as_ptr()); + if let Some(nulls) = data.nulls() { + assert_eq!( + sliced.nulls().unwrap().iter().collect::<Vec<_>>(), + nulls.iter().skip(offset).take(length).collect::<Vec<_>>() + ); + } + sliced.validate_full().unwrap(); + } + } + } + + let empty = ArrayData::new_empty(&DataType::Int32); + assert_eq!(empty.clone().sliced(0, 0), empty); + } + + #[test] + fn test_sliced_nested_struct() { + let child = ArrayData::builder(DataType::Int32) + .len(8) + .offset(1) + .add_buffer(make_i32_buffer(9)) + .nulls(Some(NullBuffer::from(vec![ + true, false, true, false, true, true, false, true, + ]))) + .build() + .unwrap(); + let inner = ArrayData::builder(DataType::Struct(Fields::from(vec![Field::new( + "x", + DataType::Int32, + true, + )]))) + .len(6) + .offset(1) + .nulls(Some(NullBuffer::from(vec![ + true, false, true, true, false, true, + ]))) + .add_child_data(child) + .build() + .unwrap(); + let data = ArrayData::builder(DataType::Struct(Fields::from(vec![Field::new( + "inner", + inner.data_type().clone(), + true, + )]))) + .len(4) + .offset(1) + .nulls(Some(NullBuffer::from(vec![false, true, false, true]))) + .add_child_data(inner) + .build() + .unwrap(); + + for (offset, length) in [(0, 4), (1, 2), (4, 0), (0, 0)] { + let owned = data.clone(); + let children = owned.child_data().as_ptr(); + let grandchildren = owned.child_data()[0].child_data().as_ptr(); + let buffers = owned.child_data()[0].child_data()[0].buffers().as_ptr(); + let sliced = owned.sliced(offset, length); + let inner = &sliced.child_data()[0]; + let child = &inner.child_data()[0]; + assert_eq!(sliced, data.slice(offset, length)); + assert_eq!(sliced.offset(), 0); + assert_eq!(inner.offset(), 0); + assert_eq!(child.offset(), 3 + offset); + assert_eq!(sliced.child_data().as_ptr(), children); + assert_eq!(inner.child_data().as_ptr(), grandchildren); + assert_eq!(child.buffers().as_ptr(), buffers); + sliced.validate_full().unwrap(); + } + + let sliced = data.sliced(1, 2).sliced(1, 1); + assert_eq!(sliced.null_count(), 1); + assert_eq!(sliced.child_data()[0].null_count(), 0); + assert_eq!(sliced.child_data()[0].child_data()[0].offset(), 5); + assert_eq!(sliced.child_data()[0].child_data()[0].null_count(), 0); + } + + #[test] + fn test_sliced_list_preserves_children() { + let child = ArrayData::builder(DataType::Int32) + .len(6) + .offset(1) + .add_buffer(make_i32_buffer(7)) + .nulls(Some(NullBuffer::from(vec![ + true, false, true, false, true, true, + ]))) + .build() + .unwrap(); + let data = ArrayData::builder(DataType::List(Arc::new(Field::new( + "item", + DataType::Int32, + true, + )))) + .len(4) + .offset(1) + .add_buffer(Buffer::from_slice_ref([0i32, 1, 2, 3, 4, 5])) + .nulls(Some(NullBuffer::from(vec![false, true, false, true]))) + .add_child_data(child) + .build() + .unwrap(); + let buffers = data.buffers().as_ptr(); + let children = data.child_data().as_ptr(); + let child_buffers = data.child_data()[0].buffers().as_ptr(); + let expected = data.slice(1, 2); + let sliced = data.sliced(1, 2); + assert_eq!(sliced, expected); + assert_eq!(sliced.offset(), 2); + assert_eq!(sliced.child_data()[0].len(), 6); + assert_eq!(sliced.child_data()[0].offset(), 1); + assert_eq!(sliced.buffers().as_ptr(), buffers); + assert_eq!(sliced.child_data().as_ptr(), children); + assert_eq!(sliced.child_data()[0].buffers().as_ptr(), child_buffers); + sliced.validate_full().unwrap(); + } + + #[test] + #[should_panic(expected = "offset + length overflow")] + fn test_sliced_panics_on_offset_length_overflow() { + ArrayData::new_null(&DataType::Int32, 4).sliced(1, usize::MAX); + } + + #[test] + #[should_panic(expected = "end <= self.len()")] + fn test_sliced_panics_on_offset_out_of_bounds() { + ArrayData::new_null(&DataType::Int32, 4).sliced(5, 0); + } + + #[test] + #[should_panic(expected = "end <= self.len()")] + fn test_sliced_panics_on_length_out_of_bounds() { + ArrayData::new_null(&DataType::Int32, 4).sliced(3, 2); + } + #[test] fn test_typed_offsets_length_overflow() { let data = ArrayData {
