use std::fmt::{self, Debug, Display, Formatter};
use std::sync::OnceLock;
use crate::enums::error::MinarrowError;
use crate::enums::shape_dim::ShapeDim;
use crate::structs::views::bitmask_view::BitmaskV;
use crate::traits::concatenate::Concatenate;
use crate::traits::print::MAX_PREVIEW;
use crate::traits::shape::Shape;
use crate::{Array, ArrayV, Bitmask, FieldArray, MaskedArray, NumericArray};
#[derive(Clone, PartialEq)]
pub struct NumericArrayV {
pub array: NumericArray,
pub offset: usize,
len: usize,
null_count: OnceLock<usize>,
}
impl NumericArrayV {
pub fn new(array: NumericArray, offset: usize, len: usize) -> Self {
assert!(
offset + len <= array.len(),
"NumericArrayView: window out of bounds (offset + len = {}, array.len = {})",
offset + len,
array.len()
);
Self {
array,
offset,
len,
null_count: OnceLock::new(),
}
}
pub fn new_nc(
array: NumericArray,
offset: usize,
len: usize,
null_count: usize,
) -> Self {
assert!(
offset + len <= array.len(),
"NumericArrayView: window out of bounds (offset + len = {}, array.len = {})",
offset + len,
array.len()
);
let lock = OnceLock::new();
let _ = lock.set(null_count); Self {
array,
offset,
len,
null_count: lock,
}
}
#[inline]
pub fn is_empty(&self) -> bool {
self.len == 0
}
#[inline]
pub fn inner_array(&self) -> Array {
Array::NumericArray(self.array.clone()) }
#[inline]
pub fn get_f64(&self, i: usize) -> Option<f64> {
if i >= self.len {
return None;
}
let phys_idx = self.offset + i;
match &self.array {
NumericArray::Int32(arr) => arr.get(phys_idx).map(|v| v as f64),
NumericArray::Int64(arr) => arr.get(phys_idx).map(|v| v as f64),
NumericArray::UInt32(arr) => arr.get(phys_idx).map(|v| v as f64),
NumericArray::UInt64(arr) => arr.get(phys_idx).map(|v| v as f64),
NumericArray::Float32(arr) => arr.get(phys_idx).map(|v| v as f64),
NumericArray::Float64(arr) => arr.get(phys_idx),
NumericArray::Null => None,
#[cfg(feature = "extended_numeric_types")]
NumericArray::Int8(arr) => arr.get(phys_idx).map(|v| v as f64),
#[cfg(feature = "extended_numeric_types")]
NumericArray::Int16(arr) => arr.get(phys_idx).map(|v| v as f64),
#[cfg(feature = "extended_numeric_types")]
NumericArray::UInt8(arr) => arr.get(phys_idx).map(|v| v as f64),
#[cfg(feature = "extended_numeric_types")]
NumericArray::UInt16(arr) => arr.get(phys_idx).map(|v| v as f64),
}
}
#[inline]
pub unsafe fn get_f64_unchecked(&self, i: usize) -> Option<f64> {
let phys_idx = self.offset + i;
match &self.array {
NumericArray::Int32(arr) => unsafe { arr.get_unchecked(phys_idx) }.map(|v| v as f64),
NumericArray::Int64(arr) => unsafe { arr.get_unchecked(phys_idx) }.map(|v| v as f64),
NumericArray::UInt32(arr) => unsafe { arr.get_unchecked(phys_idx) }.map(|v| v as f64),
NumericArray::UInt64(arr) => unsafe { arr.get_unchecked(phys_idx) }.map(|v| v as f64),
NumericArray::Float32(arr) => unsafe { arr.get_unchecked(phys_idx) }.map(|v| v as f64),
NumericArray::Float64(arr) => unsafe { arr.get_unchecked(phys_idx) },
NumericArray::Null => None,
#[cfg(feature = "extended_numeric_types")]
NumericArray::Int8(arr) => unsafe { arr.get_unchecked(phys_idx) }.map(|v| v as f64),
#[cfg(feature = "extended_numeric_types")]
NumericArray::Int16(arr) => unsafe { arr.get_unchecked(phys_idx) }.map(|v| v as f64),
#[cfg(feature = "extended_numeric_types")]
NumericArray::UInt8(arr) => unsafe { arr.get_unchecked(phys_idx) }.map(|v| v as f64),
#[cfg(feature = "extended_numeric_types")]
NumericArray::UInt16(arr) => unsafe { arr.get_unchecked(phys_idx) }.map(|v| v as f64),
}
}
#[inline]
pub fn slice(&self, offset: usize, len: usize) -> Self {
assert!(
offset + len <= self.len,
"NumericArrayView::slice: out of bounds"
);
Self {
array: self.array.clone(),
offset: self.offset + offset,
len,
null_count: OnceLock::new(),
}
}
pub fn to_numeric_array(&self) -> NumericArray {
self.inner_array().slice_clone(self.offset, self.len).num()
}
#[inline]
pub fn end(&self) -> usize {
self.offset + self.len
}
#[inline]
pub fn as_tuple(&self) -> (NumericArray, usize, usize) {
(self.array.clone(), self.offset, self.len)
}
#[inline]
pub fn as_tuple_ref(&self) -> (&NumericArray, usize, usize) {
(&self.array, self.offset, self.len)
}
#[inline]
pub fn len(&self) -> usize {
self.len
}
#[inline]
pub fn null_count(&self) -> usize {
*self
.null_count
.get_or_init(|| match self.array.null_mask() {
Some(mask) => mask.view(self.offset, self.len).count_zeros(),
None => 0,
})
}
#[inline]
pub fn null_mask_view(&self) -> Option<BitmaskV> {
self.array
.null_mask()
.map(|mask| mask.view(self.offset, self.len))
}
#[inline]
pub fn set_null_count(&self, count: usize) -> Result<(), usize> {
self.null_count.set(count).map_err(|_| count)
}
pub fn guarantee_f64(&mut self) -> (&[f64], Option<&Bitmask>, Option<usize>) {
if !matches!(&self.array, NumericArray::Float64(_)) {
let old = std::mem::take(&mut self.array);
self.array = old.cow_into_f64();
}
let NumericArray::Float64(arr) = &self.array else { unreachable!() };
let slice = &arr.data.as_slice()[self.offset..self.offset + self.len];
let mask = arr.null_mask.as_ref();
let nc = if mask.is_some() {
Some(self.null_count())
} else {
None
};
(slice, mask, nc)
}
}
impl From<NumericArray> for NumericArrayV {
fn from(array: NumericArray) -> Self {
let len = array.len();
NumericArrayV {
array,
offset: 0,
len,
null_count: OnceLock::new(),
}
}
}
impl From<FieldArray> for NumericArrayV {
fn from(field_array: FieldArray) -> Self {
match field_array.array {
Array::NumericArray(arr) => {
let len = arr.len();
NumericArrayV {
array: arr,
offset: 0,
len,
null_count: OnceLock::new(),
}
}
_ => panic!("FieldArray does not contain a NumericArray"),
}
}
}
impl From<Array> for NumericArrayV {
fn from(array: Array) -> Self {
match array {
Array::NumericArray(arr) => {
let len = arr.len();
NumericArrayV {
array: arr,
offset: 0,
len,
null_count: OnceLock::new(),
}
}
_ => panic!("Array is not a NumericArray"),
}
}
}
impl From<ArrayV> for NumericArrayV {
fn from(view: ArrayV) -> Self {
let (array, offset, len) = view.as_tuple();
match array {
Array::NumericArray(inner) => Self {
array: inner,
offset,
len,
null_count: OnceLock::new(),
},
_ => panic!("From<ArrayView>: expected NumericArray variant"),
}
}
}
impl Debug for NumericArrayV {
fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
f.debug_struct("NumericArrayView")
.field("offset", &self.offset)
.field("len", &self.len)
.field("array", &self.array)
.field("cached_null_count", &self.null_count.get())
.finish()
}
}
impl Shape for NumericArrayV {
fn shape(&self) -> ShapeDim {
ShapeDim::Rank1(self.len())
}
}
impl Concatenate for NumericArrayV {
fn concat(self, other: Self) -> Result<Self, MinarrowError> {
let self_array = self.to_numeric_array();
let other_array = other.to_numeric_array();
let concatenated = self_array.concat(other_array)?;
Ok(NumericArrayV::from(concatenated))
}
}
impl Display for NumericArrayV {
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
let dtype = match &self.array {
#[cfg(feature = "extended_numeric_types")]
NumericArray::Int8(_) => "Int8",
#[cfg(feature = "extended_numeric_types")]
NumericArray::Int16(_) => "Int16",
NumericArray::Int32(_) => "Int32",
NumericArray::Int64(_) => "Int64",
#[cfg(feature = "extended_numeric_types")]
NumericArray::UInt8(_) => "UInt8",
#[cfg(feature = "extended_numeric_types")]
NumericArray::UInt16(_) => "UInt16",
NumericArray::UInt32(_) => "UInt32",
NumericArray::UInt64(_) => "UInt64",
NumericArray::Float32(_) => "Float32",
NumericArray::Float64(_) => "Float64",
NumericArray::Null => "Null",
};
writeln!(
f,
"NumericArrayView<{dtype}> [{} rows] (offset: {}, nulls: {})",
self.len(),
self.offset,
self.null_count()
)?;
let max = self.len().min(MAX_PREVIEW);
for i in 0..max {
match self.get_f64(i) {
Some(v) => writeln!(f, " {v}")?,
None => writeln!(f, " ·")?,
}
}
if self.len() > MAX_PREVIEW {
writeln!(f, " ... ({} more)", self.len() - MAX_PREVIEW)?;
}
Ok(())
}
}
#[cfg(test)]
mod tests {
use std::sync::Arc;
use super::*;
use crate::{Array, Bitmask, IntegerArray, NumericArray, vec64};
#[test]
fn test_numeric_array_view_basic_indexing_and_slice() {
let mut arr = IntegerArray::<i32>::default();
arr.push(100);
arr.push(200);
arr.push(300);
arr.push(400);
let numeric = NumericArray::Int32(Arc::new(arr));
let view = NumericArrayV::new(numeric.clone(), 1, 2);
assert_eq!(view.len(), 2);
assert_eq!(view.offset, 1);
assert_eq!(view.get_f64(0), Some(200.0));
assert_eq!(view.get_f64(1), Some(300.0));
assert_eq!(view.get_f64(2), None);
let sub = view.slice(1, 1);
assert_eq!(sub.len(), 1);
assert_eq!(sub.get_f64(0), Some(300.0));
assert_eq!(sub.get_f64(1), None);
}
#[test]
fn test_numeric_array_view_null_count_and_cache() {
let mut arr = IntegerArray::<i32>::default();
arr.push(1);
arr.push(2);
arr.push(3);
arr.push(4);
let mut mask = Bitmask::new_set_all(4, true);
mask.set(2, false);
arr.null_mask = Some(mask);
let numeric = NumericArray::Int32(Arc::new(arr));
let view = NumericArrayV::new(numeric.clone(), 0, 4);
assert_eq!(view.null_count(), 1, "Null count should detect one null");
assert_eq!(view.null_count(), 1);
let view2 = view.slice(0, 2);
assert_eq!(view2.null_count(), 0);
let view3 = view.slice(2, 2);
assert_eq!(view3.null_count(), 1);
}
#[test]
fn test_numeric_array_view_with_supplied_null_count() {
let mut arr = IntegerArray::<i32>::default();
arr.push(5);
arr.push(6);
let numeric = NumericArray::Int32(Arc::new(arr));
let view = NumericArrayV::new_nc(numeric.clone(), 0, 2, 99);
assert_eq!(view.null_count(), 99);
assert!(view.set_null_count(101).is_err());
assert_eq!(view.null_count(), 99);
}
#[test]
fn test_numeric_array_view_to_numeric_array_and_as_tuple() {
let mut arr = IntegerArray::<i32>::default();
for v in 10..20 {
arr.push(v);
}
let numeric = NumericArray::Int32(Arc::new(arr));
let view = NumericArrayV::new(numeric.clone(), 4, 3);
let arr2 = view.to_numeric_array();
if let NumericArray::Int32(a2) = arr2 {
assert_eq!(a2.data, vec64![14, 15, 16]);
} else {
panic!("Unexpected variant");
}
let tup = view.as_tuple();
assert_eq!(&tup.0, &numeric);
assert_eq!(tup.1, 4);
assert_eq!(tup.2, 3);
}
#[test]
fn test_numeric_array_view_null_mask_view() {
let mut arr = IntegerArray::<i32>::default();
arr.push(2);
arr.push(4);
arr.push(6);
let mut mask = Bitmask::new_set_all(3, true);
mask.set(0, false);
arr.null_mask = Some(mask);
let numeric = NumericArray::Int32(Arc::new(arr));
let view = NumericArrayV::new(numeric, 1, 2);
let mask_view = view.null_mask_view().expect("Should have mask");
assert_eq!(mask_view.len(), 2);
assert!(mask_view.get(0));
assert!(mask_view.get(1));
}
#[test]
fn test_numeric_array_view_from_numeric_array_and_array() {
let mut arr = IntegerArray::<i32>::default();
arr.push(1);
arr.push(2);
let numeric = NumericArray::Int32(Arc::new(arr));
let view_from_numeric = NumericArrayV::from(numeric.clone());
assert_eq!(view_from_numeric.len(), 2);
assert_eq!(view_from_numeric.get_f64(0), Some(1.0));
let array = Array::NumericArray(numeric);
let view_from_array = NumericArrayV::from(array);
assert_eq!(view_from_array.len(), 2);
assert_eq!(view_from_array.get_f64(1), Some(2.0));
}
#[test]
#[should_panic(expected = "Array is not a NumericArray")]
fn test_numeric_array_view_from_array_panics_on_wrong_variant() {
let array = Array::Null;
let _view = NumericArrayV::from(array);
}
}