Skip to content
This repository has been archived by the owner on Feb 18, 2024. It is now read-only.

Correctly update child-offsets in GrowableUnion #1360

Merged
merged 6 commits into from
Jan 16, 2023
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 5 additions & 0 deletions src/array/growable/binary.rs
Original file line number Diff line number Diff line change
Expand Up @@ -81,6 +81,11 @@ impl<'a, O: Offset> Growable<'a> for GrowableBinary<'a, O> {
self.validity.extend_constant(additional, false);
}

#[inline]
fn len(&self) -> usize {
self.offsets.len() - 1
}

fn as_arc(&mut self) -> Arc<dyn Array> {
self.to().arced()
}
Expand Down
5 changes: 5 additions & 0 deletions src/array/growable/boolean.rs
Original file line number Diff line number Diff line change
Expand Up @@ -71,6 +71,11 @@ impl<'a> Growable<'a> for GrowableBoolean<'a> {
self.validity.extend_constant(additional, false);
}

#[inline]
fn len(&self) -> usize {
self.values.len()
}

fn as_arc(&mut self) -> Arc<dyn Array> {
Arc::new(self.to())
}
Expand Down
5 changes: 5 additions & 0 deletions src/array/growable/dictionary.rs
Original file line number Diff line number Diff line change
Expand Up @@ -132,6 +132,11 @@ impl<'a, T: DictionaryKey> Growable<'a> for GrowableDictionary<'a, T> {
);
}

#[inline]
fn len(&self) -> usize {
self.key_values.len()
}

#[inline]
fn extend_validity(&mut self, additional: usize) {
self.key_values
Expand Down
5 changes: 5 additions & 0 deletions src/array/growable/fixed_binary.rs
Original file line number Diff line number Diff line change
Expand Up @@ -78,6 +78,11 @@ impl<'a> Growable<'a> for GrowableFixedSizeBinary<'a> {
self.validity.extend_constant(additional, false);
}

#[inline]
fn len(&self) -> usize {
self.values.len() / self.size
}

fn as_arc(&mut self) -> Arc<dyn Array> {
Arc::new(self.to())
}
Expand Down
5 changes: 5 additions & 0 deletions src/array/growable/fixed_size_list.rs
Original file line number Diff line number Diff line change
Expand Up @@ -85,6 +85,11 @@ impl<'a> Growable<'a> for GrowableFixedSizeList<'a> {
self.validity.extend_constant(additional, false);
}

#[inline]
fn len(&self) -> usize {
self.values.len() / self.size
}

fn as_arc(&mut self) -> Arc<dyn Array> {
Arc::new(self.to())
}
Expand Down
5 changes: 5 additions & 0 deletions src/array/growable/list.rs
Original file line number Diff line number Diff line change
Expand Up @@ -97,6 +97,11 @@ impl<'a, O: Offset> Growable<'a> for GrowableList<'a, O> {
self.validity.extend_constant(additional, false);
}

#[inline]
fn len(&self) -> usize {
self.offsets.len() - 1
}

fn as_arc(&mut self) -> Arc<dyn Array> {
Arc::new(self.to())
}
Expand Down
3 changes: 3 additions & 0 deletions src/array/growable/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -43,6 +43,9 @@ pub trait Growable<'a> {
/// Extends this [`Growable`] with null elements, disregarding the bound arrays
fn extend_validity(&mut self, additional: usize);

/// The current length of the [`Growable`].
fn len(&self) -> usize;

/// Converts this [`Growable`] to an [`Arc<dyn Array>`], thereby finishing the mutation.
/// Self will be empty after such operation.
fn as_arc(&mut self) -> Arc<dyn Array> {
Expand Down
5 changes: 5 additions & 0 deletions src/array/growable/null.rs
Original file line number Diff line number Diff line change
Expand Up @@ -38,6 +38,11 @@ impl<'a> Growable<'a> for GrowableNull {
self.length += additional;
}

#[inline]
fn len(&self) -> usize {
self.length
}

fn as_arc(&mut self) -> Arc<dyn Array> {
Arc::new(NullArray::new(self.data_type.clone(), self.length))
}
Expand Down
5 changes: 5 additions & 0 deletions src/array/growable/primitive.rs
Original file line number Diff line number Diff line change
Expand Up @@ -82,6 +82,11 @@ impl<'a, T: NativeType> Growable<'a> for GrowablePrimitive<'a, T> {
self.validity.extend_constant(additional, false);
}

#[inline]
fn len(&self) -> usize {
self.values.len()
}

#[inline]
fn as_arc(&mut self) -> Arc<dyn Array> {
Arc::new(self.to())
Expand Down
12 changes: 12 additions & 0 deletions src/array/growable/structure.rs
Original file line number Diff line number Diff line change
Expand Up @@ -104,6 +104,18 @@ impl<'a> Growable<'a> for GrowableStruct<'a> {
self.validity.extend_constant(additional, false);
}

#[inline]
fn len(&self) -> usize {
// All children should have the same indexing, so just use the first
// one. If we don't have children, we might still have a validity
// array, so use that.
if let Some(child) = self.values.get(0) {
child.len()
} else {
self.validity.len()
}
}

fn as_arc(&mut self) -> Arc<dyn Array> {
Arc::new(self.to())
}
Expand Down
14 changes: 12 additions & 2 deletions src/array/growable/union.rs
Original file line number Diff line number Diff line change
Expand Up @@ -73,10 +73,15 @@ impl<'a> Growable<'a> for GrowableUnion<'a> {
if let Some(x) = self.offsets.as_mut() {
let offsets = &array.offsets().unwrap()[start..start + len];

x.extend(offsets);
// in a dense union, each slot has its own offset. We extend the fields accordingly.
for (&type_, &offset) in types.iter().zip(offsets.iter()) {
self.fields[type_ as usize].extend(index, offset as usize, 1);
let field = &mut self.fields[type_ as usize];
// The offset for the element that is about to be extended is the current length
// of the child field of the corresponding type. Note that this may be very
// different than the original offset from the array we are extending from as
// it is a function of the previous extensions to this child.
x.push(field.len() as i32);
field.extend(index, offset as usize, 1);
}
} else {
// in a sparse union, every field has the same length => extend all fields equally
Expand All @@ -88,6 +93,11 @@ impl<'a> Growable<'a> for GrowableUnion<'a> {

fn extend_validity(&mut self, _additional: usize) {}

#[inline]
fn len(&self) -> usize {
self.types.len()
}

fn as_arc(&mut self) -> Arc<dyn Array> {
self.to().arced()
}
Expand Down
5 changes: 5 additions & 0 deletions src/array/growable/utf8.rs
Original file line number Diff line number Diff line change
Expand Up @@ -88,6 +88,11 @@ impl<'a, O: Offset> Growable<'a> for GrowableUtf8<'a, O> {
self.validity.extend_constant(additional, false);
}

#[inline]
fn len(&self) -> usize {
self.offsets.len() - 1
}

fn as_arc(&mut self) -> Arc<dyn Array> {
Arc::new(self.to())
}
Expand Down
5 changes: 5 additions & 0 deletions tests/it/array/growable/binary.rs
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,7 @@ fn no_offsets() {
let mut a = GrowableBinary::new(vec![&array], false, 0);

a.extend(0, 1, 2);
assert_eq!(a.len(), 2);

let result: BinaryArray<i32> = a.into();

Expand All @@ -27,6 +28,7 @@ fn with_offsets() {
let mut a = GrowableBinary::new(vec![&array], false, 0);

a.extend(0, 0, 3);
assert_eq!(a.len(), 3);

let result: BinaryArray<i32> = a.into();

Expand All @@ -42,6 +44,7 @@ fn test_string_offsets() {
let mut a = GrowableBinary::new(vec![&array], false, 0);

a.extend(0, 0, 3);
assert_eq!(a.len(), 3);

let result: BinaryArray<i32> = a.into();

Expand All @@ -58,6 +61,7 @@ fn test_multiple_with_validity() {

a.extend(0, 0, 2);
a.extend(1, 0, 2);
assert_eq!(a.len(), 4);

let result: BinaryArray<i32> = a.into();

Expand All @@ -74,6 +78,7 @@ fn test_string_null_offset_validity() {

a.extend(0, 1, 2);
a.extend_validity(1);
assert_eq!(a.len(), 3);

let result: BinaryArray<i32> = a.into();

Expand Down
1 change: 1 addition & 0 deletions tests/it/array/growable/boolean.rs
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@ fn test_bool() {
let mut a = GrowableBoolean::new(vec![&array], false, 0);

a.extend(0, 1, 2);
assert_eq!(a.len(), 2);

let result: BooleanArray = a.into();

Expand Down
2 changes: 2 additions & 0 deletions tests/it/array/growable/dictionary.rs
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,7 @@ fn test_single() -> Result<()> {
let mut growable = GrowableDictionary::new(&[&array], false, 0);

growable.extend(0, 1, 2);
assert_eq!(growable.len(), 2);

let result: DictionaryArray<i32> = growable.into();

Expand Down Expand Up @@ -56,6 +57,7 @@ fn test_multi() -> Result<()> {

growable.extend(0, 1, 2);
growable.extend(1, 1, 2);
assert_eq!(growable.len(), 4);

let result: DictionaryArray<i32> = growable.into();

Expand Down
5 changes: 5 additions & 0 deletions tests/it/array/growable/fixed_binary.rs
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,7 @@ fn basic() {
let mut a = GrowableFixedSizeBinary::new(vec![&array], false, 0);

a.extend(0, 1, 2);
assert_eq!(a.len(), 2);

let result: FixedSizeBinaryArray = a.into();

Expand All @@ -30,6 +31,7 @@ fn offsets() {
let mut a = GrowableFixedSizeBinary::new(vec![&array], false, 0);

a.extend(0, 0, 3);
assert_eq!(a.len(), 3);

let result: FixedSizeBinaryArray = a.into();

Expand All @@ -46,6 +48,7 @@ fn multiple_with_validity() {

a.extend(0, 0, 2);
a.extend(1, 0, 2);
assert_eq!(a.len(), 4);

let result: FixedSizeBinaryArray = a.into();

Expand All @@ -63,6 +66,7 @@ fn null_offset_validity() {

a.extend(0, 1, 2);
a.extend_validity(1);
assert_eq!(a.len(), 3);

let result: FixedSizeBinaryArray = a.into();

Expand All @@ -81,6 +85,7 @@ fn sized_offsets() {

a.extend(0, 1, 1);
a.extend(0, 0, 1);
assert_eq!(a.len(), 2);

let result: FixedSizeBinaryArray = a.into();

Expand Down
3 changes: 3 additions & 0 deletions tests/it/array/growable/fixed_size_list.rs
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,7 @@ fn basic() {

let mut a = GrowableFixedSizeList::new(vec![&array], false, 0);
a.extend(0, 0, 1);
assert_eq!(a.len(), 1);

let result: FixedSizeListArray = a.into();

Expand All @@ -42,6 +43,7 @@ fn null_offset() {

let mut a = GrowableFixedSizeList::new(vec![&array], false, 0);
a.extend(0, 1, 1);
assert_eq!(a.len(), 1);

let result: FixedSizeListArray = a.into();

Expand Down Expand Up @@ -70,6 +72,7 @@ fn test_from_two_lists() {
let mut a = GrowableFixedSizeList::new(vec![&array_1, &array_2], false, 6);
a.extend(0, 0, 2);
a.extend(1, 1, 1);
assert_eq!(a.len(), 3);

let result: FixedSizeListArray = a.into();

Expand Down
5 changes: 5 additions & 0 deletions tests/it/array/growable/list.rs
Original file line number Diff line number Diff line change
Expand Up @@ -33,6 +33,7 @@ fn extension() {

let mut a = GrowableList::new(vec![&array_ext], false, 0);
a.extend(0, 0, 1);
assert_eq!(a.len(), 1);

let result: ListArray<i32> = a.into();
assert_eq!(array_ext.data_type(), result.data_type());
Expand All @@ -51,6 +52,7 @@ fn basic() {

let mut a = GrowableList::new(vec![&array], false, 0);
a.extend(0, 0, 1);
assert_eq!(a.len(), 1);

let result: ListArray<i32> = a.into();

Expand All @@ -72,6 +74,7 @@ fn null_offset() {

let mut a = GrowableList::new(vec![&array], false, 0);
a.extend(0, 1, 1);
assert_eq!(a.len(), 1);

let result: ListArray<i32> = a.into();

Expand All @@ -93,6 +96,7 @@ fn null_offsets() {

let mut a = GrowableList::new(vec![&array], false, 0);
a.extend(0, 1, 1);
assert_eq!(a.len(), 1);

let result: ListArray<i32> = a.into();

Expand Down Expand Up @@ -121,6 +125,7 @@ fn test_from_two_lists() {
let mut a = GrowableList::new(vec![&array_1, &array_2], false, 6);
a.extend(0, 0, 2);
a.extend(1, 1, 1);
assert_eq!(a.len(), 3);

let result: ListArray<i32> = a.into();

Expand Down
1 change: 1 addition & 0 deletions tests/it/array/growable/null.rs
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,7 @@ fn null() {

mutable.extend(0, 1, 2);
mutable.extend(1, 0, 1);
assert_eq!(mutable.len(), 3);

let result: NullArray = mutable.into();

Expand Down
5 changes: 5 additions & 0 deletions tests/it/array/growable/primitive.rs
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,7 @@ fn basics() {
let b = PrimitiveArray::<u8>::from(vec![Some(1), Some(2), Some(3)]);
let mut a = GrowablePrimitive::new(vec![&b], false, 3);
a.extend(0, 0, 2);
assert_eq!(a.len(), 2);
let result: PrimitiveArray<u8> = a.into();
let expected = PrimitiveArray::<u8>::from(vec![Some(1), Some(2)]);
assert_eq!(result, expected);
Expand All @@ -21,6 +22,7 @@ fn offset() {
let b = b.slice(1, 2);
let mut a = GrowablePrimitive::new(vec![&b], false, 2);
a.extend(0, 0, 2);
assert_eq!(a.len(), 2);
let result: PrimitiveArray<u8> = a.into();
let expected = PrimitiveArray::<u8>::from(vec![Some(2), Some(3)]);
assert_eq!(result, expected);
Expand All @@ -33,6 +35,7 @@ fn null_offset() {
let b = b.slice(1, 2);
let mut a = GrowablePrimitive::new(vec![&b], false, 2);
a.extend(0, 0, 2);
assert_eq!(a.len(), 2);
let result: PrimitiveArray<u8> = a.into();
let expected = PrimitiveArray::<u8>::from(vec![None, Some(3)]);
assert_eq!(result, expected);
Expand All @@ -46,6 +49,7 @@ fn null_offset_validity() {
a.extend(0, 0, 2);
a.extend_validity(3);
a.extend(0, 1, 1);
assert_eq!(a.len(), 6);
let result: PrimitiveArray<u8> = a.into();
let expected = PrimitiveArray::<u8>::from(&[Some(2), Some(3), None, None, None, Some(3)]);
assert_eq!(result, expected);
Expand All @@ -58,6 +62,7 @@ fn joining_arrays() {
let mut a = GrowablePrimitive::new(vec![&b, &c], false, 4);
a.extend(0, 0, 2);
a.extend(1, 1, 2);
assert_eq!(a.len(), 4);
let result: PrimitiveArray<u8> = a.into();

let expected = PrimitiveArray::<u8>::from(&[Some(1), Some(2), Some(5), Some(6)]);
Expand Down
Loading