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

Commit

Permalink
Bug fix in offset for sliced unions (#891)
Browse files Browse the repository at this point in the history
  • Loading branch information
@ncpenke
ncpenke authored Mar 7, 2022
1 parent 81bfadd commit b0774c8
Show file tree
Hide file tree
Showing 2 changed files with 182 additions and 6 deletions.
12 changes: 9 additions & 3 deletions src/array/union/mod.rs
View file
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)),
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
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(())
}

0 comments on commit b0774c8

Please sign in to comment.