From 31ce117a1d65dd30ab41de063ad28700e9b3cdaa Mon Sep 17 00:00:00 2001 From: Adam Gutglick Date: Fri, 12 Jun 2026 17:19:51 +0100 Subject: [PATCH 1/9] Json-Variant exprs Signed-off-by: Adam Gutglick --- Cargo.lock | 1 + encodings/parquet-variant/Cargo.toml | 3 +- encodings/parquet-variant/src/arrow.rs | 9 +- .../src/fns/json_to_variant.rs | 651 ++++++++++++++++++ encodings/parquet-variant/src/fns/mod.rs | 40 ++ .../src/fns/variant_to_json.rs | 464 +++++++++++++ encodings/parquet-variant/src/kernel.rs | 2 +- encodings/parquet-variant/src/lib.rs | 17 +- .../src/scalar_fn/fns/variant_get/mod.rs | 30 +- vortex-proto/proto/expr.proto | 11 + vortex-proto/src/generated/vortex.expr.rs | 14 + 11 files changed, 1224 insertions(+), 18 deletions(-) create mode 100644 encodings/parquet-variant/src/fns/json_to_variant.rs create mode 100644 encodings/parquet-variant/src/fns/mod.rs create mode 100644 encodings/parquet-variant/src/fns/variant_to_json.rs diff --git a/Cargo.lock b/Cargo.lock index a5df0c5d207..e6fc779f3b7 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -10282,6 +10282,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/arrow.rs b/encodings/parquet-variant/src/arrow.rs index 326e7f9d7ee..5be6d51739a 100644 --- a/encodings/parquet-variant/src/arrow.rs +++ b/encodings/parquet-variant/src/arrow.rs @@ -58,7 +58,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 +99,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 +115,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); diff --git a/encodings/parquet-variant/src/fns/json_to_variant.rs b/encodings/parquet-variant/src/fns/json_to_variant.rs new file mode 100644 index 00000000000..4b7e7316bac --- /dev/null +++ b/encodings/parquet-variant/src/fns/json_to_variant.rs @@ -0,0 +1,651 @@ +// 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 std::sync::Arc; + +use arrow_schema::FieldRef; +use parquet_variant_compute::ShreddedSchemaBuilder; +use parquet_variant_compute::shred_variant; +use prost::Message; +use vortex_array::ArrayRef; +use vortex_array::ExecutionCtx; +use vortex_array::arrays::ExtensionArray; +use vortex_array::arrays::extension::ExtensionArrayExt; +use vortex_array::arrow::ArrowSessionExt; +use vortex_array::dtype::DType; +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::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_json::Json; +use vortex_proto::expr as pb; +use vortex_session::VortexSession; +use vortex_session::registry::CachedId; + +use crate::ParquetVariant; +use crate::kernel::to_parquet_variant_path; + +/// Parses JSON strings into Parquet Variant values, optionally shredding fields. +/// +/// Accepts `Utf8` inputs or [`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 +/// `typed_value` child, following the [Parquet Variant shredding] rules: rows whose value does +/// not match the requested type stay readable through the residual variant value. +/// +/// # Not an inverse of `variant_to_json` +/// +/// `json_to_variant` and [`VariantToJson`](crate::VariantToJson) are lossy, normalizing +/// conversions and are NOT inverses of each other: +/// - JSON whitespace is not preserved. +/// - Object keys may be reordered: Variant metadata stores keys in sorted order, so +/// `variant_to_json` emits fields in a canonical order, not source order. +/// - Number formatting and precision change: e.g. `1.0` may render as `1`, exponent forms and +/// very large numbers are re-rendered, and floating-point values are re-encoded. +/// - Duplicate object keys are collapsed to a single entry. +/// - Unicode escape sequences are normalized (e.g. `\u0041` becomes `A`). +/// - `variant_to_json` stringifies Variant-only types — date, timestamp, UUID, binary, +/// decimal — so `json_to_variant(variant_to_json(v))` yields plain strings/numbers and loses +/// the original type information. +/// - Shredding structure is not recoverable from JSON: `variant_to_json` unshreds its input +/// first, and re-parsing produces an unshredded Variant unless a new [`ShreddingSpec`] is +/// supplied. +/// +/// [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!( + is_json_string_dtype(input_dtype), + "JsonToVariant input must be Utf8 or 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)?; + let input_nullable = input.dtype().is_nullable(); + + let storage = if input.dtype().is_utf8() { + input + } else if is_json_string_dtype(input.dtype()) { + input + .execute::(ctx)? + .storage_array() + .clone() + } else { + vortex_bail!( + "JsonToVariant input must be Utf8 or a Json extension, found {}", + input.dtype() + ); + }; + + let session = ctx.session().clone(); + let arrow_strings = session.arrow().execute_arrow(storage, 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 options.shredding().is_empty() { + arrow_variant + } else { + let mut builder = ShreddedSchemaBuilder::new(); + for (path, dtype) in 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 input_nullable { + ParquetVariant::from_arrow_variant_nullable(&arrow_variant) + } else { + ParquetVariant::from_arrow_variant(&arrow_variant) + } + } +} + +/// Returns whether `dtype` is acceptable input for [`JsonToVariant`]: `Utf8` or the [`Json`] +/// extension dtype. +fn is_json_string_dtype(dtype: &DType) -> bool { + dtype.is_utf8() + || dtype + .as_extension_opt() + .is_some_and(|ext_dtype| ext_dtype.is::()) +} + +/// 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 { + use std::sync::LazyLock; + + use vortex_array::Canonical; + use vortex_array::EmptyMetadata; + use vortex_array::IntoArray; + use vortex_array::VortexSessionExecute; + 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::Nullability; + use vortex_array::dtype::PType; + use vortex_array::dtype::extension::ExtDType; + use vortex_array::expr::Expression; + use vortex_array::expr::proto::ExprSerializeProtoExt; + use vortex_array::expr::root; + use vortex_array::expr::variant_get; + use vortex_array::session::ArraySession; + use vortex_error::vortex_bail; + + use super::*; + use crate::ParquetVariantArrayExt; + use crate::fns::json_to_variant; + + static SESSION: LazyLock = LazyLock::new(|| { + let session = VortexSession::empty().with::(); + 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 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 options_roundtrip_serialization() -> VortexResult<()> { + let specs = [ + ShreddingSpec::empty(), + ShreddingSpec::try_new([ + (VariantPath::field("a"), i64_dtype()), + ( + VariantPath::new([ + VariantPathElement::field("b"), + VariantPathElement::field("c"), + ]), + DType::Utf8(Nullability::NonNullable), + ), + ])?, + ]; + + for spec in specs { + let options = JsonToVariantOptions::new(spec); + let metadata = JsonToVariant + .serialize(&options)? + .ok_or_else(|| vortex_err!("expected metadata"))?; + let actual = JsonToVariant.deserialize(&metadata, &VortexSession::empty())?; + assert_eq!(actual, options); + } + Ok(()) + } + + #[test] + fn expression_roundtrip_serialization() -> VortexResult<()> { + let expr: Expression = json_to_variant(root(), shred_field_as_i64("a")?); + let proto = expr.serialize_proto()?; + let actual = Expression::from_proto(&proto, &SESSION)?; + + assert_eq!(actual, expr); + Ok(()) + } + + #[test] + fn options_display() -> VortexResult<()> { + assert_eq!(JsonToVariantOptions::unshredded().to_string(), ""); + assert_eq!( + JsonToVariantOptions::new(ShreddingSpec::try_new([ + (VariantPath::field("a"), i64_dtype()), + ( + VariantPath::field("b"), + DType::Utf8(Nullability::NonNullable) + ), + ])?) + .to_string(), + "$.a as i64?, $.b as utf8" + ); + Ok(()) + } + + #[test] + fn return_dtype_propagates_nullability() { + let expr = json_to_variant(root(), ShreddingSpec::empty()); + assert_eq!( + expr.return_dtype(&DType::Utf8(Nullability::NonNullable)) + .unwrap(), + DType::Variant(Nullability::NonNullable) + ); + assert_eq!( + expr.return_dtype(&DType::Utf8(Nullability::Nullable)) + .unwrap(), + DType::Variant(Nullability::Nullable) + ); + } + + #[test] + fn return_dtype_accepts_json_extension() -> VortexResult<()> { + let json_dtype = DType::Extension( + ExtDType::::try_new(EmptyMetadata, DType::Utf8(Nullability::Nullable))?.erased(), + ); + let expr = json_to_variant(root(), ShreddingSpec::empty()); + assert_eq!( + expr.return_dtype(&json_dtype)?, + DType::Variant(Nullability::Nullable) + ); + Ok(()) + } + + #[test] + fn return_dtype_rejects_non_string_input() { + let expr = json_to_variant(root(), ShreddingSpec::empty()); + let err = expr + .return_dtype(&DType::Bool(Nullability::NonNullable)) + .unwrap_err(); + assert!( + err.to_string() + .contains("JsonToVariant input must be Utf8 or a Json extension"), + "unexpected error: {err}" + ); + } + + #[test] + fn converts_utf8_json_rows() -> VortexResult<()> { + let 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 = ExtensionArray::try_new_from_vtable(Json, EmptyMetadata, storage)?.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)]) + } + + #[test] + fn null_rows_stay_null_and_json_null_becomes_variant_null() -> VortexResult<()> { + let 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); + 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() { + let 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()); + } + + #[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 shredding_produces_typed_value_child() -> VortexResult<()> { + let 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]) + ); + + // 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 = 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)); + assert_nth_scalar_is_null!(result, 1); + let mut ctx = SESSION.create_execution_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)]) + ); + Ok(()) + } + + #[test] + fn shredding_root_path_shreds_top_level_values() -> VortexResult<()> { + let 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/fns/mod.rs b/encodings/parquet-variant/src/fns/mod.rs new file mode 100644 index 00000000000..d930cb4cd80 --- /dev/null +++ b/encodings/parquet-variant/src/fns/mod.rs @@ -0,0 +1,40 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Scalar functions converting between JSON strings and Variant values. + +mod json_to_variant; +mod variant_to_json; + +pub use json_to_variant::JsonToVariant; +pub use json_to_variant::JsonToVariantOptions; +pub use json_to_variant::ShreddingSpec; +pub use variant_to_json::VariantToJson; +use vortex_array::expr::Expression; +use vortex_array::scalar_fn::EmptyOptions; +use vortex_array::scalar_fn::ScalarFnVTableExt; + +/// Creates a [`JsonToVariant`] expression that parses `child`'s JSON strings into Variant +/// values, shredding the paths selected by `shredding`. +/// +/// `child` must produce `Utf8` or [`Json`](vortex_json::Json) extension values; the result is +/// `Variant` with the input's nullability. Rows containing invalid JSON fail the expression. +/// +/// Note that this is NOT an inverse of [`variant_to_json`]: both conversions normalize their +/// input. 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]) +} + +/// Creates a [`VariantToJson`] expression that renders `child`'s Variant values as JSON +/// strings with the [`Json`](vortex_json::Json) extension dtype. +/// +/// Shredded inputs are unshredded before rendering, and the result keeps the input's +/// nullability. +/// +/// Note that this is NOT an inverse of [`json_to_variant`]: both conversions normalize their +/// input, and Variant-only types (dates, timestamps, UUIDs, binary, decimals) are rendered as +/// plain JSON strings or numbers. See [`VariantToJson`] for the full list of caveats. +pub fn variant_to_json(child: Expression) -> Expression { + VariantToJson.new_expr(EmptyOptions, [child]) +} diff --git a/encodings/parquet-variant/src/fns/variant_to_json.rs b/encodings/parquet-variant/src/fns/variant_to_json.rs new file mode 100644 index 00000000000..6f017d65fab --- /dev/null +++ b/encodings/parquet-variant/src/fns/variant_to_json.rs @@ -0,0 +1,464 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! The `vortex.variant_to_json` scalar function. + +use std::fmt; +use std::fmt::Formatter; +use std::sync::Arc; + +use arrow_schema::DataType; +use arrow_schema::Field; +use arrow_schema::Fields; +use vortex_array::ArrayRef; +use vortex_array::EmptyMetadata; +use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; +use vortex_array::arrays::ExtensionArray; +use vortex_array::arrow::FromArrowArray; +use vortex_array::dtype::DType; +use vortex_array::dtype::extension::ExtDType; +use vortex_array::expr::Expression; +use vortex_array::scalar_fn::Arity; +use vortex_array::scalar_fn::ChildName; +use vortex_array::scalar_fn::EmptyOptions; +use vortex_array::scalar_fn::ExecutionArgs; +use vortex_array::scalar_fn::ScalarFnId; +use vortex_array::scalar_fn::ScalarFnVTable; +use vortex_error::VortexResult; +use vortex_error::vortex_ensure; +use vortex_json::Json; +use vortex_session::VortexSession; +use vortex_session::registry::CachedId; + +use crate::ParquetVariant; +use crate::ParquetVariantArrayExt; +use crate::arrow::export_storage_to_target; +use crate::arrow::export_unshredded_storage_to_target; +use crate::arrow::parquet_variant_for_export; + +/// Renders Variant values as JSON strings with the [`Json`] extension dtype. +/// +/// Accepts `Variant` inputs backed by Parquet Variant storage, including shredded storage, +/// which is unshredded before rendering. Null rows stay null, while variant-null values render +/// as the JSON literal `null`. The output nullability matches the input's. +/// +/// Inputs are exported through their Parquet Variant storage, sharing the Arrow export path's +/// constraints: canonical variants whose shredded child is an object-field tree (rather than a +/// top-level typed value) cannot currently be reattached to Parquet storage and fail. +/// +/// # Not an inverse of `json_to_variant` +/// +/// [`JsonToVariant`](crate::JsonToVariant) and `variant_to_json` are lossy, normalizing +/// conversions and are NOT inverses of each other: +/// - JSON whitespace is not preserved. +/// - Object keys may be reordered: Variant metadata stores keys in sorted order, so +/// `variant_to_json` emits fields in a canonical order, not source order. +/// - Number formatting and precision change: e.g. `1.0` may render as `1`, exponent forms and +/// very large numbers are re-rendered, and floating-point values are re-encoded. +/// - Duplicate object keys are collapsed to a single entry. +/// - Unicode escape sequences are normalized (e.g. `\u0041` becomes `A`). +/// - `variant_to_json` stringifies Variant-only types — date, timestamp, UUID, binary, +/// decimal — so `json_to_variant(variant_to_json(v))` yields plain strings/numbers and loses +/// the original type information. +/// - Shredding structure is not recoverable from JSON: `variant_to_json` unshreds its input +/// first, and re-parsing produces an unshredded Variant unless a new +/// [`ShreddingSpec`](crate::ShreddingSpec) is supplied. +#[derive(Clone)] +pub struct VariantToJson; + +impl ScalarFnVTable for VariantToJson { + type Options = EmptyOptions; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("vortex.variant_to_json"); + *ID + } + + fn serialize(&self, _options: &Self::Options) -> VortexResult>> { + Ok(Some(vec![])) + } + + fn deserialize( + &self, + _metadata: &[u8], + _session: &VortexSession, + ) -> VortexResult { + Ok(EmptyOptions) + } + + 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 VariantToJson expression"), + } + } + + fn fmt_sql( + &self, + _options: &Self::Options, + expr: &Expression, + f: &mut Formatter<'_>, + ) -> fmt::Result { + write!(f, "variant_to_json(")?; + expr.child(0).fmt_sql(f)?; + write!(f, ")") + } + + fn return_dtype(&self, _options: &Self::Options, arg_dtypes: &[DType]) -> VortexResult { + let input_dtype = &arg_dtypes[0]; + vortex_ensure!( + matches!(input_dtype, DType::Variant(_)), + "VariantToJson input must be Variant, found {input_dtype}" + ); + + let storage_dtype = DType::Utf8(input_dtype.nullability()); + Ok(DType::Extension( + ExtDType::::try_new(EmptyMetadata, storage_dtype)?.erased(), + )) + } + + fn execute( + &self, + _options: &Self::Options, + args: &dyn ExecutionArgs, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + let input = args.get(0)?; + let input_nullable = input.dtype().is_nullable(); + + let parquet_array = parquet_variant_for_export(input, ctx)?; + let parquet_array = parquet_array.as_::(); + + // `parquet_variant_compute::variant_to_json` only accepts unshredded + // `STRUCT` storage, so request exactly that shape and + // unshred any typed values first. + let target_fields: Fields = vec![ + Arc::new(Field::new("metadata", DataType::Binary, false)), + Arc::new(Field::new("value", DataType::Binary, true)), + ] + .into(); + let arrow_storage = if parquet_array.typed_value_array().is_some() { + export_unshredded_storage_to_target(&parquet_array, &target_fields, ctx)? + } else { + export_storage_to_target(&parquet_array, &target_fields, ctx)? + }; + + let arrow_json = parquet_variant_compute::variant_to_json(&arrow_storage)?; + let storage = ArrayRef::from_arrow(&arrow_json, input_nullable)?; + + ExtensionArray::try_new_from_vtable(Json, EmptyMetadata, storage).map(IntoArray::into_array) + } + + fn is_null_sensitive(&self, _options: &Self::Options) -> bool { + false + } + + fn is_fallible(&self, _options: &Self::Options) -> bool { + false + } +} + +#[cfg(test)] +mod tests { + use std::sync::LazyLock; + + use chrono::NaiveDate; + use parquet_variant::Variant as PqVariant; + use parquet_variant::VariantBuilder; + use parquet_variant_compute::VariantArrayBuilder; + use vortex_array::VortexSessionExecute; + use vortex_array::accessor::ArrayAccessor; + use vortex_array::arrays::PrimitiveArray; + use vortex_array::arrays::VarBinViewArray; + use vortex_array::arrays::Variant; + use vortex_array::arrays::VariantArray; + use vortex_array::arrays::extension::ExtensionArrayExt; + use vortex_array::arrays::variant::VariantArrayExt; + use vortex_array::dtype::Nullability; + use vortex_array::dtype::PType; + use vortex_array::expr::Expression; + use vortex_array::expr::proto::ExprSerializeProtoExt; + use vortex_array::expr::root; + use vortex_array::scalar_fn::fns::variant_get::VariantPath; + use vortex_array::session::ArraySession; + use vortex_array::validity::Validity; + use vortex_error::vortex_err; + + use super::*; + use crate::ShreddingSpec; + use crate::fns::json_to_variant; + use crate::fns::variant_to_json; + + static SESSION: LazyLock = LazyLock::new(|| { + let session = VortexSession::empty().with::(); + crate::initialize(&session); + session + }); + + fn json_dtype(nullability: Nullability) -> VortexResult { + Ok(DType::Extension( + ExtDType::::try_new(EmptyMetadata, DType::Utf8(nullability))?.erased(), + )) + } + + fn execute_variant_to_json(input: ArrayRef) -> VortexResult { + input + .apply(&variant_to_json(root()))? + .execute::(&mut SESSION.create_execution_ctx()) + } + + fn json_strings(array: &ArrayRef) -> VortexResult>> { + let mut ctx = SESSION.create_execution_ctx(); + let ext = array.clone().execute::(&mut ctx)?; + let storage = ext + .storage_array() + .clone() + .execute::(&mut ctx)?; + Ok(storage.with_iterator(|iter| { + iter.map(|value| value.map(|bytes| String::from_utf8_lossy(bytes).into_owned())) + .collect() + })) + } + + fn unshredded_variant( + rows: impl IntoIterator>, + ) -> VortexResult { + let rows = rows.into_iter().collect::>(); + let mut builder = VariantArrayBuilder::new(rows.len()); + for row in rows { + builder.append_variant(row); + } + ParquetVariant::from_arrow_variant(&builder.build()) + } + + fn json_rows_to_variant( + rows: Vec>, + shredding: ShreddingSpec, + ) -> VortexResult { + let input = VarBinViewArray::from_iter_nullable_str(rows).into_array(); + input + .apply(&json_to_variant(root(), shredding))? + .execute::(&mut SESSION.create_execution_ctx()) + } + + #[test] + fn expression_roundtrip_serialization() -> VortexResult<()> { + let expr: Expression = variant_to_json(root()); + let proto = expr.serialize_proto()?; + let actual = Expression::from_proto(&proto, &SESSION)?; + + assert_eq!(actual, expr); + Ok(()) + } + + #[test] + fn return_dtype_is_json_extension() -> VortexResult<()> { + let expr = variant_to_json(root()); + assert_eq!( + expr.return_dtype(&DType::Variant(Nullability::NonNullable))?, + json_dtype(Nullability::NonNullable)? + ); + assert_eq!( + expr.return_dtype(&DType::Variant(Nullability::Nullable))?, + json_dtype(Nullability::Nullable)? + ); + Ok(()) + } + + #[test] + fn return_dtype_rejects_non_variant_input() { + let expr = variant_to_json(root()); + let err = expr + .return_dtype(&DType::Utf8(Nullability::NonNullable)) + .unwrap_err(); + assert!( + err.to_string() + .contains("VariantToJson input must be Variant"), + "unexpected error: {err}" + ); + } + + #[test] + fn formats_sql() { + let expr = variant_to_json(root()); + assert_eq!(expr.to_string(), "variant_to_json($)"); + } + + #[test] + fn renders_unshredded_values() -> VortexResult<()> { + let input = unshredded_variant([ + PqVariant::from(42i32), + PqVariant::from("hello"), + PqVariant::from(true), + PqVariant::Null, + ])?; + + let result = execute_variant_to_json(input)?; + + assert_eq!(result.dtype(), &json_dtype(Nullability::NonNullable)?); + assert_eq!( + json_strings(&result)?, + vec![ + Some("42".to_string()), + Some(r#""hello""#.to_string()), + Some("true".to_string()), + Some("null".to_string()), + ] + ); + Ok(()) + } + + #[test] + fn null_rows_stay_null_and_variant_null_renders_as_json_null() -> VortexResult<()> { + let input = + json_rows_to_variant(vec![Some("1"), None, Some("null")], ShreddingSpec::empty())?; + + let result = execute_variant_to_json(input)?; + + assert_eq!(result.dtype(), &json_dtype(Nullability::Nullable)?); + assert_eq!( + json_strings(&result)?, + vec![Some("1".to_string()), None, Some("null".to_string())] + ); + Ok(()) + } + + fn typed_value_only_variant() -> VortexResult { + let rows = [ + VariantBuilder::new().with_value(10i32).finish(), + VariantBuilder::new().with_value(20i32).finish(), + VariantBuilder::new().with_value(30i32).finish(), + ]; + let metadata = + VarBinViewArray::from_iter_bin(rows.iter().map(|(metadata, _)| metadata.as_slice())) + .into_array(); + let typed_value = PrimitiveArray::from_iter([10i32, 20, 30]).into_array(); + Ok( + ParquetVariant::try_new(Validity::NonNullable, metadata, None, Some(typed_value))? + .into_array(), + ) + } + + #[test] + fn unshreds_typed_value_only_storage() -> VortexResult<()> { + let result = execute_variant_to_json(typed_value_only_variant()?)?; + + assert_eq!( + json_strings(&result)?, + vec![ + Some("10".to_string()), + Some("20".to_string()), + Some("30".to_string()), + ] + ); + Ok(()) + } + + #[test] + fn unshreds_partially_shredded_storage() -> VortexResult<()> { + let spec = ShreddingSpec::try_new([( + VariantPath::field("a"), + DType::Primitive(PType::I64, Nullability::Nullable), + )])?; + let input = json_rows_to_variant( + vec![ + Some(r#"{"a": 1, "b": "x"}"#), + Some(r#"{"a": "not-a-number", "b": "y"}"#), + Some(r#"{"b": "z"}"#), + ], + spec, + )?; + assert!( + input.as_::().typed_value_array().is_some(), + "fixture must be shredded" + ); + + let result = execute_variant_to_json(input)?; + + assert_eq!( + json_strings(&result)?, + vec![ + Some(r#"{"a":1,"b":"x"}"#.to_string()), + Some(r#"{"a":"not-a-number","b":"y"}"#.to_string()), + Some(r#"{"b":"z"}"#.to_string()), + ] + ); + Ok(()) + } + + #[test] + fn renders_canonical_variant_with_shredded_child() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let canonical = typed_value_only_variant()? + .execute::(&mut ctx)? + .into_array(); + assert!( + canonical.as_::().shredded().is_some(), + "fixture must carry a canonical shredded child" + ); + + let result = execute_variant_to_json(canonical)?; + + assert_eq!( + json_strings(&result)?, + vec![ + Some("10".to_string()), + Some("20".to_string()), + Some("30".to_string()), + ] + ); + Ok(()) + } + + #[test] + fn json_round_trip_normalizes_whitespace_and_key_order() -> VortexResult<()> { + let input = json_rows_to_variant( + vec![Some(r#"{ "b" : 1 , "a" : 2 }"#), Some("[ 1 , 2 , 3 ]")], + ShreddingSpec::empty(), + )?; + + let result = execute_variant_to_json(input)?; + + assert_eq!( + json_strings(&result)?, + vec![ + Some(r#"{"a":2,"b":1}"#.to_string()), + Some("[1,2,3]".to_string()), + ] + ); + Ok(()) + } + + #[test] + fn variant_only_types_are_stringified_so_reparsing_loses_types() -> VortexResult<()> { + let date = + NaiveDate::from_ymd_opt(2026, 6, 11).ok_or_else(|| vortex_err!("invalid test date"))?; + let input = unshredded_variant([PqVariant::from(date)])?; + + let rendered = execute_variant_to_json(input)?; + let json = json_strings(&rendered)?; + assert_eq!(json, vec![Some(r#""2026-06-11""#.to_string())]); + + // Re-parsing the rendered JSON yields a plain string variant, not a date: the type + // information is lost, demonstrating that the conversions are not inverses. + let reparsed = rendered + .apply(&json_to_variant(root(), ShreddingSpec::empty()))? + .execute::(&mut SESSION.create_execution_ctx())?; + let mut ctx = SESSION.create_execution_ctx(); + let row0 = reparsed.execute_scalar(0, &mut ctx)?; + let value = row0 + .as_variant() + .value() + .ok_or_else(|| vortex_err!("expected non-null variant"))?; + assert_eq!( + value.as_utf8().value().map(|value| value.to_string()), + Some("2026-06-11".to_string()) + ); + Ok(()) + } +} diff --git a/encodings/parquet-variant/src/kernel.rs b/encodings/parquet-variant/src/kernel.rs index 684ce0217f4..afbf65cb26a 100644 --- a/encodings/parquet-variant/src/kernel.rs +++ b/encodings/parquet-variant/src/kernel.rs @@ -106,7 +106,7 @@ impl ExecuteParentKernel for VariantGetKernel { } } -fn to_parquet_variant_path(path: &VariantPath) -> VortexResult> { +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..2b57dadb8e1 100644 --- a/encodings/parquet-variant/src/lib.rs +++ b/encodings/parquet-variant/src/lib.rs @@ -26,6 +26,7 @@ mod array; mod arrow; +mod fns; mod kernel; mod operations; mod validity; @@ -34,16 +35,30 @@ mod vtable; use std::sync::Arc; pub use array::ParquetVariantArrayExt; +pub use fns::JsonToVariant; +pub use fns::JsonToVariantOptions; +pub use fns::ShreddingSpec; +pub use fns::VariantToJson; +pub use fns::json_to_variant; +pub use fns::variant_to_json; use vortex_array::arrow::ArrowSessionExt; +use vortex_array::scalar_fn::session::ScalarFnSessionExt; use vortex_array::session::ArraySessionExt; 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 registers the [`Json`](vortex_json::Json) extension dtype so that sessions able +/// to execute [`VariantToJson`] can serialize and deserialize its output dtype. pub fn initialize(session: &VortexSession) { session.arrays().register(ParquetVariant); kernel::initialize(session); session.arrow().register_exporter(Arc::new(ParquetVariant)); session.arrow().register_importer(Arc::new(ParquetVariant)); + vortex_json::initialize(session); + session.scalar_fns().register(JsonToVariant); + session.scalar_fns().register(VariantToJson); } 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-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 { From 9e93030a6b676b63de00b66665df2b28a87f943f Mon Sep 17 00:00:00 2001 From: Adam Gutglick Date: Mon, 15 Jun 2026 14:36:21 +0100 Subject: [PATCH 2/9] Fix thing Signed-off-by: Adam Gutglick --- encodings/parquet-variant/src/fns/mod.rs | 4 ++-- encodings/parquet-variant/src/fns/variant_to_json.rs | 10 ++++++---- 2 files changed, 8 insertions(+), 6 deletions(-) diff --git a/encodings/parquet-variant/src/fns/mod.rs b/encodings/parquet-variant/src/fns/mod.rs index d930cb4cd80..b1958a24000 100644 --- a/encodings/parquet-variant/src/fns/mod.rs +++ b/encodings/parquet-variant/src/fns/mod.rs @@ -20,7 +20,7 @@ use vortex_array::scalar_fn::ScalarFnVTableExt; /// `child` must produce `Utf8` or [`Json`](vortex_json::Json) extension values; the result is /// `Variant` with the input's nullability. Rows containing invalid JSON fail the expression. /// -/// Note that this is NOT an inverse of [`variant_to_json`]: both conversions normalize their +/// Note that this is NOT an inverse of [`variant_to_json()`]: both conversions normalize their /// input. 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]) @@ -32,7 +32,7 @@ pub fn json_to_variant(child: Expression, shredding: ShreddingSpec) -> Expressio /// Shredded inputs are unshredded before rendering, and the result keeps the input's /// nullability. /// -/// Note that this is NOT an inverse of [`json_to_variant`]: both conversions normalize their +/// Note that this is NOT an inverse of [`json_to_variant()`]: both conversions normalize their /// input, and Variant-only types (dates, timestamps, UUIDs, binary, decimals) are rendered as /// plain JSON strings or numbers. See [`VariantToJson`] for the full list of caveats. pub fn variant_to_json(child: Expression) -> Expression { diff --git a/encodings/parquet-variant/src/fns/variant_to_json.rs b/encodings/parquet-variant/src/fns/variant_to_json.rs index 6f017d65fab..23b9aa96c1b 100644 --- a/encodings/parquet-variant/src/fns/variant_to_json.rs +++ b/encodings/parquet-variant/src/fns/variant_to_json.rs @@ -157,10 +157,6 @@ impl ScalarFnVTable for VariantToJson { fn is_null_sensitive(&self, _options: &Self::Options) -> bool { false } - - fn is_fallible(&self, _options: &Self::Options) -> bool { - false - } } #[cfg(test)] @@ -289,6 +285,12 @@ mod tests { assert_eq!(expr.to_string(), "variant_to_json($)"); } + #[test] + fn is_fallible() { + let expr = variant_to_json(root()); + assert!(expr.signature().is_fallible()); + } + #[test] fn renders_unshredded_values() -> VortexResult<()> { let input = unshredded_variant([ From 5523c025c59343c8307d5c7bcf1950e60deb0614 Mon Sep 17 00:00:00 2001 From: Adam Gutglick Date: Wed, 17 Jun 2026 12:29:12 +0100 Subject: [PATCH 3/9] some fixes Signed-off-by: Adam Gutglick --- encodings/parquet-variant/src/array.rs | 175 ++++++++++++++++++ encodings/parquet-variant/src/arrow.rs | 8 +- .../src/fns/variant_to_json.rs | 120 +++++++++++- 3 files changed, 296 insertions(+), 7 deletions(-) diff --git a/encodings/parquet-variant/src/array.rs b/encodings/parquet-variant/src/array.rs index 7be0265fc96..67da0d59e1c 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 = LEGACY_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 5be6d51739a..bb409fd390d 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"; @@ -138,11 +139,16 @@ pub(crate) fn parquet_variant_for_export( 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) } diff --git a/encodings/parquet-variant/src/fns/variant_to_json.rs b/encodings/parquet-variant/src/fns/variant_to_json.rs index 23b9aa96c1b..0d411415753 100644 --- a/encodings/parquet-variant/src/fns/variant_to_json.rs +++ b/encodings/parquet-variant/src/fns/variant_to_json.rs @@ -39,13 +39,13 @@ use crate::arrow::parquet_variant_for_export; /// Renders Variant values as JSON strings with the [`Json`] extension dtype. /// -/// Accepts `Variant` inputs backed by Parquet Variant storage, including shredded storage, -/// which is unshredded before rendering. Null rows stay null, while variant-null values render -/// as the JSON literal `null`. The output nullability matches the input's. +/// Accepts `Variant` inputs backed by Parquet Variant storage, including shredded storage +/// (top-level, object, and nested fields), which is unshredded before rendering. Null rows stay +/// null, while variant-null values render as the JSON literal `null`. The output nullability +/// matches the input's. /// -/// Inputs are exported through their Parquet Variant storage, sharing the Arrow export path's -/// constraints: canonical variants whose shredded child is an object-field tree (rather than a -/// top-level typed value) cannot currently be reattached to Parquet storage and fail. +/// Inputs are exported through their Parquet Variant storage, so a `Variant` whose core storage +/// is not Parquet Variant-backed is not supported and fails. /// /// # Not an inverse of `json_to_variant` /// @@ -167,6 +167,7 @@ mod tests { use parquet_variant::Variant as PqVariant; use parquet_variant::VariantBuilder; use parquet_variant_compute::VariantArrayBuilder; + use vortex_array::Canonical; use vortex_array::VortexSessionExecute; use vortex_array::accessor::ArrayAccessor; use vortex_array::arrays::PrimitiveArray; @@ -181,8 +182,10 @@ mod tests { use vortex_array::expr::proto::ExprSerializeProtoExt; use vortex_array::expr::root; use vortex_array::scalar_fn::fns::variant_get::VariantPath; + use vortex_array::scalar_fn::fns::variant_get::VariantPathElement; use vortex_array::session::ArraySession; use vortex_array::validity::Validity; + use vortex_error::vortex_bail; use vortex_error::vortex_err; use super::*; @@ -417,6 +420,111 @@ mod tests { Ok(()) } + /// Shreds `rows` per `spec`, then canonicalizes so the typed values are lifted into a logical + /// shredded child (as a file read-back would produce). + fn canonical_shredded(rows: Vec>, spec: ShreddingSpec) -> VortexResult { + let shredded = json_rows_to_variant(rows, spec)?; + let Canonical::Variant(canonical) = + shredded.execute::(&mut SESSION.create_execution_ctx())? + else { + vortex_bail!("expected canonical variant"); + }; + Ok(canonical.into_array()) + } + + /// Rendering a canonicalized shredded variant must match rendering the same data unshredded: + /// the shredding is a storage optimization and is invisible to JSON output. + fn assert_canonical_matches_unshredded( + rows: Vec>, + spec: ShreddingSpec, + ) -> VortexResult<()> { + let unshredded = json_rows_to_variant(rows.clone(), ShreddingSpec::empty())?; + let want = json_strings(&execute_variant_to_json(unshredded)?)?; + + let canonical = canonical_shredded(rows, spec)?; + assert!( + canonical.as_::().shredded().is_some(), + "fixture must carry a canonical shredded child" + ); + let got = json_strings(&execute_variant_to_json(canonical)?)?; + + assert_eq!(got, want); + Ok(()) + } + + /// Regression for the canonicalization bug: object-field shredding lifted into a logical + /// shredded child must round-trip back to the full object, including the shredded field and + /// the non-conforming/missing-field fallbacks. + #[test] + fn renders_canonical_object_shredded_variant() -> VortexResult<()> { + let spec = ShreddingSpec::try_new([( + VariantPath::field("a"), + DType::Primitive(PType::I64, Nullability::Nullable), + )])?; + let canonical = canonical_shredded( + vec![ + Some(r#"{"a": 1, "b": "x"}"#), + Some(r#"{"a": "not-a-number", "b": "y"}"#), + Some(r#"{"b": "z"}"#), + ], + spec, + )?; + assert!( + canonical.as_::().shredded().is_some(), + "fixture must carry a canonical shredded child" + ); + + let result = execute_variant_to_json(canonical)?; + + assert_eq!( + json_strings(&result)?, + vec![ + Some(r#"{"a":1,"b":"x"}"#.to_string()), + Some(r#"{"a":"not-a-number","b":"y"}"#.to_string()), + Some(r#"{"b":"z"}"#.to_string()), + ] + ); + Ok(()) + } + + #[test] + fn canonical_object_shredding_matches_unshredded() -> VortexResult<()> { + assert_canonical_matches_unshredded( + vec![ + Some(r#"{"a": 1, "b": "x", "c": 3}"#), + Some(r#"{"a": 2, "b": "y", "c": 4}"#), + None, + ], + ShreddingSpec::try_new([ + ( + VariantPath::field("a"), + DType::Primitive(PType::I64, Nullability::Nullable), + ), + ( + VariantPath::field("c"), + DType::Primitive(PType::I64, Nullability::Nullable), + ), + ])?, + ) + } + + #[test] + fn canonical_nested_dotted_shredding_matches_unshredded() -> VortexResult<()> { + assert_canonical_matches_unshredded( + vec![ + Some(r#"{"a": {"b": 100}, "c": "keep"}"#), + Some(r#"{"a": {"b": 200}, "c": "keep2"}"#), + ], + ShreddingSpec::try_new([( + VariantPath::new([ + VariantPathElement::field("a"), + VariantPathElement::field("b"), + ]), + DType::Primitive(PType::I64, Nullability::Nullable), + )])?, + ) + } + #[test] fn json_round_trip_normalizes_whitespace_and_key_order() -> VortexResult<()> { let input = json_rows_to_variant( From a972d2206885ad0ae2b942c989702d3cbd0bde8e Mon Sep 17 00:00:00 2001 From: Adam Gutglick Date: Wed, 17 Jun 2026 14:52:03 +0100 Subject: [PATCH 4/9] Less code Signed-off-by: Adam Gutglick --- .../src/fns/json_to_variant.rs | 106 ++----------- .../src/fns/variant_to_json.rs | 140 ++++-------------- 2 files changed, 40 insertions(+), 206 deletions(-) diff --git a/encodings/parquet-variant/src/fns/json_to_variant.rs b/encodings/parquet-variant/src/fns/json_to_variant.rs index 4b7e7316bac..3e64fe2cf70 100644 --- a/encodings/parquet-variant/src/fns/json_to_variant.rs +++ b/encodings/parquet-variant/src/fns/json_to_variant.rs @@ -134,7 +134,10 @@ impl ScalarFnVTable for JsonToVariant { fn return_dtype(&self, _options: &Self::Options, arg_dtypes: &[DType]) -> VortexResult { let input_dtype = &arg_dtypes[0]; vortex_ensure!( - is_json_string_dtype(input_dtype), + input_dtype.is_utf8() + || input_dtype + .as_extension_opt() + .is_some_and(|ext_dtype| ext_dtype.is::()), "JsonToVariant input must be Utf8 or a Json extension, found {input_dtype}" ); @@ -152,7 +155,11 @@ impl ScalarFnVTable for JsonToVariant { let storage = if input.dtype().is_utf8() { input - } else if is_json_string_dtype(input.dtype()) { + } else if input + .dtype() + .as_extension_opt() + .is_some_and(|ext_dtype| ext_dtype.is::()) + { input .execute::(ctx)? .storage_array() @@ -188,15 +195,6 @@ impl ScalarFnVTable for JsonToVariant { } } -/// Returns whether `dtype` is acceptable input for [`JsonToVariant`]: `Utf8` or the [`Json`] -/// extension dtype. -fn is_json_string_dtype(dtype: &DType) -> bool { - dtype.is_utf8() - || dtype - .as_extension_opt() - .is_some_and(|ext_dtype| ext_dtype.is::()) -} - /// A list of `(path, dtype)` directives describing which Variant paths to shred and as what /// type. /// @@ -301,7 +299,6 @@ mod tests { use vortex_array::assert_nth_scalar_is_null; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; - use vortex_array::dtype::extension::ExtDType; use vortex_array::expr::Expression; use vortex_array::expr::proto::ExprSerializeProtoExt; use vortex_array::expr::root; @@ -357,33 +354,6 @@ mod tests { Ok(()) } - #[test] - fn options_roundtrip_serialization() -> VortexResult<()> { - let specs = [ - ShreddingSpec::empty(), - ShreddingSpec::try_new([ - (VariantPath::field("a"), i64_dtype()), - ( - VariantPath::new([ - VariantPathElement::field("b"), - VariantPathElement::field("c"), - ]), - DType::Utf8(Nullability::NonNullable), - ), - ])?, - ]; - - for spec in specs { - let options = JsonToVariantOptions::new(spec); - let metadata = JsonToVariant - .serialize(&options)? - .ok_or_else(|| vortex_err!("expected metadata"))?; - let actual = JsonToVariant.deserialize(&metadata, &VortexSession::empty())?; - assert_eq!(actual, options); - } - Ok(()) - } - #[test] fn expression_roundtrip_serialization() -> VortexResult<()> { let expr: Expression = json_to_variant(root(), shred_field_as_i64("a")?); @@ -394,64 +364,6 @@ mod tests { Ok(()) } - #[test] - fn options_display() -> VortexResult<()> { - assert_eq!(JsonToVariantOptions::unshredded().to_string(), ""); - assert_eq!( - JsonToVariantOptions::new(ShreddingSpec::try_new([ - (VariantPath::field("a"), i64_dtype()), - ( - VariantPath::field("b"), - DType::Utf8(Nullability::NonNullable) - ), - ])?) - .to_string(), - "$.a as i64?, $.b as utf8" - ); - Ok(()) - } - - #[test] - fn return_dtype_propagates_nullability() { - let expr = json_to_variant(root(), ShreddingSpec::empty()); - assert_eq!( - expr.return_dtype(&DType::Utf8(Nullability::NonNullable)) - .unwrap(), - DType::Variant(Nullability::NonNullable) - ); - assert_eq!( - expr.return_dtype(&DType::Utf8(Nullability::Nullable)) - .unwrap(), - DType::Variant(Nullability::Nullable) - ); - } - - #[test] - fn return_dtype_accepts_json_extension() -> VortexResult<()> { - let json_dtype = DType::Extension( - ExtDType::::try_new(EmptyMetadata, DType::Utf8(Nullability::Nullable))?.erased(), - ); - let expr = json_to_variant(root(), ShreddingSpec::empty()); - assert_eq!( - expr.return_dtype(&json_dtype)?, - DType::Variant(Nullability::Nullable) - ); - Ok(()) - } - - #[test] - fn return_dtype_rejects_non_string_input() { - let expr = json_to_variant(root(), ShreddingSpec::empty()); - let err = expr - .return_dtype(&DType::Bool(Nullability::NonNullable)) - .unwrap_err(); - assert!( - err.to_string() - .contains("JsonToVariant input must be Utf8 or a Json extension"), - "unexpected error: {err}" - ); - } - #[test] fn converts_utf8_json_rows() -> VortexResult<()> { let input = diff --git a/encodings/parquet-variant/src/fns/variant_to_json.rs b/encodings/parquet-variant/src/fns/variant_to_json.rs index 0d411415753..bd586188cea 100644 --- a/encodings/parquet-variant/src/fns/variant_to_json.rs +++ b/encodings/parquet-variant/src/fns/variant_to_json.rs @@ -173,13 +173,10 @@ mod tests { use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::VarBinViewArray; use vortex_array::arrays::Variant; - use vortex_array::arrays::VariantArray; use vortex_array::arrays::extension::ExtensionArrayExt; use vortex_array::arrays::variant::VariantArrayExt; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; - use vortex_array::expr::Expression; - use vortex_array::expr::proto::ExprSerializeProtoExt; use vortex_array::expr::root; use vortex_array::scalar_fn::fns::variant_get::VariantPath; use vortex_array::scalar_fn::fns::variant_get::VariantPathElement; @@ -205,12 +202,6 @@ mod tests { )) } - fn execute_variant_to_json(input: ArrayRef) -> VortexResult { - input - .apply(&variant_to_json(root()))? - .execute::(&mut SESSION.create_execution_ctx()) - } - fn json_strings(array: &ArrayRef) -> VortexResult>> { let mut ctx = SESSION.create_execution_ctx(); let ext = array.clone().execute::(&mut ctx)?; @@ -245,55 +236,6 @@ mod tests { .execute::(&mut SESSION.create_execution_ctx()) } - #[test] - fn expression_roundtrip_serialization() -> VortexResult<()> { - let expr: Expression = variant_to_json(root()); - let proto = expr.serialize_proto()?; - let actual = Expression::from_proto(&proto, &SESSION)?; - - assert_eq!(actual, expr); - Ok(()) - } - - #[test] - fn return_dtype_is_json_extension() -> VortexResult<()> { - let expr = variant_to_json(root()); - assert_eq!( - expr.return_dtype(&DType::Variant(Nullability::NonNullable))?, - json_dtype(Nullability::NonNullable)? - ); - assert_eq!( - expr.return_dtype(&DType::Variant(Nullability::Nullable))?, - json_dtype(Nullability::Nullable)? - ); - Ok(()) - } - - #[test] - fn return_dtype_rejects_non_variant_input() { - let expr = variant_to_json(root()); - let err = expr - .return_dtype(&DType::Utf8(Nullability::NonNullable)) - .unwrap_err(); - assert!( - err.to_string() - .contains("VariantToJson input must be Variant"), - "unexpected error: {err}" - ); - } - - #[test] - fn formats_sql() { - let expr = variant_to_json(root()); - assert_eq!(expr.to_string(), "variant_to_json($)"); - } - - #[test] - fn is_fallible() { - let expr = variant_to_json(root()); - assert!(expr.signature().is_fallible()); - } - #[test] fn renders_unshredded_values() -> VortexResult<()> { let input = unshredded_variant([ @@ -303,7 +245,9 @@ mod tests { PqVariant::Null, ])?; - let result = execute_variant_to_json(input)?; + let result = input + .apply(&variant_to_json(root()))? + .execute::(&mut SESSION.create_execution_ctx())?; assert_eq!(result.dtype(), &json_dtype(Nullability::NonNullable)?); assert_eq!( @@ -323,7 +267,9 @@ mod tests { let input = json_rows_to_variant(vec![Some("1"), None, Some("null")], ShreddingSpec::empty())?; - let result = execute_variant_to_json(input)?; + let result = input + .apply(&variant_to_json(root()))? + .execute::(&mut SESSION.create_execution_ctx())?; assert_eq!(result.dtype(), &json_dtype(Nullability::Nullable)?); assert_eq!( @@ -351,7 +297,12 @@ mod tests { #[test] fn unshreds_typed_value_only_storage() -> VortexResult<()> { - let result = execute_variant_to_json(typed_value_only_variant()?)?; + let result = { + let input = typed_value_only_variant()?; + input + .apply(&variant_to_json(root()))? + .execute::(&mut SESSION.create_execution_ctx()) + }?; assert_eq!( json_strings(&result)?, @@ -383,7 +334,9 @@ mod tests { "fixture must be shredded" ); - let result = execute_variant_to_json(input)?; + let result = input + .apply(&variant_to_json(root()))? + .execute::(&mut SESSION.create_execution_ctx())?; assert_eq!( json_strings(&result)?, @@ -396,30 +349,6 @@ mod tests { Ok(()) } - #[test] - fn renders_canonical_variant_with_shredded_child() -> VortexResult<()> { - let mut ctx = SESSION.create_execution_ctx(); - let canonical = typed_value_only_variant()? - .execute::(&mut ctx)? - .into_array(); - assert!( - canonical.as_::().shredded().is_some(), - "fixture must carry a canonical shredded child" - ); - - let result = execute_variant_to_json(canonical)?; - - assert_eq!( - json_strings(&result)?, - vec![ - Some("10".to_string()), - Some("20".to_string()), - Some("30".to_string()), - ] - ); - Ok(()) - } - /// Shreds `rows` per `spec`, then canonicalizes so the typed values are lifted into a logical /// shredded child (as a file read-back would produce). fn canonical_shredded(rows: Vec>, spec: ShreddingSpec) -> VortexResult { @@ -439,14 +368,22 @@ mod tests { spec: ShreddingSpec, ) -> VortexResult<()> { let unshredded = json_rows_to_variant(rows.clone(), ShreddingSpec::empty())?; - let want = json_strings(&execute_variant_to_json(unshredded)?)?; + let want = json_strings( + &unshredded + .apply(&variant_to_json(root()))? + .execute::(&mut SESSION.create_execution_ctx())?, + )?; let canonical = canonical_shredded(rows, spec)?; assert!( canonical.as_::().shredded().is_some(), "fixture must carry a canonical shredded child" ); - let got = json_strings(&execute_variant_to_json(canonical)?)?; + let got = json_strings( + &canonical + .apply(&variant_to_json(root()))? + .execute::(&mut SESSION.create_execution_ctx())?, + )?; assert_eq!(got, want); Ok(()) @@ -474,7 +411,9 @@ mod tests { "fixture must carry a canonical shredded child" ); - let result = execute_variant_to_json(canonical)?; + let result = canonical + .apply(&variant_to_json(root()))? + .execute::(&mut SESSION.create_execution_ctx())?; assert_eq!( json_strings(&result)?, @@ -525,32 +464,15 @@ mod tests { ) } - #[test] - fn json_round_trip_normalizes_whitespace_and_key_order() -> VortexResult<()> { - let input = json_rows_to_variant( - vec![Some(r#"{ "b" : 1 , "a" : 2 }"#), Some("[ 1 , 2 , 3 ]")], - ShreddingSpec::empty(), - )?; - - let result = execute_variant_to_json(input)?; - - assert_eq!( - json_strings(&result)?, - vec![ - Some(r#"{"a":2,"b":1}"#.to_string()), - Some("[1,2,3]".to_string()), - ] - ); - Ok(()) - } - #[test] fn variant_only_types_are_stringified_so_reparsing_loses_types() -> VortexResult<()> { let date = NaiveDate::from_ymd_opt(2026, 6, 11).ok_or_else(|| vortex_err!("invalid test date"))?; let input = unshredded_variant([PqVariant::from(date)])?; - let rendered = execute_variant_to_json(input)?; + let rendered = input + .apply(&variant_to_json(root()))? + .execute::(&mut SESSION.create_execution_ctx())?; let json = json_strings(&rendered)?; assert_eq!(json, vec![Some(r#""2026-06-11""#.to_string())]); From a6fec7a683aba72aac180fb05f2dde3fcd871511 Mon Sep 17 00:00:00 2001 From: Adam Gutglick Date: Wed, 17 Jun 2026 16:42:54 +0100 Subject: [PATCH 5/9] Fix test sessions Signed-off-by: Adam Gutglick --- encodings/parquet-variant/src/fns/json_to_variant.rs | 9 ++++++++- encodings/parquet-variant/src/fns/variant_to_json.rs | 9 ++++++++- 2 files changed, 16 insertions(+), 2 deletions(-) diff --git a/encodings/parquet-variant/src/fns/json_to_variant.rs b/encodings/parquet-variant/src/fns/json_to_variant.rs index 3e64fe2cf70..0e47f465dd5 100644 --- a/encodings/parquet-variant/src/fns/json_to_variant.rs +++ b/encodings/parquet-variant/src/fns/json_to_variant.rs @@ -295,14 +295,17 @@ mod tests { use vortex_array::arrays::VarBinViewArray; use vortex_array::arrays::struct_::StructArrayExt; use vortex_array::arrays::variant::VariantArrayExt; + use vortex_array::arrow::ArrowSession; use vortex_array::assert_arrays_eq; use vortex_array::assert_nth_scalar_is_null; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; + use vortex_array::dtype::session::DTypeSession; use vortex_array::expr::Expression; use vortex_array::expr::proto::ExprSerializeProtoExt; use vortex_array::expr::root; use vortex_array::expr::variant_get; + use vortex_array::scalar_fn::session::ScalarFnSession; use vortex_array::session::ArraySession; use vortex_error::vortex_bail; @@ -311,7 +314,11 @@ mod tests { use crate::fns::json_to_variant; static SESSION: LazyLock = LazyLock::new(|| { - let session = VortexSession::empty().with::(); + let session = VortexSession::empty() + .with::() + .with::() + .with::() + .with::(); crate::initialize(&session); session }); diff --git a/encodings/parquet-variant/src/fns/variant_to_json.rs b/encodings/parquet-variant/src/fns/variant_to_json.rs index bd586188cea..2ed09239565 100644 --- a/encodings/parquet-variant/src/fns/variant_to_json.rs +++ b/encodings/parquet-variant/src/fns/variant_to_json.rs @@ -175,11 +175,14 @@ mod tests { use vortex_array::arrays::Variant; use vortex_array::arrays::extension::ExtensionArrayExt; use vortex_array::arrays::variant::VariantArrayExt; + use vortex_array::arrow::ArrowSession; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; + use vortex_array::dtype::session::DTypeSession; use vortex_array::expr::root; use vortex_array::scalar_fn::fns::variant_get::VariantPath; use vortex_array::scalar_fn::fns::variant_get::VariantPathElement; + use vortex_array::scalar_fn::session::ScalarFnSession; use vortex_array::session::ArraySession; use vortex_array::validity::Validity; use vortex_error::vortex_bail; @@ -191,7 +194,11 @@ mod tests { use crate::fns::variant_to_json; static SESSION: LazyLock = LazyLock::new(|| { - let session = VortexSession::empty().with::(); + let session = VortexSession::empty() + .with::() + .with::() + .with::() + .with::(); crate::initialize(&session); session }); From cdb8a94e22c8eba25558712c21fce80809ef7508 Mon Sep 17 00:00:00 2001 From: Adam Gutglick Date: Fri, 19 Jun 2026 10:28:11 +0100 Subject: [PATCH 6/9] thing Signed-off-by: Adam Gutglick --- encodings/parquet-variant/src/arrow.rs | 2 +- encodings/parquet-variant/src/fns/variant_to_json.rs | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/encodings/parquet-variant/src/arrow.rs b/encodings/parquet-variant/src/arrow.rs index bb409fd390d..da4e648a138 100644 --- a/encodings/parquet-variant/src/arrow.rs +++ b/encodings/parquet-variant/src/arrow.rs @@ -259,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/fns/variant_to_json.rs b/encodings/parquet-variant/src/fns/variant_to_json.rs index 2ed09239565..c8995f80b4d 100644 --- a/encodings/parquet-variant/src/fns/variant_to_json.rs +++ b/encodings/parquet-variant/src/fns/variant_to_json.rs @@ -112,7 +112,7 @@ impl ScalarFnVTable for VariantToJson { fn return_dtype(&self, _options: &Self::Options, arg_dtypes: &[DType]) -> VortexResult { let input_dtype = &arg_dtypes[0]; vortex_ensure!( - matches!(input_dtype, DType::Variant(_)), + input_dtype.is_variant(), "VariantToJson input must be Variant, found {input_dtype}" ); From 8cb09a918329fe19910c9becd2db79a5c9fcc2fb Mon Sep 17 00:00:00 2001 From: Adam Gutglick Date: Tue, 23 Jun 2026 11:02:43 +0100 Subject: [PATCH 7/9] second pass Signed-off-by: Adam Gutglick --- Cargo.lock | 2 + .../src/fns/json_to_variant.rs | 570 ------------------ encodings/parquet-variant/src/fns/mod.rs | 40 -- .../src/fns/variant_to_json.rs | 503 ---------------- .../src/json_to_variant_tests.rs | 283 +++++++++ encodings/parquet-variant/src/kernel.rs | 96 +++ encodings/parquet-variant/src/lib.rs | 20 +- vortex-json/Cargo.toml | 7 + vortex-json/src/json_to_variant.rs | 395 ++++++++++++ vortex-json/src/lib.rs | 7 + 10 files changed, 796 insertions(+), 1127 deletions(-) delete mode 100644 encodings/parquet-variant/src/fns/json_to_variant.rs delete mode 100644 encodings/parquet-variant/src/fns/mod.rs delete mode 100644 encodings/parquet-variant/src/fns/variant_to_json.rs create mode 100644 encodings/parquet-variant/src/json_to_variant_tests.rs create mode 100644 vortex-json/src/json_to_variant.rs diff --git a/Cargo.lock b/Cargo.lock index e6fc779f3b7..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", ] diff --git a/encodings/parquet-variant/src/fns/json_to_variant.rs b/encodings/parquet-variant/src/fns/json_to_variant.rs deleted file mode 100644 index 0e47f465dd5..00000000000 --- a/encodings/parquet-variant/src/fns/json_to_variant.rs +++ /dev/null @@ -1,570 +0,0 @@ -// 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 std::sync::Arc; - -use arrow_schema::FieldRef; -use parquet_variant_compute::ShreddedSchemaBuilder; -use parquet_variant_compute::shred_variant; -use prost::Message; -use vortex_array::ArrayRef; -use vortex_array::ExecutionCtx; -use vortex_array::arrays::ExtensionArray; -use vortex_array::arrays::extension::ExtensionArrayExt; -use vortex_array::arrow::ArrowSessionExt; -use vortex_array::dtype::DType; -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::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_json::Json; -use vortex_proto::expr as pb; -use vortex_session::VortexSession; -use vortex_session::registry::CachedId; - -use crate::ParquetVariant; -use crate::kernel::to_parquet_variant_path; - -/// Parses JSON strings into Parquet Variant values, optionally shredding fields. -/// -/// Accepts `Utf8` inputs or [`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 -/// `typed_value` child, following the [Parquet Variant shredding] rules: rows whose value does -/// not match the requested type stay readable through the residual variant value. -/// -/// # Not an inverse of `variant_to_json` -/// -/// `json_to_variant` and [`VariantToJson`](crate::VariantToJson) are lossy, normalizing -/// conversions and are NOT inverses of each other: -/// - JSON whitespace is not preserved. -/// - Object keys may be reordered: Variant metadata stores keys in sorted order, so -/// `variant_to_json` emits fields in a canonical order, not source order. -/// - Number formatting and precision change: e.g. `1.0` may render as `1`, exponent forms and -/// very large numbers are re-rendered, and floating-point values are re-encoded. -/// - Duplicate object keys are collapsed to a single entry. -/// - Unicode escape sequences are normalized (e.g. `\u0041` becomes `A`). -/// - `variant_to_json` stringifies Variant-only types — date, timestamp, UUID, binary, -/// decimal — so `json_to_variant(variant_to_json(v))` yields plain strings/numbers and loses -/// the original type information. -/// - Shredding structure is not recoverable from JSON: `variant_to_json` unshreds its input -/// first, and re-parsing produces an unshredded Variant unless a new [`ShreddingSpec`] is -/// supplied. -/// -/// [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.is_utf8() - || input_dtype - .as_extension_opt() - .is_some_and(|ext_dtype| ext_dtype.is::()), - "JsonToVariant input must be Utf8 or 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)?; - let input_nullable = input.dtype().is_nullable(); - - let storage = if input.dtype().is_utf8() { - input - } else if input - .dtype() - .as_extension_opt() - .is_some_and(|ext_dtype| ext_dtype.is::()) - { - input - .execute::(ctx)? - .storage_array() - .clone() - } else { - vortex_bail!( - "JsonToVariant input must be Utf8 or a Json extension, found {}", - input.dtype() - ); - }; - - let session = ctx.session().clone(); - let arrow_strings = session.arrow().execute_arrow(storage, 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 options.shredding().is_empty() { - arrow_variant - } else { - let mut builder = ShreddedSchemaBuilder::new(); - for (path, dtype) in 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 input_nullable { - ParquetVariant::from_arrow_variant_nullable(&arrow_variant) - } else { - ParquetVariant::from_arrow_variant(&arrow_variant) - } - } -} - -/// 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 { - use std::sync::LazyLock; - - use vortex_array::Canonical; - use vortex_array::EmptyMetadata; - use vortex_array::IntoArray; - use vortex_array::VortexSessionExecute; - 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::arrow::ArrowSession; - use vortex_array::assert_arrays_eq; - use vortex_array::assert_nth_scalar_is_null; - use vortex_array::dtype::Nullability; - use vortex_array::dtype::PType; - use vortex_array::dtype::session::DTypeSession; - use vortex_array::expr::Expression; - use vortex_array::expr::proto::ExprSerializeProtoExt; - use vortex_array::expr::root; - use vortex_array::expr::variant_get; - use vortex_array::scalar_fn::session::ScalarFnSession; - use vortex_array::session::ArraySession; - use vortex_error::vortex_bail; - - use super::*; - use crate::ParquetVariantArrayExt; - use crate::fns::json_to_variant; - - static SESSION: LazyLock = LazyLock::new(|| { - let session = VortexSession::empty() - .with::() - .with::() - .with::() - .with::(); - 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 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 expression_roundtrip_serialization() -> VortexResult<()> { - let expr: Expression = json_to_variant(root(), shred_field_as_i64("a")?); - let proto = expr.serialize_proto()?; - let actual = Expression::from_proto(&proto, &SESSION)?; - - assert_eq!(actual, expr); - Ok(()) - } - - #[test] - fn converts_utf8_json_rows() -> VortexResult<()> { - let 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 = ExtensionArray::try_new_from_vtable(Json, EmptyMetadata, storage)?.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)]) - } - - #[test] - fn null_rows_stay_null_and_json_null_becomes_variant_null() -> VortexResult<()> { - let 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); - 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() { - let 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()); - } - - #[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 shredding_produces_typed_value_child() -> VortexResult<()> { - let 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]) - ); - - // 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 = 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)); - assert_nth_scalar_is_null!(result, 1); - let mut ctx = SESSION.create_execution_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)]) - ); - Ok(()) - } - - #[test] - fn shredding_root_path_shreds_top_level_values() -> VortexResult<()> { - let 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/fns/mod.rs b/encodings/parquet-variant/src/fns/mod.rs deleted file mode 100644 index b1958a24000..00000000000 --- a/encodings/parquet-variant/src/fns/mod.rs +++ /dev/null @@ -1,40 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright the Vortex contributors - -//! Scalar functions converting between JSON strings and Variant values. - -mod json_to_variant; -mod variant_to_json; - -pub use json_to_variant::JsonToVariant; -pub use json_to_variant::JsonToVariantOptions; -pub use json_to_variant::ShreddingSpec; -pub use variant_to_json::VariantToJson; -use vortex_array::expr::Expression; -use vortex_array::scalar_fn::EmptyOptions; -use vortex_array::scalar_fn::ScalarFnVTableExt; - -/// Creates a [`JsonToVariant`] expression that parses `child`'s JSON strings into Variant -/// values, shredding the paths selected by `shredding`. -/// -/// `child` must produce `Utf8` or [`Json`](vortex_json::Json) extension values; the result is -/// `Variant` with the input's nullability. Rows containing invalid JSON fail the expression. -/// -/// Note that this is NOT an inverse of [`variant_to_json()`]: both conversions normalize their -/// input. 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]) -} - -/// Creates a [`VariantToJson`] expression that renders `child`'s Variant values as JSON -/// strings with the [`Json`](vortex_json::Json) extension dtype. -/// -/// Shredded inputs are unshredded before rendering, and the result keeps the input's -/// nullability. -/// -/// Note that this is NOT an inverse of [`json_to_variant()`]: both conversions normalize their -/// input, and Variant-only types (dates, timestamps, UUIDs, binary, decimals) are rendered as -/// plain JSON strings or numbers. See [`VariantToJson`] for the full list of caveats. -pub fn variant_to_json(child: Expression) -> Expression { - VariantToJson.new_expr(EmptyOptions, [child]) -} diff --git a/encodings/parquet-variant/src/fns/variant_to_json.rs b/encodings/parquet-variant/src/fns/variant_to_json.rs deleted file mode 100644 index c8995f80b4d..00000000000 --- a/encodings/parquet-variant/src/fns/variant_to_json.rs +++ /dev/null @@ -1,503 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright the Vortex contributors - -//! The `vortex.variant_to_json` scalar function. - -use std::fmt; -use std::fmt::Formatter; -use std::sync::Arc; - -use arrow_schema::DataType; -use arrow_schema::Field; -use arrow_schema::Fields; -use vortex_array::ArrayRef; -use vortex_array::EmptyMetadata; -use vortex_array::ExecutionCtx; -use vortex_array::IntoArray; -use vortex_array::arrays::ExtensionArray; -use vortex_array::arrow::FromArrowArray; -use vortex_array::dtype::DType; -use vortex_array::dtype::extension::ExtDType; -use vortex_array::expr::Expression; -use vortex_array::scalar_fn::Arity; -use vortex_array::scalar_fn::ChildName; -use vortex_array::scalar_fn::EmptyOptions; -use vortex_array::scalar_fn::ExecutionArgs; -use vortex_array::scalar_fn::ScalarFnId; -use vortex_array::scalar_fn::ScalarFnVTable; -use vortex_error::VortexResult; -use vortex_error::vortex_ensure; -use vortex_json::Json; -use vortex_session::VortexSession; -use vortex_session::registry::CachedId; - -use crate::ParquetVariant; -use crate::ParquetVariantArrayExt; -use crate::arrow::export_storage_to_target; -use crate::arrow::export_unshredded_storage_to_target; -use crate::arrow::parquet_variant_for_export; - -/// Renders Variant values as JSON strings with the [`Json`] extension dtype. -/// -/// Accepts `Variant` inputs backed by Parquet Variant storage, including shredded storage -/// (top-level, object, and nested fields), which is unshredded before rendering. Null rows stay -/// null, while variant-null values render as the JSON literal `null`. The output nullability -/// matches the input's. -/// -/// Inputs are exported through their Parquet Variant storage, so a `Variant` whose core storage -/// is not Parquet Variant-backed is not supported and fails. -/// -/// # Not an inverse of `json_to_variant` -/// -/// [`JsonToVariant`](crate::JsonToVariant) and `variant_to_json` are lossy, normalizing -/// conversions and are NOT inverses of each other: -/// - JSON whitespace is not preserved. -/// - Object keys may be reordered: Variant metadata stores keys in sorted order, so -/// `variant_to_json` emits fields in a canonical order, not source order. -/// - Number formatting and precision change: e.g. `1.0` may render as `1`, exponent forms and -/// very large numbers are re-rendered, and floating-point values are re-encoded. -/// - Duplicate object keys are collapsed to a single entry. -/// - Unicode escape sequences are normalized (e.g. `\u0041` becomes `A`). -/// - `variant_to_json` stringifies Variant-only types — date, timestamp, UUID, binary, -/// decimal — so `json_to_variant(variant_to_json(v))` yields plain strings/numbers and loses -/// the original type information. -/// - Shredding structure is not recoverable from JSON: `variant_to_json` unshreds its input -/// first, and re-parsing produces an unshredded Variant unless a new -/// [`ShreddingSpec`](crate::ShreddingSpec) is supplied. -#[derive(Clone)] -pub struct VariantToJson; - -impl ScalarFnVTable for VariantToJson { - type Options = EmptyOptions; - - fn id(&self) -> ScalarFnId { - static ID: CachedId = CachedId::new("vortex.variant_to_json"); - *ID - } - - fn serialize(&self, _options: &Self::Options) -> VortexResult>> { - Ok(Some(vec![])) - } - - fn deserialize( - &self, - _metadata: &[u8], - _session: &VortexSession, - ) -> VortexResult { - Ok(EmptyOptions) - } - - 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 VariantToJson expression"), - } - } - - fn fmt_sql( - &self, - _options: &Self::Options, - expr: &Expression, - f: &mut Formatter<'_>, - ) -> fmt::Result { - write!(f, "variant_to_json(")?; - expr.child(0).fmt_sql(f)?; - write!(f, ")") - } - - fn return_dtype(&self, _options: &Self::Options, arg_dtypes: &[DType]) -> VortexResult { - let input_dtype = &arg_dtypes[0]; - vortex_ensure!( - input_dtype.is_variant(), - "VariantToJson input must be Variant, found {input_dtype}" - ); - - let storage_dtype = DType::Utf8(input_dtype.nullability()); - Ok(DType::Extension( - ExtDType::::try_new(EmptyMetadata, storage_dtype)?.erased(), - )) - } - - fn execute( - &self, - _options: &Self::Options, - args: &dyn ExecutionArgs, - ctx: &mut ExecutionCtx, - ) -> VortexResult { - let input = args.get(0)?; - let input_nullable = input.dtype().is_nullable(); - - let parquet_array = parquet_variant_for_export(input, ctx)?; - let parquet_array = parquet_array.as_::(); - - // `parquet_variant_compute::variant_to_json` only accepts unshredded - // `STRUCT` storage, so request exactly that shape and - // unshred any typed values first. - let target_fields: Fields = vec![ - Arc::new(Field::new("metadata", DataType::Binary, false)), - Arc::new(Field::new("value", DataType::Binary, true)), - ] - .into(); - let arrow_storage = if parquet_array.typed_value_array().is_some() { - export_unshredded_storage_to_target(&parquet_array, &target_fields, ctx)? - } else { - export_storage_to_target(&parquet_array, &target_fields, ctx)? - }; - - let arrow_json = parquet_variant_compute::variant_to_json(&arrow_storage)?; - let storage = ArrayRef::from_arrow(&arrow_json, input_nullable)?; - - ExtensionArray::try_new_from_vtable(Json, EmptyMetadata, storage).map(IntoArray::into_array) - } - - fn is_null_sensitive(&self, _options: &Self::Options) -> bool { - false - } -} - -#[cfg(test)] -mod tests { - use std::sync::LazyLock; - - use chrono::NaiveDate; - use parquet_variant::Variant as PqVariant; - use parquet_variant::VariantBuilder; - use parquet_variant_compute::VariantArrayBuilder; - use vortex_array::Canonical; - use vortex_array::VortexSessionExecute; - use vortex_array::accessor::ArrayAccessor; - use vortex_array::arrays::PrimitiveArray; - use vortex_array::arrays::VarBinViewArray; - use vortex_array::arrays::Variant; - use vortex_array::arrays::extension::ExtensionArrayExt; - use vortex_array::arrays::variant::VariantArrayExt; - use vortex_array::arrow::ArrowSession; - use vortex_array::dtype::Nullability; - use vortex_array::dtype::PType; - use vortex_array::dtype::session::DTypeSession; - use vortex_array::expr::root; - use vortex_array::scalar_fn::fns::variant_get::VariantPath; - use vortex_array::scalar_fn::fns::variant_get::VariantPathElement; - use vortex_array::scalar_fn::session::ScalarFnSession; - use vortex_array::session::ArraySession; - use vortex_array::validity::Validity; - use vortex_error::vortex_bail; - use vortex_error::vortex_err; - - use super::*; - use crate::ShreddingSpec; - use crate::fns::json_to_variant; - use crate::fns::variant_to_json; - - static SESSION: LazyLock = LazyLock::new(|| { - let session = VortexSession::empty() - .with::() - .with::() - .with::() - .with::(); - crate::initialize(&session); - session - }); - - fn json_dtype(nullability: Nullability) -> VortexResult { - Ok(DType::Extension( - ExtDType::::try_new(EmptyMetadata, DType::Utf8(nullability))?.erased(), - )) - } - - fn json_strings(array: &ArrayRef) -> VortexResult>> { - let mut ctx = SESSION.create_execution_ctx(); - let ext = array.clone().execute::(&mut ctx)?; - let storage = ext - .storage_array() - .clone() - .execute::(&mut ctx)?; - Ok(storage.with_iterator(|iter| { - iter.map(|value| value.map(|bytes| String::from_utf8_lossy(bytes).into_owned())) - .collect() - })) - } - - fn unshredded_variant( - rows: impl IntoIterator>, - ) -> VortexResult { - let rows = rows.into_iter().collect::>(); - let mut builder = VariantArrayBuilder::new(rows.len()); - for row in rows { - builder.append_variant(row); - } - ParquetVariant::from_arrow_variant(&builder.build()) - } - - fn json_rows_to_variant( - rows: Vec>, - shredding: ShreddingSpec, - ) -> VortexResult { - let input = VarBinViewArray::from_iter_nullable_str(rows).into_array(); - input - .apply(&json_to_variant(root(), shredding))? - .execute::(&mut SESSION.create_execution_ctx()) - } - - #[test] - fn renders_unshredded_values() -> VortexResult<()> { - let input = unshredded_variant([ - PqVariant::from(42i32), - PqVariant::from("hello"), - PqVariant::from(true), - PqVariant::Null, - ])?; - - let result = input - .apply(&variant_to_json(root()))? - .execute::(&mut SESSION.create_execution_ctx())?; - - assert_eq!(result.dtype(), &json_dtype(Nullability::NonNullable)?); - assert_eq!( - json_strings(&result)?, - vec![ - Some("42".to_string()), - Some(r#""hello""#.to_string()), - Some("true".to_string()), - Some("null".to_string()), - ] - ); - Ok(()) - } - - #[test] - fn null_rows_stay_null_and_variant_null_renders_as_json_null() -> VortexResult<()> { - let input = - json_rows_to_variant(vec![Some("1"), None, Some("null")], ShreddingSpec::empty())?; - - let result = input - .apply(&variant_to_json(root()))? - .execute::(&mut SESSION.create_execution_ctx())?; - - assert_eq!(result.dtype(), &json_dtype(Nullability::Nullable)?); - assert_eq!( - json_strings(&result)?, - vec![Some("1".to_string()), None, Some("null".to_string())] - ); - Ok(()) - } - - fn typed_value_only_variant() -> VortexResult { - let rows = [ - VariantBuilder::new().with_value(10i32).finish(), - VariantBuilder::new().with_value(20i32).finish(), - VariantBuilder::new().with_value(30i32).finish(), - ]; - let metadata = - VarBinViewArray::from_iter_bin(rows.iter().map(|(metadata, _)| metadata.as_slice())) - .into_array(); - let typed_value = PrimitiveArray::from_iter([10i32, 20, 30]).into_array(); - Ok( - ParquetVariant::try_new(Validity::NonNullable, metadata, None, Some(typed_value))? - .into_array(), - ) - } - - #[test] - fn unshreds_typed_value_only_storage() -> VortexResult<()> { - let result = { - let input = typed_value_only_variant()?; - input - .apply(&variant_to_json(root()))? - .execute::(&mut SESSION.create_execution_ctx()) - }?; - - assert_eq!( - json_strings(&result)?, - vec![ - Some("10".to_string()), - Some("20".to_string()), - Some("30".to_string()), - ] - ); - Ok(()) - } - - #[test] - fn unshreds_partially_shredded_storage() -> VortexResult<()> { - let spec = ShreddingSpec::try_new([( - VariantPath::field("a"), - DType::Primitive(PType::I64, Nullability::Nullable), - )])?; - let input = json_rows_to_variant( - vec![ - Some(r#"{"a": 1, "b": "x"}"#), - Some(r#"{"a": "not-a-number", "b": "y"}"#), - Some(r#"{"b": "z"}"#), - ], - spec, - )?; - assert!( - input.as_::().typed_value_array().is_some(), - "fixture must be shredded" - ); - - let result = input - .apply(&variant_to_json(root()))? - .execute::(&mut SESSION.create_execution_ctx())?; - - assert_eq!( - json_strings(&result)?, - vec![ - Some(r#"{"a":1,"b":"x"}"#.to_string()), - Some(r#"{"a":"not-a-number","b":"y"}"#.to_string()), - Some(r#"{"b":"z"}"#.to_string()), - ] - ); - Ok(()) - } - - /// Shreds `rows` per `spec`, then canonicalizes so the typed values are lifted into a logical - /// shredded child (as a file read-back would produce). - fn canonical_shredded(rows: Vec>, spec: ShreddingSpec) -> VortexResult { - let shredded = json_rows_to_variant(rows, spec)?; - let Canonical::Variant(canonical) = - shredded.execute::(&mut SESSION.create_execution_ctx())? - else { - vortex_bail!("expected canonical variant"); - }; - Ok(canonical.into_array()) - } - - /// Rendering a canonicalized shredded variant must match rendering the same data unshredded: - /// the shredding is a storage optimization and is invisible to JSON output. - fn assert_canonical_matches_unshredded( - rows: Vec>, - spec: ShreddingSpec, - ) -> VortexResult<()> { - let unshredded = json_rows_to_variant(rows.clone(), ShreddingSpec::empty())?; - let want = json_strings( - &unshredded - .apply(&variant_to_json(root()))? - .execute::(&mut SESSION.create_execution_ctx())?, - )?; - - let canonical = canonical_shredded(rows, spec)?; - assert!( - canonical.as_::().shredded().is_some(), - "fixture must carry a canonical shredded child" - ); - let got = json_strings( - &canonical - .apply(&variant_to_json(root()))? - .execute::(&mut SESSION.create_execution_ctx())?, - )?; - - assert_eq!(got, want); - Ok(()) - } - - /// Regression for the canonicalization bug: object-field shredding lifted into a logical - /// shredded child must round-trip back to the full object, including the shredded field and - /// the non-conforming/missing-field fallbacks. - #[test] - fn renders_canonical_object_shredded_variant() -> VortexResult<()> { - let spec = ShreddingSpec::try_new([( - VariantPath::field("a"), - DType::Primitive(PType::I64, Nullability::Nullable), - )])?; - let canonical = canonical_shredded( - vec![ - Some(r#"{"a": 1, "b": "x"}"#), - Some(r#"{"a": "not-a-number", "b": "y"}"#), - Some(r#"{"b": "z"}"#), - ], - spec, - )?; - assert!( - canonical.as_::().shredded().is_some(), - "fixture must carry a canonical shredded child" - ); - - let result = canonical - .apply(&variant_to_json(root()))? - .execute::(&mut SESSION.create_execution_ctx())?; - - assert_eq!( - json_strings(&result)?, - vec![ - Some(r#"{"a":1,"b":"x"}"#.to_string()), - Some(r#"{"a":"not-a-number","b":"y"}"#.to_string()), - Some(r#"{"b":"z"}"#.to_string()), - ] - ); - Ok(()) - } - - #[test] - fn canonical_object_shredding_matches_unshredded() -> VortexResult<()> { - assert_canonical_matches_unshredded( - vec![ - Some(r#"{"a": 1, "b": "x", "c": 3}"#), - Some(r#"{"a": 2, "b": "y", "c": 4}"#), - None, - ], - ShreddingSpec::try_new([ - ( - VariantPath::field("a"), - DType::Primitive(PType::I64, Nullability::Nullable), - ), - ( - VariantPath::field("c"), - DType::Primitive(PType::I64, Nullability::Nullable), - ), - ])?, - ) - } - - #[test] - fn canonical_nested_dotted_shredding_matches_unshredded() -> VortexResult<()> { - assert_canonical_matches_unshredded( - vec![ - Some(r#"{"a": {"b": 100}, "c": "keep"}"#), - Some(r#"{"a": {"b": 200}, "c": "keep2"}"#), - ], - ShreddingSpec::try_new([( - VariantPath::new([ - VariantPathElement::field("a"), - VariantPathElement::field("b"), - ]), - DType::Primitive(PType::I64, Nullability::Nullable), - )])?, - ) - } - - #[test] - fn variant_only_types_are_stringified_so_reparsing_loses_types() -> VortexResult<()> { - let date = - NaiveDate::from_ymd_opt(2026, 6, 11).ok_or_else(|| vortex_err!("invalid test date"))?; - let input = unshredded_variant([PqVariant::from(date)])?; - - let rendered = input - .apply(&variant_to_json(root()))? - .execute::(&mut SESSION.create_execution_ctx())?; - let json = json_strings(&rendered)?; - assert_eq!(json, vec![Some(r#""2026-06-11""#.to_string())]); - - // Re-parsing the rendered JSON yields a plain string variant, not a date: the type - // information is lost, demonstrating that the conversions are not inverses. - let reparsed = rendered - .apply(&json_to_variant(root(), ShreddingSpec::empty()))? - .execute::(&mut SESSION.create_execution_ctx())?; - let mut ctx = SESSION.create_execution_ctx(); - let row0 = reparsed.execute_scalar(0, &mut ctx)?; - let value = row0 - .as_variant() - .value() - .ok_or_else(|| vortex_err!("expected non-null variant"))?; - assert_eq!( - value.as_utf8().value().map(|value| value.to_string()), - Some("2026-06-11".to_string()) - ); - Ok(()) - } -} 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..7da7f9d4419 --- /dev/null +++ b/encodings/parquet-variant/src/json_to_variant_tests.rs @@ -0,0 +1,283 @@ +// 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 [`JsonToVariantKernel`](crate::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::arrow::ArrowSession; +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::dtype::session::DTypeSession; +use vortex_array::expr::root; +use vortex_array::expr::variant_get; +use vortex_array::scalar_fn::fns::variant_get::VariantPath; +use vortex_array::scalar_fn::session::ScalarFnSession; +use vortex_array::session::ArraySession; +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 = VortexSession::empty() + .with::() + .with::() + .with::() + .with::(); + 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 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 converts_utf8_json_rows() -> VortexResult<()> { + let 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 = ExtensionArray::try_new_from_vtable(Json, EmptyMetadata, storage)?.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)]) +} + +#[test] +fn dict_encoded_input_converts_each_row() -> VortexResult<()> { + // A dictionary-encoded string column exercises the dict-pushdown / canonicalization path: + // `json_to_variant` is not null-sensitive, so it pushes into the dict values (canonical + // strings) where the kernel fires; either way every row must convert correctly. + let values = 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 = + 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); + 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() { + let 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()); +} + +#[test] +fn shredding_produces_typed_value_child() -> VortexResult<()> { + let 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]) + ); + + // 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 = + 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)); + assert_nth_scalar_is_null!(result, 1); + let mut ctx = SESSION.create_execution_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)]) + ); + Ok(()) +} + +#[test] +fn shredding_root_path_shreds_top_level_values() -> VortexResult<()> { + let 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 afbf65cb26a..fd2ce9db12c 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,20 @@ 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::VarBinView; 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 +46,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 +72,12 @@ pub(crate) fn initialize(session: &VortexSession) { TakeExecuteAdaptor(ParquetVariant), ); kernels.register_execute_parent_kernel(VariantGet.id(), ParquetVariant, VariantGetKernel); + kernels.register_execute_parent_kernel(JsonToVariant.id(), VarBinView, JsonToVariantKernel); + kernels.register_execute_parent_kernel( + JsonToVariant.id(), + Extension, + JsonExtensionToVariantKernel, + ); } #[derive(Default, Debug)] @@ -106,6 +120,88 @@ impl ExecuteParentKernel for VariantGetKernel { } } +/// Performs the [`JsonToVariant`] conversion (and optional shredding) over a JSON string array. +/// +/// `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 input JSON +/// string array (any string encoding); 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 canonical `VarBinView` string input. +#[derive(Default, Debug)] +struct JsonToVariantKernel; + +impl ExecuteParentKernel for JsonToVariantKernel { + type Parent = ExactScalarFn; + + fn execute_parent( + &self, + array: ArrayView<'_, VarBinView>, + parent: ScalarFnArrayView<'_, JsonToVariant>, + child_idx: usize, + ctx: &mut ExecutionCtx, + ) -> VortexResult> { + if child_idx != 0 { + return Ok(None); + } + json_strings_to_variant(array.as_ref().clone(), parent, ctx).map(Some) + } +} + +/// Builds Parquet Variant arrays for [`JsonToVariant`] over a [`Json`] extension input. +/// +/// `JsonToVariant` also accepts `Json` extension values directly. This kernel unwraps the +/// extension's string storage and runs the same conversion as [`JsonToVariantKernel`]. 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() diff --git a/encodings/parquet-variant/src/lib.rs b/encodings/parquet-variant/src/lib.rs index 2b57dadb8e1..1bf0770333b 100644 --- a/encodings/parquet-variant/src/lib.rs +++ b/encodings/parquet-variant/src/lib.rs @@ -26,7 +26,8 @@ mod array; mod arrow; -mod fns; +#[cfg(test)] +mod json_to_variant_tests; mod kernel; mod operations; mod validity; @@ -35,30 +36,21 @@ mod vtable; use std::sync::Arc; pub use array::ParquetVariantArrayExt; -pub use fns::JsonToVariant; -pub use fns::JsonToVariantOptions; -pub use fns::ShreddingSpec; -pub use fns::VariantToJson; -pub use fns::json_to_variant; -pub use fns::variant_to_json; use vortex_array::arrow::ArrowSessionExt; -use vortex_array::scalar_fn::session::ScalarFnSessionExt; 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, Arrow extension, and scalar function support with a /// session. -/// -/// This also registers the [`Json`](vortex_json::Json) extension dtype so that sessions able -/// to execute [`VariantToJson`] can serialize and deserialize its output dtype. pub fn initialize(session: &VortexSession) { session.arrays().register(ParquetVariant); kernel::initialize(session); session.arrow().register_exporter(Arc::new(ParquetVariant)); session.arrow().register_importer(Arc::new(ParquetVariant)); - vortex_json::initialize(session); - session.scalar_fns().register(JsonToVariant); - session.scalar_fns().register(VariantToJson); } 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..f179ee4d24b --- /dev/null +++ b/vortex-json/src/json_to_variant.rs @@ -0,0 +1,395 @@ +// 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::VarBinView; +use vortex_array::arrays::VarBinViewArray; +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 `Utf8` inputs or [`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 canonical string encodings (`VarBinView` for `Utf8`, and +/// the extension encoding for a [`Json`] input). The fallback [`execute`](ScalarFnVTable::execute) +/// here only canonicalizes its input to one of those encodings 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.is_utf8() + || input_dtype + .as_extension_opt() + .is_some_and(|ext_dtype| ext_dtype.is::()), + "JsonToVariant input must be Utf8 or 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 canonical + // string encodings (`VarBinView` for `Utf8`, the extension encoding for a `Json` input). + // Reaching this fallback means no such kernel fired, so canonicalize the input to one of + // those encodings 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().is_utf8() { + if input.is::() { + return Err(no_kernel()); + } + input.execute::(ctx)?.into_array() + } else 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 Utf8 or 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 (canonical-string) 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 `Utf8` or [`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 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. Both canonical input forms are covered: a + // bare `Utf8`/`VarBinView`, and a `Json` extension. + let utf8 = VarBinViewArray::from_iter_str(["1", "2"]).into_array(); + let json_extension = ExtensionArray::try_new_from_vtable( + Json, + EmptyMetadata, + VarBinViewArray::from_iter_str(["1", "2"]).into_array(), + )? + .into_array(); + + for input in [utf8, json_extension] { + 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); } From 17007a1a1135e557fd5e2188f4679b0c3d39e44a Mon Sep 17 00:00:00 2001 From: Adam Gutglick Date: Tue, 23 Jun 2026 11:20:34 +0100 Subject: [PATCH 8/9] fix tests Signed-off-by: Adam Gutglick --- encodings/parquet-variant/src/array.rs | 2 +- .../src/json_to_variant_tests.rs | 20 +++++++------------ encodings/parquet-variant/src/lib.rs | 4 ++++ 3 files changed, 12 insertions(+), 14 deletions(-) diff --git a/encodings/parquet-variant/src/array.rs b/encodings/parquet-variant/src/array.rs index 67da0d59e1c..7e55357c78f 100644 --- a/encodings/parquet-variant/src/array.rs +++ b/encodings/parquet-variant/src/array.rs @@ -832,7 +832,7 @@ mod tests { ); // Canonicalize: the forward transform lifts `typed_value` into a logical shredded child. - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let Canonical::Variant(canonical) = original.clone().execute::(&mut ctx)? else { return Err(vortex_err!("expected canonical variant")); }; diff --git a/encodings/parquet-variant/src/json_to_variant_tests.rs b/encodings/parquet-variant/src/json_to_variant_tests.rs index 7da7f9d4419..7b00610d82c 100644 --- a/encodings/parquet-variant/src/json_to_variant_tests.rs +++ b/encodings/parquet-variant/src/json_to_variant_tests.rs @@ -21,18 +21,14 @@ 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::arrow::ArrowSession; 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::dtype::session::DTypeSession; use vortex_array::expr::root; use vortex_array::expr::variant_get; use vortex_array::scalar_fn::fns::variant_get::VariantPath; -use vortex_array::scalar_fn::session::ScalarFnSession; -use vortex_array::session::ArraySession; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_err; @@ -45,11 +41,7 @@ use crate::ShreddingSpec; use crate::json_to_variant; static SESSION: LazyLock = LazyLock::new(|| { - let session = VortexSession::empty() - .with::() - .with::() - .with::() - .with::(); + let session = vortex_array::array_session(); crate::initialize(&session); session }); @@ -158,7 +150,7 @@ fn null_rows_stay_null_and_json_null_becomes_variant_null() -> VortexResult<()> 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); + 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)); @@ -213,7 +205,8 @@ fn shredding_produces_typed_value_child() -> VortexResult<()> { .execute::(&mut ctx)?; assert_arrays_eq!( typed, - PrimitiveArray::from_option_iter([Some(1i64), Some(2), None, None]) + PrimitiveArray::from_option_iter([Some(1i64), Some(2), None, None]), + &mut ctx ); // Mismatched rows keep their original value through the variant fallback. @@ -241,8 +234,8 @@ fn shredding_preserves_null_rows() -> VortexResult<()> { let result = execute_json_to_variant(input, shred_field_as_i64("a")?)?; assert_eq!(result.dtype(), &DType::Variant(Nullability::Nullable)); - assert_nth_scalar_is_null!(result, 1); let mut ctx = SESSION.create_execution_ctx(); + assert_nth_scalar_is_null!(result, 1, &mut ctx); let typed = result .apply(&variant_get( root(), @@ -252,7 +245,8 @@ fn shredding_preserves_null_rows() -> VortexResult<()> { .execute::(&mut ctx)?; assert_arrays_eq!( typed, - PrimitiveArray::from_option_iter([Some(1i64), None, Some(3)]) + PrimitiveArray::from_option_iter([Some(1i64), None, Some(3)]), + &mut ctx ); Ok(()) } diff --git a/encodings/parquet-variant/src/lib.rs b/encodings/parquet-variant/src/lib.rs index 1bf0770333b..05e93c5d08c 100644 --- a/encodings/parquet-variant/src/lib.rs +++ b/encodings/parquet-variant/src/lib.rs @@ -48,7 +48,11 @@ pub use vtable::ParquetVariantArray; /// 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)); From 8318ee1f60d27458f75c97ea176e0cb00a615cbd Mon Sep 17 00:00:00 2001 From: Adam Gutglick Date: Tue, 23 Jun 2026 16:06:22 +0100 Subject: [PATCH 9/9] Only support direct JSON inputs for expressions Signed-off-by: Adam Gutglick --- .../src/json_to_variant_tests.rs | 68 ++++++++---- encodings/parquet-variant/src/kernel.rs | 36 ++----- vortex-json/src/json_to_variant.rs | 101 +++++++++--------- 3 files changed, 103 insertions(+), 102 deletions(-) diff --git a/encodings/parquet-variant/src/json_to_variant_tests.rs b/encodings/parquet-variant/src/json_to_variant_tests.rs index 7b00610d82c..cd3b2f035ee 100644 --- a/encodings/parquet-variant/src/json_to_variant_tests.rs +++ b/encodings/parquet-variant/src/json_to_variant_tests.rs @@ -4,8 +4,8 @@ //! 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 [`JsonToVariantKernel`](crate::kernel) registered here, so these end-to-end tests live -//! in `vortex-parquet-variant` where that kernel is registered. +//! 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; @@ -54,6 +54,10 @@ 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 @@ -82,8 +86,21 @@ fn assert_variant_i64_rows(array: &ArrayRef, expected: &[Option]) -> Vortex } #[test] -fn converts_utf8_json_rows() -> VortexResult<()> { - let input = VarBinViewArray::from_iter_str([r#"{"a": 1}"#, "2", r#"{"a": 3}"#]).into_array(); +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())?; @@ -117,7 +134,7 @@ fn converts_utf8_json_rows() -> VortexResult<()> { #[test] fn converts_json_extension_input() -> VortexResult<()> { let storage = VarBinViewArray::from_iter_str(["1", "2"]).into_array(); - let input = ExtensionArray::try_new_from_vtable(Json, EmptyMetadata, storage)?.into_array(); + let input = json_input(storage)?; let result = execute_json_to_variant(input, ShreddingSpec::empty())?; @@ -127,10 +144,10 @@ fn converts_json_extension_input() -> VortexResult<()> { #[test] fn dict_encoded_input_converts_each_row() -> VortexResult<()> { - // A dictionary-encoded string column exercises the dict-pushdown / canonicalization path: + // 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 - // strings) where the kernel fires; either way every row must convert correctly. - let values = VarBinViewArray::from_iter_str(["1", "2"]).into_array(); + // 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(); @@ -142,8 +159,9 @@ fn dict_encoded_input_converts_each_row() -> VortexResult<()> { #[test] fn null_rows_stay_null_and_json_null_becomes_variant_null() -> VortexResult<()> { - let input = - VarBinViewArray::from_iter_nullable_str([Some("1"), None, Some("null")]).into_array(); + 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())?; @@ -158,22 +176,26 @@ fn null_rows_stay_null_and_json_null_becomes_variant_null() -> VortexResult<()> } #[test] -fn invalid_json_errors() { - let input = VarBinViewArray::from_iter_str([r#"{"a": 1}"#, r#"{"a":"#]).into_array(); +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 = 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 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")?)?; @@ -227,9 +249,10 @@ fn shredding_produces_typed_value_child() -> VortexResult<()> { #[test] fn shredding_preserves_null_rows() -> VortexResult<()> { - let input = + let input = json_input( VarBinViewArray::from_iter_nullable_str([Some(r#"{"a": 1}"#), None, Some(r#"{"a": 3}"#)]) - .into_array(); + .into_array(), + )?; let result = execute_json_to_variant(input, shred_field_as_i64("a")?)?; @@ -253,7 +276,8 @@ fn shredding_preserves_null_rows() -> VortexResult<()> { #[test] fn shredding_root_path_shreds_top_level_values() -> VortexResult<()> { - let input = VarBinViewArray::from_iter_str(["1", "2", r#""not-a-number""#]).into_array(); + 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)?; diff --git a/encodings/parquet-variant/src/kernel.rs b/encodings/parquet-variant/src/kernel.rs index fd2ce9db12c..88a4b8931ab 100644 --- a/encodings/parquet-variant/src/kernel.rs +++ b/encodings/parquet-variant/src/kernel.rs @@ -24,7 +24,6 @@ use vortex_array::arrays::Dict; use vortex_array::arrays::Extension; use vortex_array::arrays::Filter; use vortex_array::arrays::Slice; -use vortex_array::arrays::VarBinView; use vortex_array::arrays::dict::TakeExecute; use vortex_array::arrays::dict::TakeExecuteAdaptor; use vortex_array::arrays::extension::ExtensionArrayExt; @@ -72,7 +71,6 @@ pub(crate) fn initialize(session: &VortexSession) { TakeExecuteAdaptor(ParquetVariant), ); kernels.register_execute_parent_kernel(VariantGet.id(), ParquetVariant, VariantGetKernel); - kernels.register_execute_parent_kernel(JsonToVariant.id(), VarBinView, JsonToVariantKernel); kernels.register_execute_parent_kernel( JsonToVariant.id(), Extension, @@ -120,13 +118,13 @@ impl ExecuteParentKernel for VariantGetKernel { } } -/// Performs the [`JsonToVariant`] conversion (and optional shredding) over a JSON string array. +/// 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 input JSON -/// string array (any string encoding); it is executed to Arrow and parsed. Nullability of the -/// result follows `parent.dtype()`, which equals the input's nullability. +/// `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>, @@ -156,32 +154,10 @@ fn json_strings_to_variant( } } -/// Builds Parquet Variant arrays for [`JsonToVariant`] over canonical `VarBinView` string input. -#[derive(Default, Debug)] -struct JsonToVariantKernel; - -impl ExecuteParentKernel for JsonToVariantKernel { - type Parent = ExactScalarFn; - - fn execute_parent( - &self, - array: ArrayView<'_, VarBinView>, - parent: ScalarFnArrayView<'_, JsonToVariant>, - child_idx: usize, - ctx: &mut ExecutionCtx, - ) -> VortexResult> { - if child_idx != 0 { - return Ok(None); - } - json_strings_to_variant(array.as_ref().clone(), parent, ctx).map(Some) - } -} - /// Builds Parquet Variant arrays for [`JsonToVariant`] over a [`Json`] extension input. /// -/// `JsonToVariant` also accepts `Json` extension values directly. This kernel unwraps the -/// extension's string storage and runs the same conversion as [`JsonToVariantKernel`]. It is keyed -/// on the shared extension encoding, so it declines any non-`Json` extension. +/// 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; diff --git a/vortex-json/src/json_to_variant.rs b/vortex-json/src/json_to_variant.rs index f179ee4d24b..d0febf4e646 100644 --- a/vortex-json/src/json_to_variant.rs +++ b/vortex-json/src/json_to_variant.rs @@ -13,8 +13,6 @@ use vortex_array::ExecutionCtx; use vortex_array::IntoArray; use vortex_array::arrays::Extension; use vortex_array::arrays::ExtensionArray; -use vortex_array::arrays::VarBinView; -use vortex_array::arrays::VarBinViewArray; use vortex_array::arrays::scalar_fn::ScalarFnFactoryExt; use vortex_array::dtype::DType; use vortex_array::expr::Expression; @@ -38,9 +36,9 @@ use crate::Json; /// Parses JSON strings into Variant values, optionally shredding fields. /// -/// Accepts `Utf8` inputs or [`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. +/// 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 @@ -50,10 +48,10 @@ use crate::Json; /// /// 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 canonical string encodings (`VarBinView` for `Utf8`, and -/// the extension encoding for a [`Json`] input). The fallback [`execute`](ScalarFnVTable::execute) -/// here only canonicalizes its input to one of those encodings and re-dispatches so that kernel -/// runs; it errors if no Variant encoding is registered with the session. +/// `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 /// @@ -135,11 +133,10 @@ impl ScalarFnVTable for JsonToVariant { fn return_dtype(&self, _options: &Self::Options, arg_dtypes: &[DType]) -> VortexResult { let input_dtype = &arg_dtypes[0]; vortex_ensure!( - input_dtype.is_utf8() - || input_dtype - .as_extension_opt() - .is_some_and(|ext_dtype| ext_dtype.is::()), - "JsonToVariant input must be Utf8 or a Json extension, found {input_dtype}" + 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())) @@ -154,11 +151,11 @@ impl ScalarFnVTable for JsonToVariant { 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 canonical - // string encodings (`VarBinView` for `Utf8`, the extension encoding for a `Json` input). - // Reaching this fallback means no such kernel fired, so canonicalize the input to one of - // those encodings and re-dispatch. If the input is already canonical here, no kernel is - // registered for it; bail with a clear error rather than looping. + // 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 \ @@ -166,12 +163,7 @@ impl ScalarFnVTable for JsonToVariant { ) }; - let canonical = if input.dtype().is_utf8() { - if input.is::() { - return Err(no_kernel()); - } - input.execute::(ctx)?.into_array() - } else if input + let canonical = if input .dtype() .as_extension_opt() .is_some_and(|ext_dtype| ext_dtype.is::()) @@ -182,7 +174,7 @@ impl ScalarFnVTable for JsonToVariant { input.execute::(ctx)?.into_array() } else { vortex_bail!( - "JsonToVariant input must be Utf8 or a Json extension, found {}", + "JsonToVariant input must be a Json extension, found {}", input.dtype() ); }; @@ -195,7 +187,7 @@ impl ScalarFnVTable for JsonToVariant { // `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 (canonical-string) values, where the kernel fires directly. + // into their JSON extension values, where the kernel fires directly. false } } @@ -203,8 +195,8 @@ impl ScalarFnVTable for JsonToVariant { /// Creates a [`JsonToVariant`] expression that parses `child`'s JSON strings into Variant /// values, shredding the paths selected by `shredding`. /// -/// `child` must produce `Utf8` or [`Json`] extension values; the result is `Variant` with the -/// input's nullability. Rows containing invalid JSON fail the expression. +/// `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. @@ -360,36 +352,45 @@ mod tests { 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. Both canonical input forms are covered: a - // bare `Utf8`/`VarBinView`, and a `Json` extension. - let utf8 = VarBinViewArray::from_iter_str(["1", "2"]).into_array(); - let json_extension = ExtensionArray::try_new_from_vtable( + // 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(); - for input in [utf8, json_extension] { - 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}" - ); - } + 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(()) } }