[Variant] Minor code cleanups

parquet-variant-compute/src/variant_get.rs

@@ -135,7 +135,8 @@ fn shredded_get_path(
     let shred_basic_variant =
         |target: VariantArray, path: VariantPath<'_>, as_field: Option<&Field>| {
             let as_type = as_field.map(|f| f.data_type());
-            let mut builder = make_variant_to_arrow_row_builder(path, as_type, cast_options)?;
+            let mut builder =
+                make_variant_to_arrow_row_builder(path, as_type, cast_options, target.len())?;
             for i in 0..target.len() {
                 if target.is_null(i) {
                     builder.append_null()?;

parquet-variant-compute/src/variant_to_arrow.rs

@@ -76,7 +76,7 @@ impl<'a> VariantToArrowRowBuilder<'a> {
         }
     }
 
-    pub fn finish(&mut self) -> Result<ArrayRef> {
+    pub fn finish(self) -> Result<ArrayRef> {
         use VariantToArrowRowBuilder::*;
         match self {
             Int8(b) => b.finish(),
@@ -97,19 +97,41 @@ pub(crate) fn make_variant_to_arrow_row_builder<'a>(
     path: VariantPath<'a>,
     data_type: Option<&'a DataType>,
     cast_options: &'a CastOptions,
+    capacity: usize,
 ) -> Result<VariantToArrowRowBuilder<'a>> {
     use VariantToArrowRowBuilder::*;
 
     let mut builder = match data_type {
         // If no data type was requested, build an unshredded VariantArray.
-        None => BinaryVariant(VariantToBinaryVariantArrowRowBuilder::new(16)),
-        Some(DataType::Int8) => Int8(VariantToPrimitiveArrowRowBuilder::new(cast_options)),
-        Some(DataType::Int16) => Int16(VariantToPrimitiveArrowRowBuilder::new(cast_options)),
-        Some(DataType::Int32) => Int32(VariantToPrimitiveArrowRowBuilder::new(cast_options)),
-        Some(DataType::Int64) => Int64(VariantToPrimitiveArrowRowBuilder::new(cast_options)),
-        Some(DataType::Float16) => Float16(VariantToPrimitiveArrowRowBuilder::new(cast_options)),
-        Some(DataType::Float32) => Float32(VariantToPrimitiveArrowRowBuilder::new(cast_options)),
-        Some(DataType::Float64) => Float64(VariantToPrimitiveArrowRowBuilder::new(cast_options)),
+        None => BinaryVariant(VariantToBinaryVariantArrowRowBuilder::new(capacity)),
+        Some(DataType::Int8) => Int8(VariantToPrimitiveArrowRowBuilder::new(
+            cast_options,
+            capacity,
+        )),
+        Some(DataType::Int16) => Int16(VariantToPrimitiveArrowRowBuilder::new(
+            cast_options,
+            capacity,
+        )),
+        Some(DataType::Int32) => Int32(VariantToPrimitiveArrowRowBuilder::new(
+            cast_options,
+            capacity,
+        )),
+        Some(DataType::Int64) => Int64(VariantToPrimitiveArrowRowBuilder::new(
+            cast_options,
+            capacity,
+        )),
+        Some(DataType::Float16) => Float16(VariantToPrimitiveArrowRowBuilder::new(
+            cast_options,
+            capacity,
+        )),
+        Some(DataType::Float32) => Float32(VariantToPrimitiveArrowRowBuilder::new(
+            cast_options,
+            capacity,
+        )),
+        Some(DataType::Float64) => Float64(VariantToPrimitiveArrowRowBuilder::new(
+            cast_options,
+            capacity,
+        )),
         _ => {
             return Err(ArrowError::NotYetImplemented(format!(
                 "variant_get with path={:?} and data_type={:?} not yet implemented",
@@ -150,7 +172,7 @@ impl<'a> VariantPathRowBuilder<'a> {
         }
     }
 
-    fn finish(&mut self) -> Result<ArrayRef> {
+    fn finish(self) -> Result<ArrayRef> {
         self.builder.finish()
     }
 }
@@ -180,9 +202,9 @@ pub(crate) struct VariantToPrimitiveArrowRowBuilder<'a, T: ArrowPrimitiveType> {
 }
 
 impl<'a, T: ArrowPrimitiveType> VariantToPrimitiveArrowRowBuilder<'a, T> {
-    fn new(cast_options: &'a CastOptions<'a>) -> Self {
+    fn new(cast_options: &'a CastOptions<'a>, capacity: usize) -> Self {
         Self {
-            builder: PrimitiveBuilder::<T>::new(),
+            builder: PrimitiveBuilder::<T>::with_capacity(capacity),
             cast_options,
         }
     }
@@ -217,7 +239,7 @@ where
         }
     }
 
-    fn finish(&mut self) -> Result<ArrayRef> {
+    fn finish(mut self) -> Result<ArrayRef> {
         Ok(Arc::new(self.builder.finish()))
     }
 }
@@ -253,9 +275,7 @@ impl VariantToBinaryVariantArrowRowBuilder {
         Ok(true)
     }
 
-    fn finish(&mut self) -> Result<ArrayRef> {
-        // VariantArrayBuilder::build takes ownership, so we need to replace it
-        let builder = std::mem::replace(&mut self.builder, VariantArrayBuilder::new(0));
-        Ok(Arc::new(builder.build()))
+    fn finish(self) -> Result<ArrayRef> {
+        Ok(Arc::new(self.builder.build()))
     }
 }