This is an automated email from the ASF dual-hosted git repository.
jayzhan211 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 9500647b24 fix(variant): accept Dictionary and RunEndEncoded metadata
in VariantArray (#10810)
9500647b24 is described below
commit 9500647b24ae6c9d271daa9c7d0487deb8849a16
Author: Peter Lee <[email protected]>
AuthorDate: Fri Aug 28 06:40:58 2026 +0800
fix(variant): accept Dictionary and RunEndEncoded metadata in VariantArray
(#10810)
# Which issue does this PR close?
- Closes #10802.
# Rationale for this change
The canonical `arrow.parquet.variant` extension permits Dictionary and
RunEndEncoded metadata, but `VariantArray` only accepted unencoded
binary arrays. Downstream users therefore had to decode valid metadata
eagerly before using Variant compute kernels.
# What changes are included in this PR?
- Accept Dictionary and RunEndEncoded metadata with Binary, LargeBinary,
or BinaryView logical values.
- Resolve encoded values and logical nulls in the shared metadata byte
accessor.
- Route `variant_to_json` through `VariantArray` so it shares the same
validation and row access.
- Add coverage for Dictionary/REE metadata, nulls, slices,
`variant_get`, and `variant_to_json`.
# Are these changes tested?
- `cargo test -p parquet-variant-compute --lib` — 364 passed
- `cargo test -p parquet-variant-compute --doc` — 13 passed
- `cargo clippy -p parquet-variant-compute --all-targets -- -D warnings`
- `cargo fmt --all -- --check`
- `git diff --check`
# Are there any user-facing changes?
Yes. `VariantArray` and its compute kernels now accept canonical
Dictionary- and RunEndEncoded metadata. There are no public API
signature changes or breaking changes.
# AI-assisted development
Codex assisted with tracing callers, implementing the focused change,
and drafting tests. I reviewed the resulting code and verified it with
the checks above.
---------
Co-authored-by: Claude Fable 5 <[email protected]>
---
parquet-variant-compute/src/to_json.rs | 62 +++---------
parquet-variant-compute/src/variant_array.rs | 143 ++++++++++++++++++++++++---
2 files changed, 138 insertions(+), 67 deletions(-)
diff --git a/parquet-variant-compute/src/to_json.rs
b/parquet-variant-compute/src/to_json.rs
index efe50c2a09..0a173fb5e9 100644
--- a/parquet-variant-compute/src/to_json.rs
+++ b/parquet-variant-compute/src/to_json.rs
@@ -15,73 +15,33 @@
// specific language governing permissions and limitations
// under the License.
-//! Module for transforming a batch of Variants represented as
-//! STRUCT<metadata: BINARY, value: BINARY> into a batch of JSON strings.
+//! Module for transforming a batch of Variants into JSON strings.
-use arrow::array::{Array, ArrayRef, BinaryArray, BooleanBufferBuilder,
StringArray, StructArray};
+use crate::VariantArray;
+use arrow::array::{ArrayRef, BooleanBufferBuilder, StringArray};
use arrow::buffer::{Buffer, NullBuffer, OffsetBuffer, ScalarBuffer};
-use arrow::datatypes::DataType;
use arrow_schema::ArrowError;
-use parquet_variant::Variant;
use parquet_variant_json::VariantToJson;
-/// Transform a batch of Variant represented as STRUCT<metadata: BINARY,
value: BINARY> to a batch
-/// of JSON strings where nulls are preserved. The JSON strings in the input
must be valid.
+/// Transform a [`VariantArray`] to a batch of JSON strings where nulls are
preserved.
pub fn variant_to_json(input: &ArrayRef) -> Result<StringArray, ArrowError> {
- let struct_array = input
- .as_any()
- .downcast_ref::<StructArray>()
- .ok_or_else(|| ArrowError::CastError("Expected StructArray as
input".into()))?;
-
- // Validate field types
- let data_type = struct_array.data_type();
- match data_type {
- DataType::Struct(inner_fields) => {
- if inner_fields.len() != 2
- || inner_fields[0].data_type() != &DataType::Binary
- || inner_fields[1].data_type() != &DataType::Binary
- {
- return Err(ArrowError::CastError(
- "Expected struct with two binary fields".into(),
- ));
- }
- }
- _ => {
- return Err(ArrowError::CastError(
- "Expected StructArray with known fields".into(),
- ));
- }
- }
-
- let metadata_array = struct_array
- .column(0)
- .as_any()
- .downcast_ref::<BinaryArray>()
- .ok_or_else(|| ArrowError::CastError("Expected BinaryArray for
'metadata'".into()))?;
-
- let value_array = struct_array
- .column(1)
- .as_any()
- .downcast_ref::<BinaryArray>()
- .ok_or_else(|| ArrowError::CastError("Expected BinaryArray for
'value'".into()))?;
+ let array = VariantArray::try_new(input)?;
// Zero-copy builder
// The size per JSON string is assumed to be 128 bytes. If this holds
true, resizing could be
// minimized for performance.
- let mut json_buffer: Vec<u8> = Vec::with_capacity(struct_array.len() *
128);
- let mut offsets: Vec<i32> = Vec::with_capacity(struct_array.len() + 1);
- let mut validity = BooleanBufferBuilder::new(struct_array.len());
+ let mut json_buffer: Vec<u8> = Vec::with_capacity(array.len() * 128);
+ let mut offsets: Vec<i32> = Vec::with_capacity(array.len() + 1);
+ let mut validity = BooleanBufferBuilder::new(array.len());
let mut current_offset: i32 = 0;
offsets.push(current_offset);
- for i in 0..struct_array.len() {
- if struct_array.is_null(i) {
+ for i in 0..array.len() {
+ if array.is_null(i) {
validity.append(false);
offsets.push(current_offset);
} else {
- let metadata = metadata_array.value(i);
- let value = value_array.value(i);
- let variant = Variant::new(metadata, value);
+ let variant = array.try_value(i)?;
let start_len = json_buffer.len();
variant.to_json(&mut json_buffer)?;
let written = (json_buffer.len() - start_len) as i32;
diff --git a/parquet-variant-compute/src/variant_array.rs
b/parquet-variant-compute/src/variant_array.rs
index b8a0894719..4dd3b1aab0 100644
--- a/parquet-variant-compute/src/variant_array.rs
+++ b/parquet-variant-compute/src/variant_array.rs
@@ -22,7 +22,10 @@ use crate::type_conversion::{
generic_conversion_single_value,
generic_conversion_single_value_with_result,
primitive_conversion_single_value,
};
-use arrow::array::{Array, ArrayRef, AsArray, StructArray, new_null_array};
+use arrow::array::{
+ Array, ArrayRef, AsArray, StructArray, downcast_dictionary_array,
downcast_run_array,
+ new_null_array,
+};
use arrow::buffer::NullBuffer;
use arrow::compute::cast;
use arrow::datatypes::{
@@ -41,12 +44,30 @@ use parquet_variant::{
use std::borrow::Cow;
use std::sync::Arc;
-/// Returns the raw bytes at the given index from a binary-like array, return
`None` if the array isn't binary-like.
+/// Returns the logical bytes at the given index from a binary-like array,
resolving dictionary
+/// and run-end encodings. Returns `None` for nulls or if the logical values
aren't binary-like.
pub(crate) fn binary_array_value(array: &dyn Array, index: usize) ->
Option<&[u8]> {
+ if array.is_null(index) {
+ return None;
+ }
match array.data_type() {
DataType::Binary => Some(array.as_binary::<i32>().value(index)),
DataType::LargeBinary => Some(array.as_binary::<i64>().value(index)),
DataType::BinaryView => Some(array.as_binary_view().value(index)),
+ DataType::Dictionary(..) => downcast_dictionary_array! {
+ array => {
+ let index = array.key(index)?;
+ binary_array_value(array.values().as_ref(), index)
+ },
+ _ => unreachable!(),
+ },
+ DataType::RunEndEncoded(..) => downcast_run_array! {
+ array => {
+ let index = array.get_physical_index(index);
+ binary_array_value(array.values().as_ref(), index)
+ },
+ _ => unreachable!(),
+ },
_ => None,
}
}
@@ -83,6 +104,25 @@ pub(crate) fn validate_binary_array(array: &dyn Array,
field_name: &str) -> Resu
}
}
+/// Validates that a metadata array has binary-like logical values.
+fn validate_metadata_array(array: &dyn Array) -> Result<()> {
+ let is_binary = |data_type: &DataType| {
+ matches!(
+ data_type,
+ DataType::Binary | DataType::LargeBinary | DataType::BinaryView
+ )
+ };
+ match array.data_type() {
+ data_type if is_binary(data_type) => Ok(()),
+ DataType::Dictionary(_, values) if is_binary(values) => Ok(()),
+ DataType::RunEndEncoded(_, values) if is_binary(values.data_type()) =>
Ok(()),
+ _ => Err(ArrowError::InvalidArgumentError(format!(
+ "VariantArray 'metadata' field must be Binary, LargeBinary,
BinaryView, or a Dictionary or RunEndEncoded array of one of those types, got
{}",
+ array.data_type()
+ ))),
+ }
+}
+
/// Arrow Variant [`ExtensionType`].
///
/// Represents the canonical Arrow Extension Type for storing variants.
@@ -263,7 +303,7 @@ pub struct VariantArray {
/// Reference to the underlying StructArray
inner: StructArray,
- /// The metadata column of this variant (Binary, LargeBinary, or
BinaryView)
+ /// The metadata column of this variant (Binary, LargeBinary, or
BinaryView), possibly encoded
metadata: ArrayRef,
/// how is this variant array shredded?
@@ -285,7 +325,7 @@ impl VariantArray {
/// # Requirements of the `StructArray`
///
/// 1. A required field named `metadata` which is binary, large_binary, or
- /// binary_view
+ /// binary_view, optionally dictionary or run-end-encoded
///
/// 2. A required field named `value` that is binary, large_binary, or
/// binary_view
@@ -293,10 +333,6 @@ impl VariantArray {
/// 3. An optional field named `typed_value` which can be any primitive
type
/// or be a list, large_list, list_view or struct
///
- /// NOTE: It is also permissible for the metadata field to be
- /// Dictionary-Encoded, preferably (but not required) with an index type of
- /// int8.
- ///
pub fn try_new(inner: &dyn Array) -> Result<Self> {
// Canonicalize shredded typed_value fields (e.g. decimal narrowing)
let inner = canonicalize_shredded_types(inner)?;
@@ -309,13 +345,13 @@ impl VariantArray {
// Note the specification allows for any order so we must search by
name
- // Ensure the StructArray has a metadata field that is a binary type
+ // Ensure the StructArray has a metadata field with binary-like
logical values
let Some(metadata_col) = inner.column_by_name("metadata") else {
return Err(ArrowError::InvalidArgumentError(
"Invalid VariantArray: StructArray must contain a 'metadata'
field".to_string(),
));
};
- validate_binary_array(metadata_col.as_ref(), "metadata")?;
+ validate_metadata_array(metadata_col.as_ref())?;
let shredding_state = ShreddingState::try_from(inner)?;
@@ -1187,7 +1223,7 @@ fn typed_value_to_variant(typed_value: &ArrayRef, index:
usize) -> Result<Varian
/// Canonicalize shredded typed_value fields (e.g. decimal narrowing) and
/// verify that all data types in the struct are legal for a variant array.
fn canonicalize_shredded_types(array: &dyn Array) -> Result<ArrayRef> {
- let new_type = canonicalize_and_verify_data_type(array.data_type())?;
+ let new_type = canonicalize_and_verify_data_type_impl(array.data_type(),
true)?;
if let Cow::Borrowed(_) = new_type
&& let Some(array) = array.as_struct_opt()
{
@@ -1200,6 +1236,13 @@ fn canonicalize_shredded_types(array: &dyn Array) ->
Result<ArrayRef> {
/// appear in a (possibly shredded) variant array. It also narrows decimal
types to the smallest
/// valid precision (e.g. Decimal128 -> Decimal32 when the precision fits).
fn canonicalize_and_verify_data_type(data_type: &DataType) -> Result<Cow<'_,
DataType>> {
+ canonicalize_and_verify_data_type_impl(data_type, false)
+}
+
+fn canonicalize_and_verify_data_type_impl(
+ data_type: &DataType,
+ skip_top_level_metadata: bool,
+) -> Result<Cow<'_, DataType>> {
use DataType::*;
// helper macros
@@ -1279,6 +1322,9 @@ fn canonicalize_and_verify_data_type(data_type:
&DataType) -> Result<Cow<'_, Dat
// of the data type. Even if some fields change, the others are
shallow arc clones.
let mut new_fields = std::collections::HashMap::new();
for (i, field) in fields.iter().enumerate() {
+ if skip_top_level_metadata && field.name() == "metadata" {
+ continue;
+ }
if let Cow::Owned(new_field) =
canonicalize_and_verify_field(field)? {
new_fields.insert(i, new_field);
}
@@ -1346,14 +1392,15 @@ impl ShreddedVariantFieldArray {
#[cfg(test)]
mod test {
- use crate::VariantArrayBuilder;
+ use crate::{GetOptions, VariantArrayBuilder, json_to_variant, variant_get,
variant_to_json};
use std::str::FromStr;
use super::*;
use arrow::array::{
- BinaryArray, BinaryViewArray, Decimal32Array, Decimal64Array,
Decimal128Array,
- FixedSizeBinaryArray, Int32Array, Int64Array, LargeBinaryArray,
LargeListArray,
- LargeListViewArray, ListArray, ListViewArray, Time64MicrosecondArray,
+ BinaryArray, BinaryDictionaryBuilder, BinaryRunBuilder,
BinaryViewArray, Decimal32Array,
+ Decimal64Array, Decimal128Array, FixedSizeBinaryArray, Int8Array,
Int32Array, Int64Array,
+ LargeBinaryArray, LargeListArray, LargeListViewArray, ListArray,
ListViewArray,
+ StringArray, Time64MicrosecondArray,
};
use arrow::buffer::{OffsetBuffer, ScalarBuffer};
use arrow_schema::{Field, Fields};
@@ -1429,10 +1476,74 @@ mod test {
let err = VariantArray::try_new(&array);
assert_eq!(
err.unwrap_err().to_string(),
- "Invalid argument error: VariantArray 'metadata' field must be
Binary, LargeBinary, or BinaryView, got Int32"
+ "Invalid argument error: VariantArray 'metadata' field must be
Binary, LargeBinary, BinaryView, or a Dictionary or RunEndEncoded array of one
of those types, got Int32"
);
}
+ #[test]
+ fn encoded_metadata_supports_nulls_slices_and_variant_get() {
+ let json: ArrayRef = Arc::new(StringArray::from(vec![
+ Some(r#"{"a":0}"#),
+ Some(r#"{"a":1}"#),
+ None,
+ Some(r#"{"b":3}"#),
+ Some(r#"{"b":4}"#),
+ ]));
+ let baseline = json_to_variant(&json).unwrap();
+ let metadata = baseline.metadata_column().as_binary_view();
+ let metadata_a = metadata.value(0);
+ let metadata_b = metadata.value(3);
+
+ let logical_metadata = [
+ Some(metadata_a),
+ Some(metadata_a),
+ None,
+ Some(metadata_b),
+ Some(metadata_b),
+ ];
+
+ let mut dictionary = BinaryDictionaryBuilder::<Int8Type>::new();
+ dictionary.extend(logical_metadata);
+ let dictionary: ArrayRef = Arc::new(dictionary.finish());
+
+ let mut ree = BinaryRunBuilder::<Int16Type>::new();
+ ree.extend(logical_metadata);
+ let run_end_encoded: ArrayRef = Arc::new(ree.finish());
+
+ for metadata in [dictionary, run_end_encoded] {
+ assert_eq!(binary_array_value(metadata.as_ref(), 2), None);
+ let fields = Fields::from(vec![
+ Field::new("metadata", metadata.data_type().clone(), false),
+ Field::new("value",
baseline.value_column().data_type().clone(), false),
+ ]);
+ let input = StructArray::try_new(
+ fields,
+ vec![metadata, baseline.value_column().clone()],
+ baseline.nulls().cloned(),
+ )
+ .unwrap()
+ .slice(1, 3);
+
+ let variant = VariantArray::try_new(&input).unwrap();
+ assert_eq!(variant.value(0), baseline.value(1));
+ assert!(variant.is_null(1));
+ assert_eq!(variant.value(2), baseline.value(3));
+
+ let input: ArrayRef = Arc::new(input);
+ let options = GetOptions::new_with_path("b".try_into().unwrap())
+ .with_as_type(Some(Arc::new(Field::new("b", DataType::Int8,
true))));
+ let result = variant_get(&input, options).unwrap();
+ assert_eq!(
+ result.as_primitive::<Int8Type>(),
+ &Int8Array::from(vec![None, None, Some(3)])
+ );
+ assert_eq!(
+ variant_to_json(&input).unwrap(),
+ StringArray::from(vec![Some(r#"{"a":1}"#), None,
Some(r#"{"b":3}"#)])
+ );
+ }
+ }
+
#[test]
fn invalid_value_field_type() {
let fields = Fields::from(vec