diff --git a/parquet-variant-compute/src/variant_get.rs b/parquet-variant-compute/src/variant_get.rs index 774da0e72e8c..38c577564dea 100644 --- a/parquet-variant-compute/src/variant_get.rs +++ b/parquet-variant-compute/src/variant_get.rs @@ -15,17 +15,20 @@ // specific language governing permissions and limitations // under the License. use arrow::{ - array::{self, Array, ArrayRef, StructArray, make_array}, + array::{ + self, Array, ArrayRef, GenericListArray, GenericListViewArray, ListLikeArray, StructArray, + UInt64Array, make_array, + }, buffer::NullBuffer, - compute::CastOptions, + compute::{CastOptions, take}, datatypes::Field, error::Result, }; use arrow_schema::{ArrowError, DataType, FieldRef}; use parquet_variant::{VariantPath, VariantPathElement}; +use crate::ShreddingState; use crate::VariantArray; -use crate::variant_array::ShreddingState; use crate::variant_to_arrow::make_variant_to_arrow_row_builder; use arrow::array::AsArray; @@ -43,12 +46,70 @@ pub(crate) enum ShreddedPathStep { NotShredded, } +/// Build the next shredding state by taking one list-like element (at `index`) per input row. +/// +fn take_list_like_index_as_shredding_state( + typed_value: &dyn Array, + index: usize, +) -> Result> { + let list_array = typed_value.as_any().downcast_ref::().ok_or_else(|| { + ArrowError::ComputeError(format!( + "Expected array type '{}' while handling list-like path step, got '{}'", + std::any::type_name::(), + typed_value.data_type() + )) + })?; + + let values = list_array.values(); + + let Some(struct_array) = values.as_struct_opt() else { + return Ok(None); + }; + let shredding_state = ShreddingState::try_from(struct_array)?; + + let value_array = shredding_state.value_field(); + let typed_array = shredding_state.typed_value_field(); + + // If list elements have neither typed nor fallback value, this path step is missing. + if value_array.is_none() && typed_array.is_none() { + return Ok(None); + } + + let mut take_indices = Vec::with_capacity(list_array.len()); + for row in 0..list_array.len() { + let row_range = list_array.element_range(row); + let take_index = (index < row_range.len()).then(|| (row_range.start + index) as u64); + take_indices.push(take_index); + } + + let index_array = UInt64Array::from(take_indices); + + // Gather both typed and fallback values at the requested element index. + let taken_value = value_array + .map(|value| take(value, &index_array, None)) + .transpose()?; + let taken_typed = typed_array + .map(|typed| take(typed, &index_array, None)) + .transpose()?; + + Ok(Some(ShreddingState::new(taken_value, taken_typed))) +} + /// Given a shredded variant field -- a `(value?, typed_value?)` pair -- try to take one path step /// deeper. For a `VariantPathElement::Field`, if there is no `typed_value` at this level, if /// `typed_value` is not a struct, or if the requested field name does not exist, traversal returns /// a missing-path step (`Missing` or `NotShredded` depending on whether `value` exists). /// -/// TODO: Support `VariantPathElement::Index`? It wouldn't be easy, and maybe not even possible. +/// Safe-cast behavior (`cast_options.safe = true`): +/// - Type mismatch during path traversal (for example field access on non-struct, index access on +/// non-list) returns [`ShreddedPathStep::Missing`] or [`ShreddedPathStep::NotShredded`], allowing +/// the caller to continue with null/fallback semantics. +/// - List index out-of-bounds produces nulls for the corresponding rows. +/// +/// Unsafe-cast behavior (`cast_options.safe = false`): +/// - Field access on non-struct returns [`ArrowError::CastError`]. +/// - List index path steps follow JSONPath semantics and return missing/null for non-list or +/// out-of-bounds rows. pub(crate) fn follow_shredded_path_element( shredding_state: &ShreddingState, path_element: &VariantPathElement<'_>, @@ -69,7 +130,7 @@ pub(crate) fn follow_shredded_path_element( VariantPathElement::Field { name } => { // Try to step into the requested field name of a struct. // First, try to downcast to StructArray - let Some(struct_array) = typed_value.as_any().downcast_ref::() else { + let Some(struct_array) = typed_value.as_struct_opt() else { // Object field path step follows JSONPath semantics and returns missing path step (NotShredded/Missing) on non-struct path return Ok(missing_path_step()); }; @@ -93,12 +154,30 @@ pub(crate) fn follow_shredded_path_element( let state = ShreddingState::try_from(struct_array)?; Ok(ShreddedPathStep::Success(state)) } - VariantPathElement::Index { .. } => { - // TODO: Support array indexing. Among other things, it will require slicing not - // only the array we have here, but also the corresponding metadata and null masks. - Err(ArrowError::NotYetImplemented( - "Pathing into shredded variant array index".into(), - )) + VariantPathElement::Index { index } => { + let state = match typed_value.data_type() { + DataType::List(_) => take_list_like_index_as_shredding_state::< + GenericListArray, + >(typed_value.as_ref(), *index)?, + DataType::LargeList(_) => take_list_like_index_as_shredding_state::< + GenericListArray, + >(typed_value.as_ref(), *index)?, + DataType::ListView(_) => take_list_like_index_as_shredding_state::< + GenericListViewArray, + >(typed_value.as_ref(), *index)?, + DataType::LargeListView(_) => take_list_like_index_as_shredding_state::< + GenericListViewArray, + >(typed_value.as_ref(), *index)?, + _ => { + // JSONPath semantics: indexing a non-list yields no match. + return Ok(missing_path_step()); + } + }; + + match state { + Some(state) => Ok(ShreddedPathStep::Success(state)), + None => Ok(missing_path_step()), + } } } } @@ -356,7 +435,8 @@ mod test { use super::{GetOptions, variant_get}; use crate::variant_array::{ShreddedVariantFieldArray, StructArrayBuilder}; use crate::{ - VariantArray, VariantArrayBuilder, cast_to_variant, json_to_variant, shred_variant, + ShreddedSchemaBuilder, VariantArray, VariantArrayBuilder, cast_to_variant, json_to_variant, + shred_variant, }; use arrow::array::{ Array, ArrayRef, AsArray, BinaryArray, BinaryViewArray, BooleanArray, Date32Array, @@ -448,7 +528,7 @@ mod test { fn get_primitive_variant_inside_object_of_list() { single_variant_get_test( r#"{"some_field": [1234]}"#, - VariantPath::try_from("some_field").unwrap().join(0), + VariantPath::try_from("some_field[0]").unwrap(), "1234", ); } @@ -1741,6 +1821,7 @@ mod test { Some(nulls), )) } + /// This test manually constructs a shredded variant array representing objects /// like {"x": 1, "y": "foo"} and {"x": 42} and tests extracting the "x" field /// as VariantArray using variant_get. @@ -1777,6 +1858,210 @@ mod test { assert_eq!(&result, &expected); } + type ShreddedListLikeArrayGen = fn() -> ArrayRef; + type ShreddedListLikeCase = (&'static str, ShreddedListLikeArrayGen); + + fn shredded_list_like_cases() -> [ShreddedListLikeCase; 4] { + [ + ("list", shredded_list_variant_array), + ("large_list", shredded_large_list_variant_array), + ("list_view", shredded_list_view_variant_array), + ("large_list_view", shredded_large_list_view_variant_array), + ] + } + + #[test] + fn test_shredded_list_like_index_access_from_value_field() { + let options = GetOptions::new_with_path(VariantPath::from(1)); + + for (case, array_gen) in shredded_list_like_cases() { + let array = array_gen(); + let result = variant_get(&array, options.clone()).unwrap(); + let result_variant = VariantArray::try_new(&result).unwrap(); + + assert_eq!(result_variant.value(0), Variant::from("drama"), "{case}"); + assert_eq!(result_variant.value(1).as_int64(), Some(123), "{case}"); + } + } + + #[test] + fn test_shredded_list_like_index_out_of_bounds_unsafe_cast_returns_null() { + let options = + GetOptions::new_with_path(VariantPath::from(10)).with_cast_options(CastOptions { + safe: false, + ..Default::default() + }); + + for (case, array_gen) in shredded_list_like_cases() { + let result = variant_get(&array_gen(), options.clone()).unwrap(); + let result_variant = VariantArray::try_new(&result).unwrap(); + assert_eq!(result_variant.value(0), Variant::Null, "{case}"); + assert_eq!(result_variant.value(1), Variant::Null, "{case}"); + } + } + + /// Test extracting shredded list-like field with type conversion. + #[test] + fn test_shredded_list_like_as_string() { + let field = Field::new("typed_value", DataType::Utf8, false); + let options = GetOptions::new_with_path(VariantPath::from(0)) + .with_as_type(Some(FieldRef::from(field))); + let expected: ArrayRef = Arc::new(StringArray::from(vec![Some("comedy"), Some("horror")])); + + for (case, array_gen) in shredded_list_like_cases() { + let result = variant_get(&array_gen(), options.clone()).unwrap(); + assert_eq!(&result, &expected, "{case}"); + } + } + + #[test] + fn test_shredded_list_like_index_access_from_value_field_as_int64() { + let field = Field::new("typed_value", DataType::Int64, true); + let options = GetOptions::new_with_path(VariantPath::from(1)) + .with_as_type(Some(FieldRef::from(field))); + let expected: ArrayRef = Arc::new(Int64Array::from(vec![None, Some(123)])); + + for (case, array_gen) in shredded_list_like_cases() { + let result = variant_get(&array_gen(), options.clone()).unwrap(); + // "drama" -> NULL, 123 -> 123. + assert_eq!(&result, &expected, "{case}"); + } + } + + #[test] + fn test_shredded_list_in_struct_index_access() { + let array = shredded_struct_with_list_variant_array(); + let options = GetOptions::new_with_path(VariantPath::try_from("a[1]").unwrap()); + let result = variant_get(&array, options).unwrap(); + let result_variant = VariantArray::try_new(&result).unwrap(); + + assert_eq!(result_variant.value(0), Variant::from("drama")); + assert_eq!(result_variant.value(1).as_int64(), Some(123)); + } + + #[test] + fn test_shredded_struct_in_list_field_access() { + let array = shredded_list_of_struct_variant_array(); + let field = Field::new("x", DataType::Int32, true); + let path = VariantPath::from(0).join("x"); + let options = GetOptions::new_with_path(path).with_as_type(Some(FieldRef::from(field))); + let result = variant_get(&array, options).unwrap(); + + let expected: ArrayRef = Arc::new(Int32Array::from(vec![Some(1), Some(3)])); + assert_eq!(&result, &expected); + } + + #[test] + fn test_shredded_list_of_lists_index_access() { + let array = shredded_list_of_lists_variant_array(); + let path = VariantPath::from(0).join(1); + + let result = variant_get(&array, GetOptions::new_with_path(path.clone())).unwrap(); + let result_variant = VariantArray::try_new(&result).unwrap(); + assert_eq!(result_variant.value(0), Variant::from("b")); + assert_eq!(result_variant.value(1).as_int64(), Some(123)); + + let field = Field::new("typed_value", DataType::Int64, true); + let casted = variant_get( + &array, + GetOptions::new_with_path(path).with_as_type(Some(FieldRef::from(field))), + ) + .unwrap(); + let expected: ArrayRef = Arc::new(Int64Array::from(vec![None, Some(123)])); + assert_eq!(&casted, &expected); + } + + /// Helper to create a shredded list-like variant array used by list index tests. + /// + /// Rows: + /// 1. `["comedy", "drama"]` (fully shred-able as `Utf8`) + /// 2. `["horror", 123]` (partially shredded, with fallback for the numeric element) + fn shredded_list_like_variant_array(list_schema: DataType) -> ArrayRef { + let json_rows: ArrayRef = Arc::new(StringArray::from(vec![ + Some(r#"["comedy", "drama"]"#), + Some(r#"["horror", 123]"#), + ])); + let input = json_to_variant(&json_rows).unwrap(); + + let shredded = shred_variant(&input, &list_schema).unwrap(); + ArrayRef::from(shredded) + } + + fn shredded_list_of_lists_variant_array() -> ArrayRef { + let json_rows: ArrayRef = Arc::new(StringArray::from(vec![ + Some(r#"[["a", "b"], ["c", "d"]]"#), + Some(r#"[["x", 123], ["y", "z"]]"#), + ])); + let input = json_to_variant(&json_rows).unwrap(); + + let inner_list = DataType::List(Arc::new(Field::new("item", DataType::Utf8, true))); + let outer_list = DataType::List(Arc::new(Field::new("item", inner_list, true))); + let shredded = shred_variant(&input, &outer_list).unwrap(); + ArrayRef::from(shredded) + } + + fn shredded_list_variant_array() -> ArrayRef { + shredded_list_like_variant_array(DataType::List(Arc::new(Field::new( + "item", + DataType::Utf8, + true, + )))) + } + + fn shredded_large_list_variant_array() -> ArrayRef { + shredded_list_like_variant_array(DataType::LargeList(Arc::new(Field::new( + "item", + DataType::Utf8, + true, + )))) + } + + fn shredded_list_view_variant_array() -> ArrayRef { + shredded_list_like_variant_array(DataType::ListView(Arc::new(Field::new( + "item", + DataType::Utf8, + true, + )))) + } + + fn shredded_large_list_view_variant_array() -> ArrayRef { + shredded_list_like_variant_array(DataType::LargeListView(Arc::new(Field::new( + "item", + DataType::Utf8, + true, + )))) + } + + fn shredded_struct_with_list_variant_array() -> ArrayRef { + let json_rows: ArrayRef = Arc::new(StringArray::from(vec![ + Some(r#"{"a": ["comedy", "drama"]}"#), + Some(r#"{"a": ["horror", 123]}"#), + ])); + let input = json_to_variant(&json_rows).unwrap(); + + let list_schema = DataType::List(Arc::new(Field::new("item", DataType::Utf8, true))); + let shredding_schema = ShreddedSchemaBuilder::default() + .with_path("a", &list_schema) + .unwrap() + .build(); + let shredded = shred_variant(&input, &shredding_schema).unwrap(); + ArrayRef::from(shredded) + } + + fn shredded_list_of_struct_variant_array() -> ArrayRef { + let json_rows: ArrayRef = Arc::new(StringArray::from(vec![ + Some(r#"[{"x": 1}, {"x": 2}]"#), + Some(r#"[{"x": 3}, {"y": 4}]"#), + ])); + let input = json_to_variant(&json_rows).unwrap(); + + let struct_type = + DataType::Struct(Fields::from(vec![Field::new("x", DataType::Int32, true)])); + let list_schema = DataType::List(Arc::new(Field::new("item", struct_type, true))); + let shredded = shred_variant(&input, &list_schema).unwrap(); + ArrayRef::from(shredded) + } + /// Helper function to create a shredded variant array representing objects /// /// This creates an array that represents: @@ -2494,6 +2779,33 @@ mod test { } } + #[test] + fn test_strict_cast_options_index_on_non_list_returns_null() { + use arrow::compute::CastOptions; + use arrow::datatypes::{DataType, Field}; + use parquet_variant::VariantPath; + use std::sync::Arc; + + // Use existing test data that has Int32 typed_value at the top level. + let variant_array = perfectly_shredded_int32_variant_array(); + let options = GetOptions { + path: VariantPath::from(0), + as_type: Some(Arc::new(Field::new("result", DataType::Int32, true))), + cast_options: CastOptions { + safe: false, + ..Default::default() + }, + }; + + let variant_array_ref: Arc = variant_array.clone(); + let result = variant_get(&variant_array_ref, options).unwrap(); + + assert_eq!(result.len(), 3); + assert!(result.is_null(0)); + assert!(result.is_null(1)); + assert!(result.is_null(2)); + } + #[test] fn test_error_message_boolean_type_display() { let mut builder = VariantArrayBuilder::new(1); diff --git a/parquet-variant/src/path.rs b/parquet-variant/src/path.rs index 8e68d9efadf2..0837e49f8abf 100644 --- a/parquet-variant/src/path.rs +++ b/parquet-variant/src/path.rs @@ -346,5 +346,14 @@ mod tests { err.to_string(), "Parser error: Invalid token in bracket request: `abc`. Expected a quoted string or a number(e.g., `['field']` or `[123]`)" ); + + // Out-of-range integer indexes are invalid path tokens. + let too_large_index = (usize::MAX as u128) + 1; + let err = VariantPath::try_from(format!("[{too_large_index}]").as_str()).unwrap_err(); + assert!( + err.to_string() + .contains("Parser error: Invalid token in bracket request"), + "{err}" + ); } }