use std::sync::Arc;
use arrow2::{
array::*,
buffer::Buffer,
datatypes::*,
error::Result,
scalar::{PrimitiveScalar, Utf8Scalar},
};
#[test]
fn sparse_debug() -> 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 = vec![0, 0, 1].into();
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, None);
assert_eq!(format!("{:?}", array), "UnionArray[1, None, c]");
Ok(())
}
#[test]
fn dense_debug() -> 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 = vec![0, 0, 1].into();
let fields = vec![
Arc::new(Int32Array::from(&[Some(1), None, Some(2)])) as Arc<dyn Array>,
Arc::new(Utf8Array::<i32>::from(&[Some("c")])) as Arc<dyn Array>,
];
let offsets = Some(vec![0, 1, 0].into());
let array = UnionArray::from_data(data_type, types, fields, offsets);
assert_eq!(format!("{:?}", array), "UnionArray[1, None, c]");
Ok(())
}
#[test]
fn 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), 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.clone(), types, fields.clone(), None);
let result = array.slice(1, 2);
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(())
}