This is an automated email from the ASF dual-hosted git repository.
mbrobbel pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/arrow-rs.git
The following commit(s) were added to refs/heads/main by this push:
new 89b4b13eb5 [Variant] Minor code cleanups (#8356)
89b4b13eb5 is described below
commit 89b4b13eb5a9c69147883016918d9bac357f14d1
Author: Ryan Johnson <[email protected]>
AuthorDate: Tue Sep 16 07:34:44 2025 -0600
[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
---
parquet-variant-compute/src/variant_get.rs | 3 +-
parquet-variant-compute/src/variant_to_arrow.rs | 54 +++++++++++++++++--------
2 files changed, 39 insertions(+), 18 deletions(-)
diff --git a/parquet-variant-compute/src/variant_get.rs
b/parquet-variant-compute/src/variant_get.rs
index 3ac6d2be6c..44c3ebbbc0 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 4deeaffe4e..60f74e365d 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<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()))
}
}
