diff --git a/parquet-variant-compute/src/variant_get.rs b/parquet-variant-compute/src/variant_get.rs index 3ac6d2be6c72..44c3ebbbc02e 100644 --- a/parquet-variant-compute/src/variant_get.rs +++ b/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()?; diff --git a/parquet-variant-compute/src/variant_to_arrow.rs b/parquet-variant-compute/src/variant_to_arrow.rs index 4deeaffe4e5b..60f74e365dd4 100644 --- a/parquet-variant-compute/src/variant_to_arrow.rs +++ b/parquet-variant-compute/src/variant_to_arrow.rs @@ -76,7 +76,7 @@ impl<'a> VariantToArrowRowBuilder<'a> { } } - pub fn finish(&mut self) -> Result { + pub fn finish(self) -> Result { 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> { 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 { + fn finish(self) -> Result { 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::::new(), + builder: PrimitiveBuilder::::with_capacity(capacity), cast_options, } } @@ -217,7 +239,7 @@ where } } - fn finish(&mut self) -> Result { + fn finish(mut self) -> Result { Ok(Arc::new(self.builder.finish())) } } @@ -253,9 +275,7 @@ impl VariantToBinaryVariantArrowRowBuilder { Ok(true) } - fn finish(&mut self) -> Result { - // 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 { + Ok(Arc::new(self.builder.build())) } }