refactor: consume_opt

aljazerzen · Feb 19, 2024 · f7865a4 · f7865a4
1 parent 25b3bc9
commit f7865a4
Show file tree

Hide file tree

Showing 2 changed files with 28 additions and 21 deletions.
diff --git a/connector_arrow/src/util/row_writer.rs b/connector_arrow/src/util/row_writer.rs
@@ -1,6 +1,6 @@
 use std::any::Any;
 
-use arrow::array::{ArrayBuilder, ArrayRef};
+use arrow::array::{ArrayBuilder, ArrayRef, FixedSizeBinaryBuilder};
 use arrow::datatypes::*;
 use arrow::record_batch::RecordBatch;
 
@@ -155,11 +155,12 @@ macro_rules! impl_consume_ty {
                         .expect(concat!("bad cast to ", stringify!($Builder)))
                         .append_value(value);
                 }
-                fn consume_opt(&mut self, value: Option<<$ArrTy as ArrowType>::Native>) {
+
+                fn consume_null(&mut self) {
                     self.next_builder()
                         .downcast_mut::<arrow::array::builder::$Builder>()
                         .expect(concat!("bad cast to ", stringify!($Builder)))
-                        .append_option(value);
+                        .append_null();
                 }
             }
         )+
@@ -210,7 +211,8 @@ impl ConsumeTy<NullType> for ArrowRowWriter {
     fn consume(&mut self, _: ()) {
         self.next_builder();
     }
-    fn consume_opt(&mut self, _: Option<()>) {
+
+    fn consume_null(&mut self) {
         self.next_builder();
     }
 }
@@ -225,11 +227,12 @@ macro_rules! impl_consume_ref_ty {
                         .expect(concat!("bad cast to ", stringify!($Builder)))
                         .append_value(&value);
                 }
-                fn consume_opt(&mut self, value: Option<<$ArrTy as ArrowType>::Native>) {
+
+                fn consume_null(&mut self) {
                     self.next_builder()
                         .downcast_mut::<arrow::array::builder::$Builder>()
                         .expect(concat!("bad cast to ", stringify!($Builder)))
-                        .append_option(value);
+                        .append_null();
                 }
             }
         )+
@@ -251,15 +254,11 @@ impl ConsumeTy<FixedSizeBinaryType> for ArrowRowWriter {
             .append_value(&value)
             .unwrap();
     }
-    fn consume_opt(&mut self, value: Option<<FixedSizeBinaryType as ArrowType>::Native>) {
-        let builder = self
-            .next_builder()
-            .downcast_mut::<arrow::array::builder::FixedSizeBinaryBuilder>()
-            .expect(concat!("bad cast to ", stringify!(FixedSizeBinaryBuilder)));
-        if let Some(value) = value {
-            builder.append_value(value).unwrap();
-        } else {
-            builder.append_null();
-        }
+
+    fn consume_null(&mut self) {
+        self.next_builder()
+            .downcast_mut::<FixedSizeBinaryBuilder>()
+            .expect(concat!("bad cast to ", stringify!($Builder)))
+            .append_null();
     }
 }
diff --git a/connector_arrow/src/util/transport.rs b/connector_arrow/src/util/transport.rs
@@ -16,9 +16,15 @@ macro_rules! impl_transport_match {
             }
         } else {
             match $f.data_type() {
-                Null => ConsumeTy::<NullType>::consume_opt($c, Some(())),
+                Null => ConsumeTy::<NullType>::consume($c, ()),
                 $(
-                    $Pat => ConsumeTy::<$ArrTy>::consume_opt($c, ProduceTy::<$ArrTy>::produce_opt($p)?),
+                    $Pat => {
+                        if let Some(v) = ProduceTy::<$ArrTy>::produce_opt($p)? {
+                            ConsumeTy::<$ArrTy>::consume($c, v)
+                        } else {
+                            ConsumeTy::<$ArrTy>::consume_null($c)
+                        }
+                    },
                 )*
                 _ => todo!("unimplemented transport of {:?}", $f.data_type()),
             }
@@ -175,7 +181,8 @@ pub trait Consume:
 /// Ability to consume a value of an an arrow type
 pub trait ConsumeTy<T: ArrowType> {
     fn consume(&mut self, value: T::Native);
-    fn consume_opt(&mut self, value: Option<T::Native>);
+
+    fn consume_null(&mut self);
 }
 
 pub mod print {
@@ -192,8 +199,9 @@ pub mod print {
         fn consume(&mut self, value: T::Native) {
             println!("{}: {value:?}", std::any::type_name::<T>());
         }
-        fn consume_opt(&mut self, value: Option<T::Native>) {
-            println!("{}: {value:?}", std::any::type_name::<T>());
+
+        fn consume_null(&mut self) {
+            println!("{}: null", std::any::type_name::<T>());
         }
     }
 }