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

Bug fix in offset for sliced unions #891

Merged
merged 1 commit into from
Mar 7, 2022
Merged

Bug fix in offset for sliced unions #891

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
12 changes: 9 additions & 3 deletions src/array/union/mod.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -191,7 +191,10 @@ impl UnionArray {
fields: self.fields.clone(),
fields_hash: self.fields_hash.clone(),
types: self.types.clone().slice(offset, length),
offsets: self.offsets.clone(),
offsets: self
.offsets
.clone()
.map(|offsets| offsets.slice(offset, length)),
Copy link
Contributor Author

@ncpenke ncpenke Mar 7, 2022 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@jorgecarleitao For dense unions wasn't sure if slicing the offsets or using (index + offset) in field_slot was more correct.

offset: self.offset + offset,
}
}
Expand All @@ -208,7 +211,10 @@ impl UnionArray {
fields: self.fields.clone(),
fields_hash: self.fields_hash.clone(),
types: self.types.clone().slice_unchecked(offset, length),
offsets: self.offsets.clone(),
offsets: self
.offsets
.clone()
.map(|offsets| offsets.slice_unchecked(offset, length)),
offset: self.offset + offset,
}
}
Expand Down Expand Up @@ -249,7 +255,7 @@ impl UnionArray {
self.offsets()
.as_ref()
.map(|x| x[index] as usize)
.unwrap_or(index)
.unwrap_or(index + self.offset)
}

/// Returns the index and slot of the field to select from `self.fields`.
Expand Down
176 changes: 173 additions & 3 deletions tests/it/array/union.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,6 +1,12 @@
use std::sync::Arc;

use arrow2::{array::*, buffer::Buffer, datatypes::*, error::Result};
use arrow2::{
array::*,
buffer::Buffer,
datatypes::*,
error::Result,
scalar::{PrimitiveScalar, Utf8Scalar},
};

#[test]
fn debug() -> Result<()> {
Expand Down Expand Up @@ -39,9 +45,173 @@ fn slice() -> Result<()> {

let result = array.slice(1, 2);

let types = Buffer::from_slice([0, 1]);
let expected = UnionArray::from_data(data_type, types, fields, None);
let sliced_types = Buffer::from_slice([0, 1]);
let sliced_fields = vec![
Arc::new(Int32Array::from(&[None, Some(2)])) as Arc<dyn Array>,
Arc::new(Utf8Array::<i32>::from(&[Some("b"), Some("c")])) as Arc<dyn Array>,
];
let expected = UnionArray::from_data(data_type, sliced_types, sliced_fields, None);

assert_eq!(expected, result);
Ok(())
}

#[test]
fn iter_sparse() -> Result<()> {
let fields = vec![
Field::new("a", DataType::Int32, true),
Field::new("b", DataType::Utf8, true),
];
let data_type = DataType::Union(fields, None, UnionMode::Sparse);
let types = Buffer::from_slice([0, 0, 1]);
let fields = vec![
Arc::new(Int32Array::from(&[Some(1), None, Some(2)])) as Arc<dyn Array>,
Arc::new(Utf8Array::<i32>::from(&[Some("a"), Some("b"), Some("c")])) as Arc<dyn Array>,
];

let array = UnionArray::from_data(data_type, types, fields.clone(), None);
let mut iter = array.iter();

assert_eq!(
iter.next()
.unwrap()
.as_any()
.downcast_ref::<PrimitiveScalar<i32>>()
.unwrap()
.value(),
Some(1)
);
assert_eq!(
iter.next()
.unwrap()
.as_any()
.downcast_ref::<PrimitiveScalar<i32>>()
.unwrap()
.value(),
None
);
assert_eq!(
iter.next()
.unwrap()
.as_any()
.downcast_ref::<Utf8Scalar<i32>>()
.unwrap()
.value(),
Some("c")
);
assert_eq!(iter.next(), None);

Ok(())
}

#[test]
fn iter_dense() -> Result<()> {
let fields = vec![
Field::new("a", DataType::Int32, true),
Field::new("b", DataType::Utf8, true),
];
let data_type = DataType::Union(fields, None, UnionMode::Dense);
let types = Buffer::from_slice([0, 0, 1]);
let offsets = Buffer::<i32>::from_slice([0, 1, 0]);
let fields = vec![
Arc::new(Int32Array::from(&[Some(1), None])) as Arc<dyn Array>,
Arc::new(Utf8Array::<i32>::from(&[Some("c")])) as Arc<dyn Array>,
];

let array = UnionArray::from_data(data_type, types, fields.clone(), Some(offsets));
let mut iter = array.iter();

assert_eq!(
iter.next()
.unwrap()
.as_any()
.downcast_ref::<PrimitiveScalar<i32>>()
.unwrap()
.value(),
Some(1)
);
assert_eq!(
iter.next()
.unwrap()
.as_any()
.downcast_ref::<PrimitiveScalar<i32>>()
.unwrap()
.value(),
None
);
assert_eq!(
iter.next()
.unwrap()
.as_any()
.downcast_ref::<Utf8Scalar<i32>>()
.unwrap()
.value(),
Some("c")
);
assert_eq!(iter.next(), None);

Ok(())
}

#[test]
fn iter_sparse_slice() -> Result<()> {
let fields = vec![
Field::new("a", DataType::Int32, true),
Field::new("b", DataType::Utf8, true),
];
let data_type = DataType::Union(fields, None, UnionMode::Sparse);
let types = Buffer::from_slice([0, 0, 1]);
let fields = vec![
Arc::new(Int32Array::from(&[Some(1), Some(3), Some(2)])) as Arc<dyn Array>,
Arc::new(Utf8Array::<i32>::from(&[Some("a"), Some("b"), Some("c")])) as Arc<dyn Array>,
];

let array = UnionArray::from_data(data_type, types, fields.clone(), None);
let array_slice = array.slice(1, 1);
let mut iter = array_slice.iter();

assert_eq!(
iter.next()
.unwrap()
.as_any()
.downcast_ref::<PrimitiveScalar<i32>>()
.unwrap()
.value(),
Some(3)
);
assert_eq!(iter.next(), None);

Ok(())
}

#[test]
fn iter_dense_slice() -> Result<()> {
let fields = vec![
Field::new("a", DataType::Int32, true),
Field::new("b", DataType::Utf8, true),
];
let data_type = DataType::Union(fields, None, UnionMode::Dense);
let types = Buffer::from_slice([0, 0, 1]);
let offsets = Buffer::<i32>::from_slice([0, 1, 0]);
let fields = vec![
Arc::new(Int32Array::from(&[Some(1), Some(3)])) as Arc<dyn Array>,
Arc::new(Utf8Array::<i32>::from(&[Some("c")])) as Arc<dyn Array>,
];

let array = UnionArray::from_data(data_type, types, fields.clone(), Some(offsets));
let array_slice = array.slice(1, 1);
let mut iter = array_slice.iter();

assert_eq!(
iter.next()
.unwrap()
.as_any()
.downcast_ref::<PrimitiveScalar<i32>>()
.unwrap()
.value(),
Some(3)
);
assert_eq!(iter.next(), None);

Ok(())
}