handled all tinyInt cases and added tests

spaceandtimelabs · Sep 28, 2024 · 4ab705f · 4ab705f
1 parent 7fef8d8
commit 4ab705f
Show file tree

Hide file tree

Showing 7 changed files with 409 additions and 1 deletion.
diff --git a/crates/proof-of-sql/src/base/database/column.rs b/crates/proof-of-sql/src/base/database/column.rs
@@ -107,6 +107,9 @@ impl<'a, S: Scalar> Column<'a, S> {
             LiteralValue::Boolean(value) => {
                 Column::Boolean(alloc.alloc_slice_fill_copy(length, *value))
             }
+            LiteralValue::TinyInt(value) => {
+                Column::TinyInt(alloc.alloc_slice_fill_copy(length, *value))
+            }
             LiteralValue::SmallInt(value) => {
                 Column::SmallInt(alloc.alloc_slice_fill_copy(length, *value))
             }

diff --git a/crates/proof-of-sql/src/base/database/literal_value.rs b/crates/proof-of-sql/src/base/database/literal_value.rs
@@ -13,6 +13,8 @@ use serde::{Deserialize, Serialize};
 pub enum LiteralValue<S: Scalar> {
     /// Boolean literals
     Boolean(bool),
+    /// i8 literals
+    TinyInt(i8),
     /// i16 literals
     SmallInt(i16),
     /// i32 literals
@@ -41,6 +43,7 @@ impl<S: Scalar> LiteralValue<S> {
     pub fn column_type(&self) -> ColumnType {
         match self {
             Self::Boolean(_) => ColumnType::Boolean,
+            Self::TinyInt(_) => ColumnType::TinyInt,
             Self::SmallInt(_) => ColumnType::SmallInt,
             Self::Int(_) => ColumnType::Int,
             Self::BigInt(_) => ColumnType::BigInt,
@@ -56,6 +59,7 @@ impl<S: Scalar> LiteralValue<S> {
     pub(crate) fn to_scalar(&self) -> S {
         match self {
             Self::Boolean(b) => b.into(),
+            Self::TinyInt(i) => i.into(),
             Self::SmallInt(i) => i.into(),
             Self::Int(i) => i.into(),
             Self::BigInt(i) => i.into(),

diff --git a/crates/proof-of-sql/src/base/database/owned_and_arrow_conversions.rs b/crates/proof-of-sql/src/base/database/owned_and_arrow_conversions.rs
@@ -148,6 +148,14 @@ impl<S: Scalar> TryFrom<&ArrayRef> for OwnedColumn<S> {
                     .collect::<Option<Vec<bool>>>()
                     .ok_or(OwnedArrowConversionError::NullNotSupportedYet)?,
             )),
+            DataType::Int8 => Ok(Self::TinyInt(
+                value
+                    .as_any()
+                    .downcast_ref::<Int8Array>()
+                    .unwrap()
+                    .values()
+                    .to_vec(),
+            )),
             DataType::Int16 => Ok(Self::SmallInt(
                 value
                     .as_any()

diff --git a/crates/proof-of-sql/src/base/database/owned_column.rs b/crates/proof-of-sql/src/base/database/owned_column.rs
@@ -232,6 +232,15 @@ impl<S: Scalar> OwnedColumn<S> {
         Self::try_from_scalars(&scalars, column_type)
     }
 
+    #[cfg(test)]
+    /// Returns an iterator over the raw data of the column
+    /// assuming the underlying type is [i8], panicking if it is not.
+    pub fn i8_iter(&self) -> impl Iterator<Item = &i8> {
+        match self {
+            OwnedColumn::TinyInt(col) => col.iter(),
+            _ => panic!("Expected TinyInt column"),
+        }
+    }
     #[cfg(test)]
     /// Returns an iterator over the raw data of the column
     /// assuming the underlying type is [i16], panicking if it is not.