flesh out columnar batch

Author: wjones127

kernel/src/columnar_batch.rs

@@ -1,42 +1,42 @@
-use std::sync::Arc;
-
 use crate::{
     schema::{DataType, Schema, StructField},
     DeltaResult,
 };
 
 pub trait ColumnarBatch {
+    type Column: ColumnVector;
+
     /// Get the schema of the batch.
     fn schema(&self) -> DeltaResult<Schema>;
 
     /// Get the column at the specified index.
-    fn column(&self, index: usize) -> &dyn ColumnVector;
+    fn column(&self, index: usize) -> Self::Column;
 
     /// Number of rows in the batch.
     fn size(&self) -> usize;
 
-    fn with_column(
-        &self,
-        index: usize,
-        field: StructField,
-        column: Arc<dyn ColumnVector>,
-    ) -> DeltaResult<Self>
-    where
-        Self: Sized;
+    // fn with_column(
+    //     &self,
+    //     index: usize,
+    //     field: StructField,
+    //     column: &dyn ColumnVector,
+    // ) -> DeltaResult<Self>
+    // where
+    //     Self: Sized;
 
     fn with_deleted_column_at(&self, index: usize) -> DeltaResult<Self>
     where
         Self: Sized;
 
-    fn with_schema(&self, schema: Schema) -> DeltaResult<Self>
-    where
-        Self: Sized;
+    // fn with_schema(&self, schema: Schema) -> DeltaResult<Self>
+    // where
+    //     Self: Sized;
 
     fn slice(&self, offset: usize, length: usize) -> DeltaResult<Self>
     where
         Self: Sized;
 
-    fn rows(&self) -> Box<dyn Iterator<Item = &dyn Row>>;
+    fn rows(&self) -> Box<dyn Iterator<Item = Box<dyn Row<Column = Self::Column>>>>;
 }
 
 // TODO: do all these methods do bounds checking? Should we offer alternative
@@ -50,43 +50,54 @@ pub trait ColumnVector {
     fn is_null(&self, i: usize) -> bool;
     fn get_i32(&self, i: usize) -> DeltaResult<Option<i32>>;
     fn get_string(&self, i: usize) -> DeltaResult<Option<&str>>;
-    fn get_struct(&self, i: usize) -> DeltaResult<Option<&dyn Row>>;
-    fn get_array(&self, i: usize) -> DeltaResult<Option<&dyn ArrayValue>>;
-    fn get_map(&self, i: usize) -> DeltaResult<Option<&dyn MapValue>>;
+    fn get_struct(&self, i: usize) -> DeltaResult<Option<Box<dyn Row<Column = Self>>>>;
+    fn get_array(&self, i: usize) -> DeltaResult<Option<Box<dyn ArrayValue<Column = Self>>>>;
+    fn get_map(&self, i: usize) -> DeltaResult<Option<Box<dyn MapValue<Column = Self>>>>;
 }
 
 pub trait Row {
+    type Column: ColumnVector;
+
     fn schema(&self) -> DeltaResult<Schema>;
     fn is_null(&self, i: usize) -> bool;
     fn get_i32(&self, i: usize) -> DeltaResult<Option<i32>>;
     fn get_string(&self, i: usize) -> DeltaResult<Option<&str>>;
-    fn get_struct(&self, i: usize) -> DeltaResult<Option<&dyn Row>>;
-    fn get_array(&self, i: usize) -> DeltaResult<Option<&dyn ArrayValue>>;
-    fn get_map(&self, i: usize) -> DeltaResult<Option<&dyn MapValue>>;
+    fn get_struct(&self, i: usize) -> DeltaResult<Option<Box<dyn Row<Column = Self::Column>>>>;
+    fn get_array(
+        &self,
+        i: usize,
+    ) -> DeltaResult<Option<Box<dyn ArrayValue<Column = Self::Column>>>>;
+    fn get_map(&self, i: usize) -> DeltaResult<Option<Box<dyn MapValue<Column = Self::Column>>>>;
 }
 
 // Based on: https://github.com/delta-io/delta/pull/2087
 
 pub trait ArrayValue {
+    type Column: ColumnVector;
+
     /// Return the number of elements in the array
     fn size(&self) -> usize;
 
     /// Get the elements in the array
-    fn elements(&self) -> &dyn ColumnVector;
+    fn elements(&self) -> Self::Column;
 }
 
 pub trait MapValue {
+    type Column: ColumnVector;
+
     /// Return the number of elements in the map
     fn size(&self) -> usize;
 
     /// Get the keys in the map
-    fn keys(&self) -> &dyn ColumnVector;
+    fn keys(&self) -> Self::Column;
 
     /// Get the values in the map
-    fn values(&self) -> &dyn ColumnVector;
+    fn values(&self) -> Self::Column;
 }
 
-mod arrow {
+pub mod arrow {
+    use std::sync::Arc;
+
     use arrow_array::cast::AsArray;
     use arrow_array::types::Int32Type;
     use arrow_array::Array as ArrowArray;
@@ -99,33 +110,67 @@ mod arrow {
     use super::*;
 
     impl ColumnarBatch for RecordBatch {
+        type Column = Arc<dyn ArrowArray>;
+
         fn schema(&self) -> DeltaResult<Schema> {
             Ok(self.schema().as_ref().try_into()?)
         }
 
-        fn column(&self, index: usize) -> &dyn ColumnVector {
-            &self.column(index).as_ref() as &dyn ColumnVector
+        fn column(&self, index: usize) -> Self::Column {
+            self.column(index).clone()
+        }
+
+        fn size(&self) -> usize {
+            self.num_rows()
+        }
+
+        fn with_deleted_column_at(&self, index: usize) -> DeltaResult<Self>
+        where
+            Self: Sized,
+        {
+            let indices = (0..self.num_columns())
+                .filter(|i| *i != index)
+                .collect::<Vec<_>>();
+            RecordBatch::project(&self, &indices).map_err(|err| Error::Arrow(err))
+        }
+
+        fn slice(&self, offset: usize, length: usize) -> DeltaResult<Self>
+        where
+            Self: Sized,
+        {
+            Ok(RecordBatch::slice(self, offset, length))
+        }
+
+        fn rows(&self) -> Box<dyn Iterator<Item = Box<dyn Row<Column = Self::Column>>>> {
+            let batch = self.clone();
+            Box::new((0..self.size()).into_iter().map(move |i| {
+                let row = Box::new(ArrowRow {
+                    batch: batch.clone(),
+                    row_index: i,
+                });
+                row as Box<dyn Row<Column = Self::Column>>
+            }))
         }
     }
 
-    impl ColumnVector for dyn ArrowArray {
+    impl ColumnVector for Arc<dyn ArrowArray> {
         fn data_type(&self) -> DeltaResult<DataType> {
-            Ok(self.data_type().try_into()?)
+            Ok(self.as_ref().data_type().try_into()?)
         }
 
         fn size(&self) -> usize {
             self.len()
         }
 
         fn is_null(&self, i: usize) -> bool {
-            self.is_null(i)
+            self.as_ref().is_null(i)
         }
 
         /// Get the i32 value at the specified index.
         ///
         /// This will panic if the column is not boolean or if the index is out of bounds.
         fn get_i32(&self, i: usize) -> DeltaResult<Option<i32>> {
-            if self.is_null(i) {
+            if self.as_ref().is_null(i) {
                 Ok(None)
             } else {
                 Ok(Some(self.as_primitive::<Int32Type>().value(i)))
@@ -136,10 +181,10 @@ mod arrow {
         ///
         /// This will panic if the column is not string or if the index is out of bounds.
         fn get_string(&self, i: usize) -> DeltaResult<Option<&str>> {
-            if self.is_null(i) {
+            if self.as_ref().is_null(i) {
                 Ok(None)
             } else {
-                match self.data_type() {
+                match self.as_ref().data_type() {
                     ArrowDataType::Utf8 => Ok(Some(self.as_string::<i32>().value(i))),
                     ArrowDataType::LargeUtf8 => Ok(Some(self.as_string::<i64>().value(i))),
                     _ => panic!("get_string called on non-string column"),
@@ -150,79 +195,94 @@ mod arrow {
         /// Get the struct value at the specified index.
         ///
         /// This will panic if the column is not struct or if the index is out of bounds.
-        fn get_struct(&self, i: usize) -> DeltaResult<Option<&dyn Row>> {
-            if self.is_null(i) {
+        fn get_struct(&self, i: usize) -> DeltaResult<Option<Box<dyn Row<Column = Self>>>> {
+            if self.as_ref().is_null(i) {
                 Ok(None)
             } else {
                 let batch = self
                     .as_struct_opt()
-                    .expect("get_struct called on non-struct column");
+                    .expect("get_struct called on non-struct column")
+                    .clone();
                 let row = ArrowRow {
                     batch,
                     row_index: i,
                 };
-                Ok(Some(&row))
+                Ok(Some(Box::new(row)))
             }
         }
 
         /// Get the array value at the specified index.
         ///
         /// This will panic if the column is not array or if the index is out of bounds.
-        fn get_array(&self, i: usize) -> DeltaResult<Option<&dyn ArrayValue>> {
-            if self.is_null(i) {
+        fn get_array(&self, i: usize) -> DeltaResult<Option<Box<dyn ArrayValue<Column = Self>>>> {
+            if self.as_ref().is_null(i) {
                 Ok(None)
             } else {
-                let batch = self
-                    .as_list_opt()
-                    .expect("get_array called on non-array column");
-                let array = ArrowArray {
-                    array: batch.value(i),
+                let sub_array = match self.as_ref().data_type() {
+                    ArrowDataType::List(_) => self.as_list_opt::<i32>().unwrap().value(i),
+                    ArrowDataType::LargeList(_) => self.as_list_opt::<i32>().unwrap().value(i),
+                    _ => panic!("get_array called on non-array column"),
                 };
-                Ok(Some(&array))
+                Ok(Some(Box::new(ArrowArraySlice(sub_array))))
+            }
+        }
+
+        fn get_map(&self, i: usize) -> DeltaResult<Option<Box<dyn MapValue<Column = Self>>>> {
+            if self.as_ref().is_null(i) {
+                Ok(None)
+            } else {
+                let arr = self.as_map().value(i);
+                let map_array = ArrowMapValue {
+                    keys: arr.column(0).clone(),
+                    values: arr.column(1).clone(),
+                };
+                Ok(Some(Box::new(map_array)))
             }
         }
     }
 
-    trait ArrowTabular {
-        fn schema(&self) -> &ArrowSchema;
-        fn column(&self, index: usize) -> &dyn ArrowArray;
+    pub trait ArrowTabular {
+        fn schema(&self) -> Arc<ArrowSchema>;
+        fn column(&self, index: usize) -> &Arc<dyn ArrowArray>;
     }
 
     impl ArrowTabular for RecordBatch {
-        fn schema(&self) -> &ArrowSchema {
-            self.schema().as_ref()
+        fn schema(&self) -> Arc<ArrowSchema> {
+            self.schema()
         }
 
-        fn column(&self, index: usize) -> &dyn ArrowArray {
-            self.column(index).as_ref()
+        fn column(&self, index: usize) -> &Arc<dyn ArrowArray> {
+            &self.column(index)
         }
     }
 
     impl ArrowTabular for StructArray {
-        fn schema(&self) -> &ArrowSchema {
-            self.schema()
+        fn schema(&self) -> Arc<ArrowSchema> {
+            Arc::new(ArrowSchema::new(self.fields().clone()))
         }
 
-        fn column(&self, index: usize) -> &dyn ArrowArray {
-            self.column(index).as_ref()
+        fn column(&self, index: usize) -> &Arc<dyn ArrowArray> {
+            &self.column(index)
         }
     }
 
     /// A reference to a row in a RecordBatch or StructArray.
-    pub struct ArrowRow<'a, T: ArrowTabular> {
-        batch: &'a T,
+    pub struct ArrowRow<T: ArrowTabular> {
+        batch: T,
         row_index: usize,
     }
 
-    impl<'a, T: ArrowTabular> Row for ArrowRow<'a, T> {
+    impl<T: ArrowTabular> Row for ArrowRow<T> {
+        type Column = Arc<dyn ArrowArray>;
+
         fn schema(&self) -> DeltaResult<Schema> {
-            ArrowTabular::schema(self.batch)
+            ArrowTabular::schema(&self.batch)
                 .try_into()
                 .map_err(|err| Error::Arrow(err))
         }
 
         fn is_null(&self, i: usize) -> bool {
-            self.batch.column(i).is_null(self.row_index)
+            ArrowArray::is_null(self.batch.column(i), self.row_index)
         }
 
         fn get_i32(&self, i: usize) -> DeltaResult<Option<i32>> {
@@ -233,24 +293,57 @@ mod arrow {
             self.batch.column(i).get_string(self.row_index)
         }
 
-        fn get_struct(&self, i: usize) -> DeltaResult<Option<&dyn Row>> {
+        fn get_struct(&self, i: usize) -> DeltaResult<Option<Box<dyn Row<Column = Self::Column>>>> {
             self.batch.column(i).get_struct(self.row_index)
         }
 
-        fn get_array(&self, i: usize) -> DeltaResult<Option<&dyn ArrayValue>> {
+        fn get_array(
+            &self,
+            i: usize,
+        ) -> DeltaResult<Option<Box<dyn ArrayValue<Column = Self::Column>>>> {
             self.batch.column(i).get_array(self.row_index)
         }
 
-        fn get_map(&self, i: usize) -> DeltaResult<Option<&dyn MapValue>> {
+        fn get_map(
+            &self,
+            i: usize,
+        ) -> DeltaResult<Option<Box<dyn MapValue<Column = Self::Column>>>> {
             self.batch.column(i).get_map(self.row_index)
         }
     }
 
-    pub struct ArrowArraySlice<'a> {
-        batch: &'a dyn ArrowArray,
-        offset: usize,
-        length: usize,
+    pub struct ArrowArraySlice(Arc<dyn ArrowArray>);
+
+    impl ArrayValue for ArrowArraySlice {
+        type Column = Arc<dyn ArrowArray>;
+
+        fn size(&self) -> usize {
+            self.0.len()
+        }
+
+        fn elements(&self) -> Self::Column {
+            self.0.clone()
+        }
+    }
+
+    pub struct ArrowMapValue {
+        keys: Arc<dyn ArrowArray>,
+        values: Arc<dyn ArrowArray>,
     }
 
-    
+    impl MapValue for ArrowMapValue {
+        type Column = Arc<dyn ArrowArray>;
+
+        fn size(&self) -> usize {
+            self.keys.len()
+        }
+
+        fn keys(&self) -> Self::Column {
+            self.keys.clone()
+        }
+
+        fn values(&self) -> Self::Column {
+            self.values.clone()
+        }
+    }
 }