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

52 changes: 52 additions & 0 deletions src/scalar/fixed_size_binary.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,52 @@
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 {
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, DataType::FixedSizeBinary(x.len()));
illumination-k marked this conversation as resolved.
Show resolved Hide resolved
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
}
}
69 changes: 69 additions & 0 deletions src/scalar/fixed_size_list.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,69 @@
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: Arc<dyn Array>,
is_valid: bool,
Copy link
Owner

Choose a reason for hiding this comment

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

What do you think about

Suggested change
values: Arc<dyn Array>,
is_valid: bool,
values: Option<Arc<dyn Array>>,

and remove is_valid?

Copy link
Contributor Author

Choose a reason for hiding this comment

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

I referred to ListScalar, but I think Option<Arc<dyn Array>> is better and value method should return Option<&Arc<dyn Array>> for consistency of other scalar.

StructScalar and ListScalar also have is_valid. Should we use Option and remove is_valid for them? (Maybe another PR)

data_type: DataType,
}

impl PartialEq for FixedSizeListScalar {
fn eq(&self, other: &Self) -> bool {
(self.data_type == other.data_type)
&& (self.is_valid == other.is_valid)
&& ((!self.is_valid) | (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 (is_valid, values) = match values {
Some(values) => {
assert_eq!(inner_data_type, values.data_type());
assert_eq!(size, values.len());
(true, values)
}
None => (false, new_empty_array(inner_data_type.clone()).into()),
};
Self {
values,
is_valid,
data_type,
}
}

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

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

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

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();
}
45 changes: 45 additions & 0 deletions tests/it/scalar/fixed_size_list.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,45 @@
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().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