Skip to content

Commit

Permalink
Changed methods to slice arrays (jorgecarleitao#1396)
Browse files Browse the repository at this point in the history
  • Loading branch information
jorgecarleitao authored and ritchie46 committed Mar 29, 2023
1 parent ee1df36 commit 6137e16
Show file tree
Hide file tree
Showing 71 changed files with 480 additions and 437 deletions.
9 changes: 6 additions & 3 deletions benches/bitmap_ops.rs
Original file line number Diff line number Diff line change
Expand Up @@ -26,7 +26,8 @@ fn add_benchmark(c: &mut Criterion) {
&format!("bitmap count zeros 85% slice 2^{log2_size}"),
|b| {
b.iter(|| {
let r = bitmap.clone().slice(offset, len);
let mut r = bitmap.clone();
r.slice(offset, len);
assert!(r.unset_bits() > 0);
})
},
Expand All @@ -39,13 +40,15 @@ fn add_benchmark(c: &mut Criterion) {
&format!("bitmap count zeros 51% slice 2^{log2_size}"),
|b| {
b.iter(|| {
let r = bitmap.clone().slice(offset, len);
let mut r = bitmap.clone();
r.slice(offset, len);
assert!(r.unset_bits() > 0);
})
},
);

let bitmap1 = bitmap.clone().slice(1, size - 1);
let mut bitmap1 = bitmap.clone();
bitmap1.slice(1, size - 1);
c.bench_function(&format!("bitmap not 2^{log2_size}"), |b| {
b.iter(|| {
let r = !&bitmap1;
Expand Down
2 changes: 1 addition & 1 deletion guide/src/low_level.md
Original file line number Diff line number Diff line change
Expand Up @@ -34,7 +34,7 @@ let x = vec![1u32, 2, 3];
let x: Buffer<u32> = x.into();
assert_eq!(x.as_slice(), &[1u32, 2, 3]);

let x = x.slice(1, 2); // O(1)
let x = x.sliced(1, 2); // O(1)
assert_eq!(x.as_slice(), &[2, 3]);
# }
```
Expand Down
43 changes: 19 additions & 24 deletions src/array/binary/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -175,41 +175,34 @@ impl<O: Offset> BinaryArray<O> {
self.validity.as_ref()
}

/// Creates a new [`BinaryArray`] by slicing this [`BinaryArray`].
/// Slices this [`BinaryArray`].
/// # Implementation
/// This function is `O(1)`: all data will be shared between both arrays.
/// This function is `O(1)`.
/// # Panics
/// iff `offset + length > self.len()`.
#[must_use]
pub fn slice(&self, offset: usize, length: usize) -> Self {
pub fn slice(&mut self, offset: usize, length: usize) {
assert!(
offset + length <= self.len(),
"the offset of the new Buffer cannot exceed the existing length"
);
unsafe { self.slice_unchecked(offset, length) }
}

/// Creates a new [`BinaryArray`] by slicing this [`BinaryArray`].
/// Slices this [`BinaryArray`].
/// # Implementation
/// This function is `O(1)`: all data will be shared between both arrays.
/// This function is `O(1)`.
/// # Safety
/// The caller must ensure that `offset + length <= self.len()`.
#[must_use]
pub unsafe fn slice_unchecked(&self, offset: usize, length: usize) -> Self {
let validity = self
.validity
.clone()
.map(|bitmap| bitmap.slice_unchecked(offset, length))
.and_then(|bitmap| (bitmap.unset_bits() > 0).then(|| bitmap));
let offsets = self.offsets.clone().slice_unchecked(offset, length + 1);
Self {
data_type: self.data_type.clone(),
offsets,
values: self.values.clone(),
validity,
}
pub unsafe fn slice_unchecked(&mut self, offset: usize, length: usize) {
self.validity.as_mut().and_then(|bitmap| {
bitmap.slice_unchecked(offset, length);
(bitmap.unset_bits() > 0).then(|| bitmap)
});
self.offsets.slice_unchecked(offset, length + 1);
}

impl_sliced!();

/// Boxes self into a [`Box<dyn Array>`].
pub fn boxed(self) -> Box<dyn Array> {
Box::new(self)
Expand Down Expand Up @@ -440,12 +433,14 @@ impl<O: Offset> Array for BinaryArray<O> {
self.validity.as_ref()
}

fn slice(&self, offset: usize, length: usize) -> Box<dyn Array> {
Box::new(self.slice(offset, length))
fn slice(&mut self, offset: usize, length: usize) {
self.slice(offset, length)
}
unsafe fn slice_unchecked(&self, offset: usize, length: usize) -> Box<dyn Array> {
Box::new(self.slice_unchecked(offset, length))

unsafe fn slice_unchecked(&mut self, offset: usize, length: usize) {
self.slice_unchecked(offset, length)
}

fn with_validity(&self, validity: Option<Bitmap>) -> Box<dyn Array> {
Box::new(self.clone().with_validity(validity))
}
Expand Down
37 changes: 17 additions & 20 deletions src/array/boolean/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -143,41 +143,36 @@ impl BooleanArray {
self.values.get_bit_unchecked(i)
}

/// Returns a slice of this [`BooleanArray`].
/// Slices this [`BooleanArray`].
/// # Implementation
/// This operation is `O(1)` as it amounts to increase up to two ref counts.
/// # Panic
/// This function panics iff `offset + length > self.len()`.
#[inline]
#[must_use]
pub fn slice(&self, offset: usize, length: usize) -> Self {
pub fn slice(&mut self, offset: usize, length: usize) {
assert!(
offset + length <= self.len(),
"the offset of the new Buffer cannot exceed the existing length"
);
unsafe { self.slice_unchecked(offset, length) }
}

/// Returns a slice of this [`BooleanArray`].
/// Slices this [`BooleanArray`].
/// # Implementation
/// This operation is `O(1)` as it amounts to increase two ref counts.
/// # Safety
/// The caller must ensure that `offset + length <= self.len()`.
#[inline]
#[must_use]
pub unsafe fn slice_unchecked(&self, offset: usize, length: usize) -> Self {
let validity = self
.validity
.clone()
.map(|bitmap| bitmap.slice_unchecked(offset, length))
.and_then(|bitmap| (bitmap.unset_bits() > 0).then(|| bitmap));
Self {
data_type: self.data_type.clone(),
values: self.values.clone().slice_unchecked(offset, length),
validity,
}
pub unsafe fn slice_unchecked(&mut self, offset: usize, length: usize) {
self.validity.as_mut().and_then(|bitmap| {
bitmap.slice_unchecked(offset, length);
(bitmap.unset_bits() > 0).then(|| bitmap)
});
self.values.slice_unchecked(offset, length);
}

impl_sliced!();

/// Returns this [`BooleanArray`] with a new validity.
/// # Panic
/// This function panics iff `validity.len() != self.len()`.
Expand Down Expand Up @@ -400,13 +395,15 @@ impl Array for BooleanArray {
}

#[inline]
fn slice(&self, offset: usize, length: usize) -> Box<dyn Array> {
Box::new(self.slice(offset, length))
fn slice(&mut self, offset: usize, length: usize) {
self.slice(offset, length)
}

#[inline]
unsafe fn slice_unchecked(&self, offset: usize, length: usize) -> Box<dyn Array> {
Box::new(self.slice_unchecked(offset, length))
unsafe fn slice_unchecked(&mut self, offset: usize, length: usize) {
self.slice_unchecked(offset, length)
}

fn with_validity(&self, validity: Option<Bitmap>) -> Box<dyn Array> {
Box::new(self.clone().with_validity(validity))
}
Expand Down
30 changes: 12 additions & 18 deletions src/array/dictionary/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -291,28 +291,22 @@ impl<K: DictionaryKey> DictionaryArray<K> {
DataType::Dictionary(K::KEY_TYPE, Box::new(values_datatype), false)
}

/// Creates a new [`DictionaryArray`] by slicing the existing [`DictionaryArray`].
/// Slices this [`DictionaryArray`].
/// # Panics
/// iff `offset + length > self.len()`.
pub fn slice(&self, offset: usize, length: usize) -> Self {
Self {
data_type: self.data_type.clone(),
keys: self.keys.clone().slice(offset, length),
values: self.values.clone(),
}
pub fn slice(&mut self, offset: usize, length: usize) {
self.keys.slice(offset, length);
}

/// Creates a new [`DictionaryArray`] by slicing the existing [`DictionaryArray`].
/// Slices this [`DictionaryArray`].
/// # Safety
/// Safe iff `offset + length <= self.len()`.
pub unsafe fn slice_unchecked(&self, offset: usize, length: usize) -> Self {
Self {
data_type: self.data_type.clone(),
keys: self.keys.clone().slice_unchecked(offset, length),
values: self.values.clone(),
}
pub unsafe fn slice_unchecked(&mut self, offset: usize, length: usize) {
self.keys.slice_unchecked(offset, length);
}

impl_sliced!();

/// Returns this [`DictionaryArray`] with a new validity.
/// # Panic
/// This function panics iff `validity.len() != self.len()`.
Expand Down Expand Up @@ -437,12 +431,12 @@ impl<K: DictionaryKey> Array for DictionaryArray<K> {
self.keys.validity()
}

fn slice(&self, offset: usize, length: usize) -> Box<dyn Array> {
Box::new(self.slice(offset, length))
fn slice(&mut self, offset: usize, length: usize) {
self.slice(offset, length)
}

unsafe fn slice_unchecked(&self, offset: usize, length: usize) -> Box<dyn Array> {
Box::new(self.slice_unchecked(offset, length))
unsafe fn slice_unchecked(&mut self, offset: usize, length: usize) {
self.slice_unchecked(offset, length)
}

fn with_validity(&self, validity: Option<Bitmap>) -> Box<dyn Array> {
Expand Down
6 changes: 3 additions & 3 deletions src/array/equal/struct_.rs
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,7 @@ pub(super) fn equal(lhs: &StructArray, rhs: &StructArray) -> bool {
l_validity.iter().zip(r_validity.iter()).enumerate().all(
|(i, (lhs_is_valid, rhs_is_valid))| {
if lhs_is_valid && rhs_is_valid {
lhs.slice(i, 1) == rhs.slice(i, 1)
lhs.sliced(i, 1) == rhs.sliced(i, 1)
} else {
lhs_is_valid == rhs_is_valid
}
Expand All @@ -27,7 +27,7 @@ pub(super) fn equal(lhs: &StructArray, rhs: &StructArray) -> bool {
.all(|(lhs, rhs)| {
l_validity.iter().enumerate().all(|(i, lhs_is_valid)| {
if lhs_is_valid {
lhs.slice(i, 1) == rhs.slice(i, 1)
lhs.sliced(i, 1) == rhs.sliced(i, 1)
} else {
// rhs is always valid => different
false
Expand All @@ -42,7 +42,7 @@ pub(super) fn equal(lhs: &StructArray, rhs: &StructArray) -> bool {
.all(|(lhs, rhs)| {
r_validity.iter().enumerate().all(|(i, rhs_is_valid)| {
if rhs_is_valid {
lhs.slice(i, 1) == rhs.slice(i, 1)
lhs.sliced(i, 1) == rhs.sliced(i, 1)
} else {
// lhs is always valid => different
false
Expand Down
43 changes: 17 additions & 26 deletions src/array/fixed_size_binary/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -97,44 +97,35 @@ impl FixedSizeBinaryArray {

// must use
impl FixedSizeBinaryArray {
/// Returns a slice of this [`FixedSizeBinaryArray`].
/// Slices this [`FixedSizeBinaryArray`].
/// # Implementation
/// This operation is `O(1)` as it amounts to increase 3 ref counts.
/// This operation is `O(1)`.
/// # Panics
/// panics iff `offset + length > self.len()`
#[must_use]
pub fn slice(&self, offset: usize, length: usize) -> Self {
pub fn slice(&mut self, offset: usize, length: usize) {
assert!(
offset + length <= self.len(),
"the offset of the new Buffer cannot exceed the existing length"
);
unsafe { self.slice_unchecked(offset, length) }
}

/// Returns a slice of this [`FixedSizeBinaryArray`].
/// Slices this [`FixedSizeBinaryArray`].
/// # Implementation
/// This operation is `O(1)` as it amounts to increase 3 ref counts.
/// This operation is `O(1)`.
/// # Safety
/// The caller must ensure that `offset + length <= self.len()`.
#[must_use]
pub unsafe fn slice_unchecked(&self, offset: usize, length: usize) -> Self {
let validity = self
.validity
.clone()
.map(|bitmap| bitmap.slice_unchecked(offset, length))
.and_then(|bitmap| (bitmap.unset_bits() > 0).then(|| bitmap));
let values = self
.values
.clone()
pub unsafe fn slice_unchecked(&mut self, offset: usize, length: usize) {
self.validity.as_mut().and_then(|bitmap| {
bitmap.slice_unchecked(offset, length);
(bitmap.unset_bits() > 0).then(|| bitmap)
});
self.values
.slice_unchecked(offset * self.size, length * self.size);
Self {
data_type: self.data_type.clone(),
size: self.size,
values,
validity,
}
}

impl_sliced!();

/// Returns this [`FixedSizeBinaryArray`] with a new validity.
/// # Panic
/// This function panics iff `validity.len() != self.len()`.
Expand Down Expand Up @@ -267,12 +258,12 @@ impl Array for FixedSizeBinaryArray {
self.validity.as_ref()
}

fn slice(&self, offset: usize, length: usize) -> Box<dyn Array> {
Box::new(self.slice(offset, length))
fn slice(&mut self, offset: usize, length: usize) {
self.slice(offset, length)
}

unsafe fn slice_unchecked(&self, offset: usize, length: usize) -> Box<dyn Array> {
Box::new(self.slice_unchecked(offset, length))
unsafe fn slice_unchecked(&mut self, offset: usize, length: usize) {
self.slice_unchecked(offset, length)
}

fn with_validity(&self, validity: Option<Bitmap>) -> Box<dyn Array> {
Expand Down
Loading

0 comments on commit 6137e16

Please sign in to comment.