Skip to content

Commit 24b0d83

Browse files
committed
make slice work for nested types
revert changes made in ARROW-11394 See commit 9f96561 Only slice into structs
1 parent f1fb2b1 commit 24b0d83

File tree

3 files changed

+36
-23
lines changed

3 files changed

+36
-23
lines changed

arrow/src/array/array_struct.rs

Lines changed: 1 addition & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -84,12 +84,7 @@ impl From<ArrayData> for StructArray {
8484
fn from(data: ArrayData) -> Self {
8585
let mut boxed_fields = vec![];
8686
for cd in data.child_data() {
87-
let child_data = if data.offset() != 0 || data.len() != cd.len() {
88-
cd.slice(data.offset(), data.len())
89-
} else {
90-
cd.clone()
91-
};
92-
boxed_fields.push(make_array(child_data));
87+
boxed_fields.push(make_array(cd.clone()));
9388
}
9489
Self { data, boxed_fields }
9590
}

arrow/src/array/data.rs

Lines changed: 30 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -383,15 +383,36 @@ impl ArrayData {
383383
pub fn slice(&self, offset: usize, length: usize) -> ArrayData {
384384
assert!((offset + length) <= self.len());
385385

386-
let mut new_data = self.clone();
387-
388-
new_data.len = length;
389-
new_data.offset = offset + self.offset;
390-
391-
new_data.null_count =
392-
count_nulls(new_data.null_buffer(), new_data.offset, new_data.len);
393-
394-
new_data
386+
if let DataType::Struct(_) = self.data_type() {
387+
// Slice into children
388+
let new_offset = self.offset + offset;
389+
let new_data = ArrayData {
390+
data_type: self.data_type().clone(),
391+
len: length,
392+
null_count: count_nulls(self.null_buffer(), new_offset, length),
393+
offset: new_offset,
394+
buffers: self.buffers.clone(),
395+
// Slice child data, to propagate offsets down to them
396+
child_data: self
397+
.child_data()
398+
.iter()
399+
.map(|data| data.slice(offset, length))
400+
.collect(),
401+
null_bitmap: self.null_bitmap().clone(),
402+
};
403+
404+
new_data
405+
} else {
406+
let mut new_data = self.clone();
407+
408+
new_data.len = length;
409+
new_data.offset = offset + self.offset;
410+
411+
new_data.null_count =
412+
count_nulls(new_data.null_buffer(), new_data.offset, new_data.len);
413+
414+
new_data
415+
}
395416
}
396417

397418
/// Returns the `buffer` as a slice of type `T` starting at self.offset

arrow/src/array/transform/structure.rs

Lines changed: 5 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -26,13 +26,10 @@ pub(super) fn build_extend(array: &ArrayData) -> Extend {
2626
index: usize,
2727
start: usize,
2828
len: usize| {
29-
mutable.child_data.iter_mut().for_each(|child| {
30-
child.extend(
31-
index,
32-
array.offset() + start,
33-
array.offset() + start + len,
34-
)
35-
})
29+
mutable
30+
.child_data
31+
.iter_mut()
32+
.for_each(|child| child.extend(index, start, start + len))
3633
},
3734
)
3835
} else {
@@ -41,7 +38,7 @@ pub(super) fn build_extend(array: &ArrayData) -> Extend {
4138
index: usize,
4239
start: usize,
4340
len: usize| {
44-
(array.offset() + start..array.offset() + start + len).for_each(|i| {
41+
(start..start + len).for_each(|i| {
4542
if array.is_valid(i) {
4643
mutable
4744
.child_data

0 commit comments

Comments
 (0)