diff --git a/Cargo.lock b/Cargo.lock index a5df0c5d207..b46a5aa28ec 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -10167,8 +10167,10 @@ version = "0.1.0" dependencies = [ "arrow-array", "arrow-schema", + "prost 0.14.4", "vortex-array", "vortex-error", + "vortex-proto", "vortex-session", ] @@ -10282,6 +10284,7 @@ dependencies = [ "vortex-error", "vortex-file", "vortex-io", + "vortex-json", "vortex-layout", "vortex-mask", "vortex-proto", diff --git a/encodings/parquet-variant/Cargo.toml b/encodings/parquet-variant/Cargo.toml index 6b3900abb27..d0c0a0d418e 100644 --- a/encodings/parquet-variant/Cargo.toml +++ b/encodings/parquet-variant/Cargo.toml @@ -27,8 +27,9 @@ prost = { workspace = true } vortex-array = { workspace = true } vortex-buffer = { workspace = true } vortex-error = { workspace = true } +vortex-json = { workspace = true } vortex-mask = { workspace = true } -vortex-proto = { workspace = true } +vortex-proto = { workspace = true, features = ["expr"] } vortex-session = { workspace = true } [dev-dependencies] diff --git a/encodings/parquet-variant/src/array.rs b/encodings/parquet-variant/src/array.rs index 7be0265fc96..7e55357c78f 100644 --- a/encodings/parquet-variant/src/array.rs +++ b/encodings/parquet-variant/src/array.rs @@ -19,9 +19,11 @@ use vortex_array::arrays::List; use vortex_array::arrays::ListArray; use vortex_array::arrays::Struct; use vortex_array::arrays::StructArray; +use vortex_array::arrays::Variant; use vortex_array::arrays::VariantArray; use vortex_array::arrays::list::ListArrayExt; use vortex_array::arrays::struct_::StructArrayExt; +use vortex_array::arrays::variant::VariantArrayExt; #[expect( deprecated, reason = "TODO(aduffy): figure out what to do with Parquet Variant" @@ -42,6 +44,7 @@ use vortex_array::vtable::validity_to_child; use vortex_buffer::BitBuffer; use vortex_error::VortexExpect; use vortex_error::VortexResult; +use vortex_error::vortex_err; use crate::ParquetVariant; use crate::ParquetVariantArray; @@ -307,6 +310,111 @@ fn inferred_shredded_field_validity( Ok(Validity::from_mask(validity, Nullability::Nullable)) } +/// Reconstructs a Parquet `typed_value` tree from the storage-agnostic canonical shredded tree. +/// +/// This is the inverse of [`logical_shredded_from_parquet_typed_value`]: canonicalization strips +/// the Parquet `value`/`typed_value` wrapper shells out of the shredded tree (representing +/// partially shredded fields as nested [`VariantArray`]s), and this re-adds them. The result is a +/// valid Parquet `typed_value` that can be reattached to a [`ParquetVariant`] and serialized to +/// Arrow, where [`parquet_variant_compute::unshred_variant`] merges it back with the residual +/// `value`. +/// +/// `forward` then `inverse` is structure-preserving for the shapes canonicalization produces; +/// fields the forward transform omitted (Parquet wrapper shells with no `typed_value`) are served +/// from the residual `value` and are intentionally not reconstructed here. +pub(crate) fn parquet_typed_value_from_logical_shredded( + shredded: ArrayRef, + ctx: &mut ExecutionCtx, +) -> VortexResult { + if let Some(list_array) = shredded.as_opt::() { + let elements = parquet_shredded_field_from_logical(list_array.elements().clone(), ctx)?; + return Ok(ListArray::try_new( + elements, + list_array.offsets().clone(), + list_array.list_validity(), + )? + .into_array()); + } + + let Some(struct_array) = shredded.as_opt::() else { + // A bare typed leaf (a fully shredded scalar) is already a valid Parquet `typed_value`. + return Ok(shredded); + }; + + let mut names = Vec::with_capacity(struct_array.names().len()); + let mut fields = Vec::with_capacity(struct_array.names().len()); + for (name, field) in struct_array + .names() + .iter() + .zip(struct_array.iter_unmasked_fields()) + { + names.push(FieldName::from(name.as_ref())); + fields.push(parquet_shredded_field_from_logical(field.clone(), ctx)?); + } + + Ok(StructArray::try_new( + FieldNames::from_iter(names), + fields, + struct_array.len(), + struct_array.validity()?, + )? + .into_array()) +} + +/// Reconstructs one Parquet shredded field shell (`{value?, typed_value}`) from its canonical +/// representation, the inverse of [`logical_shredded_from_parquet_field`]. +fn parquet_shredded_field_from_logical( + logical_field: ArrayRef, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let len = logical_field.len(); + + // Partially shredded fields canonicalize to a nested Variant whose core storage holds the + // residual `value` and whose own shredded tree holds the typed children. + if let Some(variant) = logical_field.as_opt::() { + let core = variant + .core_storage() + .as_opt::() + .ok_or_else(|| { + vortex_err!( + "cannot rebuild Parquet shredded field: nested Variant lacks Parquet Variant core storage" + ) + })?; + let value = core.value_array().cloned().ok_or_else(|| { + vortex_err!("cannot rebuild Parquet shredded field: partially shredded Variant has no residual value") + })?; + let typed_value = variant + .shredded() + .cloned() + .map(|shredded| parquet_typed_value_from_logical_shredded(shredded, ctx)) + .transpose()?; + + let mut names = vec![FieldName::from("value")]; + let mut fields = vec![value]; + if let Some(typed_value) = typed_value { + names.push(FieldName::from("typed_value")); + fields.push(typed_value); + } + return Ok(StructArray::try_new( + FieldNames::from_iter(names), + fields, + len, + Validity::NonNullable, + )? + .into_array()); + } + + // Fully shredded field: rebuild its typed subtree and wrap it in a `typed_value`-only shell. + let typed_value = parquet_typed_value_from_logical_shredded(logical_field, ctx)?; + Ok(StructArray::try_new( + FieldNames::from_iter([FieldName::from("typed_value")]), + vec![typed_value], + len, + Validity::NonNullable, + )? + .into_array()) +} + /// Accessors and Arrow conversion for Parquet Variant storage arrays. pub trait ParquetVariantArrayExt: TypedArrayRef { /// Returns the non-nullable Parquet Variant metadata child. @@ -390,6 +498,7 @@ mod tests { use arrow_array::Array as _; use arrow_array::ArrayRef as ArrowArrayRef; use arrow_array::Int32Array; + use arrow_array::StringArray; use arrow_array::StructArray; use arrow_array::builder::BinaryViewBuilder; use arrow_buffer::NullBuffer; @@ -397,12 +506,17 @@ mod tests { use arrow_schema::Field; use arrow_schema::Fields; use parquet_variant::Variant as PqVariant; + use parquet_variant_compute::ShreddedSchemaBuilder; use parquet_variant_compute::VariantArray as ArrowVariantArray; use parquet_variant_compute::VariantArrayBuilder; + use parquet_variant_compute::json_to_variant; + use parquet_variant_compute::shred_variant; + use vortex_array::Canonical; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::VarBinViewArray; + use vortex_array::arrays::variant::VariantArrayExt; use vortex_array::assert_arrays_eq; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; @@ -414,6 +528,7 @@ mod tests { use crate::ParquetVariant; use crate::array::ParquetVariantArrayExt; + use crate::array::parquet_typed_value_from_logical_shredded; static SESSION: LazyLock = LazyLock::new(|| { let session = vortex_array::array_session(); @@ -690,4 +805,64 @@ mod tests { assert_arrow_variant_storage_roundtrip(struct_array) } + + /// `parquet_typed_value_from_logical_shredded` must invert the wrapper-stripping that + /// canonicalization performs: an object-shredded Parquet variant, once canonicalized and then + /// rebuilt, must produce the same per-row values as the original. + #[test] + fn parquet_typed_value_inverse_roundtrips_object_shredding() -> VortexResult<()> { + // Shred `$.a` as Int32 over conforming, non-conforming, and missing-field rows. + let json: ArrowArrayRef = Arc::new(StringArray::from(vec![ + Some(r#"{"a":1,"b":"x"}"#), + Some(r#"{"a":"not-a-number","b":"y"}"#), + Some(r#"{"b":"z"}"#), + ])); + let shredding = ShreddedSchemaBuilder::new() + .with_path("a", &DataType::Int32)? + .build(); + let shredded = shred_variant(&json_to_variant(&json)?, &shredding)?; + let original = ParquetVariant::from_arrow_variant(&shredded)?; + assert!( + original + .as_opt::() + .ok_or_else(|| vortex_err!("expected parquet variant"))? + .typed_value_array() + .is_some(), + "fixture must be shredded" + ); + + // Canonicalize: the forward transform lifts `typed_value` into a logical shredded child. + let mut ctx = SESSION.create_execution_ctx(); + let Canonical::Variant(canonical) = original.clone().execute::(&mut ctx)? else { + return Err(vortex_err!("expected canonical variant")); + }; + let core = canonical + .core_storage() + .as_opt::() + .ok_or_else(|| vortex_err!("expected parquet variant core storage"))?; + let logical = canonical + .shredded() + .ok_or_else(|| vortex_err!("expected canonical shredded child"))? + .clone(); + + // Inverse transform: rebuild a Parquet `typed_value` and reattach it. + let rebuilt_typed_value = parquet_typed_value_from_logical_shredded(logical, &mut ctx)?; + let rebuilt = ParquetVariant::try_new( + ParquetVariantArrayExt::validity(&core), + core.metadata_array().clone(), + core.value_array().cloned(), + Some(rebuilt_typed_value), + )? + .into_array(); + + assert_eq!(rebuilt.len(), original.len()); + for idx in 0..original.len() { + assert_eq!( + rebuilt.execute_scalar(idx, &mut ctx)?, + original.execute_scalar(idx, &mut ctx)?, + "row {idx}" + ); + } + Ok(()) + } } diff --git a/encodings/parquet-variant/src/arrow.rs b/encodings/parquet-variant/src/arrow.rs index 326e7f9d7ee..da4e648a138 100644 --- a/encodings/parquet-variant/src/arrow.rs +++ b/encodings/parquet-variant/src/arrow.rs @@ -36,6 +36,7 @@ use vortex_session::registry::Id; use crate::ParquetVariant; use crate::ParquetVariantArrayExt; +use crate::array::parquet_typed_value_from_logical_shredded; /// Arrow canonical extension name for Parquet Variant storage. const PARQUET_VARIANT_ARROW_EXTENSION_NAME: &str = "arrow.parquet.variant"; @@ -58,7 +59,7 @@ fn parquet_variant_storage_request(fields: &Fields) -> Option<(bool, bool)> { (has_metadata && (has_value || has_typed_value)).then_some((has_value, has_typed_value)) } -fn export_storage_to_target( +pub(crate) fn export_storage_to_target( parquet_array: &T, target_fields: &Fields, ctx: &mut ExecutionCtx, @@ -99,7 +100,7 @@ fn export_storage_to_target( )?)) } -fn export_unshredded_storage_to_target( +pub(crate) fn export_unshredded_storage_to_target( parquet_array: &T, target_fields: &Fields, ctx: &mut ExecutionCtx, @@ -115,7 +116,10 @@ fn export_unshredded_storage_to_target( export_storage_to_target(&unshredded_parquet, target_fields, ctx) } -fn parquet_variant_for_export(array: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { +pub(crate) fn parquet_variant_for_export( + array: ArrayRef, + ctx: &mut ExecutionCtx, +) -> VortexResult { let executed = array.execute_until::(ctx)?; if executed.is::() { return Ok(executed); @@ -135,11 +139,16 @@ fn parquet_variant_for_export(array: ArrayRef, ctx: &mut ExecutionCtx) -> Vortex return Ok(core_storage); }; + // The canonical shredded child has had its Parquet `value`/`typed_value` wrapper shells + // stripped; rebuild them so the reattached `typed_value` is valid Parquet storage that + // `to_arrow` and `unshred_variant` can consume. + let typed_value = parquet_typed_value_from_logical_shredded(shredded.clone(), ctx)?; + ParquetVariant::try_new( ParquetVariantArrayExt::validity(&parquet_core), parquet_core.metadata_array().clone(), parquet_core.value_array().cloned(), - Some(shredded.clone()), + Some(typed_value), ) .map(IntoArray::into_array) } @@ -250,7 +259,7 @@ impl ArrowImportVTable for ParquetVariant { field: &Field, dtype: &DType, ) -> VortexResult { - if !matches!(dtype, DType::Variant(_)) + if !dtype.is_variant() || field .metadata() .get(EXTENSION_TYPE_NAME_KEY) diff --git a/encodings/parquet-variant/src/json_to_variant_tests.rs b/encodings/parquet-variant/src/json_to_variant_tests.rs new file mode 100644 index 00000000000..cd3b2f035ee --- /dev/null +++ b/encodings/parquet-variant/src/json_to_variant_tests.rs @@ -0,0 +1,301 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Execution tests for the `vortex.json_to_variant` scalar function. +//! +//! The function definition lives in `vortex-json`; the JSON->Variant construction is performed +//! by the execute-parent kernel registered here, so these end-to-end tests live in +//! `vortex-parquet-variant` where that kernel is registered. + +use std::sync::LazyLock; + +use vortex_array::ArrayRef; +use vortex_array::Canonical; +use vortex_array::EmptyMetadata; +use vortex_array::IntoArray; +use vortex_array::VortexSessionExecute; +use vortex_array::arrays::DictArray; +use vortex_array::arrays::ExtensionArray; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::arrays::StructArray; +use vortex_array::arrays::VarBinViewArray; +use vortex_array::arrays::struct_::StructArrayExt; +use vortex_array::arrays::variant::VariantArrayExt; +use vortex_array::assert_arrays_eq; +use vortex_array::assert_nth_scalar_is_null; +use vortex_array::dtype::DType; +use vortex_array::dtype::Nullability; +use vortex_array::dtype::PType; +use vortex_array::expr::root; +use vortex_array::expr::variant_get; +use vortex_array::scalar_fn::fns::variant_get::VariantPath; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_err; +use vortex_json::Json; +use vortex_session::VortexSession; + +use crate::ParquetVariant; +use crate::ParquetVariantArrayExt; +use crate::ShreddingSpec; +use crate::json_to_variant; + +static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + crate::initialize(&session); + session +}); + +fn i64_dtype() -> DType { + DType::Primitive(PType::I64, Nullability::Nullable) +} + +fn shred_field_as_i64(field: &str) -> VortexResult { + ShreddingSpec::try_new([(VariantPath::field(field), i64_dtype())]) +} + +fn json_input(storage: ArrayRef) -> VortexResult { + Ok(ExtensionArray::try_new_from_vtable(Json, EmptyMetadata, storage)?.into_array()) +} + +fn execute_json_to_variant(input: ArrayRef, shredding: ShreddingSpec) -> VortexResult { + let expr = json_to_variant(root(), shredding); + input + .apply(&expr)? + .execute::(&mut SESSION.create_execution_ctx()) +} + +fn assert_variant_i64_rows(array: &ArrayRef, expected: &[Option]) -> VortexResult<()> { + assert_eq!(array.len(), expected.len()); + let mut ctx = SESSION.create_execution_ctx(); + for (idx, expected) in expected.iter().enumerate() { + let scalar = array.execute_scalar(idx, &mut ctx)?; + let variant = scalar.as_variant(); + match expected { + Some(expected) => { + let value = variant + .value() + .ok_or_else(|| vortex_err!("expected non-null variant at row {idx}"))? + .cast(&i64_dtype())?; + assert_eq!(value.as_primitive().typed_value::(), Some(*expected)); + } + None => assert!(scalar.is_null(), "expected null row {idx}"), + } + } + Ok(()) +} + +#[test] +fn rejects_bare_utf8_input() { + let input = VarBinViewArray::from_iter_str(["1", "2"]).into_array(); + + let err = execute_json_to_variant(input, ShreddingSpec::empty()).unwrap_err(); + assert!( + err.to_string().contains("Json extension"), + "unexpected error: {err}" + ); +} + +#[test] +fn converts_json_extension_rows() -> VortexResult<()> { + let input = json_input( + VarBinViewArray::from_iter_str([r#"{"a": 1}"#, "2", r#"{"a": 3}"#]).into_array(), + )?; + + let result = execute_json_to_variant(input, ShreddingSpec::empty())?; + + assert_eq!(result.dtype(), &DType::Variant(Nullability::NonNullable)); + let mut ctx = SESSION.create_execution_ctx(); + let row0 = result.execute_scalar(0, &mut ctx)?; + let object = row0 + .as_variant() + .value() + .ok_or_else(|| vortex_err!("expected non-null variant"))?; + let field = object + .as_struct() + .field("a") + .ok_or_else(|| vortex_err!("expected field a"))? + .as_variant() + .value() + .ok_or_else(|| vortex_err!("expected non-null field a"))? + .cast(&i64_dtype())?; + assert_eq!(field.as_primitive().typed_value::(), Some(1)); + + let row1 = result.execute_scalar(1, &mut ctx)?; + let value = row1 + .as_variant() + .value() + .ok_or_else(|| vortex_err!("expected non-null variant"))? + .cast(&i64_dtype())?; + assert_eq!(value.as_primitive().typed_value::(), Some(2)); + Ok(()) +} + +#[test] +fn converts_json_extension_input() -> VortexResult<()> { + let storage = VarBinViewArray::from_iter_str(["1", "2"]).into_array(); + let input = json_input(storage)?; + + let result = execute_json_to_variant(input, ShreddingSpec::empty())?; + + assert_eq!(result.dtype(), &DType::Variant(Nullability::NonNullable)); + assert_variant_i64_rows(&result, &[Some(1), Some(2)]) +} + +#[test] +fn dict_encoded_input_converts_each_row() -> VortexResult<()> { + // A dictionary-encoded JSON column exercises the dict-pushdown / canonicalization path: + // `json_to_variant` is not null-sensitive, so it pushes into the dict values (canonical + // JSON extension values) where the kernel fires; either way every row must convert correctly. + let values = json_input(VarBinViewArray::from_iter_str(["1", "2"]).into_array())?; + let codes = PrimitiveArray::from_iter([0u8, 1, 0, 1, 0]).into_array(); + let input = DictArray::try_new(codes, values)?.into_array(); + + let result = execute_json_to_variant(input, ShreddingSpec::empty())?; + + assert_eq!(result.dtype(), &DType::Variant(Nullability::NonNullable)); + assert_variant_i64_rows(&result, &[Some(1), Some(2), Some(1), Some(2), Some(1)]) +} + +#[test] +fn null_rows_stay_null_and_json_null_becomes_variant_null() -> VortexResult<()> { + let input = json_input( + VarBinViewArray::from_iter_nullable_str([Some("1"), None, Some("null")]).into_array(), + )?; + + let result = execute_json_to_variant(input, ShreddingSpec::empty())?; + + assert_eq!(result.dtype(), &DType::Variant(Nullability::Nullable)); + let mut ctx = SESSION.create_execution_ctx(); + assert!(!result.execute_scalar(0, &mut ctx)?.is_null()); + assert_nth_scalar_is_null!(result, 1, &mut ctx); + let row2 = result.execute_scalar(2, &mut ctx)?; + assert!(!row2.is_null(), "JSON null must not be a row null"); + assert_eq!(row2.as_variant().is_variant_null(), Some(true)); + Ok(()) +} + +#[test] +fn invalid_json_errors() -> VortexResult<()> { + let input = + json_input(VarBinViewArray::from_iter_str([r#"{"a": 1}"#, r#"{"a":"#]).into_array())?; + + let err = execute_json_to_variant(input, ShreddingSpec::empty()).unwrap_err(); + assert!(!err.to_string().is_empty()); + Ok(()) +} + +#[test] +fn shredding_produces_typed_value_child() -> VortexResult<()> { + let input = json_input( + VarBinViewArray::from_iter_str([ + r#"{"a": 1, "b": "x"}"#, + r#"{"a": 2, "b": "y"}"#, + r#"{"a": "not-a-number", "b": "z"}"#, + r#"{"b": "missing-a"}"#, + ]) + .into_array(), + )?; + + let result = execute_json_to_variant(input, shred_field_as_i64("a")?)?; + + assert!( + result.as_::().typed_value_array().is_some(), + "expected shredded typed_value child" + ); + + // The canonical form must expose field `a` through the shredded tree. + let mut ctx = SESSION.create_execution_ctx(); + let Canonical::Variant(canonical) = result.clone().execute::(&mut ctx)? else { + vortex_bail!("expected canonical variant array"); + }; + let shredded = canonical + .shredded() + .ok_or_else(|| vortex_err!("expected canonical shredded child"))? + .clone() + .execute::(&mut ctx)?; + assert!(shredded.unmasked_field_by_name_opt("a").is_some()); + + // Typed extraction must serve shredded rows and fall back for mismatched rows. + let typed = result + .clone() + .apply(&variant_get( + root(), + VariantPath::field("a"), + Some(i64_dtype()), + ))? + .execute::(&mut ctx)?; + assert_arrays_eq!( + typed, + PrimitiveArray::from_option_iter([Some(1i64), Some(2), None, None]), + &mut ctx + ); + + // Mismatched rows keep their original value through the variant fallback. + let untyped = result + .apply(&variant_get( + root(), + VariantPath::field("a"), + Some(DType::Utf8(Nullability::Nullable)), + ))? + .execute::(&mut ctx)?; + let row2 = untyped.execute_scalar(2, &mut ctx)?; + assert_eq!( + row2.as_utf8().value().map(|value| value.to_string()), + Some("not-a-number".to_string()) + ); + Ok(()) +} + +#[test] +fn shredding_preserves_null_rows() -> VortexResult<()> { + let input = json_input( + VarBinViewArray::from_iter_nullable_str([Some(r#"{"a": 1}"#), None, Some(r#"{"a": 3}"#)]) + .into_array(), + )?; + + let result = execute_json_to_variant(input, shred_field_as_i64("a")?)?; + + assert_eq!(result.dtype(), &DType::Variant(Nullability::Nullable)); + let mut ctx = SESSION.create_execution_ctx(); + assert_nth_scalar_is_null!(result, 1, &mut ctx); + let typed = result + .apply(&variant_get( + root(), + VariantPath::field("a"), + Some(i64_dtype()), + ))? + .execute::(&mut ctx)?; + assert_arrays_eq!( + typed, + PrimitiveArray::from_option_iter([Some(1i64), None, Some(3)]), + &mut ctx + ); + Ok(()) +} + +#[test] +fn shredding_root_path_shreds_top_level_values() -> VortexResult<()> { + let input = + json_input(VarBinViewArray::from_iter_str(["1", "2", r#""not-a-number""#]).into_array())?; + let spec = ShreddingSpec::try_new([(VariantPath::root(), i64_dtype())])?; + + let result = execute_json_to_variant(input, spec)?; + + assert!( + result.as_::().typed_value_array().is_some(), + "expected shredded typed_value child" + ); + assert_variant_i64_rows(&result.slice(0..2)?, &[Some(1), Some(2)])?; + let mut ctx = SESSION.create_execution_ctx(); + let row2 = result.execute_scalar(2, &mut ctx)?; + let value = row2 + .as_variant() + .value() + .ok_or_else(|| vortex_err!("expected non-null variant"))?; + assert_eq!( + value.as_utf8().value().map(|value| value.to_string()), + Some("not-a-number".to_string()) + ); + Ok(()) +} diff --git a/encodings/parquet-variant/src/kernel.rs b/encodings/parquet-variant/src/kernel.rs index 684ce0217f4..88a4b8931ab 100644 --- a/encodings/parquet-variant/src/kernel.rs +++ b/encodings/parquet-variant/src/kernel.rs @@ -11,7 +11,9 @@ use arrow_schema::FieldRef; use parquet_variant::VariantPath as PqVariantPath; use parquet_variant::VariantPathElement as PqVariantPathElement; use parquet_variant_compute::GetOptions; +use parquet_variant_compute::ShreddedSchemaBuilder; use parquet_variant_compute::VariantArray as ArrowVariantArray; +use parquet_variant_compute::shred_variant; use parquet_variant_compute::variant_get as arrow_variant_get; use vortex_array::ArrayRef; use vortex_array::ArrayVTable; @@ -19,16 +21,19 @@ use vortex_array::ArrayView; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; use vortex_array::arrays::Dict; +use vortex_array::arrays::Extension; use vortex_array::arrays::Filter; use vortex_array::arrays::Slice; use vortex_array::arrays::dict::TakeExecute; use vortex_array::arrays::dict::TakeExecuteAdaptor; +use vortex_array::arrays::extension::ExtensionArrayExt; use vortex_array::arrays::filter::FilterExecuteAdaptor; use vortex_array::arrays::filter::FilterKernel; use vortex_array::arrays::scalar_fn::ExactScalarFn; use vortex_array::arrays::scalar_fn::ScalarFnArrayView; use vortex_array::arrays::slice::SliceExecuteAdaptor; use vortex_array::arrays::slice::SliceKernel; +use vortex_array::arrow::ArrowSessionExt; use vortex_array::arrow::FromArrowArray; use vortex_array::dtype::DType; use vortex_array::kernel::ExecuteParentKernel; @@ -40,6 +45,8 @@ use vortex_array::scalar_fn::fns::variant_get::VariantPathElement; use vortex_error::VortexResult; use vortex_error::vortex_ensure_eq; use vortex_error::vortex_err; +use vortex_json::Json; +use vortex_json::JsonToVariant; use vortex_mask::Mask; use vortex_session::VortexSession; @@ -64,6 +71,11 @@ pub(crate) fn initialize(session: &VortexSession) { TakeExecuteAdaptor(ParquetVariant), ); kernels.register_execute_parent_kernel(VariantGet.id(), ParquetVariant, VariantGetKernel); + kernels.register_execute_parent_kernel( + JsonToVariant.id(), + Extension, + JsonExtensionToVariantKernel, + ); } #[derive(Default, Debug)] @@ -106,7 +118,67 @@ impl ExecuteParentKernel for VariantGetKernel { } } -fn to_parquet_variant_path(path: &VariantPath) -> VortexResult> { +/// Performs the [`JsonToVariant`] conversion (and optional shredding) over JSON string storage. +/// +/// `JsonToVariant`'s definition lives in `vortex-json`; the registered `execute_parent` kernels +/// delegate here to do the actual JSON parsing and optional shredding using +/// `parquet_variant_compute`, producing a [`ParquetVariant`] array. `strings` is the JSON +/// extension's storage array; it is executed to Arrow and parsed. Nullability of the result follows +/// `parent.dtype()`, which equals the input's nullability. +fn json_strings_to_variant( + strings: ArrayRef, + parent: ScalarFnArrayView<'_, JsonToVariant>, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let nullable = parent.dtype().is_nullable(); + let session = ctx.session().clone(); + let arrow_strings = session.arrow().execute_arrow(strings, None, ctx)?; + // Any row that fails to parse as JSON fails the whole conversion. + let arrow_variant = parquet_variant_compute::json_to_variant(&arrow_strings)?; + + let arrow_variant = if parent.options.shredding().is_empty() { + arrow_variant + } else { + let mut builder = ShreddedSchemaBuilder::new(); + for (path, dtype) in parent.options.shredding().fields() { + let field: FieldRef = Arc::new(session.arrow().to_arrow_field("shredded", dtype)?); + builder = builder.with_path(to_parquet_variant_path(path)?, field)?; + } + shred_variant(&arrow_variant, &builder.build())? + }; + + if nullable { + ParquetVariant::from_arrow_variant_nullable(&arrow_variant) + } else { + ParquetVariant::from_arrow_variant(&arrow_variant) + } +} + +/// Builds Parquet Variant arrays for [`JsonToVariant`] over a [`Json`] extension input. +/// +/// This kernel unwraps the extension's string storage and runs the JSON conversion. It is keyed on +/// the shared extension encoding, so it declines any non-`Json` extension. +#[derive(Default, Debug)] +struct JsonExtensionToVariantKernel; + +impl ExecuteParentKernel for JsonExtensionToVariantKernel { + type Parent = ExactScalarFn; + + fn execute_parent( + &self, + array: ArrayView<'_, Extension>, + parent: ScalarFnArrayView<'_, JsonToVariant>, + child_idx: usize, + ctx: &mut ExecutionCtx, + ) -> VortexResult> { + if child_idx != 0 || !array.ext_dtype().is::() { + return Ok(None); + } + json_strings_to_variant(array.storage_array().clone(), parent, ctx).map(Some) + } +} + +pub(crate) fn to_parquet_variant_path(path: &VariantPath) -> VortexResult> { path.elements() .iter() .map(|element| match element { diff --git a/encodings/parquet-variant/src/lib.rs b/encodings/parquet-variant/src/lib.rs index 0076af48d22..05e93c5d08c 100644 --- a/encodings/parquet-variant/src/lib.rs +++ b/encodings/parquet-variant/src/lib.rs @@ -26,6 +26,8 @@ mod array; mod arrow; +#[cfg(test)] +mod json_to_variant_tests; mod kernel; mod operations; mod validity; @@ -36,12 +38,21 @@ use std::sync::Arc; pub use array::ParquetVariantArrayExt; use vortex_array::arrow::ArrowSessionExt; use vortex_array::session::ArraySessionExt; +pub use vortex_json::JsonToVariant; +pub use vortex_json::JsonToVariantOptions; +pub use vortex_json::ShreddingSpec; +pub use vortex_json::json_to_variant; use vortex_session::VortexSession; pub use vtable::ParquetVariant; pub use vtable::ParquetVariantArray; -/// Register Parquet Variant array and Arrow extension support with a session. +/// Register Parquet Variant array, Arrow extension, and scalar function support with a +/// session. +/// +/// This also initializes [`vortex_json`], registering the `Json` extension dtype and the +/// `json_to_variant` scalar function whose execution this crate provides. pub fn initialize(session: &VortexSession) { + vortex_json::initialize(session); session.arrays().register(ParquetVariant); kernel::initialize(session); session.arrow().register_exporter(Arc::new(ParquetVariant)); diff --git a/vortex-array/src/scalar_fn/fns/variant_get/mod.rs b/vortex-array/src/scalar_fn/fns/variant_get/mod.rs index 179264ba033..3e50f5f5c30 100644 --- a/vortex-array/src/scalar_fn/fns/variant_get/mod.rs +++ b/vortex-array/src/scalar_fn/fns/variant_get/mod.rs @@ -55,16 +55,7 @@ impl ScalarFnVTable for VariantGet { .path() .elements() .iter() - .map(|element| match element { - VariantPathElement::Field(name) => pb::VariantPathElement { - element: Some(variant_path_element::Element::Field( - name.as_ref().to_string(), - )), - }, - VariantPathElement::Index(index) => pb::VariantPathElement { - element: Some(variant_path_element::Element::Index(*index)), - }, - }) + .map(VariantPathElement::to_proto) .collect(); let dtype = options.dtype().map(TryInto::try_into).transpose()?; @@ -355,15 +346,30 @@ impl VariantPathElement { Self::Index(index) } - fn from_proto(value: pb::VariantPathElement) -> VortexResult { + /// Decodes a path element from its protobuf representation. + pub fn from_proto(value: pb::VariantPathElement) -> VortexResult { match value .element - .ok_or_else(|| vortex_err!("VariantGet path element missing value"))? + .ok_or_else(|| vortex_err!("Variant path element missing value"))? { variant_path_element::Element::Field(field) => Ok(Self::field(field)), variant_path_element::Element::Index(index) => Ok(Self::index(index)), } } + + /// Encodes this path element into its protobuf representation. + pub fn to_proto(&self) -> pb::VariantPathElement { + match self { + VariantPathElement::Field(name) => pb::VariantPathElement { + element: Some(variant_path_element::Element::Field( + name.as_ref().to_string(), + )), + }, + VariantPathElement::Index(index) => pb::VariantPathElement { + element: Some(variant_path_element::Element::Index(*index)), + }, + } + } } impl From for VariantPathElement { diff --git a/vortex-json/Cargo.toml b/vortex-json/Cargo.toml index 4afd224f6aa..71eb015b5ee 100644 --- a/vortex-json/Cargo.toml +++ b/vortex-json/Cargo.toml @@ -19,6 +19,13 @@ workspace = true [dependencies] arrow-array = { workspace = true } arrow-schema = { workspace = true, features = ["canonical_extension_types"] } +prost = { workspace = true } vortex-array = { workspace = true, default-features = false } vortex-error = { workspace = true, default-features = false } +vortex-proto = { workspace = true, features = ["expr"] } vortex-session = { workspace = true } + +[dev-dependencies] +vortex-array = { workspace = true, default-features = false, features = [ + "_test-harness", +] } diff --git a/vortex-json/src/json_to_variant.rs b/vortex-json/src/json_to_variant.rs new file mode 100644 index 00000000000..d0febf4e646 --- /dev/null +++ b/vortex-json/src/json_to_variant.rs @@ -0,0 +1,396 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! The `vortex.json_to_variant` scalar function. + +use std::fmt; +use std::fmt::Display; +use std::fmt::Formatter; + +use prost::Message; +use vortex_array::ArrayRef; +use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; +use vortex_array::arrays::Extension; +use vortex_array::arrays::ExtensionArray; +use vortex_array::arrays::scalar_fn::ScalarFnFactoryExt; +use vortex_array::dtype::DType; +use vortex_array::expr::Expression; +use vortex_array::scalar_fn::Arity; +use vortex_array::scalar_fn::ChildName; +use vortex_array::scalar_fn::ExecutionArgs; +use vortex_array::scalar_fn::ScalarFnId; +use vortex_array::scalar_fn::ScalarFnVTable; +use vortex_array::scalar_fn::ScalarFnVTableExt; +use vortex_array::scalar_fn::fns::variant_get::VariantPath; +use vortex_array::scalar_fn::fns::variant_get::VariantPathElement; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_ensure; +use vortex_error::vortex_err; +use vortex_proto::expr as pb; +use vortex_session::VortexSession; +use vortex_session::registry::CachedId; + +use crate::Json; + +/// Parses JSON strings into Variant values, optionally shredding fields. +/// +/// Accepts [`Json`] extension inputs and returns `Variant` values with the input's nullability. +/// Null rows stay null, the JSON literal `null` becomes a variant-null value, and any row that +/// fails to parse as JSON fails the whole expression. +/// +/// A non-empty [`ShreddingSpec`] additionally shreds the selected paths into a typed shredded +/// child, following the [Parquet Variant shredding] rules: rows whose value does not match the +/// requested type stay readable through the residual variant value. +/// +/// # Execution +/// +/// Building a Variant requires a concrete Variant encoding, so this function does not perform the +/// conversion itself. The Variant encoding registered with the session supplies it as an +/// `execute_parent` kernel keyed on the extension encoding for a [`Json`] input. The fallback +/// [`execute`](ScalarFnVTable::execute) here only canonicalizes the input to that encoding and +/// re-dispatches so that kernel runs; it errors if no Variant encoding is registered with the +/// session. +/// +/// # Normalization +/// +/// `json_to_variant` is a lossy, normalizing conversion: the parsed Variant does not round-trip +/// back to the exact source JSON text. +/// - JSON whitespace is not preserved. +/// - Object keys are stored in Variant metadata in sorted order, not source order. +/// - Number representations are normalized (e.g. `1.0` and `1` may parse to the same value; +/// exponent forms and very large numbers are re-encoded). +/// - Duplicate object keys are collapsed to a single entry. +/// - Unicode escape sequences are normalized (e.g. `A` becomes `A`). +/// +/// [Parquet Variant shredding]: https://github.com/apache/parquet-format/blob/master/VariantShredding.md +#[derive(Clone)] +pub struct JsonToVariant; + +impl ScalarFnVTable for JsonToVariant { + type Options = JsonToVariantOptions; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("vortex.json_to_variant"); + *ID + } + + fn serialize(&self, options: &Self::Options) -> VortexResult>> { + let shredding = options + .shredding() + .fields() + .iter() + .map(|(path, dtype)| { + Ok(pb::ShreddingSpecField { + path: path + .elements() + .iter() + .map(VariantPathElement::to_proto) + .collect(), + dtype: Some(dtype.try_into()?), + }) + }) + .collect::>()?; + + Ok(Some(pb::JsonToVariantOpts { shredding }.encode_to_vec())) + } + + fn deserialize(&self, metadata: &[u8], session: &VortexSession) -> VortexResult { + let opts = pb::JsonToVariantOpts::decode(metadata)?; + let fields = opts + .shredding + .into_iter() + .map(|field| { + let path = field + .path + .into_iter() + .map(VariantPathElement::from_proto) + .collect::>()?; + let dtype = field + .dtype + .as_ref() + .ok_or_else(|| vortex_err!("ShreddingSpecField missing dtype")) + .and_then(|dtype| DType::from_proto(dtype, session))?; + Ok((path, dtype)) + }) + .collect::>>()?; + + Ok(JsonToVariantOptions::new(ShreddingSpec::try_new(fields)?)) + } + + fn arity(&self, _options: &Self::Options) -> Arity { + Arity::Exact(1) + } + + fn child_name(&self, _options: &Self::Options, child_idx: usize) -> ChildName { + match child_idx { + 0 => ChildName::from("input"), + _ => unreachable!("Invalid child index {child_idx} for JsonToVariant expression"), + } + } + + fn return_dtype(&self, _options: &Self::Options, arg_dtypes: &[DType]) -> VortexResult { + let input_dtype = &arg_dtypes[0]; + vortex_ensure!( + input_dtype + .as_extension_opt() + .is_some_and(|ext_dtype| ext_dtype.is::()), + "JsonToVariant input must be a Json extension, found {input_dtype}" + ); + + Ok(DType::Variant(input_dtype.nullability())) + } + + fn execute( + &self, + options: &Self::Options, + args: &dyn ExecutionArgs, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + let input = args.get(0)?; + + // This function does not build Variants itself: the Variant encoding registered with the + // session supplies the conversion as an `execute_parent` kernel keyed on the extension + // encoding for a `Json` input. Reaching this fallback means no such kernel fired, so + // canonicalize the input to that encoding and re-dispatch. If the input is already + // canonical here, no kernel is registered for it; bail with a clear error rather than + // looping. + let no_kernel = || { + vortex_err!( + "json_to_variant requires a registered Variant encoding to build Variant values \ + from JSON, but none is registered with this session" + ) + }; + + let canonical = if input + .dtype() + .as_extension_opt() + .is_some_and(|ext_dtype| ext_dtype.is::()) + { + if input.is::() { + return Err(no_kernel()); + } + input.execute::(ctx)?.into_array() + } else { + vortex_bail!( + "JsonToVariant input must be a Json extension, found {}", + input.dtype() + ); + }; + + self.try_new_array(canonical.len(), options.clone(), [canonical])? + .execute::(ctx) + } + + fn is_null_sensitive(&self, _options: &Self::Options) -> bool { + // `json_to_variant` maps null rows to null rows and the JSON literal `null` to a + // variant-null value independently of which other rows are null, so it commutes with + // validity masking. Marking it not null-sensitive also lets it push through dictionaries + // into their JSON extension values, where the kernel fires directly. + false + } +} + +/// Creates a [`JsonToVariant`] expression that parses `child`'s JSON strings into Variant +/// values, shredding the paths selected by `shredding`. +/// +/// `child` must produce [`Json`] extension values; the result is `Variant` with the input's +/// nullability. Rows containing invalid JSON fail the expression. +/// +/// Note that this is a lossy, normalizing conversion. See [`JsonToVariant`] for the full list of +/// caveats. +pub fn json_to_variant(child: Expression, shredding: ShreddingSpec) -> Expression { + JsonToVariant.new_expr(JsonToVariantOptions::new(shredding), [child]) +} + +/// A list of `(path, dtype)` directives describing which Variant paths to shred and as what +/// type. +/// +/// Paths must contain only object-field elements; list index elements are rejected because +/// Parquet Variant shredding schemas cannot express list element shredding yet. The root path +/// (`$`) shreds the top-level value itself. When entries overlap (e.g. `$.a` and `$.a.b`), +/// later entries overwrite earlier ones. +#[derive(Clone, Debug, Default, PartialEq, Eq, Hash)] +pub struct ShreddingSpec(Vec<(VariantPath, DType)>); + +impl ShreddingSpec { + /// Creates an empty spec, meaning no shredding is performed. + pub fn empty() -> Self { + Self::default() + } + + /// Creates a spec from `(path, dtype)` directives. + /// + /// # Errors + /// + /// Returns an error if any path contains a list index element. + pub fn try_new(fields: impl IntoIterator) -> VortexResult { + let fields = fields.into_iter().collect::>(); + for (path, _) in &fields { + vortex_ensure!( + path.elements() + .iter() + .all(|element| matches!(element, VariantPathElement::Field(_))), + "ShreddingSpec paths must only contain object fields, found list index in {path}" + ); + } + Ok(Self(fields)) + } + + /// Returns the `(path, dtype)` directives. + pub fn fields(&self) -> &[(VariantPath, DType)] { + &self.0 + } + + /// Returns whether this spec contains no directives. + pub fn is_empty(&self) -> bool { + self.0.is_empty() + } +} + +impl Display for ShreddingSpec { + fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { + for (idx, (path, dtype)) in self.0.iter().enumerate() { + if idx > 0 { + write!(f, ", ")?; + } + write!(f, "{path} as {dtype}")?; + } + Ok(()) + } +} + +/// Options for [`JsonToVariant`]. +#[derive(Clone, Debug, PartialEq, Eq, Hash)] +pub struct JsonToVariantOptions { + /// The paths to shred into typed storage, if any. + shredding: ShreddingSpec, +} + +impl JsonToVariantOptions { + /// Creates options that shred the paths selected by `shredding`. + pub fn new(shredding: ShreddingSpec) -> Self { + Self { shredding } + } + + /// Creates options that perform no shredding. + pub fn unshredded() -> Self { + Self::new(ShreddingSpec::empty()) + } + + /// Returns the shredding spec. + pub fn shredding(&self) -> &ShreddingSpec { + &self.shredding + } +} + +impl Display for JsonToVariantOptions { + fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { + self.shredding.fmt(f) + } +} + +#[cfg(test)] +mod tests { + #![allow(clippy::missing_docs_in_private_items)] + + use vortex_array::ArrayRef; + use vortex_array::EmptyMetadata; + use vortex_array::IntoArray; + use vortex_array::VortexSessionExecute; + use vortex_array::arrays::VarBinViewArray; + use vortex_array::dtype::Nullability; + use vortex_array::dtype::PType; + use vortex_array::dtype::session::DTypeSession; + use vortex_array::expr::proto::ExprSerializeProtoExt; + use vortex_array::expr::root; + use vortex_array::scalar_fn::session::ScalarFnSession; + use vortex_array::scalar_fn::session::ScalarFnSessionExt; + use vortex_array::session::ArraySession; + + use super::*; + + fn i64_dtype() -> DType { + DType::Primitive(PType::I64, Nullability::Nullable) + } + + /// A session that knows the `JsonToVariant` definition but has no Variant encoding registered, + /// so executing the function exercises the fallback that errors when no kernel is present. + fn session() -> VortexSession { + let session = VortexSession::empty() + .with::() + .with::() + .with::(); + session.scalar_fns().register(JsonToVariant); + session + } + + #[test] + fn shredding_spec_rejects_index_paths() { + let err = ShreddingSpec::try_new([( + VariantPath::new([VariantPathElement::field("a"), VariantPathElement::index(0)]), + i64_dtype(), + )]) + .unwrap_err(); + + assert!( + err.to_string() + .contains("ShreddingSpec paths must only contain object fields"), + "unexpected error: {err}" + ); + } + + #[test] + fn expression_roundtrip_serialization() -> VortexResult<()> { + let spec = ShreddingSpec::try_new([(VariantPath::field("a"), i64_dtype())])?; + let expr: Expression = json_to_variant(root(), spec); + let proto = expr.serialize_proto()?; + let actual = Expression::from_proto(&proto, &session())?; + + assert_eq!(actual, expr); + Ok(()) + } + + #[test] + fn utf8_input_is_rejected() { + let input = VarBinViewArray::from_iter_str(["1", "2"]).into_array(); + let err = JsonToVariant + .try_new_array(input.len(), JsonToVariantOptions::unshredded(), [input]) + .unwrap_err(); + + assert!( + err.to_string().contains("Json extension"), + "unexpected error: {err}" + ); + } + + #[test] + fn execute_without_variant_kernel_errors() -> VortexResult<()> { + // With no Variant encoding registered, executing over an already-canonical input must + // surface a clear error rather than looping. + let input = ExtensionArray::try_new_from_vtable( + Json, + EmptyMetadata, + VarBinViewArray::from_iter_str(["1", "2"]).into_array(), + )? + .into_array(); + + let dtype = input.dtype().clone(); + let array = JsonToVariant.try_new_array( + input.len(), + JsonToVariantOptions::unshredded(), + [input], + )?; + + let err = array + .execute::(&mut session().create_execution_ctx()) + .unwrap_err(); + + assert!( + err.to_string().contains("Variant encoding"), + "unexpected error for {dtype}: {err}" + ); + Ok(()) + } +} diff --git a/vortex-json/src/lib.rs b/vortex-json/src/lib.rs index fc7b0b1c9b3..21f86ca9a7c 100644 --- a/vortex-json/src/lib.rs +++ b/vortex-json/src/lib.rs @@ -11,12 +11,18 @@ mod arrow; mod dtype; +mod json_to_variant; use std::sync::Arc; pub use dtype::Json; +pub use json_to_variant::JsonToVariant; +pub use json_to_variant::JsonToVariantOptions; +pub use json_to_variant::ShreddingSpec; +pub use json_to_variant::json_to_variant; use vortex_array::arrow::ArrowSessionExt; use vortex_array::dtype::session::DTypeSessionExt; +use vortex_array::scalar_fn::session::ScalarFnSessionExt; use vortex_session::VortexSession; /// Register JSON extension support with a session. @@ -24,4 +30,5 @@ pub fn initialize(session: &VortexSession) { session.dtypes().register(Json); session.arrow().register_exporter(Arc::new(Json)); session.arrow().register_importer(Arc::new(Json)); + session.scalar_fns().register(JsonToVariant); } diff --git a/vortex-proto/proto/expr.proto b/vortex-proto/proto/expr.proto index 22cef069da0..00ffaac433c 100644 --- a/vortex-proto/proto/expr.proto +++ b/vortex-proto/proto/expr.proto @@ -61,6 +61,17 @@ message VariantPathElement { } } +// Options for `vortex.json_to_variant` +message JsonToVariantOpts { + repeated ShreddingSpecField shredding = 1; +} + +// One (path, dtype) shredding directive for `vortex.json_to_variant`. +message ShreddingSpecField { + repeated VariantPathElement path = 1; + vortex.dtype.DType dtype = 2; +} + // Options for `vortex.binary` message BinaryOpts { BinaryOp op = 1; diff --git a/vortex-proto/src/generated/vortex.expr.rs b/vortex-proto/src/generated/vortex.expr.rs index 0ae2cbe7368..a44328623e3 100644 --- a/vortex-proto/src/generated/vortex.expr.rs +++ b/vortex-proto/src/generated/vortex.expr.rs @@ -69,6 +69,20 @@ pub mod variant_path_element { Index(u64), } } +/// Options for `vortex.json_to_variant` +#[derive(Clone, PartialEq, ::prost::Message)] +pub struct JsonToVariantOpts { + #[prost(message, repeated, tag = "1")] + pub shredding: ::prost::alloc::vec::Vec, +} +/// One (path, dtype) shredding directive for `vortex.json_to_variant`. +#[derive(Clone, PartialEq, ::prost::Message)] +pub struct ShreddingSpecField { + #[prost(message, repeated, tag = "1")] + pub path: ::prost::alloc::vec::Vec, + #[prost(message, optional, tag = "2")] + pub dtype: ::core::option::Option, +} /// Options for `vortex.binary` #[derive(Clone, Copy, PartialEq, Eq, Hash, ::prost::Message)] pub struct BinaryOpts {