This is an automated email from the ASF dual-hosted git repository.
alamb 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 d2f1611253 [Variant] Add `variant_get` access as `Variant` (#9681)
d2f1611253 is described below
commit d2f16112531f5a559a2141036a7ddd3ac867c324
Author: Konstantin Tarasov <[email protected]>
AuthorDate: Tue Jun 23 15:56:58 2026 -0400
[Variant] Add `variant_get` access as `Variant` (#9681)
# 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 #8154.
# Rationale for this change
Check issue
<!--
Why are you proposing this change? If this is already explained clearly
in the issue then this section is not needed.
Explaining clearly why changes are proposed helps reviewers understand
your changes and offer better suggestions for fixes.
-->
# What changes are included in this PR?
- Added support for `variant_get` access as `Variant`
- Added unit tests
<!--
There is no need to duplicate the description in the issue here but it
is sometimes worth providing a summary of the individual changes in this
PR.
-->
# Are these changes tested?
Yes, added unit tests
<!--
We typically require tests for all PRs in order to:
1. Prevent the code from being accidentally broken by subsequent changes
2. Serve as another way to document the expected behavior of the code
If tests are not included in your PR, please explain why (for example,
are they covered by existing tests)?
-->
# Are there any user-facing changes?
Can now extract unshredded `VariantArray` with `variant_get`
<!--
If there are user-facing changes then we may require documentation to be
updated before approving the PR.
If there are any breaking changes to public APIs, please call them out.
-->
---------
Co-authored-by: Ryan Johnson <[email protected]>
---
parquet-variant-compute/src/variant_get.rs | 186 ++++++++++++++++++++++++-----
1 file changed, 159 insertions(+), 27 deletions(-)
diff --git a/parquet-variant-compute/src/variant_get.rs
b/parquet-variant-compute/src/variant_get.rs
index 0caef50759..fc01ca8d3b 100644
--- a/parquet-variant-compute/src/variant_get.rs
+++ b/parquet-variant-compute/src/variant_get.rs
@@ -28,8 +28,8 @@ use arrow_schema::{ArrowError, DataType, FieldRef};
use parquet_variant::{VariantPath, VariantPathElement};
use crate::ShreddingState;
-use crate::VariantArray;
use crate::variant_to_arrow::make_variant_to_arrow_row_builder;
+use crate::{VariantArray, VariantType, unshred_variant};
use arrow::array::AsArray;
use std::sync::Arc;
@@ -201,10 +201,46 @@ fn shredded_get_path(
VariantArray::from_parts(metadata, value, typed_value,
accumulated_nulls)
};
- // Helper that shreds a VariantArray to a specific type.
+ // Helper that extracts the value at `path` and casts it to the requested
type, or returns it as
+ // an unshredded binary variant when `Variant` output is requested.
let shred_basic_variant =
|target: VariantArray, path: VariantPath<'_>, as_field:
Option<&Field>| {
- let as_type = as_field.map(|f| f.data_type());
+ // A `VariantType` extension on `as_field` requests `Variant`
output: return an
+ // unshredded binary variant instead of casting to a concrete
Arrow type.
+ let requested_variant =
+
as_field.is_some_and(Field::has_valid_extension_type::<VariantType>);
+
+ // A `typed_value` in that field requests shredded output -- a
`VariantArray` with
+ // `typed_value` columns. We produce only unshredded variant
output. Shredded output is
+ // tracked in https://github.com/apache/arrow-rs/issues/8153.
Reject such a request
+ // instead of silently dropping the shredding it asked for.
+ if requested_variant && requested_field_is_shredded(as_field) {
+ return Err(ArrowError::NotYetImplemented(
+ "variant_get with shredded `Variant` output is not yet
supported".to_string(),
+ ));
+ }
+
+ // Collapse any shredding back to binary. Only the `NotShredded`
step below passes a
+ // non-empty `path`, and there `target` is already a plain `value`
column (no
+ // `typed_value`) -- so `unshred_variant` hits its clone
fast-path, with nothing deeper
+ // to shred. The builder then walks any remaining path per-row,
emitting variant output
+ // because `as_type` is `None`.
+ let target = if requested_variant {
+ unshred_variant(&target)?
+ } else {
+ target
+ };
+
+ // Path exhausted, variant requested: return the target directly.
+ if requested_variant && path.is_empty() {
+ return Ok(ArrayRef::from(target));
+ }
+
+ let as_type = if requested_variant {
+ None
+ } else {
+ as_field.map(|f| f.data_type())
+ };
let mut builder = make_variant_to_arrow_row_builder(
target.metadata_column(),
path,
@@ -250,6 +286,14 @@ fn shredded_get_path(
}
ShreddedPathStep::Missing => {
let num_rows = input.len();
+ if
as_field.is_some_and(Field::has_valid_extension_type::<VariantType>) {
+ let all_nulls =
Some(arrow::buffer::NullBuffer::from(vec![false; num_rows]));
+ // Propagating metadata is not necessary for an all-NULL
array, but is cheaper than constructing
+ // a new empty metadata array. (n * 3 bytes vs Arc bump)
+ let metadata = input.metadata_column().clone();
+ let arr = VariantArray::from_parts(metadata, None, None,
all_nulls);
+ return Ok(ArrayRef::from(arr));
+ }
let arr = match as_field.map(|f| f.data_type()) {
Some(data_type) => array::new_null_array(data_type,
num_rows),
None => Arc::new(array::NullArray::new(num_rows)) as _,
@@ -293,36 +337,43 @@ fn shredded_get_path(
//
// For shredded/partially-shredded targets (`typed_value` present),
recurse into each field
// separately to take advantage of deeper shredding in child fields.
- if let DataType::Struct(fields) = as_field.data_type() {
- if target.typed_value_column().is_none() {
- return shred_basic_variant(target, VariantPath::default(),
Some(as_field));
- }
-
- let children = fields
- .iter()
- .map(|field| {
- shredded_get_path(
- &target,
- &[VariantPathElement::from(field.name().as_str())],
- Some(field),
- cast_options,
- )
- })
- .collect::<Result<Vec<_>>>()?;
-
- let struct_nulls = target.nulls().cloned();
+ if !as_field.has_valid_extension_type::<VariantType>() {
+ if let DataType::Struct(fields) = as_field.data_type() {
+ if target.typed_value_column().is_none() {
+ return shred_basic_variant(target, VariantPath::default(),
Some(as_field));
+ }
- return Ok(Arc::new(StructArray::try_new(
- fields.clone(),
- children,
- struct_nulls,
- )?));
+ let children = fields
+ .iter()
+ .map(|field| {
+ let path =
&[VariantPathElement::from(field.name().as_str())];
+ shredded_get_path(&target, path, Some(field), cast_options)
+ })
+ .collect::<Result<Vec<_>>>()?;
+
+ return Ok(Arc::new(StructArray::try_new(
+ fields.clone(),
+ children,
+ target.nulls().cloned(),
+ )?));
+ }
}
// Not a struct, so directly shred the variant as the requested type
shred_basic_variant(target, VariantPath::default(), Some(as_field))
}
+/// Returns true if `as_field` requests *shredded* `Variant` output.
+///
+/// Its struct carries a `typed_value` field naming the type to shred to.
+/// A plain variant request has only `metadata` and `value`.
+fn requested_field_is_shredded(as_field: Option<&Field>) -> bool {
+ as_field.is_some_and(|f| match f.data_type() {
+ DataType::Struct(fields) => fields.iter().any(|field| field.name() ==
"typed_value"),
+ _ => false,
+ })
+}
+
fn try_perfect_shredding(variant_array: &VariantArray, as_field: &Field) ->
Option<ArrayRef> {
// Try to return the typed value directly when we have a perfect shredding
match.
if matches!(as_field.data_type(), DataType::Struct(_)) {
@@ -432,7 +483,7 @@ mod test {
use std::str::FromStr;
use std::sync::Arc;
- use super::{GetOptions, variant_get};
+ use super::{GetOptions, requested_field_is_shredded, variant_get};
use crate::variant_array::{ShreddedVariantFieldArray, StructArrayBuilder};
use crate::{
ShreddedSchemaBuilder, VariantArray, VariantArrayBuilder,
cast_to_variant, json_to_variant,
@@ -452,6 +503,7 @@ mod test {
use arrow::datatypes::DataType::{Int16, Int32, Int64};
use arrow::datatypes::i256;
use arrow::util::display::FormatOptions;
+ use arrow_schema::ArrowError;
use arrow_schema::DataType::{Boolean, Float32, Float64, Int8};
use arrow_schema::{DataType, Field, FieldRef, Fields, IntervalUnit,
TimeUnit};
use chrono::DateTime;
@@ -2339,6 +2391,86 @@ mod test {
println!("Nested path 'a.x' result: {:?}", result);
}
+ #[test]
+ fn test_variant_get_as_variant_from_unshredded_input() {
+ let (unshredded, _) = create_variant_get_as_variant_test_data();
+ let unshredded_field =
VariantArray::try_new(&unshredded).unwrap().field("result");
+
assert_variant_field_extraction_returns_unshredded_variant(&unshredded,
&unshredded_field);
+ }
+
+ #[test]
+ fn test_variant_get_as_variant_from_shredded_input() {
+ let (unshredded, shredded) = create_variant_get_as_variant_test_data();
+ let unshredded_field =
VariantArray::try_new(&unshredded).unwrap().field("result");
+ assert_variant_field_extraction_returns_unshredded_variant(&shredded,
&unshredded_field);
+ }
+
+ #[test]
+ fn test_variant_get_as_shredded_variant_is_not_yet_supported() {
+ let (_, shredded) = create_variant_get_as_variant_test_data();
+ // Deriving the request field from the shredded array yields a
`VariantType` field whose
+ // struct carries a `typed_value` -- a request to shred the output.
That is unsupported
+ // (https://github.com/apache/arrow-rs/issues/8153) and must error,
not silently return a
+ // plain binary variant.
+ let shredded_field =
VariantArray::try_new(&shredded).unwrap().field("result");
+ assert!(requested_field_is_shredded(Some(&shredded_field)));
+
+ let options =
GetOptions::new_with_path(VariantPath::try_from("field_name").unwrap())
+ .with_as_type(Some(FieldRef::from(shredded_field)));
+ let err = variant_get(&shredded, options).unwrap_err();
+ assert!(
+ matches!(err, ArrowError::NotYetImplemented(_)),
+ "expected NotYetImplemented, got {err:?}"
+ );
+ }
+
+ fn create_variant_get_as_variant_test_data() -> (ArrayRef, ArrayRef) {
+ let input_json: ArrayRef = Arc::new(StringArray::from(vec![
+ Some(r#"{"field_name": {"k": 100000}}"#),
+ Some(r#"{"field_name": {"k": "s"}}"#),
+ ]));
+
+ let unshredded = ArrayRef::from(json_to_variant(&input_json).unwrap());
+ let unshredded_variant = VariantArray::try_new(&unshredded).unwrap();
+
+ let as_type = DataType::Struct(Fields::from(vec![Field::new(
+ "field_name",
+ DataType::Struct(Fields::from(vec![Field::new("k",
DataType::Int32, true)])),
+ true,
+ )]));
+ let shredded = ArrayRef::from(shred_variant(&unshredded_variant,
&as_type).unwrap());
+
+ (unshredded, shredded)
+ }
+
+ fn assert_variant_field_extraction_returns_unshredded_variant(
+ input: &ArrayRef,
+ variant_field: &Field,
+ ) {
+ let options =
GetOptions::new_with_path(VariantPath::try_from("field_name").unwrap())
+ .with_as_type(Some(FieldRef::from(variant_field.clone())));
+
+ let result = variant_get(input, options).unwrap();
+ let result_variant = VariantArray::try_new(&result).unwrap();
+
+ assert!(result_variant.typed_value_column().is_none());
+ assert!(result_variant.value_column().is_some());
+
+ let expected_json: ArrayRef = Arc::new(StringArray::from(vec![
+ Some(r#"{"k":100000}"#),
+ Some(r#"{"k":"s"}"#),
+ ]));
+ let expected = json_to_variant(&expected_json).unwrap();
+
+ assert_eq!(result_variant.len(), expected.len());
+ for i in 0..result_variant.len() {
+ assert_eq!(result_variant.is_null(i), expected.is_null(i));
+ if !result_variant.is_null(i) {
+ assert_eq!(result_variant.value(i), expected.value(i));
+ }
+ }
+ }
+
/// Create test data for depth 0 (direct field access)
/// [{"x": 42}, {"x": "foo"}, {"y": 10}]
fn create_depth_0_test_data() -> ArrayRef {
