Skip to content

Commit c194e54

Browse files
authored
[Variant] Add unshredded Struct fast-path for variant_get(..., Struct) (#9597)
# 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 #9596. # 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? Reuse `shred_basic_variant` as a fast path for unshredded `Struct` handling in `variant_get(..., Struct)` <!-- 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 two unit tests to establish safe mode behavior. <!-- 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? <!-- 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. -->
1 parent 7f307c0 commit c194e54

File tree

1 file changed

+88
-2
lines changed

1 file changed

+88
-2
lines changed

parquet-variant-compute/src/variant_get.rs

Lines changed: 88 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -213,9 +213,20 @@ fn shredded_get_path(
213213
return Ok(shredded);
214214
}
215215

216-
// Structs are special. Recurse into each field separately, hoping to follow the shredding even
217-
// further, and build up the final struct from those individually shredded results.
216+
// Structs are special.
217+
//
218+
// For fully unshredded targets (`typed_value` absent), delegate to the row builder so we
219+
// preserve struct-level cast semantics:
220+
// - safe mode: non-object rows become NULL structs
221+
// - strict mode: non-object rows raise a cast error
222+
//
223+
// For shredded/partially-shredded targets (`typed_value` present), recurse into each field
224+
// separately to take advantage of deeper shredding in child fields.
218225
if let DataType::Struct(fields) = as_field.data_type() {
226+
if target.typed_value_field().is_none() {
227+
return shred_basic_variant(target, VariantPath::default(), Some(as_field));
228+
}
229+
219230
let children = fields
220231
.iter()
221232
.map(|field| {
@@ -3111,6 +3122,81 @@ mod test {
31113122
assert_eq!(inner_values.value(1), 100);
31123123
}
31133124

3125+
#[test]
3126+
fn test_unshredded_struct_safe_cast_non_object_rows_are_null() {
3127+
let json_strings = vec![r#"{"a": 1, "b": 2}"#, "123", "{}"];
3128+
let string_array: Arc<dyn Array> = Arc::new(StringArray::from(json_strings));
3129+
let variant_array_ref = ArrayRef::from(json_to_variant(&string_array).unwrap());
3130+
3131+
let struct_fields = Fields::from(vec![
3132+
Field::new("a", DataType::Int32, true),
3133+
Field::new("b", DataType::Int32, true),
3134+
]);
3135+
let options = GetOptions {
3136+
path: VariantPath::default(),
3137+
as_type: Some(Arc::new(Field::new(
3138+
"result",
3139+
DataType::Struct(struct_fields),
3140+
true,
3141+
))),
3142+
cast_options: CastOptions::default(),
3143+
};
3144+
3145+
let result = variant_get(&variant_array_ref, options).unwrap();
3146+
let struct_result = result.as_struct();
3147+
let field_a = struct_result
3148+
.column(0)
3149+
.as_primitive::<arrow::datatypes::Int32Type>();
3150+
let field_b = struct_result
3151+
.column(1)
3152+
.as_primitive::<arrow::datatypes::Int32Type>();
3153+
3154+
// Row 0 is an object, so the struct row is valid with extracted fields.
3155+
assert!(!struct_result.is_null(0));
3156+
assert_eq!(field_a.value(0), 1);
3157+
assert_eq!(field_b.value(0), 2);
3158+
3159+
// Row 1 is a scalar, so safe struct cast should produce a NULL struct row.
3160+
assert!(struct_result.is_null(1));
3161+
assert!(field_a.is_null(1));
3162+
assert!(field_b.is_null(1));
3163+
3164+
// Row 2 is an empty object, so the struct row is valid with missing fields as NULL.
3165+
assert!(!struct_result.is_null(2));
3166+
assert!(field_a.is_null(2));
3167+
assert!(field_b.is_null(2));
3168+
}
3169+
3170+
#[test]
3171+
fn test_unshredded_struct_strict_cast_non_object_errors() {
3172+
let json_strings = vec![r#"{"a": 1, "b": 2}"#, "123"];
3173+
let string_array: Arc<dyn Array> = Arc::new(StringArray::from(json_strings));
3174+
let variant_array_ref = ArrayRef::from(json_to_variant(&string_array).unwrap());
3175+
3176+
let struct_fields = Fields::from(vec![
3177+
Field::new("a", DataType::Int32, true),
3178+
Field::new("b", DataType::Int32, true),
3179+
]);
3180+
let options = GetOptions {
3181+
path: VariantPath::default(),
3182+
as_type: Some(Arc::new(Field::new(
3183+
"result",
3184+
DataType::Struct(struct_fields),
3185+
true,
3186+
))),
3187+
cast_options: CastOptions {
3188+
safe: false,
3189+
..Default::default()
3190+
},
3191+
};
3192+
3193+
let err = variant_get(&variant_array_ref, options).unwrap_err();
3194+
assert!(
3195+
err.to_string()
3196+
.contains("Failed to extract struct from variant")
3197+
);
3198+
}
3199+
31143200
/// Create comprehensive shredded variant with diverse null patterns and empty objects
31153201
/// Rows: normal values, top-level null, missing field a, missing field b, empty object
31163202
fn create_comprehensive_shredded_variant() -> ArrayRef {

0 commit comments

Comments
 (0)