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

Added FixedSizeListScalar and FixedSizeBinaryScalar #786

59 changes: 59 additions & 0 deletions src/scalar/fixed_size_binary.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,59 @@
use crate::datatypes::DataType;

use super::Scalar;

#[derive(Debug, Clone, PartialEq)]
/// The [`Scalar`] implementation of fixed size binary ([`Option<Box<[u8]>>`]).
pub struct FixedSizeBinaryScalar {
value: Option<Box<[u8]>>,
data_type: DataType,
}

impl FixedSizeBinaryScalar {
/// Returns a new [`FixedSizeBinaryScalar`].
/// # Panics
/// iff
/// * the `data_type` is not `FixedSizeBinary`
/// * the size of child binary is not equal
#[inline]
pub fn new<P: Into<Vec<u8>>>(data_type: DataType, value: Option<P>) -> Self {
assert_eq!(
data_type.to_physical_type(),
crate::datatypes::PhysicalType::FixedSizeBinary
);
Self {
illumination-k marked this conversation as resolved.
Show resolved Hide resolved
value: value.map(|x| {
let x: Vec<u8> = x.into();
assert_eq!(
data_type.to_logical_type(),
&DataType::FixedSizeBinary(x.len())
);
x.into_boxed_slice()
}),
data_type,
}
}

/// Its value
#[inline]
pub fn value(&self) -> Option<&[u8]> {
self.value.as_ref().map(|x| x.as_ref())
}
}

impl Scalar for FixedSizeBinaryScalar {
#[inline]
fn as_any(&self) -> &dyn std::any::Any {
self
}

#[inline]
fn is_valid(&self) -> bool {
self.value.is_some()
}

#[inline]
fn data_type(&self) -> &DataType {
&self.data_type
}
}
61 changes: 61 additions & 0 deletions src/scalar/fixed_size_list.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,61 @@
use std::any::Any;
use std::sync::Arc;

use crate::{array::*, datatypes::DataType};

use super::Scalar;

/// The scalar equivalent of [`FixedSizeListArray`]. Like [`FixedSizeListArray`], this struct holds a dynamically-typed
/// [`Array`]. The only difference is that this has only one element.
#[derive(Debug, Clone)]
pub struct FixedSizeListScalar {
values: Option<Arc<dyn Array>>,
data_type: DataType,
}

impl PartialEq for FixedSizeListScalar {
fn eq(&self, other: &Self) -> bool {
(self.data_type == other.data_type)
&& (self.values.is_some() == other.values.is_some())
&& ((self.values.is_none()) | (self.values.as_ref() == other.values.as_ref()))
}
}

impl FixedSizeListScalar {
/// returns a new [`FixedSizeListScalar`]
/// # Panics
/// iff
/// * the `data_type` is not `FixedSizeList`
/// * the child of the `data_type` is not equal to the `values`
/// * the size of child array is not equal
#[inline]
pub fn new(data_type: DataType, values: Option<Arc<dyn Array>>) -> Self {
let (field, size) = FixedSizeListArray::get_child_and_size(&data_type);
let inner_data_type = field.data_type();
let values = values.map(|x| {
assert_eq!(inner_data_type, x.data_type());
assert_eq!(size, x.len());
x
});
Self { values, data_type }
}

/// The values of the [`FixedSizeListScalar`]
pub fn values(&self) -> Option<&Arc<dyn Array>> {
self.values.as_ref()
}
}

impl Scalar for FixedSizeListScalar {
fn as_any(&self) -> &dyn Any {
self
}

fn is_valid(&self) -> bool {
self.values.is_some()
}

fn data_type(&self) -> &DataType {
&self.data_type
}
}
27 changes: 25 additions & 2 deletions src/scalar/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,10 @@ mod null;
pub use null::*;
mod struct_;
pub use struct_::*;
mod fixed_size_list;
pub use fixed_size_list::*;
mod fixed_size_binary;
pub use fixed_size_binary::*;

/// Trait object declaring an optional value with a [`DataType`].
/// This strait is often used in APIs that accept multiple scalar types.
Expand Down Expand Up @@ -120,8 +124,27 @@ pub fn new_scalar(array: &dyn Array, index: usize) -> Box<dyn Scalar> {
Box::new(StructScalar::new(array.data_type().clone(), None))
}
}
FixedSizeBinary => todo!(),
FixedSizeList => todo!(),
FixedSizeBinary => {
let array = array
.as_any()
.downcast_ref::<FixedSizeBinaryArray>()
.unwrap();
let value = if array.is_valid(index) {
Some(array.value(index))
} else {
None
};
Box::new(FixedSizeBinaryScalar::new(array.data_type().clone(), value))
}
FixedSizeList => {
let array = array.as_any().downcast_ref::<FixedSizeListArray>().unwrap();
let value = if array.is_valid(index) {
Some(array.value(index).into())
} else {
None
};
Box::new(FixedSizeListScalar::new(array.data_type().clone(), value))
}
Union | Map => todo!(),
Dictionary(key_type) => match_integer_type!(key_type, |$T| {
let array = array
Expand Down
28 changes: 28 additions & 0 deletions tests/it/scalar/fixed_size_binary.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,28 @@
use arrow2::{
datatypes::DataType,
scalar::{FixedSizeBinaryScalar, Scalar},
};

#[allow(clippy::eq_op)]
#[test]
fn equal() {
let a = FixedSizeBinaryScalar::new(DataType::FixedSizeBinary(1), Some("a"));
let b = FixedSizeBinaryScalar::new(DataType::FixedSizeBinary(1), None::<&str>);
assert_eq!(a, a);
assert_eq!(b, b);
assert!(a != b);
let b = FixedSizeBinaryScalar::new(DataType::FixedSizeBinary(1), Some("b"));
assert!(a != b);
assert_eq!(b, b);
}

#[test]
fn basics() {
let a = FixedSizeBinaryScalar::new(DataType::FixedSizeBinary(1), Some("a"));

assert_eq!(a.value(), Some(b"a".as_ref()));
assert_eq!(a.data_type(), &DataType::FixedSizeBinary(1));
assert!(a.is_valid());

let _: &dyn std::any::Any = a.as_any();
}
48 changes: 48 additions & 0 deletions tests/it/scalar/fixed_size_list.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,48 @@
use std::sync::Arc;

use arrow2::{
array::{Array, BooleanArray},
datatypes::{DataType, Field},
scalar::{FixedSizeListScalar, Scalar},
};

#[allow(clippy::eq_op)]
#[test]
fn equal() {
let dt = DataType::FixedSizeList(Box::new(Field::new("a", DataType::Boolean, true)), 2);
let a = FixedSizeListScalar::new(
dt.clone(),
Some(Arc::new(BooleanArray::from_slice([true, false])) as Arc<dyn Array>),
);

let b = FixedSizeListScalar::new(dt.clone(), None);

assert_eq!(a, a);
assert_eq!(b, b);
assert!(a != b);

let b = FixedSizeListScalar::new(
dt,
Some(Arc::new(BooleanArray::from_slice([true, true])) as Arc<dyn Array>),
);
assert!(a != b);
assert_eq!(b, b);
}

#[test]
fn basics() {
let dt = DataType::FixedSizeList(Box::new(Field::new("a", DataType::Boolean, true)), 2);
let a = FixedSizeListScalar::new(
dt.clone(),
Some(Arc::new(BooleanArray::from_slice([true, false])) as Arc<dyn Array>),
);

assert_eq!(
BooleanArray::from_slice([true, false]),
a.values().unwrap().as_ref()
);
assert_eq!(a.data_type(), &dt);
assert!(a.is_valid());

let _: &dyn std::any::Any = a.as_any();
}
2 changes: 2 additions & 0 deletions tests/it/scalar/mod.rs
Original file line number Diff line number Diff line change
@@ -1,5 +1,7 @@
mod binary;
mod boolean;
mod fixed_size_binary;
mod fixed_size_list;
mod list;
mod null;
mod primitive;
Expand Down