diff --git a/parquet-variant-compute/src/cast_to_variant.rs b/parquet-variant-compute/src/cast_to_variant.rs index b6c968b0678d..1b26ffe07d2a 100644 --- a/parquet-variant-compute/src/cast_to_variant.rs +++ b/parquet-variant-compute/src/cast_to_variant.rs @@ -58,6 +58,13 @@ pub fn cast_to_variant_with_options( input: &dyn Array, options: &CastOptions, ) -> Result { + // Fast path: any all-null input maps to an all-null VariantArray. + if input.null_count() == input.len() { + let mut array_builder = VariantArrayBuilder::new(input.len()); + array_builder.append_nulls(input.len()); + return Ok(array_builder.build()); + } + // Create row builder for the input array type let mut row_builder = make_arrow_to_variant_row_builder(input.data_type(), input, options)?; diff --git a/parquet-variant-compute/src/variant_array_builder.rs b/parquet-variant-compute/src/variant_array_builder.rs index 86ece0010042..f6b95c881af4 100644 --- a/parquet-variant-compute/src/variant_array_builder.rs +++ b/parquet-variant-compute/src/variant_array_builder.rs @@ -156,6 +156,18 @@ impl VariantArrayBuilder { self.value_offsets.push(self.value_builder.offset()); } + /// Appends `n` null rows to the builder. + pub fn append_nulls(&mut self, n: usize) { + self.nulls.append_n_nulls(n); + // The subfields are expected to be non-nullable according to the parquet variant spec. + let metadata_offset = self.metadata_builder.offset(); + let value_offset = self.value_builder.offset(); + self.metadata_offsets + .extend(std::iter::repeat_n(metadata_offset, n)); + self.value_offsets + .extend(std::iter::repeat_n(value_offset, n)); + } + /// Append the [`Variant`] to the builder as the next row pub fn append_variant(&mut self, variant: Variant) { ValueBuilder::append_variant(self.parent_state(), variant); @@ -526,6 +538,24 @@ mod test { assert_eq!(list.len(), 2); } + #[test] + fn test_variant_array_builder_append_nulls() { + let mut builder = VariantArrayBuilder::new(6); + builder.append_variant(Variant::from(1i32)); + builder.append_nulls(0); // should be a no-op + builder.append_nulls(3); + builder.append_variant(Variant::from(2i32)); + + let variant_array = builder.build(); + + assert_eq!(variant_array.len(), 5); + assert_eq!(variant_array.value(0), Variant::from(1i32)); + assert!(variant_array.is_null(1)); + assert!(variant_array.is_null(2)); + assert!(variant_array.is_null(3)); + assert_eq!(variant_array.value(4), Variant::from(2i32)); + } + #[test] fn test_extend_variant_array_builder() { let mut b = VariantArrayBuilder::new(3);