Serde

spiraldb · Apr 8, 2024 · 182678d · 182678d
1 parent 36d2dd2
commit 182678d
Show file tree

Hide file tree

Showing 10 changed files with 265 additions and 9 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/vortex-array2/Cargo.toml b/vortex-array2/Cargo.toml
@@ -15,6 +15,7 @@ arrow-buffer = { workspace = true }
 flatbuffers = { workspace = true }
 flexbuffers = { workspace = true }
 half = { workspace = true }
+humansize = { workspace = true }
 paste = { workspace = true }
 serde = { workspace = true, features = ["derive"] }
 vortex-array = { path = "../vortex-array", features = ["serde"] }

diff --git a/vortex-array2/src/array/bool/mod.rs b/vortex-array2/src/array/bool/mod.rs
@@ -8,6 +8,7 @@ use vortex_schema::DType;
 use crate::impl_encoding;
 use crate::validity::Validity;
 use crate::validity::{ArrayValidity, ValidityMetadata};
+use crate::visitor::{AcceptArrayVisitor, ArrayVisitor};
 use crate::ArrayMetadata;
 use crate::{ArrayData, TypedArrayData};
 use crate::{ArrayView, ToArrayData};
@@ -109,6 +110,13 @@ impl ToArrayData for BoolArray<'_> {
     }
 }
 
+impl AcceptArrayVisitor for BoolArray<'_> {
+    fn accept(&self, visitor: &mut dyn ArrayVisitor) -> VortexResult<()> {
+        visitor.visit_buffer(self.buffer())?;
+        visitor.visit_validity(self.validity())
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use crate::array::bool::BoolData;

diff --git a/vortex-array2/src/array/primitive/mod.rs b/vortex-array2/src/array/primitive/mod.rs
@@ -1,13 +1,14 @@
 mod compute;
 
-use arrow_buffer::{Buffer, ScalarBuffer};
+use arrow_buffer::{ArrowNativeType, Buffer, ScalarBuffer};
 use serde::{Deserialize, Serialize};
 use vortex::ptype::{NativePType, PType};
 use vortex_error::VortexResult;
 use vortex_schema::DType;
 
 use crate::impl_encoding;
 use crate::validity::{ArrayValidity, Validity, ValidityMetadata};
+use crate::visitor::{AcceptArrayVisitor, ArrayVisitor};
 use crate::ArrayMetadata;
 use crate::{ArrayData, TypedArrayData};
 use crate::{ArrayView, ToArrayData};
@@ -71,6 +72,10 @@ impl PrimitiveData {
         .unwrap()
         .try_into()
     }
+
+    pub fn from_vec<T: NativePType + ArrowNativeType>(values: Vec<T>) -> Self {
+        Self::try_new(ScalarBuffer::from(values), Validity::NonNullable).unwrap()
+    }
 }
 
 impl ArrayTrait for PrimitiveArray<'_> {
@@ -94,3 +99,10 @@ impl ToArrayData for PrimitiveArray<'_> {
         todo!()
     }
 }
+
+impl AcceptArrayVisitor for PrimitiveArray<'_> {
+    fn accept(&self, visitor: &mut dyn ArrayVisitor) -> VortexResult<()> {
+        visitor.visit_buffer(self.buffer())?;
+        visitor.visit_validity(self.validity())
+    }
+}
diff --git a/vortex-array2/src/array/ree/mod.rs b/vortex-array2/src/array/ree/mod.rs
@@ -6,6 +6,7 @@ use vortex_schema::DType;
 
 use crate::impl_encoding;
 use crate::validity::ArrayValidity;
+use crate::visitor::{AcceptArrayVisitor, ArrayVisitor};
 use crate::{Array, ArrayMetadata};
 use crate::{ArrayData, TypedArrayData};
 use crate::{ArrayView, ToArrayData};
@@ -22,6 +23,17 @@ pub struct REEArray<'a> {
     dtype: &'a DType,
     values: Array<'a>,
     run_ends: Array<'a>,
+    length: usize,
+}
+
+impl REEArray<'_> {
+    pub fn values(&self) -> &Array {
+        &self.values
+    }
+
+    pub fn run_ends(&self) -> &Array {
+        &self.run_ends
+    }
 }
 
 impl REEData {
@@ -55,6 +67,7 @@ impl<'v> TryFromArrayParts<'v, REEMetadata> for REEArray<'v> {
             run_ends: parts
                 .child(1, &metadata.ends_dtype)
                 .ok_or_else(|| vortex_err!("REEArray missing run_ends"))?,
+            length: metadata.length,
         })
     }
 }
@@ -65,7 +78,7 @@ impl ArrayTrait for REEArray<'_> {
     }
 
     fn len(&self) -> usize {
-        todo!()
+        self.length
     }
 }
 
@@ -80,3 +93,10 @@ impl ToArrayData for REEArray<'_> {
         todo!()
     }
 }
+
+impl AcceptArrayVisitor for REEArray<'_> {
+    fn accept(&self, visitor: &mut dyn ArrayVisitor) -> VortexResult<()> {
+        visitor.visit_array("values", self.values())?;
+        visitor.visit_array("run_ends", self.run_ends())
+    }
+}
diff --git a/vortex-array2/src/encoding.rs b/vortex-array2/src/encoding.rs
@@ -33,10 +33,10 @@ impl Debug for dyn ArrayEncoding + '_ {
 }
 
 impl dyn ArrayEncoding {
-    pub(crate) fn with_view<'v, R, F: Fn(&dyn ArrayTrait) -> R>(
+    pub(crate) fn with_view<'v, R, F: FnMut(&dyn ArrayTrait) -> R>(
         &self,
         view: &'v ArrayView<'v>,
-        f: F,
+        mut f: F,
     ) -> R {
         let mut result = None;
 
@@ -52,7 +52,11 @@ impl dyn ArrayEncoding {
         result.unwrap()
     }
 
-    pub(crate) fn with_data<R, F: Fn(&dyn ArrayTrait) -> R>(&self, data: &ArrayData, f: F) -> R {
+    pub(crate) fn with_data<R, F: FnMut(&dyn ArrayTrait) -> R>(
+        &self,
+        data: &ArrayData,
+        mut f: F,
+    ) -> R {
         let mut result = None;
 
         // Unwrap the result. This is safe since we validate that encoding against the

diff --git a/vortex-array2/src/lib.rs b/vortex-array2/src/lib.rs
@@ -7,10 +7,12 @@ mod data;
 pub mod encoding;
 mod implementation;
 mod metadata;
+mod tree;
 mod validity;
 mod view;
+mod visitor;
 
-use std::fmt::Debug;
+use std::fmt::{Debug, Display, Formatter};
 
 use arrow_buffer::Buffer;
 pub use context::*;
@@ -22,7 +24,9 @@ use vortex_error::VortexResult;
 use vortex_schema::DType;
 
 use crate::compute::ArrayCompute;
+use crate::encoding::EncodingRef;
 use crate::validity::ArrayValidity;
+use crate::visitor::{AcceptArrayVisitor, ArrayVisitor};
 
 #[derive(Debug, Clone)]
 pub enum Array<'v> {
@@ -32,13 +36,25 @@ pub enum Array<'v> {
 }
 
 impl Array<'_> {
+    pub fn encoding(&self) -> EncodingRef {
+        match self {
+            Array::Data(d) => d.encoding(),
+            Array::DataRef(d) => d.encoding(),
+            Array::View(v) => v.encoding(),
+        }
+    }
+
     pub fn dtype(&self) -> &DType {
         match self {
             Array::Data(d) => d.dtype(),
             Array::DataRef(d) => d.dtype(),
             Array::View(v) => v.dtype(),
         }
     }
+
+    pub fn len(&self) -> usize {
+        self.with_array(|a| a.len())
+    }
 }
 
 pub trait ToArray {
@@ -54,7 +70,7 @@ pub trait ToArrayData {
 }
 
 pub trait WithArray {
-    fn with_array<R, F: Fn(&dyn ArrayTrait) -> R>(&self, f: F) -> R;
+    fn with_array<R, F: FnMut(&dyn ArrayTrait) -> R>(&self, f: F) -> R;
 }
 
 pub trait ArrayParts<'a> {
@@ -68,7 +84,7 @@ pub trait TryFromArrayParts<'v, M: ArrayMetadata>: Sized + 'v {
 }
 
 /// Collects together the behaviour of an array.
-pub trait ArrayTrait: ArrayCompute + ArrayValidity + ToArrayData {
+pub trait ArrayTrait: ArrayCompute + ArrayValidity + AcceptArrayVisitor + ToArrayData {
     fn dtype(&self) -> &DType;
 
     fn len(&self) -> usize;
@@ -77,6 +93,25 @@ pub trait ArrayTrait: ArrayCompute + ArrayValidity + ToArrayData {
         // TODO(ngates): remove this default impl to encourage explicit implementation
         self.len() == 0
     }
+
+    fn nbytes(&self) -> usize {
+        let mut visitor = NBytesVisitor(0);
+        self.accept(&mut visitor).unwrap();
+        visitor.0
+    }
+}
+
+struct NBytesVisitor(usize);
+impl ArrayVisitor for NBytesVisitor {
+    fn visit_array(&mut self, _name: &str, array: &Array) -> VortexResult<()> {
+        self.0 += array.with_array(|a| a.nbytes());
+        Ok(())
+    }
+
+    fn visit_buffer(&mut self, buffer: &Buffer) -> VortexResult<()> {
+        self.0 += buffer.len();
+        Ok(())
+    }
 }
 
 impl ToArrayData for Array<'_> {
@@ -90,11 +125,29 @@ impl ToArrayData for Array<'_> {
 }
 
 impl WithArray for Array<'_> {
-    fn with_array<R, F: Fn(&dyn ArrayTrait) -> R>(&self, f: F) -> R {
+    fn with_array<R, F: FnMut(&dyn ArrayTrait) -> R>(&self, f: F) -> R {
         match self {
             Array::Data(d) => d.encoding().with_data(d, f),
             Array::DataRef(d) => d.encoding().with_data(d, f),
             Array::View(v) => v.encoding().with_view(v, f),
         }
     }
 }
+
+impl Display for Array<'_> {
+    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+        let prefix = match self {
+            Array::Data(_) => "",
+            Array::DataRef(_) => "&",
+            Array::View(_) => "$",
+        };
+        write!(
+            f,
+            "{}{}({}, len={})",
+            prefix,
+            self.encoding().id(),
+            self.dtype(),
+            self.len()
+        )
+    }
+}