use crate::{
bitmap::Bitmap,
datatypes::{DataType, Field},
error::Error,
};
use std::sync::Arc;
use super::{new_empty_array, new_null_array, Array};
mod ffi;
pub(super) mod fmt;
mod iterator;
pub use iterator::*;
mod mutable;
pub use mutable::*;
#[derive(Clone)]
pub struct FixedSizeListArray {
size: usize, data_type: DataType,
values: Box<dyn Array>,
validity: Option<Bitmap>,
}
impl FixedSizeListArray {
pub fn try_new(
data_type: DataType,
values: Box<dyn Array>,
validity: Option<Bitmap>,
) -> Result<Self, Error> {
let (child, size) = Self::try_child_and_size(&data_type)?;
let child_data_type = &child.data_type;
let values_data_type = values.data_type();
if child_data_type != values_data_type {
return Err(Error::oos(
format!("FixedSizeListArray's child's DataType must match. However, the expected DataType is {child_data_type:?} while it got {values_data_type:?}."),
));
}
if values.len() % size != 0 {
return Err(Error::oos(format!(
"values (of len {}) must be a multiple of size ({}) in FixedSizeListArray.",
values.len(),
size
)));
}
let len = values.len() / size;
if validity
.as_ref()
.map_or(false, |validity| validity.len() != len)
{
return Err(Error::oos(
"validity mask length must be equal to the number of values divided by size",
));
}
Ok(Self {
size,
data_type,
values,
validity,
})
}
pub fn new(data_type: DataType, values: Box<dyn Array>, validity: Option<Bitmap>) -> Self {
Self::try_new(data_type, values, validity).unwrap()
}
pub const fn size(&self) -> usize {
self.size
}
pub fn new_empty(data_type: DataType) -> Self {
let values = new_empty_array(Self::get_child_and_size(&data_type).0.data_type().clone());
Self::new(data_type, values, None)
}
pub fn new_null(data_type: DataType, length: usize) -> Self {
let (field, size) = Self::get_child_and_size(&data_type);
let values = new_null_array(field.data_type().clone(), length * size);
Self::new(data_type, values, Some(Bitmap::new_zeroed(length)))
}
pub fn boxed(self) -> Box<dyn Array> {
Box::new(self)
}
pub fn arced(self) -> Arc<dyn Array> {
Arc::new(self)
}
}
impl FixedSizeListArray {
#[must_use]
pub fn slice(&self, offset: usize, length: usize) -> Self {
assert!(
offset + length <= self.len(),
"the offset of the new Buffer cannot exceed the existing length"
);
unsafe { self.slice_unchecked(offset, length) }
}
#[must_use]
pub unsafe fn slice_unchecked(&self, offset: usize, length: usize) -> Self {
let validity = self
.validity
.clone()
.map(|bitmap| bitmap.slice_unchecked(offset, length))
.and_then(|bitmap| (bitmap.unset_bits() > 0).then(|| bitmap));
let values = self
.values
.clone()
.slice_unchecked(offset * self.size, length * self.size);
Self {
data_type: self.data_type.clone(),
size: self.size,
values,
validity,
}
}
#[must_use]
pub fn with_validity(mut self, validity: Option<Bitmap>) -> Self {
self.set_validity(validity);
self
}
pub fn set_validity(&mut self, validity: Option<Bitmap>) {
if matches!(&validity, Some(bitmap) if bitmap.len() != self.len()) {
panic!("validity must be equal to the array's length")
}
self.validity = validity;
}
}
impl FixedSizeListArray {
#[inline]
pub fn len(&self) -> usize {
self.values.len() / self.size
}
#[inline]
pub fn validity(&self) -> Option<&Bitmap> {
self.validity.as_ref()
}
pub fn values(&self) -> &Box<dyn Array> {
&self.values
}
#[inline]
pub fn value(&self, i: usize) -> Box<dyn Array> {
self.values.slice(i * self.size, self.size)
}
#[inline]
pub unsafe fn value_unchecked(&self, i: usize) -> Box<dyn Array> {
self.values.slice_unchecked(i * self.size, self.size)
}
}
impl FixedSizeListArray {
pub(crate) fn try_child_and_size(data_type: &DataType) -> Result<(&Field, usize), Error> {
match data_type.to_logical_type() {
DataType::FixedSizeList(child, size) => {
if *size == 0 {
return Err(Error::oos("FixedSizeBinaryArray expects a positive size"));
}
Ok((child.as_ref(), *size))
}
_ => Err(Error::oos(
"FixedSizeListArray expects DataType::FixedSizeList",
)),
}
}
pub(crate) fn get_child_and_size(data_type: &DataType) -> (&Field, usize) {
Self::try_child_and_size(data_type).unwrap()
}
pub fn default_datatype(data_type: DataType, size: usize) -> DataType {
let field = Box::new(Field::new("item", data_type, true));
DataType::FixedSizeList(field, size)
}
}
impl Array for FixedSizeListArray {
#[inline]
fn as_any(&self) -> &dyn std::any::Any {
self
}
#[inline]
fn as_any_mut(&mut self) -> &mut dyn std::any::Any {
self
}
#[inline]
fn len(&self) -> usize {
self.len()
}
#[inline]
fn data_type(&self) -> &DataType {
&self.data_type
}
fn validity(&self) -> Option<&Bitmap> {
self.validity.as_ref()
}
fn slice(&self, offset: usize, length: usize) -> Box<dyn Array> {
Box::new(self.slice(offset, length))
}
unsafe fn slice_unchecked(&self, offset: usize, length: usize) -> Box<dyn Array> {
Box::new(self.slice_unchecked(offset, length))
}
fn with_validity(&self, validity: Option<Bitmap>) -> Box<dyn Array> {
Box::new(self.clone().with_validity(validity))
}
fn to_boxed(&self) -> Box<dyn Array> {
Box::new(self.clone())
}
}