Use ArrayBuilder as return value of take

chmp · Sep 1, 2024 · 1aef722 · 1aef722
1 parent dd70221
commit 1aef722
Show file tree

Hide file tree

Showing 21 changed files with 118 additions and 109 deletions.
diff --git a/serde_arrow/src/internal/serialization/array_builder.rs b/serde_arrow/src/internal/serialization/array_builder.rs
@@ -106,43 +106,8 @@ impl ArrayBuilder {
 }
 
 impl ArrayBuilder {
-    /// Take the contained array builder, while leaving structure intact
-    // TODO: use ArrayBuilder as return type for the impls and use dispatch here
     pub fn take(&mut self) -> ArrayBuilder {
-        match self {
-            Self::Null(builder) => Self::Null(builder.take()),
-            Self::Bool(builder) => Self::Bool(builder.take()),
-            Self::I8(builder) => Self::I8(builder.take()),
-            Self::I16(builder) => Self::I16(builder.take()),
-            Self::I32(builder) => Self::I32(builder.take()),
-            Self::I64(builder) => Self::I64(builder.take()),
-            Self::U8(builder) => Self::U8(builder.take()),
-            Self::U16(builder) => Self::U16(builder.take()),
-            Self::U32(builder) => Self::U32(builder.take()),
-            Self::U64(builder) => Self::U64(builder.take()),
-            Self::F16(builder) => Self::F16(builder.take()),
-            Self::F32(builder) => Self::F32(builder.take()),
-            Self::F64(builder) => Self::F64(builder.take()),
-            Self::Date32(builder) => Self::Date32(builder.take()),
-            Self::Date64(builder) => Self::Date64(builder.take()),
-            Self::Time32(builder) => Self::Time32(builder.take()),
-            Self::Time64(builder) => Self::Time64(builder.take()),
-            Self::Duration(builder) => Self::Duration(builder.take()),
-            Self::Decimal128(builder) => Self::Decimal128(builder.take()),
-            Self::Utf8(builder) => Self::Utf8(builder.take()),
-            Self::LargeUtf8(builder) => Self::LargeUtf8(builder.take()),
-            Self::List(builder) => Self::List(builder.take()),
-            Self::LargeList(builder) => Self::LargeList(builder.take()),
-            Self::FixedSizedList(builder) => Self::FixedSizedList(builder.take()),
-            Self::Binary(builder) => Self::Binary(builder.take()),
-            Self::LargeBinary(builder) => Self::LargeBinary(builder.take()),
-            Self::FixedSizeBinary(builder) => Self::FixedSizeBinary(builder.take()),
-            Self::Struct(builder) => Self::Struct(builder.take()),
-            Self::Map(builder) => Self::Map(builder.take()),
-            Self::DictionaryUtf8(builder) => Self::DictionaryUtf8(builder.take()),
-            Self::Union(builder) => Self::Union(builder.take()),
-            Self::UnknownVariant(builder) => Self::UnknownVariant(builder.take()),
-        }
+        dispatch!(self, Self(builder) => builder.take())
     }
 }
 

diff --git a/serde_arrow/src/internal/serialization/binary_builder.rs b/serde_arrow/src/internal/serialization/binary_builder.rs
@@ -11,7 +11,7 @@ use crate::internal::{
     },
 };
 
-use super::simple_serializer::SimpleSerializer;
+use super::{array_builder::ArrayBuilder, simple_serializer::SimpleSerializer};
 
 #[derive(Debug, Clone)]
 
@@ -28,7 +28,7 @@ impl<O: Offset> BinaryBuilder<O> {
         }
     }
 
-    pub fn take(&mut self) -> Self {
+    pub fn take_self(&mut self) -> Self {
         Self {
             path: self.path.clone(),
             array: self.array.take(),
@@ -41,12 +41,20 @@ impl<O: Offset> BinaryBuilder<O> {
 }
 
 impl BinaryBuilder<i32> {
+    pub fn take(&mut self) -> ArrayBuilder {
+        ArrayBuilder::Binary(self.take_self())
+    }
+
     pub fn into_array(self) -> Result<Array> {
         Ok(Array::Binary(self.array))
     }
 }
 
 impl BinaryBuilder<i64> {
+    pub fn take(&mut self) -> ArrayBuilder {
+        ArrayBuilder::LargeBinary(self.take_self())
+    }
+
     pub fn into_array(self) -> Result<Array> {
         Ok(Array::LargeBinary(self.array))
     }

diff --git a/serde_arrow/src/internal/serialization/bool_builder.rs b/serde_arrow/src/internal/serialization/bool_builder.rs
@@ -9,7 +9,7 @@ use crate::internal::{
     },
 };
 
-use super::simple_serializer::SimpleSerializer;
+use super::{array_builder::ArrayBuilder, simple_serializer::SimpleSerializer};
 
 #[derive(Debug, Clone)]
 pub struct BoolBuilder {
@@ -29,15 +29,15 @@ impl BoolBuilder {
         }
     }
 
-    pub fn take(&mut self) -> Self {
-        Self {
+    pub fn take(&mut self) -> ArrayBuilder {
+        ArrayBuilder::Bool(Self {
             path: self.path.clone(),
             array: BooleanArray {
                 len: std::mem::take(&mut self.array.len),
                 validity: self.array.validity.as_mut().map(std::mem::take),
                 values: std::mem::take(&mut self.array.values),
             },
-        }
+        })
     }
 
     pub fn is_nullable(&self) -> bool {

diff --git a/serde_arrow/src/internal/serialization/date32_builder.rs b/serde_arrow/src/internal/serialization/date32_builder.rs
@@ -11,7 +11,7 @@ use crate::internal::{
     },
 };
 
-use super::simple_serializer::SimpleSerializer;
+use super::{array_builder::ArrayBuilder, simple_serializer::SimpleSerializer};
 
 #[derive(Debug, Clone)]
 pub struct Date32Builder {
@@ -27,11 +27,11 @@ impl Date32Builder {
         }
     }
 
-    pub fn take(&mut self) -> Self {
-        Self {
+    pub fn take(&mut self) -> ArrayBuilder {
+        ArrayBuilder::Date32(Self {
             path: self.path.clone(),
             array: self.array.take(),
-        }
+        })
     }
 
     pub fn is_nullable(&self) -> bool {

diff --git a/serde_arrow/src/internal/serialization/date64_builder.rs b/serde_arrow/src/internal/serialization/date64_builder.rs
@@ -9,7 +9,7 @@ use crate::internal::{
     },
 };
 
-use super::simple_serializer::SimpleSerializer;
+use super::{array_builder::ArrayBuilder, simple_serializer::SimpleSerializer};
 
 #[derive(Debug, Clone)]
 pub struct Date64Builder {
@@ -34,13 +34,13 @@ impl Date64Builder {
         }
     }
 
-    pub fn take(&mut self) -> Self {
-        Self {
+    pub fn take(&mut self) -> ArrayBuilder {
+        ArrayBuilder::Date64(Self {
             path: self.path.clone(),
             meta: self.meta.clone(),
             utc: self.utc,
             array: self.array.take(),
-        }
+        })
     }
 
     pub fn is_nullable(&self) -> bool {

diff --git a/serde_arrow/src/internal/serialization/decimal_builder.rs b/serde_arrow/src/internal/serialization/decimal_builder.rs
@@ -10,7 +10,7 @@ use crate::internal::{
     },
 };
 
-use super::simple_serializer::SimpleSerializer;
+use super::{array_builder::ArrayBuilder, simple_serializer::SimpleSerializer};
 
 #[derive(Debug, Clone)]
 pub struct DecimalBuilder {
@@ -36,16 +36,16 @@ impl DecimalBuilder {
         }
     }
 
-    pub fn take(&mut self) -> Self {
-        Self {
+    pub fn take(&mut self) -> ArrayBuilder {
+        ArrayBuilder::Decimal128(Self {
             path: self.path.clone(),
             precision: self.precision,
             scale: self.scale,
             f32_factor: self.f32_factor,
             f64_factor: self.f64_factor,
             parser: self.parser,
             array: self.array.take(),
-        }
+        })
     }
 
     pub fn is_nullable(&self) -> bool {

diff --git a/serde_arrow/src/internal/serialization/dictionary_utf8_builder.rs b/serde_arrow/src/internal/serialization/dictionary_utf8_builder.rs
@@ -28,13 +28,13 @@ impl DictionaryUtf8Builder {
         }
     }
 
-    pub fn take(&mut self) -> Self {
-        Self {
+    pub fn take(&mut self) -> ArrayBuilder {
+        ArrayBuilder::DictionaryUtf8(Self {
             path: self.path.clone(),
             indices: Box::new(self.indices.take()),
             values: Box::new(self.values.take()),
             index: std::mem::take(&mut self.index),
-        }
+        })
     }
 
     pub fn is_nullable(&self) -> bool {

diff --git a/serde_arrow/src/internal/serialization/duration_builder.rs b/serde_arrow/src/internal/serialization/duration_builder.rs
@@ -9,7 +9,7 @@ use crate::internal::{
     },
 };
 
-use super::simple_serializer::SimpleSerializer;
+use super::{array_builder::ArrayBuilder, simple_serializer::SimpleSerializer};
 
 #[derive(Debug, Clone)]
 pub struct DurationBuilder {
@@ -27,12 +27,12 @@ impl DurationBuilder {
         }
     }
 
-    pub fn take(&mut self) -> Self {
-        Self {
+    pub fn take(&mut self) -> ArrayBuilder {
+        ArrayBuilder::Duration(Self {
             path: self.path.clone(),
             unit: self.unit,
             array: self.array.take(),
-        }
+        })
     }
 
     pub fn is_nullable(&self) -> bool {

diff --git a/serde_arrow/src/internal/serialization/fixed_size_binary_builder.rs b/serde_arrow/src/internal/serialization/fixed_size_binary_builder.rs
@@ -11,7 +11,7 @@ use crate::internal::{
     },
 };
 
-use super::simple_serializer::SimpleSerializer;
+use super::{array_builder::ArrayBuilder, simple_serializer::SimpleSerializer};
 
 #[derive(Debug, Clone)]
 
@@ -34,14 +34,14 @@ impl FixedSizeBinaryBuilder {
         }
     }
 
-    pub fn take(&mut self) -> Self {
-        Self {
+    pub fn take(&mut self) -> ArrayBuilder {
+        ArrayBuilder::FixedSizeBinary(Self {
             path: self.path.clone(),
             seq: self.seq.take(),
             buffer: std::mem::take(&mut self.buffer),
             current_n: std::mem::take(&mut self.current_n),
             n: self.n,
-        }
+        })
     }
 
     pub fn is_nullable(&self) -> bool {

diff --git a/serde_arrow/src/internal/serialization/fixed_size_list_builder.rs b/serde_arrow/src/internal/serialization/fixed_size_list_builder.rs
@@ -42,15 +42,15 @@ impl FixedSizeListBuilder {
         }
     }
 
-    pub fn take(&mut self) -> Self {
-        Self {
+    pub fn take(&mut self) -> ArrayBuilder {
+        ArrayBuilder::FixedSizedList(Self {
             path: self.path.clone(),
             seq: self.seq.take(),
             meta: self.meta.clone(),
             n: self.n,
             current_count: std::mem::take(&mut self.current_count),
             element: Box::new(self.element.take()),
-        }
+        })
     }
 
     pub fn is_nullable(&self) -> bool {

diff --git a/serde_arrow/src/internal/serialization/float_builder.rs b/serde_arrow/src/internal/serialization/float_builder.rs
@@ -11,7 +11,7 @@ use crate::internal::{
     },
 };
 
-use super::simple_serializer::SimpleSerializer;
+use super::{array_builder::ArrayBuilder, simple_serializer::SimpleSerializer};
 
 #[derive(Debug, Clone)]
 pub struct FloatBuilder<I> {
@@ -27,7 +27,7 @@ impl<F: Default + 'static> FloatBuilder<F> {
         }
     }
 
-    pub fn take(&mut self) -> Self {
+    pub fn take_self(&mut self) -> Self {
         Self {
             path: self.path.clone(),
             array: self.array.take(),
@@ -40,18 +40,22 @@ impl<F: Default + 'static> FloatBuilder<F> {
 }
 
 macro_rules! impl_into_array {
-    ($ty:ty, $var:ident) => {
+    ($ty:ty, $builder_var:ident, $array_var:ident) => {
         impl FloatBuilder<$ty> {
+            pub fn take(&mut self) -> ArrayBuilder {
+                ArrayBuilder::$builder_var(self.take_self())
+            }
+
             pub fn into_array(self) -> Result<Array> {
-                Ok(Array::$var(self.array))
+                Ok(Array::$array_var(self.array))
             }
         }
     };
 }
 
-impl_into_array!(f16, Float16);
-impl_into_array!(f32, Float32);
-impl_into_array!(f64, Float64);
+impl_into_array!(f16, F16, Float16);
+impl_into_array!(f32, F32, Float32);
+impl_into_array!(f64, F64, Float64);
 
 impl<F> Context for FloatBuilder<F> {
     fn annotations(&self) -> BTreeMap<String, String> {

diff --git a/serde_arrow/src/internal/serialization/int_builder.rs b/serde_arrow/src/internal/serialization/int_builder.rs
@@ -9,7 +9,7 @@ use crate::internal::{
     },
 };
 
-use super::simple_serializer::SimpleSerializer;
+use super::{array_builder::ArrayBuilder, simple_serializer::SimpleSerializer};
 
 #[derive(Debug, Clone)]
 pub struct IntBuilder<I> {
@@ -25,7 +25,7 @@ impl<I: Default + 'static> IntBuilder<I> {
         }
     }
 
-    pub fn take(&mut self) -> Self {
+    pub fn take_self(&mut self) -> Self {
         Self {
             path: self.path.clone(),
             array: self.array.take(),
@@ -38,23 +38,27 @@ impl<I: Default + 'static> IntBuilder<I> {
 }
 
 macro_rules! impl_into_array {
-    ($ty:ty, $var:ident) => {
+    ($ty:ty, $builder_var: ident, $array_var:ident) => {
         impl IntBuilder<$ty> {
+            pub fn take(&mut self) -> ArrayBuilder {
+                ArrayBuilder::$builder_var(self.take_self())
+            }
+
             pub fn into_array(self) -> Result<Array> {
-                Ok(Array::$var(self.array))
+                Ok(Array::$array_var(self.array))
             }
         }
     };
 }
 
-impl_into_array!(i8, Int8);
-impl_into_array!(i16, Int16);
-impl_into_array!(i32, Int32);
-impl_into_array!(i64, Int64);
-impl_into_array!(u8, UInt8);
-impl_into_array!(u16, UInt16);
-impl_into_array!(u32, UInt32);
-impl_into_array!(u64, UInt64);
+impl_into_array!(i8, I8, Int8);
+impl_into_array!(i16, I16, Int16);
+impl_into_array!(i32, I32, Int32);
+impl_into_array!(i64, I64, Int64);
+impl_into_array!(u8, U8, UInt8);
+impl_into_array!(u16, U16, UInt16);
+impl_into_array!(u32, U32, UInt32);
+impl_into_array!(u64, U64, UInt64);
 
 impl<I> Context for IntBuilder<I> {
     fn annotations(&self) -> BTreeMap<String, String> {