diff --git a/src/scalar/fixed_size_binary.rs b/src/scalar/fixed_size_binary.rs new file mode 100644 index 00000000000..c09466f9eb3 --- /dev/null +++ b/src/scalar/fixed_size_binary.rs @@ -0,0 +1,59 @@ +use crate::datatypes::DataType; + +use super::Scalar; + +#[derive(Debug, Clone, PartialEq)] +/// The [`Scalar`] implementation of fixed size binary ([`Option>`]). +pub struct FixedSizeBinaryScalar { + value: Option>, + 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>>(data_type: DataType, value: Option

) -> Self { + assert_eq!( + data_type.to_physical_type(), + crate::datatypes::PhysicalType::FixedSizeBinary + ); + Self { + value: value.map(|x| { + let x: Vec = 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 + } +} diff --git a/src/scalar/fixed_size_list.rs b/src/scalar/fixed_size_list.rs new file mode 100644 index 00000000000..fd41e2100d5 --- /dev/null +++ b/src/scalar/fixed_size_list.rs @@ -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>, + 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>) -> 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> { + 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 + } +} diff --git a/src/scalar/mod.rs b/src/scalar/mod.rs index 9a2ca3fecf5..4910bb67c29 100644 --- a/src/scalar/mod.rs +++ b/src/scalar/mod.rs @@ -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. @@ -120,8 +124,27 @@ pub fn new_scalar(array: &dyn Array, index: usize) -> Box { Box::new(StructScalar::new(array.data_type().clone(), None)) } } - FixedSizeBinary => todo!(), - FixedSizeList => todo!(), + FixedSizeBinary => { + let array = array + .as_any() + .downcast_ref::() + .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::().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 diff --git a/tests/it/scalar/fixed_size_binary.rs b/tests/it/scalar/fixed_size_binary.rs new file mode 100644 index 00000000000..3962c390180 --- /dev/null +++ b/tests/it/scalar/fixed_size_binary.rs @@ -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(); +} diff --git a/tests/it/scalar/fixed_size_list.rs b/tests/it/scalar/fixed_size_list.rs new file mode 100644 index 00000000000..a76b7aca38d --- /dev/null +++ b/tests/it/scalar/fixed_size_list.rs @@ -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), + ); + + 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), + ); + 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), + ); + + 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(); +} diff --git a/tests/it/scalar/mod.rs b/tests/it/scalar/mod.rs index 64a35e0edb4..2e7d105d7fd 100644 --- a/tests/it/scalar/mod.rs +++ b/tests/it/scalar/mod.rs @@ -1,5 +1,7 @@ mod binary; mod boolean; +mod fixed_size_binary; +mod fixed_size_list; mod list; mod null; mod primitive;