Merge pull request #10 from pbower/shape

feat: Add Shape support

Author: pbower

.gitignore

@@ -19,3 +19,5 @@ target
 #  and can be added to the global gitignore or merged into this file.  For a more nuclear
 #  option (not recommended) you can uncomment the following to ignore the entire idea folder.
 #.idea/
+
+TODO.md

src/enums/shape_dim.rs

@@ -0,0 +1,60 @@
+//! # ShapeDim Enum Module
+//!
+//! Companion to [crate::traits::shape::Shape];
+//!
+//! Contains all supported `Shape` variants.
+
+use crate::traits::shape::Shape;
+
+/// Recursively-describable dimensional rank for any `Value`.
+#[derive(Clone, PartialEq, Eq)]
+pub enum ShapeDim {
+    /// Rank-0 - must always be `1`
+    Rank0(usize),
+
+    /// Array row count
+    Rank1(usize),
+
+    /// Relational table with row/column counts.
+    Rank2 { rows: usize, cols: usize },
+
+    /// 3d object
+    Rank3 { x: usize, y: usize, z: usize },
+
+    /// 4d Object
+    Rank4 {
+        a: usize,
+        b: usize,
+        c: usize,
+        d: usize,
+    },
+
+    /// N-dimensional tensor.
+    RankN(Vec<usize>),
+
+    /// Dictionary shape
+    Dictionary {
+        // Number of keys
+        n_keys: usize,
+        // Number of values for each key
+        n_values: Vec<usize>,
+    },
+
+    /// Heterogeneous ordered collection.
+    /// Covers lists, tuples, cubes (with varying row-counts) and user-custom chunked types.
+    ///
+    /// The order is significant; for tuples it is the fixed arity order.
+    Collection(Vec<ShapeDim>),
+
+    /// Shape could not be determined.
+    Unknown,
+}
+
+
+/// Implement `Shape` for `ShapeDim` so recursive calls like `item.shape_3d()`
+/// compile when iterating `Collection(Vec<ShapeDim>)`.
+impl Shape for ShapeDim {
+    fn shape(&self) -> ShapeDim {
+        self.clone()
+    }
+}

src/lib.rs

@@ -100,6 +100,7 @@ pub mod enums {
         pub mod temporal_array;
     }
     pub mod operators;
+    pub mod shape_dim;
 }
 
 /// Contains SIMD-accelerated kernels for the 'essentials' that are highly coupled to this crate
@@ -179,6 +180,7 @@ pub mod traits {
     pub mod print;
     pub mod type_unions;
     pub mod custom_value;
+    pub mod shape;
 }
 
 pub mod aliases;

src/structs/bitmask.rs

@@ -18,6 +18,8 @@ use std::fmt::{Debug, Display, Formatter, Result as FmtResult};
 use std::ops::{BitAnd, BitOr, Deref, DerefMut, Index, Not};
 
 use crate::structs::vec64::Vec64;
+use crate::traits::shape::Shape;
+use crate::enums::shape_dim::ShapeDim;
 use crate::{BitmaskV, Buffer, Length, Offset};
 
 /// TODO: Move bitmask kernels here
@@ -731,6 +733,11 @@ impl Display for Bitmask {
     }
 }
 
+impl Shape for Bitmask {
+    fn shape(&self) -> ShapeDim {
+        ShapeDim::Rank1(self.len())
+    }
+}
 
 #[cfg(test)]
 mod tests {

src/structs/chunked/super_array.rs

@@ -13,6 +13,8 @@ use std::iter::FromIterator;
 #[cfg(feature = "views")]
 use std::sync::Arc;
 
+use crate::traits::shape::Shape;
+use crate::enums::shape_dim::ShapeDim;
 #[cfg(feature = "views")]
 use crate::ArrayV;
 #[cfg(feature = "views")]
@@ -719,6 +721,12 @@ impl From<FieldArray> for SuperArray {
     }
 }
 
+impl Shape for SuperArray {
+    fn shape(&self) -> ShapeDim {
+        ShapeDim::Rank1(self.len())
+    }
+}
+
 impl Display for SuperArray {
     fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
         writeln!(

src/structs/chunked/super_table.rs

@@ -27,6 +27,8 @@ use std::sync::Arc;
 use crate::structs::field::Field;
 use crate::structs::field_array::FieldArray;
 use crate::structs::table::Table;
+use crate::traits::shape::Shape;
+use crate::enums::shape_dim::ShapeDim;
 #[cfg(feature = "views")]
 use crate::{SuperTableV, TableV};
 
@@ -154,6 +156,12 @@ impl SuperTable {
     pub fn n_cols(&self) -> usize {
         self.schema.len()
     }
+
+    #[inline]
+    pub fn n_rows(&self) -> usize {
+        self.n_rows
+    }
+
     #[inline]
     pub fn n_batches(&self) -> usize {
         self.batches.len()
@@ -223,6 +231,7 @@ impl SuperTable {
             name
         }
     }
+
 }
 
 impl Default for SuperTable {
@@ -238,6 +247,12 @@ impl FromIterator<Table> for SuperTable {
     }
 }
 
+impl Shape for SuperTable {
+    fn shape(&self) -> ShapeDim {
+        ShapeDim::Rank2 { rows: self.n_rows(), cols: self.n_cols() }
+    }
+}
+
 impl Display for SuperTable {
     fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
         writeln!(

src/structs/cube.rs

@@ -32,6 +32,8 @@ use super::field_array::FieldArray;
 #[cfg(feature = "views")]
 use crate::aliases::CubeV;
 use crate::ffi::arrow_dtype::ArrowType;
+use crate::traits::shape::Shape;
+use crate::enums::shape_dim::ShapeDim;
 use crate::{Field, Table};
 #[cfg(feature = "views")]
 use crate::TableV;
@@ -64,6 +66,7 @@ pub struct Cube {
     pub tables: Vec<Table>,
     /// Number of rows in each table
     pub n_rows: Vec<usize>,
+    
     /// Cube name
     pub name: String,
     // Third-dimensional index column names
@@ -160,6 +163,11 @@ impl Cube {
         self.tables.len()
     }
 
+    /// Returns the number of rows
+    pub fn n_rows(&self) -> Vec<usize> {
+        self.n_rows.clone()
+    }
+
     /// Returns the number of columns.
     pub fn n_cols(&self) -> usize {
         self.tables[0].n_cols()
@@ -439,6 +447,13 @@ impl IntoIterator for Cube {
     }
 }
 
+impl Shape for Cube {
+    fn shape(&self) -> ShapeDim {
+        ShapeDim::Collection(self.tables.iter().map(|t| t.shape()).collect())
+    }
+}
+
+
 #[cfg(test)]
 mod tests {
     use super::*;

src/structs/field_array.rs

@@ -21,6 +21,8 @@ use polars::series::Series;
 #[cfg(feature = "views")]
 use crate::aliases::FieldAVT;
 use crate::ffi::arrow_dtype::ArrowType;
+use crate::traits::shape::Shape;
+use crate::enums::shape_dim::ShapeDim;
 use crate::{Array, Field};
 
 
@@ -245,6 +247,12 @@ impl Display for FieldArray {
     }
 }
 
+impl Shape for FieldArray {
+    fn shape(&self) -> ShapeDim {
+        ShapeDim::Rank1(self.len())
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;