Skip to content

Commit 231ae9b

Browse files
tustvoldtustvold
authored
Return Buffers from ArrayData::buffers instead of slice (#1799) (#3783)
* Return Buffers from ArrayData::buffers instead of slice (#1799) * Clippy
1 parent eff058f commit 231ae9b

File tree

10 files changed

+132
-31
lines changed

10 files changed

+132
-31
lines changed

‎arrow-arith/src/boolean.rs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -353,7 +353,7 @@ pub fn not(left: &BooleanArray) -> Result<BooleanArray, ArrowError> {
353353
let data = left.data_ref();
354354
let null_bit_buffer = data.nulls().map(|b| b.inner().sliced());
355355

356-
let values = buffer_unary_not(&data.buffers()[0], left_offset, len);
356+
let values = buffer_unary_not(data.buffers()[0], left_offset, len);
357357

358358
let data = unsafe {
359359
ArrayData::new_unchecked(

‎arrow-array/src/array/primitive_array.rs

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1234,7 +1234,7 @@ mod tests {
12341234
fn test_primitive_array_from_vec() {
12351235
let buf = Buffer::from_slice_ref([0, 1, 2, 3, 4]);
12361236
let arr = Int32Array::from(vec![0, 1, 2, 3, 4]);
1237-
assert_eq!(buf, arr.data.buffers()[0]);
1237+
assert_eq!(buf, *arr.data.buffers()[0]);
12381238
assert_eq!(5, arr.len());
12391239
assert_eq!(0, arr.offset());
12401240
assert_eq!(0, arr.null_count());
@@ -1740,7 +1740,7 @@ mod tests {
17401740
.build()
17411741
.unwrap();
17421742
let arr = Int32Array::from(data);
1743-
assert_eq!(buf2, arr.data.buffers()[0]);
1743+
assert_eq!(buf2, *arr.data.buffers()[0]);
17441744
assert_eq!(5, arr.len());
17451745
assert_eq!(0, arr.null_count());
17461746
for i in 0..3 {

‎arrow-array/src/array/union_array.rs

Lines changed: 15 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -409,7 +409,7 @@ mod tests {
409409

410410
// Check type ids
411411
assert_eq!(
412-
union.data().buffers()[0],
412+
*union.data().buffers()[0],
413413
Buffer::from_slice_ref(&expected_type_ids)
414414
);
415415
for (i, id) in expected_type_ids.iter().enumerate() {
@@ -418,7 +418,7 @@ mod tests {
418418

419419
// Check offsets
420420
assert_eq!(
421-
union.data().buffers()[1],
421+
*union.data().buffers()[1],
422422
Buffer::from_slice_ref(&expected_value_offsets)
423423
);
424424
for (i, id) in expected_value_offsets.iter().enumerate() {
@@ -427,15 +427,15 @@ mod tests {
427427

428428
// Check data
429429
assert_eq!(
430-
union.data().child_data()[0].buffers()[0],
430+
*union.data().child_data()[0].buffers()[0],
431431
Buffer::from_slice_ref([1_i32, 4, 6])
432432
);
433433
assert_eq!(
434-
union.data().child_data()[1].buffers()[0],
434+
*union.data().child_data()[1].buffers()[0],
435435
Buffer::from_slice_ref([2_i32, 7])
436436
);
437437
assert_eq!(
438-
union.data().child_data()[2].buffers()[0],
438+
*union.data().child_data()[2].buffers()[0],
439439
Buffer::from_slice_ref([3_i32, 5]),
440440
);
441441

@@ -467,7 +467,7 @@ mod tests {
467467

468468
// Check type ids
469469
assert_eq!(
470-
union.data().buffers()[0],
470+
*union.data().buffers()[0],
471471
Buffer::from_slice_ref(&expected_type_ids)
472472
);
473473
for (i, id) in expected_type_ids.iter().enumerate() {
@@ -476,7 +476,7 @@ mod tests {
476476

477477
// Check offsets
478478
assert_eq!(
479-
union.data().buffers()[1],
479+
*union.data().buffers()[1],
480480
Buffer::from_slice_ref(&expected_value_offsets)
481481
);
482482
for (i, id) in expected_value_offsets.iter().enumerate() {
@@ -660,15 +660,15 @@ mod tests {
660660
.unwrap();
661661

662662
// Check type ids
663-
assert_eq!(Buffer::from_slice_ref(type_ids), array.data().buffers()[0]);
663+
assert_eq!(Buffer::from_slice_ref(type_ids), *array.data().buffers()[0]);
664664
for (i, id) in type_ids.iter().enumerate() {
665665
assert_eq!(id, &array.type_id(i));
666666
}
667667

668668
// Check offsets
669669
assert_eq!(
670670
Buffer::from_slice_ref(value_offsets),
671-
array.data().buffers()[1]
671+
*array.data().buffers()[1]
672672
);
673673
for (i, id) in value_offsets.iter().enumerate() {
674674
assert_eq!(id, &array.value_offset(i));
@@ -736,7 +736,7 @@ mod tests {
736736
// Check type ids
737737
assert_eq!(
738738
Buffer::from_slice_ref(&expected_type_ids),
739-
union.data().buffers()[0]
739+
*union.data().buffers()[0]
740740
);
741741
for (i, id) in expected_type_ids.iter().enumerate() {
742742
assert_eq!(id, &union.type_id(i));
@@ -747,16 +747,16 @@ mod tests {
747747

748748
// Check data
749749
assert_eq!(
750-
union.data().child_data()[0].buffers()[0],
750+
*union.data().child_data()[0].buffers()[0],
751751
Buffer::from_slice_ref([1_i32, 0, 0, 4, 0, 6, 0]),
752752
);
753753
assert_eq!(
754754
Buffer::from_slice_ref([0_i32, 2_i32, 0, 0, 0, 0, 7]),
755-
union.data().child_data()[1].buffers()[0]
755+
*union.data().child_data()[1].buffers()[0]
756756
);
757757
assert_eq!(
758758
Buffer::from_slice_ref([0_i32, 0, 3_i32, 0, 5, 0, 0]),
759-
union.data().child_data()[2].buffers()[0]
759+
*union.data().child_data()[2].buffers()[0]
760760
);
761761

762762
assert_eq!(expected_array_values.len(), union.len());
@@ -785,7 +785,7 @@ mod tests {
785785
// Check type ids
786786
assert_eq!(
787787
Buffer::from_slice_ref(&expected_type_ids),
788-
union.data().buffers()[0]
788+
*union.data().buffers()[0]
789789
);
790790
for (i, id) in expected_type_ids.iter().enumerate() {
791791
assert_eq!(id, &union.type_id(i));
@@ -847,7 +847,7 @@ mod tests {
847847
// Check type ids
848848
assert_eq!(
849849
Buffer::from_slice_ref(&expected_type_ids),
850-
union.data().buffers()[0]
850+
*union.data().buffers()[0]
851851
);
852852
for (i, id) in expected_type_ids.iter().enumerate() {
853853
assert_eq!(id, &union.type_id(i));

‎arrow-csv/src/reader/mod.rs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2482,7 +2482,7 @@ mod tests {
24822482
for v in *values {
24832483
t.update(v, None)
24842484
}
2485-
assert_eq!(&t.get(), expected, "{:?}", values)
2485+
assert_eq!(&t.get(), expected, "{values:?}")
24862486
}
24872487
}
24882488
}

‎arrow-data/src/data/buffers.rs

Lines changed: 97 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,97 @@
1+
// Licensed to the Apache Software Foundation (ASF) under one
2+
// or more contributor license agreements. See the NOTICE file
3+
// distributed with this work for additional information
4+
// regarding copyright ownership. The ASF licenses this file
5+
// to you under the Apache License, Version 2.0 (the
6+
// "License"); you may not use this file except in compliance
7+
// with the License. You may obtain a copy of the License at
8+
//
9+
// http://www.apache.org/licenses/LICENSE-2.0
10+
//
11+
// Unless required by applicable law or agreed to in writing,
12+
// software distributed under the License is distributed on an
13+
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14+
// KIND, either express or implied. See the License for the
15+
// specific language governing permissions and limitations
16+
// under the License.
17+
18+
use arrow_buffer::Buffer;
19+
use std::iter::Chain;
20+
use std::ops::Index;
21+
22+
/// A collection of [`Buffer`]
23+
#[derive(Debug, Default, Copy, Clone, Eq, PartialEq)]
24+
pub struct Buffers<'a>([Option<&'a Buffer>; 2]);
25+
26+
impl<'a> Buffers<'a> {
27+
/// Temporary will be removed once ArrayData does not store `Vec<Buffer>` directly (#3769)
28+
#[inline]
29+
pub(crate) fn from_slice(a: &'a [Buffer]) -> Self {
30+
match a.len() {
31+
0 => Self([None, None]),
32+
1 => Self([Some(&a[0]), None]),
33+
_ => Self([Some(&a[0]), Some(&a[1])]),
34+
}
35+
}
36+
37+
/// Returns the number of [`Buffer`] in this collection
38+
#[inline]
39+
pub fn len(&self) -> usize {
40+
self.0[0].is_some() as usize + self.0[1].is_some() as usize
41+
}
42+
43+
/// Returns `true` if this collection is empty
44+
#[inline]
45+
pub fn is_empty(&self) -> bool {
46+
self.0[0].is_none() && self.0[1].is_none()
47+
}
48+
49+
#[inline]
50+
pub fn iter(&self) -> IntoIter<'a> {
51+
self.into_iter()
52+
}
53+
54+
/// Converts this [`Buffers`] to a `Vec<Buffer>`
55+
#[inline]
56+
pub fn to_vec(&self) -> Vec<Buffer> {
57+
self.iter().cloned().collect()
58+
}
59+
}
60+
61+
impl<'a> Index<usize> for Buffers<'a> {
62+
type Output = &'a Buffer;
63+
64+
#[inline]
65+
fn index(&self, index: usize) -> &Self::Output {
66+
self.0[index].as_ref().unwrap()
67+
}
68+
}
69+
70+
impl<'a> IntoIterator for Buffers<'a> {
71+
type Item = &'a Buffer;
72+
type IntoIter = IntoIter<'a>;
73+
74+
#[inline]
75+
fn into_iter(self) -> Self::IntoIter {
76+
IntoIter(self.0[0].into_iter().chain(self.0[1].into_iter()))
77+
}
78+
}
79+
80+
type OptionIter<'a> = std::option::IntoIter<&'a Buffer>;
81+
82+
/// [`Iterator`] for [`Buffers`]
83+
pub struct IntoIter<'a>(Chain<OptionIter<'a>, OptionIter<'a>>);
84+
85+
impl<'a> Iterator for IntoIter<'a> {
86+
type Item = &'a Buffer;
87+
88+
#[inline]
89+
fn next(&mut self) -> Option<Self::Item> {
90+
self.0.next()
91+
}
92+
93+
#[inline]
94+
fn size_hint(&self) -> (usize, Option<usize>) {
95+
self.0.size_hint()
96+
}
97+
}

‎arrow-data/src/data/mod.rs

Lines changed: 7 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -30,6 +30,9 @@ use std::sync::Arc;
3030

3131
use crate::equal;
3232

33+
mod buffers;
34+
pub use buffers::*;
35+
3336
#[allow(unused)] // Private until ready (#1176)
3437
mod bytes;
3538
#[allow(unused)] // Private until ready (#1176)
@@ -371,9 +374,10 @@ impl ArrayData {
371374
&self.data_type
372375
}
373376

374-
/// Returns a slice of the [`Buffer`]s that hold the data.
375-
pub fn buffers(&self) -> &[Buffer] {
376-
&self.buffers[..]
377+
/// Returns the [`Buffers`] storing data for this [`ArrayData`]
378+
pub fn buffers(&self) -> Buffers<'_> {
379+
// In future ArrayData won't store data contiguously as `Vec<Buffer>` (#1799)
380+
Buffers::from_slice(&self.buffers)
377381
}
378382

379383
/// Returns a slice of children [`ArrayData`]. This will be non

‎arrow-json/src/reader.rs

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -2210,7 +2210,7 @@ mod tests {
22102210
.unwrap();
22112211
// test that the list offsets are correct
22122212
assert_eq!(
2213-
cc.data().buffers()[0],
2213+
*cc.data().buffers()[0],
22142214
Buffer::from_slice_ref([0i32, 2, 2, 4, 5])
22152215
);
22162216
let cc = as_boolean_array(cc.values());
@@ -2230,7 +2230,7 @@ mod tests {
22302230
.unwrap();
22312231
// test that the list offsets are correct
22322232
assert_eq!(
2233-
dd.data().buffers()[0],
2233+
*dd.data().buffers()[0],
22342234
Buffer::from_slice_ref([0i32, 1, 1, 2, 6])
22352235
);
22362236

@@ -2374,7 +2374,7 @@ mod tests {
23742374
let read: &ListArray = read.as_any().downcast_ref::<ListArray>().unwrap();
23752375
let expected = expected.as_any().downcast_ref::<ListArray>().unwrap();
23762376
assert_eq!(
2377-
read.data().buffers()[0],
2377+
*read.data().buffers()[0],
23782378
Buffer::from_slice_ref([0i32, 2, 3, 6, 6, 6, 7])
23792379
);
23802380
// compare list null buffers

‎arrow-select/src/filter.rs

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -82,7 +82,7 @@ impl<'a> IndexIterator<'a> {
8282
fn new(filter: &'a BooleanArray, remaining: usize) -> Self {
8383
assert_eq!(filter.null_count(), 0);
8484
let data = filter.data();
85-
let iter = BitIndexIterator::new(&data.buffers()[0], data.offset(), data.len());
85+
let iter = BitIndexIterator::new(data.buffers()[0], data.offset(), data.len());
8686
Self { remaining, iter }
8787
}
8888
}
@@ -470,7 +470,7 @@ fn filter_boolean(values: &BooleanArray, predicate: &FilterPredicate) -> Boolean
470470
assert_eq!(data.buffers().len(), 1);
471471
assert_eq!(data.child_data().len(), 0);
472472

473-
let values = filter_bits(&data.buffers()[0], data.offset(), predicate);
473+
let values = filter_bits(data.buffers()[0], data.offset(), predicate);
474474

475475
let mut builder = ArrayDataBuilder::new(DataType::Boolean)
476476
.len(predicate.count)
@@ -572,7 +572,7 @@ where
572572

573573
Self {
574574
src_offsets: array.value_offsets(),
575-
src_values: &array.data().buffers()[1],
575+
src_values: array.data().buffers()[1],
576576
dst_offsets,
577577
dst_values,
578578
cur_offset,

‎arrow-select/src/nullif.rs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -56,7 +56,7 @@ pub fn nullif(left: &dyn Array, right: &BooleanArray) -> Result<ArrayRef, ArrowE
5656
let (right, r_offset) = match right_data.nulls() {
5757
Some(nulls) => (
5858
buffer_bin_and(
59-
&right_data.buffers()[0],
59+
right_data.buffers()[0],
6060
right_data.offset(),
6161
nulls.buffer(),
6262
nulls.offset(),

‎arrow/src/ffi.rs

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1158,7 +1158,7 @@ mod tests {
11581158
// Check type ids
11591159
assert_eq!(
11601160
Buffer::from_slice_ref(&expected_type_ids),
1161-
array.data().buffers()[0]
1161+
*array.data().buffers()[0]
11621162
);
11631163
for (i, id) in expected_type_ids.iter().enumerate() {
11641164
assert_eq!(id, &array.type_id(i));
@@ -1222,7 +1222,7 @@ mod tests {
12221222
// Check type ids
12231223
assert_eq!(
12241224
Buffer::from_slice_ref(&expected_type_ids),
1225-
array.data().buffers()[0]
1225+
*array.data().buffers()[0]
12261226
);
12271227
for (i, id) in expected_type_ids.iter().enumerate() {
12281228
assert_eq!(id, &array.type_id(i));

0 commit comments

Comments
 (0)