use crate::{
bitmap::Bitmap,
datatypes::{DataType, Field},
error::Error,
offset::OffsetsBuffer,
};
use super::{new_empty_array, specification::try_check_offsets_bounds, Array};
#[cfg(feature = "arrow")]
mod data;
mod ffi;
pub(super) mod fmt;
mod iterator;
#[derive(Clone)]
pub struct MapArray {
data_type: DataType,
offsets: OffsetsBuffer<i32>,
field: Box<dyn Array>,
validity: Option<Bitmap>,
}
impl MapArray {
pub fn try_new(
data_type: DataType,
offsets: OffsetsBuffer<i32>,
field: Box<dyn Array>,
validity: Option<Bitmap>,
) -> Result<Self, Error> {
try_check_offsets_bounds(&offsets, field.len())?;
let inner_field = Self::try_get_field(&data_type)?;
if let DataType::Struct(inner) = inner_field.data_type() {
if inner.len() != 2 {
return Err(Error::InvalidArgumentError(
"MapArray's inner `Struct` must have 2 fields (keys and maps)".to_string(),
));
}
} else {
return Err(Error::InvalidArgumentError(
"MapArray expects `DataType::Struct` as its inner logical type".to_string(),
));
}
if field.data_type() != inner_field.data_type() {
return Err(Error::InvalidArgumentError(
"MapArray expects `field.data_type` to match its inner DataType".to_string(),
));
}
if validity
.as_ref()
.map_or(false, |validity| validity.len() != offsets.len_proxy())
{
return Err(Error::oos(
"validity mask length must match the number of values",
));
}
Ok(Self {
data_type,
field,
offsets,
validity,
})
}
pub fn new(
data_type: DataType,
offsets: OffsetsBuffer<i32>,
field: Box<dyn Array>,
validity: Option<Bitmap>,
) -> Self {
Self::try_new(data_type, offsets, field, validity).unwrap()
}
pub fn new_null(data_type: DataType, length: usize) -> Self {
let field = new_empty_array(Self::get_field(&data_type).data_type().clone());
Self::new(
data_type,
vec![0i32; 1 + length].try_into().unwrap(),
field,
Some(Bitmap::new_zeroed(length)),
)
}
pub fn new_empty(data_type: DataType) -> Self {
let field = new_empty_array(Self::get_field(&data_type).data_type().clone());
Self::new(data_type, OffsetsBuffer::default(), field, None)
}
}
impl MapArray {
pub fn slice(&mut self, offset: usize, length: usize) {
assert!(
offset + length <= self.len(),
"the offset of the new Buffer cannot exceed the existing length"
);
unsafe { self.slice_unchecked(offset, length) }
}
#[inline]
pub unsafe fn slice_unchecked(&mut self, offset: usize, length: usize) {
self.validity.as_mut().and_then(|bitmap| {
bitmap.slice_unchecked(offset, length);
(bitmap.unset_bits() > 0).then(|| bitmap)
});
self.offsets.slice_unchecked(offset, length + 1);
}
impl_sliced!();
impl_mut_validity!();
impl_into_array!();
pub(crate) fn try_get_field(data_type: &DataType) -> Result<&Field, Error> {
if let DataType::Map(field, _) = data_type.to_logical_type() {
Ok(field.as_ref())
} else {
Err(Error::oos(
"The data_type's logical type must be DataType::Map",
))
}
}
pub(crate) fn get_field(data_type: &DataType) -> &Field {
Self::try_get_field(data_type).unwrap()
}
}
impl MapArray {
#[inline]
pub fn len(&self) -> usize {
self.offsets.len_proxy()
}
#[inline]
pub fn offsets(&self) -> &OffsetsBuffer<i32> {
&self.offsets
}
#[inline]
pub fn field(&self) -> &Box<dyn Array> {
&self.field
}
#[inline]
pub fn value(&self, i: usize) -> Box<dyn Array> {
assert!(i < self.len());
unsafe { self.value_unchecked(i) }
}
#[inline]
pub unsafe fn value_unchecked(&self, i: usize) -> Box<dyn Array> {
let (start, end) = self.offsets.start_end_unchecked(i);
let length = end - start;
self.field.sliced_unchecked(start, length)
}
}
impl Array for MapArray {
impl_common_array!();
fn validity(&self) -> Option<&Bitmap> {
self.validity.as_ref()
}
#[inline]
fn with_validity(&self, validity: Option<Bitmap>) -> Box<dyn Array> {
Box::new(self.clone().with_validity(validity))
}
}