diff --git a/parquet-variant-compute/src/variant_get.rs b/parquet-variant-compute/src/variant_get.rs index 0caef50759b7..fc01ca8d3b20 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::); + + // 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::) { + 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::>>()?; - - let struct_nulls = target.nulls().cloned(); + if !as_field.has_valid_extension_type::() { + 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::>>()?; + + 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 { // 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 {