vortex_json/
json_to_variant.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4//! The `vortex.json_to_variant` scalar function.
5
6use std::fmt;
7use std::fmt::Display;
8use std::fmt::Formatter;
9
10use prost::Message;
11use vortex_array::ArrayRef;
12use vortex_array::ExecutionCtx;
13use vortex_array::IntoArray;
14use vortex_array::arrays::Extension;
15use vortex_array::arrays::ExtensionArray;
16use vortex_array::arrays::scalar_fn::ScalarFnFactoryExt;
17use vortex_array::dtype::DType;
18use vortex_array::expr::Expression;
19use vortex_array::scalar_fn::Arity;
20use vortex_array::scalar_fn::ChildName;
21use vortex_array::scalar_fn::ExecutionArgs;
22use vortex_array::scalar_fn::ScalarFnId;
23use vortex_array::scalar_fn::ScalarFnVTable;
24use vortex_array::scalar_fn::ScalarFnVTableExt;
25use vortex_array::scalar_fn::fns::variant_get::VariantPath;
26use vortex_array::scalar_fn::fns::variant_get::VariantPathElement;
27use vortex_error::VortexResult;
28use vortex_error::vortex_bail;
29use vortex_error::vortex_ensure;
30use vortex_error::vortex_err;
31use vortex_proto::expr as pb;
32use vortex_session::VortexSession;
33use vortex_session::registry::CachedId;
34
35use crate::Json;
36
37/// Parses JSON strings into Variant values, optionally shredding fields.
38///
39/// Accepts [`Json`] extension inputs and returns `Variant` values with the input's nullability.
40/// Null rows stay null, the JSON literal `null` becomes a variant-null value, and any row that
41/// fails to parse as JSON fails the whole expression.
42///
43/// A non-empty [`ShreddingSpec`] additionally shreds the selected paths into a typed shredded
44/// child, following the [Parquet Variant shredding] rules: rows whose value does not match the
45/// requested type stay readable through the residual variant value.
46///
47/// # Execution
48///
49/// Building a Variant requires a concrete Variant encoding, so this function does not perform the
50/// conversion itself. The Variant encoding registered with the session supplies it as an
51/// `execute_parent` kernel keyed on the extension encoding for a [`Json`] input. The fallback
52/// [`execute`](ScalarFnVTable::execute) here only canonicalizes the input to that encoding and
53/// re-dispatches so that kernel runs; it errors if no Variant encoding is registered with the
54/// session.
55///
56/// # Normalization
57///
58/// `json_to_variant` is a lossy, normalizing conversion: the parsed Variant does not round-trip
59/// back to the exact source JSON text.
60/// - JSON whitespace is not preserved.
61/// - Object keys are stored in Variant metadata in sorted order, not source order.
62/// - Number representations are normalized (e.g. `1.0` and `1` may parse to the same value;
63///   exponent forms and very large numbers are re-encoded).
64/// - Duplicate object keys are collapsed to a single entry.
65/// - Unicode escape sequences are normalized (e.g. `A` becomes `A`).
66///
67/// [Parquet Variant shredding]: https://github.com/apache/parquet-format/blob/master/VariantShredding.md
68#[derive(Clone)]
69pub struct JsonToVariant;
70
71impl ScalarFnVTable for JsonToVariant {
72    type Options = JsonToVariantOptions;
73
74    fn id(&self) -> ScalarFnId {
75        static ID: CachedId = CachedId::new("vortex.json_to_variant");
76        *ID
77    }
78
79    fn serialize(&self, options: &Self::Options) -> VortexResult<Option<Vec<u8>>> {
80        let shredding = options
81            .shredding()
82            .fields()
83            .iter()
84            .map(|(path, dtype)| {
85                Ok(pb::ShreddingSpecField {
86                    path: path
87                        .elements()
88                        .iter()
89                        .map(VariantPathElement::to_proto)
90                        .collect(),
91                    dtype: Some(dtype.try_into()?),
92                })
93            })
94            .collect::<VortexResult<_>>()?;
95
96        Ok(Some(pb::JsonToVariantOpts { shredding }.encode_to_vec()))
97    }
98
99    fn deserialize(&self, metadata: &[u8], session: &VortexSession) -> VortexResult<Self::Options> {
100        let opts = pb::JsonToVariantOpts::decode(metadata)?;
101        let fields = opts
102            .shredding
103            .into_iter()
104            .map(|field| {
105                let path = field
106                    .path
107                    .into_iter()
108                    .map(VariantPathElement::from_proto)
109                    .collect::<VortexResult<VariantPath>>()?;
110                let dtype = field
111                    .dtype
112                    .as_ref()
113                    .ok_or_else(|| vortex_err!("ShreddingSpecField missing dtype"))
114                    .and_then(|dtype| DType::from_proto(dtype, session))?;
115                Ok((path, dtype))
116            })
117            .collect::<VortexResult<Vec<_>>>()?;
118
119        Ok(JsonToVariantOptions::new(ShreddingSpec::try_new(fields)?))
120    }
121
122    fn arity(&self, _options: &Self::Options) -> Arity {
123        Arity::Exact(1)
124    }
125
126    fn child_name(&self, _options: &Self::Options, child_idx: usize) -> ChildName {
127        match child_idx {
128            0 => ChildName::from("input"),
129            _ => unreachable!("Invalid child index {child_idx} for JsonToVariant expression"),
130        }
131    }
132
133    fn return_dtype(&self, _options: &Self::Options, arg_dtypes: &[DType]) -> VortexResult<DType> {
134        let input_dtype = &arg_dtypes[0];
135        vortex_ensure!(
136            input_dtype
137                .as_extension_opt()
138                .is_some_and(|ext_dtype| ext_dtype.is::<Json>()),
139            "JsonToVariant input must be a Json extension, found {input_dtype}"
140        );
141
142        Ok(DType::Variant(input_dtype.nullability()))
143    }
144
145    fn execute(
146        &self,
147        options: &Self::Options,
148        args: &dyn ExecutionArgs,
149        ctx: &mut ExecutionCtx,
150    ) -> VortexResult<ArrayRef> {
151        let input = args.get(0)?;
152
153        // This function does not build Variants itself: the Variant encoding registered with the
154        // session supplies the conversion as an `execute_parent` kernel keyed on the extension
155        // encoding for a `Json` input. Reaching this fallback means no such kernel fired, so
156        // canonicalize the input to that encoding and re-dispatch. If the input is already
157        // canonical here, no kernel is registered for it; bail with a clear error rather than
158        // looping.
159        let no_kernel = || {
160            vortex_err!(
161                "json_to_variant requires a registered Variant encoding to build Variant values \
162                 from JSON, but none is registered with this session"
163            )
164        };
165
166        let canonical = if input
167            .dtype()
168            .as_extension_opt()
169            .is_some_and(|ext_dtype| ext_dtype.is::<Json>())
170        {
171            if input.is::<Extension>() {
172                return Err(no_kernel());
173            }
174            input.execute::<ExtensionArray>(ctx)?.into_array()
175        } else {
176            vortex_bail!(
177                "JsonToVariant input must be a Json extension, found {}",
178                input.dtype()
179            );
180        };
181
182        self.try_new_array(canonical.len(), options.clone(), [canonical])?
183            .execute::<ArrayRef>(ctx)
184    }
185
186    fn is_null_sensitive(&self, _options: &Self::Options) -> bool {
187        // `json_to_variant` maps null rows to null rows and the JSON literal `null` to a
188        // variant-null value independently of which other rows are null, so it commutes with
189        // validity masking. Marking it not null-sensitive also lets it push through dictionaries
190        // into their JSON extension values, where the kernel fires directly.
191        false
192    }
193}
194
195/// Creates a [`JsonToVariant`] expression that parses `child`'s JSON strings into Variant
196/// values, shredding the paths selected by `shredding`.
197///
198/// `child` must produce [`Json`] extension values; the result is `Variant` with the input's
199/// nullability. Rows containing invalid JSON fail the expression.
200///
201/// Note that this is a lossy, normalizing conversion. See [`JsonToVariant`] for the full list of
202/// caveats.
203pub fn json_to_variant(child: Expression, shredding: ShreddingSpec) -> Expression {
204    JsonToVariant.new_expr(JsonToVariantOptions::new(shredding), [child])
205}
206
207/// A list of `(path, dtype)` directives describing which Variant paths to shred and as what
208/// type.
209///
210/// Paths must contain only object-field elements; list index elements are rejected because
211/// Parquet Variant shredding schemas cannot express list element shredding yet. The root path
212/// (`$`) shreds the top-level value itself. When entries overlap (e.g. `$.a` and `$.a.b`),
213/// later entries overwrite earlier ones.
214#[derive(Clone, Debug, Default, PartialEq, Eq, Hash)]
215pub struct ShreddingSpec(Vec<(VariantPath, DType)>);
216
217impl ShreddingSpec {
218    /// Creates an empty spec, meaning no shredding is performed.
219    pub fn empty() -> Self {
220        Self::default()
221    }
222
223    /// Creates a spec from `(path, dtype)` directives.
224    ///
225    /// # Errors
226    ///
227    /// Returns an error if any path contains a list index element.
228    pub fn try_new(fields: impl IntoIterator<Item = (VariantPath, DType)>) -> VortexResult<Self> {
229        let fields = fields.into_iter().collect::<Vec<_>>();
230        for (path, _) in &fields {
231            vortex_ensure!(
232                path.elements()
233                    .iter()
234                    .all(|element| matches!(element, VariantPathElement::Field(_))),
235                "ShreddingSpec paths must only contain object fields, found list index in {path}"
236            );
237        }
238        Ok(Self(fields))
239    }
240
241    /// Returns the `(path, dtype)` directives.
242    pub fn fields(&self) -> &[(VariantPath, DType)] {
243        &self.0
244    }
245
246    /// Returns whether this spec contains no directives.
247    pub fn is_empty(&self) -> bool {
248        self.0.is_empty()
249    }
250}
251
252impl Display for ShreddingSpec {
253    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
254        for (idx, (path, dtype)) in self.0.iter().enumerate() {
255            if idx > 0 {
256                write!(f, ", ")?;
257            }
258            write!(f, "{path} as {dtype}")?;
259        }
260        Ok(())
261    }
262}
263
264/// Options for [`JsonToVariant`].
265#[derive(Clone, Debug, PartialEq, Eq, Hash)]
266pub struct JsonToVariantOptions {
267    /// The paths to shred into typed storage, if any.
268    shredding: ShreddingSpec,
269}
270
271impl JsonToVariantOptions {
272    /// Creates options that shred the paths selected by `shredding`.
273    pub fn new(shredding: ShreddingSpec) -> Self {
274        Self { shredding }
275    }
276
277    /// Creates options that perform no shredding.
278    pub fn unshredded() -> Self {
279        Self::new(ShreddingSpec::empty())
280    }
281
282    /// Returns the shredding spec.
283    pub fn shredding(&self) -> &ShreddingSpec {
284        &self.shredding
285    }
286}
287
288impl Display for JsonToVariantOptions {
289    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
290        self.shredding.fmt(f)
291    }
292}
293
294#[cfg(test)]
295mod tests {
296    #![allow(clippy::missing_docs_in_private_items)]
297
298    use vortex_array::ArrayRef;
299    use vortex_array::EmptyMetadata;
300    use vortex_array::IntoArray;
301    use vortex_array::VortexSessionExecute;
302    use vortex_array::arrays::VarBinViewArray;
303    use vortex_array::dtype::Nullability;
304    use vortex_array::dtype::PType;
305    use vortex_array::dtype::session::DTypeSession;
306    use vortex_array::expr::proto::ExprSerializeProtoExt;
307    use vortex_array::expr::root;
308    use vortex_array::scalar_fn::session::ScalarFnSession;
309    use vortex_array::scalar_fn::session::ScalarFnSessionExt;
310    use vortex_array::session::ArraySession;
311
312    use super::*;
313
314    fn i64_dtype() -> DType {
315        DType::Primitive(PType::I64, Nullability::Nullable)
316    }
317
318    /// A session that knows the `JsonToVariant` definition but has no Variant encoding registered,
319    /// so executing the function exercises the fallback that errors when no kernel is present.
320    fn session() -> VortexSession {
321        let session = VortexSession::empty()
322            .with::<ArraySession>()
323            .with::<DTypeSession>()
324            .with::<ScalarFnSession>();
325        session.scalar_fns().register(JsonToVariant);
326        session
327    }
328
329    #[test]
330    fn shredding_spec_rejects_index_paths() {
331        let err = ShreddingSpec::try_new([(
332            VariantPath::new([VariantPathElement::field("a"), VariantPathElement::index(0)]),
333            i64_dtype(),
334        )])
335        .unwrap_err();
336
337        assert!(
338            err.to_string()
339                .contains("ShreddingSpec paths must only contain object fields"),
340            "unexpected error: {err}"
341        );
342    }
343
344    #[test]
345    fn expression_roundtrip_serialization() -> VortexResult<()> {
346        let spec = ShreddingSpec::try_new([(VariantPath::field("a"), i64_dtype())])?;
347        let expr: Expression = json_to_variant(root(), spec);
348        let proto = expr.serialize_proto()?;
349        let actual = Expression::from_proto(&proto, &session())?;
350
351        assert_eq!(actual, expr);
352        Ok(())
353    }
354
355    #[test]
356    fn utf8_input_is_rejected() {
357        let input = VarBinViewArray::from_iter_str(["1", "2"]).into_array();
358        let err = JsonToVariant
359            .try_new_array(input.len(), JsonToVariantOptions::unshredded(), [input])
360            .unwrap_err();
361
362        assert!(
363            err.to_string().contains("Json extension"),
364            "unexpected error: {err}"
365        );
366    }
367
368    #[test]
369    fn execute_without_variant_kernel_errors() -> VortexResult<()> {
370        // With no Variant encoding registered, executing over an already-canonical input must
371        // surface a clear error rather than looping.
372        let input = ExtensionArray::try_new_from_vtable(
373            Json,
374            EmptyMetadata,
375            VarBinViewArray::from_iter_str(["1", "2"]).into_array(),
376        )?
377        .into_array();
378
379        let dtype = input.dtype().clone();
380        let array = JsonToVariant.try_new_array(
381            input.len(),
382            JsonToVariantOptions::unshredded(),
383            [input],
384        )?;
385
386        let err = array
387            .execute::<ArrayRef>(&mut session().create_execution_ctx())
388            .unwrap_err();
389
390        assert!(
391            err.to_string().contains("Variant encoding"),
392            "unexpected error for {dtype}: {err}"
393        );
394        Ok(())
395    }
396}
vortex_json/json_to_variant.rs

vortex_json/
json_to_variant.rs
