Skip to content

Commit 4778077

Browse files
authored
[Variant] Add VariantArrayBuilder::append_nulls API (#9685)
# Which issue does this PR close? - Closes #9684. # Rationale for this change Check issue # What changes are included in this PR? - Add `VariantArrayBuilder::append_nulls` - Add unit test - add fast paths for `cast_to_variant_with_options` and `string_array_to_variant` # Are these changes tested? Yes, added a unit test # Are there any user-facing changes? A new public API
1 parent a42c475 commit 4778077

File tree

2 files changed

+37
-0
lines changed

2 files changed

+37
-0
lines changed

parquet-variant-compute/src/cast_to_variant.rs

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -58,6 +58,13 @@ pub fn cast_to_variant_with_options(
5858
input: &dyn Array,
5959
options: &CastOptions,
6060
) -> Result<VariantArray, ArrowError> {
61+
// Fast path: any all-null input maps to an all-null VariantArray.
62+
if input.null_count() == input.len() {
63+
let mut array_builder = VariantArrayBuilder::new(input.len());
64+
array_builder.append_nulls(input.len());
65+
return Ok(array_builder.build());
66+
}
67+
6168
// Create row builder for the input array type
6269
let mut row_builder = make_arrow_to_variant_row_builder(input.data_type(), input, options)?;
6370

parquet-variant-compute/src/variant_array_builder.rs

Lines changed: 30 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -156,6 +156,18 @@ impl VariantArrayBuilder {
156156
self.value_offsets.push(self.value_builder.offset());
157157
}
158158

159+
/// Appends `n` null rows to the builder.
160+
pub fn append_nulls(&mut self, n: usize) {
161+
self.nulls.append_n_nulls(n);
162+
// The subfields are expected to be non-nullable according to the parquet variant spec.
163+
let metadata_offset = self.metadata_builder.offset();
164+
let value_offset = self.value_builder.offset();
165+
self.metadata_offsets
166+
.extend(std::iter::repeat_n(metadata_offset, n));
167+
self.value_offsets
168+
.extend(std::iter::repeat_n(value_offset, n));
169+
}
170+
159171
/// Append the [`Variant`] to the builder as the next row
160172
pub fn append_variant(&mut self, variant: Variant) {
161173
ValueBuilder::append_variant(self.parent_state(), variant);
@@ -526,6 +538,24 @@ mod test {
526538
assert_eq!(list.len(), 2);
527539
}
528540

541+
#[test]
542+
fn test_variant_array_builder_append_nulls() {
543+
let mut builder = VariantArrayBuilder::new(6);
544+
builder.append_variant(Variant::from(1i32));
545+
builder.append_nulls(0); // should be a no-op
546+
builder.append_nulls(3);
547+
builder.append_variant(Variant::from(2i32));
548+
549+
let variant_array = builder.build();
550+
551+
assert_eq!(variant_array.len(), 5);
552+
assert_eq!(variant_array.value(0), Variant::from(1i32));
553+
assert!(variant_array.is_null(1));
554+
assert!(variant_array.is_null(2));
555+
assert!(variant_array.is_null(3));
556+
assert_eq!(variant_array.value(4), Variant::from(2i32));
557+
}
558+
529559
#[test]
530560
fn test_extend_variant_array_builder() {
531561
let mut b = VariantArrayBuilder::new(3);

0 commit comments

Comments
 (0)