[Variant] Minor code cleanups (#8356)

# Which issue does this PR close?

We generally require a GitHub issue to be filed for all bug fixes and
enhancements and this helps us generate change logs for our releases.
You can link an issue to this PR using the GitHub syntax.

- Closes #NNN.

# Rationale for this change

Small code cleanups that accumulated during other work.

# What changes are included in this PR?

Now that `VariantToArrowRowBuilder` is an enum instead of a dyn trait,
its `finish` method can take `self` instead of `&mut self`, which
simplifies both its semantics and its call sites.

Additionally, pass the input length to the row builders so they can
reserve capacity accordingly.

# Are these changes tested?

Existing unit tests cover this code

# Are there any user-facing changes?

No

Author: scovich

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()))
     }
 }