vortex_array/scalar_fn/fns/variant_get/
mod.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use std::fmt;
5use std::fmt::Display;
6use std::fmt::Formatter;
7
8use prost::Message;
9use vortex_error::VortexResult;
10use vortex_error::vortex_ensure;
11use vortex_error::vortex_err;
12use vortex_proto::expr as pb;
13use vortex_proto::expr::variant_path_element;
14use vortex_session::VortexSession;
15use vortex_session::registry::CachedId;
16use vortex_utils::aliases::StringEscape;
17
18use crate::ArrayRef;
19use crate::ExecutionCtx;
20use crate::IntoArray;
21use crate::arrays::ChunkedArray;
22use crate::arrays::ConstantArray;
23use crate::arrays::ScalarFnArray;
24use crate::arrays::VariantArray;
25use crate::builders::builder_with_capacity_in;
26use crate::dtype::DType;
27use crate::dtype::FieldName;
28use crate::dtype::Nullability;
29use crate::expr::display::ExprDisplay;
30use crate::scalar::Scalar;
31use crate::scalar_fn::Arity;
32use crate::scalar_fn::ChildName;
33use crate::scalar_fn::ExecutionArgs;
34use crate::scalar_fn::ScalarFnId;
35use crate::scalar_fn::ScalarFnVTable;
36use crate::scalar_fn::ScalarFnVTableExt;
37
38/// Extracts a field/index path from Variant values.
39///
40/// Missing paths, type mismatches while traversing, and failed casts produce nulls. Without a
41/// requested dtype, results are returned as nullable Variant values; with one, results are cast to
42/// that dtype with nullable nullability. Encodings may serve perfectly shredded paths directly,
43/// but must fall back to the core Variant value for paths not represented by shredded storage.
44#[derive(Clone)]
45pub struct VariantGet;
46
47impl VariantGet {
48    /// Creates a lazy extraction from Variant `input`.
49    ///
50    /// # Errors
51    ///
52    /// Returns an error if `input` is not Variant data.
53    pub fn try_new(input: ArrayRef, options: VariantGetOptions) -> VortexResult<ScalarFnArray> {
54        ScalarFnArray::try_new(VariantGet.bind(options), vec![input])
55    }
56}
57
58impl ScalarFnVTable for VariantGet {
59    type Options = VariantGetOptions;
60
61    fn id(&self) -> ScalarFnId {
62        static ID: CachedId = CachedId::new("vortex.variant_get");
63        *ID
64    }
65
66    fn serialize(&self, options: &Self::Options) -> VortexResult<Option<Vec<u8>>> {
67        let path = options
68            .path()
69            .elements()
70            .iter()
71            .map(VariantPathElement::to_proto)
72            .collect();
73        let dtype = options.dtype().map(TryInto::try_into).transpose()?;
74
75        Ok(Some(pb::VariantGetOpts { path, dtype }.encode_to_vec()))
76    }
77
78    fn deserialize(&self, metadata: &[u8], session: &VortexSession) -> VortexResult<Self::Options> {
79        let opts = pb::VariantGetOpts::decode(metadata)?;
80        let path = opts
81            .path
82            .into_iter()
83            .map(VariantPathElement::from_proto)
84            .collect::<VortexResult<_>>()?;
85        let dtype = opts
86            .dtype
87            .as_ref()
88            .map(|dtype| DType::from_proto(dtype, session))
89            .transpose()?;
90
91        Ok(VariantGetOptions::new(path, dtype))
92    }
93
94    fn arity(&self, _options: &Self::Options) -> Arity {
95        Arity::Exact(1)
96    }
97
98    fn child_name(&self, _options: &Self::Options, child_idx: usize) -> ChildName {
99        match child_idx {
100            0 => ChildName::from("input"),
101            _ => unreachable!("Invalid child index {child_idx} for VariantGet expression"),
102        }
103    }
104
105    fn fmt_sql(
106        &self,
107        options: &Self::Options,
108        expr: &dyn ExprDisplay,
109        f: &mut Formatter<'_>,
110    ) -> fmt::Result {
111        write!(f, "variant_get(")?;
112        Display::fmt(expr.display_child(0), f)?;
113        let path = options.path().to_string();
114        write!(f, ", \"{}\"", StringEscape(&path))?;
115        if let Some(dtype) = options.dtype() {
116            write!(f, ", {dtype}")?;
117        }
118        write!(f, ")")
119    }
120
121    fn return_dtype(&self, options: &Self::Options, arg_dtypes: &[DType]) -> VortexResult<DType> {
122        let input_dtype = &arg_dtypes[0];
123        vortex_ensure!(
124            matches!(input_dtype, DType::Variant(_)),
125            "VariantGet input must be Variant, found {input_dtype}"
126        );
127
128        // Missing paths, traversal mismatches, and cast failures all produce nulls.
129        Ok(options
130            .dtype()
131            .map_or(DType::Variant(Nullability::Nullable), DType::as_nullable))
132    }
133
134    fn execute(
135        &self,
136        options: &Self::Options,
137        args: &dyn ExecutionArgs,
138        ctx: &mut ExecutionCtx,
139    ) -> VortexResult<ArrayRef> {
140        let input = args.get(0)?;
141        // Missing paths, traversal mismatches, and cast failures all produce nulls.
142        let dtype = options
143            .dtype()
144            .map_or(DType::Variant(Nullability::Nullable), DType::as_nullable);
145
146        if !dtype.is_variant() {
147            let mut builder = builder_with_capacity_in(ctx.allocator(), &dtype, input.len());
148            for idx in 0..input.len() {
149                let scalar = input.execute_scalar(idx, ctx)?;
150                let output = variant_get_scalar(&scalar, options, &dtype)?;
151                builder.append_scalar(&output)?;
152            }
153
154            return Ok(builder.finish());
155        }
156
157        // TODO(variant): replace this with a Variant builder once one exists.
158        // Chunked<Variant> canonicalizes to VariantArray, so this row-wise fallback is safe.
159        let mut chunks = Vec::with_capacity(input.len());
160
161        for idx in 0..input.len() {
162            let scalar = input.execute_scalar(idx, ctx)?;
163            let output = variant_get_scalar(&scalar, options, &dtype)?;
164            chunks.push(ConstantArray::new(output, 1).into_array());
165        }
166
167        let array = ChunkedArray::try_new(chunks, dtype)?.into_array();
168        VariantArray::try_new(array, None).map(|array| array.into_array())
169    }
170
171    fn is_strict(&self, _options: &Self::Options) -> bool {
172        true
173    }
174}
175
176fn variant_get_scalar(
177    scalar: &Scalar,
178    options: &VariantGetOptions,
179    output_dtype: &DType,
180) -> VortexResult<Scalar> {
181    let Some(value) = variant_path_scalar(scalar, options.path().elements())? else {
182        return Ok(Scalar::null(output_dtype.clone()));
183    };
184
185    if options.dtype().is_none_or(DType::is_variant) {
186        return Scalar::variant(value).cast(output_dtype);
187    }
188
189    if value.is_null() {
190        return Ok(Scalar::null(output_dtype.clone()));
191    }
192
193    value
194        .cast(output_dtype)
195        .or_else(|_| Ok(Scalar::null(output_dtype.clone())))
196}
197
198fn variant_path_scalar(
199    scalar: &Scalar,
200    path: &[VariantPathElement],
201) -> VortexResult<Option<Scalar>> {
202    let mut current = match variant_payload(scalar.clone()) {
203        Some(value) => value,
204        None => return Ok(None),
205    };
206
207    for element in path {
208        current = match variant_payload(current) {
209            Some(value) => value,
210            None => return Ok(None),
211        };
212
213        if current.is_null() {
214            return Ok(None);
215        }
216
217        current = match element {
218            VariantPathElement::Field(name) => {
219                let Some(struct_scalar) = current.as_struct_opt() else {
220                    return Ok(None);
221                };
222                if struct_scalar.is_null() {
223                    return Ok(None);
224                }
225                let Some(field) = struct_scalar.field(name.as_ref()) else {
226                    return Ok(None);
227                };
228                field
229            }
230            VariantPathElement::Index(index) => {
231                let Ok(index) = usize::try_from(*index) else {
232                    return Ok(None);
233                };
234                let Some(list_scalar) = current.as_list_opt() else {
235                    return Ok(None);
236                };
237                let Some(element) = list_scalar.element(index) else {
238                    return Ok(None);
239                };
240                element
241            }
242        };
243    }
244
245    Ok(variant_payload(current))
246}
247
248fn variant_payload(scalar: Scalar) -> Option<Scalar> {
249    if scalar.dtype().is_variant() {
250        scalar.as_variant().value().cloned()
251    } else {
252        Some(scalar)
253    }
254}
255
256/// Options for [`VariantGet`].
257#[derive(Clone, Debug, PartialEq, Eq, Hash)]
258pub struct VariantGetOptions {
259    path: VariantPath,
260    dtype: Option<DType>,
261}
262
263impl VariantGetOptions {
264    /// Creates options for extracting `path`, returning Variant values when `dtype` is `None`.
265    pub fn new(path: VariantPath, dtype: Option<DType>) -> Self {
266        Self { path, dtype }
267    }
268
269    /// Returns the path to extract.
270    pub fn path(&self) -> &VariantPath {
271        &self.path
272    }
273
274    /// Returns the requested output dtype, if any.
275    pub fn dtype(&self) -> Option<&DType> {
276        self.dtype.as_ref()
277    }
278}
279
280impl Display for VariantGetOptions {
281    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
282        write!(f, "{}", self.path)?;
283        if let Some(dtype) = &self.dtype {
284            write!(f, " as {dtype}")?;
285        }
286        Ok(())
287    }
288}
289
290/// A strict Variant path made from object fields and list indexes.
291#[derive(Clone, Debug, Default, PartialEq, Eq, Hash)]
292pub struct VariantPath(Vec<VariantPathElement>);
293
294impl VariantPath {
295    /// Creates a path from explicit elements.
296    pub fn new(elements: impl IntoIterator<Item = VariantPathElement>) -> Self {
297        Self(elements.into_iter().collect())
298    }
299
300    /// Creates the root path.
301    pub fn root() -> Self {
302        Self::default()
303    }
304
305    /// Creates a path containing one object field.
306    pub fn field(field: impl Into<FieldName>) -> Self {
307        Self(vec![VariantPathElement::field(field)])
308    }
309
310    /// Returns the path elements.
311    pub fn elements(&self) -> &[VariantPathElement] {
312        &self.0
313    }
314
315    /// Returns whether this path references the root Variant value.
316    pub fn is_root(&self) -> bool {
317        self.0.is_empty()
318    }
319}
320
321impl FromIterator<VariantPathElement> for VariantPath {
322    fn from_iter<T: IntoIterator<Item = VariantPathElement>>(iter: T) -> Self {
323        Self(iter.into_iter().collect())
324    }
325}
326
327impl From<VariantPathElement> for VariantPath {
328    fn from(value: VariantPathElement) -> Self {
329        Self(vec![value])
330    }
331}
332
333impl Display for VariantPath {
334    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
335        write!(f, "$")?;
336        for element in self.elements() {
337            match element {
338                VariantPathElement::Field(name) => write!(f, ".{name}")?,
339                VariantPathElement::Index(index) => write!(f, "[{index}]")?,
340            }
341        }
342        Ok(())
343    }
344}
345
346/// A single field or index step in a [`VariantPath`].
347#[derive(Clone, Debug, PartialEq, Eq, Hash)]
348pub enum VariantPathElement {
349    /// Select an object field by name.
350    Field(FieldName),
351    /// Select a list element by zero-based index.
352    Index(u64),
353}
354
355impl VariantPathElement {
356    /// Creates an object-field path element.
357    pub fn field(field: impl Into<FieldName>) -> Self {
358        Self::Field(field.into())
359    }
360
361    /// Creates a list-index path element.
362    pub fn index(index: u64) -> Self {
363        Self::Index(index)
364    }
365
366    /// Decodes a path element from its protobuf representation.
367    pub fn from_proto(value: pb::VariantPathElement) -> VortexResult<Self> {
368        match value
369            .element
370            .ok_or_else(|| vortex_err!("Variant path element missing value"))?
371        {
372            variant_path_element::Element::Field(field) => Ok(Self::field(field)),
373            variant_path_element::Element::Index(index) => Ok(Self::index(index)),
374        }
375    }
376
377    /// Encodes this path element into its protobuf representation.
378    pub fn to_proto(&self) -> pb::VariantPathElement {
379        match self {
380            VariantPathElement::Field(name) => pb::VariantPathElement {
381                element: Some(variant_path_element::Element::Field(
382                    name.as_ref().to_string(),
383                )),
384            },
385            VariantPathElement::Index(index) => pb::VariantPathElement {
386                element: Some(variant_path_element::Element::Index(*index)),
387            },
388        }
389    }
390}
391
392impl From<FieldName> for VariantPathElement {
393    fn from(value: FieldName) -> Self {
394        Self::field(value)
395    }
396}
397
398impl From<&str> for VariantPathElement {
399    fn from(value: &str) -> Self {
400        Self::field(value)
401    }
402}
403
404impl From<u64> for VariantPathElement {
405    fn from(value: u64) -> Self {
406        Self::index(value)
407    }
408}
409
410#[cfg(test)]
411mod tests {
412    use vortex_error::VortexResult;
413    use vortex_error::vortex_bail;
414    use vortex_error::vortex_ensure;
415    use vortex_error::vortex_err;
416    use vortex_session::VortexSession;
417
418    use crate::ArrayRef;
419    use crate::Canonical;
420    use crate::IntoArray;
421    use crate::VortexSessionExecute;
422    use crate::array_session;
423    use crate::arrays::Chunked;
424    use crate::arrays::ChunkedArray;
425    use crate::arrays::ConstantArray;
426    use crate::arrays::PrimitiveArray;
427    use crate::arrays::VariantArray;
428    use crate::arrays::variant::VariantArraySlotsExt;
429    use crate::assert_arrays_eq;
430    use crate::assert_nth_scalar_is_null;
431    use crate::dtype::DType;
432    use crate::dtype::FieldName;
433    use crate::dtype::FieldNames;
434    use crate::dtype::Nullability;
435    use crate::dtype::PType;
436    use crate::dtype::StructFields;
437    use crate::expr::Expression;
438    use crate::expr::proto::ExprSerializeProtoExt;
439    use crate::expr::root;
440    use crate::expr::variant_get;
441    use crate::scalar::Scalar;
442    use crate::scalar::ScalarValue;
443    use crate::scalar_fn::ScalarFnVTable;
444    use crate::scalar_fn::fns::variant_get::VariantGet;
445    use crate::scalar_fn::fns::variant_get::VariantGetOptions;
446    use crate::scalar_fn::fns::variant_get::VariantPath;
447    use crate::scalar_fn::fns::variant_get::VariantPathElement;
448
449    fn variant_object(fields: impl IntoIterator<Item = (&'static str, Scalar)>) -> Scalar {
450        let fields = fields.into_iter().collect::<Vec<_>>();
451        let names = FieldNames::from_iter(fields.iter().map(|(name, _)| FieldName::from(*name)));
452        let dtypes = vec![DType::Variant(Nullability::NonNullable); fields.len()];
453        let values = fields
454            .into_iter()
455            .map(|(_, value)| Scalar::variant(value).into_value())
456            .collect();
457        Scalar::try_new(
458            DType::Struct(StructFields::new(names, dtypes), Nullability::NonNullable),
459            Some(ScalarValue::Tuple(values)),
460        )
461        .unwrap()
462    }
463
464    fn variant_rows(rows: impl IntoIterator<Item = Scalar>) -> VortexResult<ArrayRef> {
465        let dtype = DType::Variant(Nullability::Nullable);
466        let chunks = rows
467            .into_iter()
468            .map(|row| ConstantArray::new(row.cast(&dtype).unwrap(), 1).into_array());
469        ChunkedArray::try_new(chunks, dtype.clone()).map(|array| array.into_array())
470    }
471
472    /// Test-only syntax for keeping `variant_get` cases compact without committing
473    /// to a public string grammar yet.
474    fn parse_path(path: &str) -> VortexResult<VariantPath> {
475        if path.is_empty() || path == "$" {
476            return Ok(VariantPath::root());
477        }
478
479        let mut elements = Vec::new();
480        let mut pos = usize::from(path.as_bytes().first() == Some(&b'$'));
481        if pos == 1
482            && path
483                .as_bytes()
484                .get(pos)
485                .is_some_and(|byte| !matches!(byte, b'.' | b'['))
486        {
487            vortex_bail!("Invalid Variant path {path:?}: expected '.' or '[' after '$'");
488        }
489
490        while pos < path.len() {
491            match path.as_bytes()[pos] {
492                b'.' => {
493                    pos += 1;
494                    let (field, next_pos) = parse_field(path, pos)?;
495                    elements.push(VariantPathElement::field(field));
496                    pos = next_pos;
497                }
498                b'[' => {
499                    let (index, next_pos) = parse_index(path, pos + 1)?;
500                    elements.push(VariantPathElement::index(index));
501                    pos = next_pos;
502                }
503                _ if pos == 0 => {
504                    let (field, next_pos) = parse_field(path, pos)?;
505                    elements.push(VariantPathElement::field(field));
506                    pos = next_pos;
507                }
508                _ => {
509                    vortex_bail!("Invalid Variant path {path:?}: expected '.', '[', or end of path")
510                }
511            }
512        }
513
514        Ok(VariantPath::new(elements))
515    }
516
517    fn parse_field(path: &str, start: usize) -> VortexResult<(&str, usize)> {
518        let mut pos = start;
519        while path
520            .as_bytes()
521            .get(pos)
522            .is_some_and(|byte| byte.is_ascii_alphanumeric() || *byte == b'_')
523        {
524            pos += 1;
525        }
526        vortex_ensure!(
527            pos > start,
528            "Invalid Variant path {path:?}: expected field name"
529        );
530        Ok((&path[start..pos], pos))
531    }
532
533    fn parse_index(path: &str, start: usize) -> VortexResult<(u64, usize)> {
534        let mut pos = start;
535        while path
536            .as_bytes()
537            .get(pos)
538            .is_some_and(|byte| byte.is_ascii_digit())
539        {
540            pos += 1;
541        }
542        vortex_ensure!(
543            pos > start,
544            "Invalid Variant path {path:?}: expected list index"
545        );
546        vortex_ensure!(
547            path.as_bytes().get(pos) == Some(&b']'),
548            "Invalid Variant path {path:?}: expected closing ']'"
549        );
550        let index = path[start..pos]
551            .parse()
552            .map_err(|_| vortex_err!("Invalid Variant path {path:?}: list index is too large"))?;
553        Ok((index, pos + 1))
554    }
555
556    fn execute_variant_get(
557        array: ArrayRef,
558        path: &str,
559        dtype: Option<DType>,
560    ) -> VortexResult<ArrayRef> {
561        let expr = variant_get(root(), parse_path(path)?, dtype);
562        array
563            .apply(&expr)?
564            .execute::<ArrayRef>(&mut array_session().create_execution_ctx())
565    }
566
567    #[test]
568    fn variant_get_path_parse_and_display() {
569        let path = parse_path("$.data[1].a").unwrap();
570        assert_eq!(
571            path.elements(),
572            &[
573                VariantPathElement::field("data"),
574                VariantPathElement::index(1),
575                VariantPathElement::field("a")
576            ]
577        );
578        assert_eq!(path.to_string(), "$.data[1].a");
579
580        let bare_path = parse_path("data[2]").unwrap();
581        assert_eq!(bare_path.to_string(), "$.data[2]");
582        assert!(parse_path("$.").is_err());
583        assert!(parse_path("$data").is_err());
584        assert!(parse_path("$.data[-1]").is_err());
585    }
586
587    #[test]
588    fn variant_get_return_dtype_is_nullable_variant_without_requested_dtype() {
589        let expr = variant_get(root(), VariantPath::field("data"), None);
590        let dtype = expr
591            .return_dtype(&DType::Variant(Nullability::NonNullable))
592            .unwrap();
593
594        assert_eq!(dtype, DType::Variant(Nullability::Nullable));
595    }
596
597    #[test]
598    fn variant_get_return_dtype_makes_requested_dtype_nullable() {
599        let requested = DType::Primitive(PType::I64, Nullability::NonNullable);
600        let expr = variant_get(root(), VariantPath::field("data"), Some(requested));
601        let dtype = expr
602            .return_dtype(&DType::Variant(Nullability::NonNullable))
603            .unwrap();
604
605        assert_eq!(dtype, DType::Primitive(PType::I64, Nullability::Nullable));
606    }
607
608    #[test]
609    fn variant_get_rejects_non_variant_input() {
610        let expr = variant_get(root(), VariantPath::field("data"), None);
611        let err = expr
612            .return_dtype(&DType::Utf8(Nullability::NonNullable))
613            .unwrap_err();
614
615        assert!(err.to_string().contains("VariantGet input must be Variant"));
616    }
617
618    #[test]
619    fn variant_get_formats_sql() {
620        let expr = variant_get(
621            root(),
622            parse_path("$.data[1].a").unwrap(),
623            Some(DType::Utf8(Nullability::NonNullable)),
624        );
625
626        assert_eq!(expr.to_string(), "variant_get($, \"$.data[1].a\", utf8)");
627    }
628
629    #[test]
630    fn variant_get_options_roundtrip_serialization() {
631        let options = VariantGetOptions::new(
632            VariantPath::new([
633                VariantPathElement::field("data"),
634                VariantPathElement::index(1),
635                VariantPathElement::field("a"),
636            ]),
637            Some(DType::Primitive(PType::I32, Nullability::NonNullable)),
638        );
639        let metadata = VariantGet.serialize(&options).unwrap().unwrap();
640        let actual = VariantGet
641            .deserialize(&metadata, &VortexSession::empty())
642            .unwrap();
643
644        assert_eq!(actual, options);
645    }
646
647    #[test]
648    fn variant_get_expression_roundtrip_serialization() {
649        let expr: Expression = variant_get(
650            root(),
651            parse_path("$.data[1].a").unwrap(),
652            Some(DType::Primitive(PType::I32, Nullability::NonNullable)),
653        );
654        let proto = expr.serialize_proto().unwrap();
655        let actual = Expression::from_proto(&proto, &array_session()).unwrap();
656
657        assert_eq!(actual, expr);
658    }
659
660    #[test]
661    fn variant_get_generic_fallback_extracts_field_and_list_index() -> VortexResult<()> {
662        let items = Scalar::list(
663            DType::Variant(Nullability::NonNullable),
664            vec![
665                Scalar::variant(Scalar::primitive(10i32, Nullability::NonNullable)),
666                Scalar::variant(Scalar::primitive(20i32, Nullability::NonNullable)),
667            ],
668            Nullability::NonNullable,
669        );
670        let array = variant_rows([
671            Scalar::variant(variant_object([("items", items)])),
672            Scalar::variant(variant_object([(
673                "items",
674                Scalar::list_empty(
675                    DType::Variant(Nullability::NonNullable).into(),
676                    Nullability::NonNullable,
677                ),
678            )])),
679            Scalar::variant(variant_object([(
680                "items",
681                Scalar::list(
682                    DType::Variant(Nullability::NonNullable),
683                    vec![
684                        Scalar::variant(Scalar::utf8("x", Nullability::NonNullable)),
685                        Scalar::variant(Scalar::utf8("wrong", Nullability::NonNullable)),
686                    ],
687                    Nullability::NonNullable,
688                ),
689            )])),
690        ])?;
691
692        let result = execute_variant_get(
693            array,
694            "$.items[1]",
695            Some(DType::Primitive(PType::I32, Nullability::NonNullable)),
696        )?;
697        let mut ctx = array_session().create_execution_ctx();
698
699        assert_arrays_eq!(
700            result,
701            PrimitiveArray::from_option_iter([Some(20i32), None, None]),
702            &mut ctx
703        );
704        Ok(())
705    }
706
707    #[test]
708    fn variant_get_reads_chunked_variant_input() -> VortexResult<()> {
709        let array = variant_rows([
710            Scalar::variant(variant_object([(
711                "a",
712                Scalar::primitive(10i32, Nullability::NonNullable),
713            )])),
714            Scalar::variant(variant_object([(
715                "b",
716                Scalar::primitive(20i32, Nullability::NonNullable),
717            )])),
718            Scalar::variant(variant_object([(
719                "a",
720                Scalar::primitive(30i32, Nullability::NonNullable),
721            )])),
722            Scalar::null(DType::Variant(Nullability::Nullable)),
723        ])?;
724        assert!(array.is::<Chunked>());
725
726        let result = execute_variant_get(
727            array,
728            "$.a",
729            Some(DType::Primitive(PType::I32, Nullability::NonNullable)),
730        )?;
731        let mut ctx = array_session().create_execution_ctx();
732
733        assert_arrays_eq!(
734            result,
735            PrimitiveArray::from_option_iter([Some(10i32), None, Some(30), None]),
736            &mut ctx
737        );
738        Ok(())
739    }
740
741    #[test]
742    fn variant_get_fallback_typed_output_is_contiguous() -> VortexResult<()> {
743        let array = variant_rows([
744            Scalar::variant(variant_object([(
745                "a",
746                Scalar::primitive(10i32, Nullability::NonNullable),
747            )])),
748            Scalar::variant(variant_object([(
749                "a",
750                Scalar::primitive(20i32, Nullability::NonNullable),
751            )])),
752            Scalar::variant(variant_object([(
753                "b",
754                Scalar::primitive(30i32, Nullability::NonNullable),
755            )])),
756        ])?;
757
758        let result = execute_variant_get(
759            array,
760            "$.a",
761            Some(DType::Primitive(PType::I32, Nullability::NonNullable)),
762        )?;
763
764        assert!(!result.is::<Chunked>());
765        let mut ctx = array_session().create_execution_ctx();
766        assert_arrays_eq!(
767            result,
768            PrimitiveArray::from_option_iter([Some(10i32), Some(20), None]),
769            &mut ctx
770        );
771        Ok(())
772    }
773
774    #[test]
775    fn variant_get_generic_fallback_preserves_variant_null() -> VortexResult<()> {
776        let array = variant_rows([
777            Scalar::variant(variant_object([(
778                "a",
779                Scalar::utf8("ok", Nullability::NonNullable),
780            )])),
781            Scalar::null(DType::Variant(Nullability::Nullable)),
782            Scalar::variant(variant_object([("a", Scalar::null(DType::Null))])),
783            Scalar::variant(variant_object([(
784                "b",
785                Scalar::primitive(2i32, Nullability::NonNullable),
786            )])),
787        ])?;
788
789        let result = execute_variant_get(array, "$.a", None)?;
790
791        let mut ctx = array_session().create_execution_ctx();
792        let row0 = result.execute_scalar(0, &mut ctx)?;
793        assert_eq!(
794            row0.as_variant()
795                .value()
796                .and_then(|value| value.as_utf8().value())
797                .map(|value| value.as_str()),
798            Some("ok")
799        );
800        assert_nth_scalar_is_null!(result, 1, &mut ctx);
801        assert_eq!(
802            result
803                .execute_scalar(2, &mut ctx)?
804                .as_variant()
805                .is_variant_null(),
806            Some(true)
807        );
808        assert_nth_scalar_is_null!(result, 3, &mut ctx);
809        Ok(())
810    }
811
812    #[test]
813    fn variant_get_fallback_variant_output_canonicalizes() -> VortexResult<()> {
814        let array = variant_rows([
815            Scalar::variant(variant_object([(
816                "a",
817                Scalar::primitive(10i32, Nullability::NonNullable),
818            )])),
819            Scalar::variant(variant_object([(
820                "a",
821                Scalar::primitive(20i32, Nullability::NonNullable),
822            )])),
823        ])?;
824
825        let result = execute_variant_get(array, "$.a", None)?;
826        let variant = result
827            .clone()
828            .execute::<VariantArray>(&mut array_session().create_execution_ctx())?;
829        let canonical = result.execute::<Canonical>(&mut array_session().create_execution_ctx())?;
830        let Canonical::Variant(canonical_variant) = canonical else {
831            vortex_bail!("expected Variant canonical array");
832        };
833
834        assert_eq!(variant.len(), 2);
835        assert_eq!(canonical_variant.len(), 2);
836        assert_eq!(variant.core_storage().dtype(), variant.dtype());
837        assert_eq!(variant.core_storage().len(), variant.len());
838
839        let mut ctx = array_session().create_execution_ctx();
840        for (idx, expected) in [10i32, 20].into_iter().enumerate() {
841            let scalar = variant.execute_scalar(idx, &mut ctx)?;
842            let actual = scalar
843                .as_variant()
844                .value()
845                .and_then(|value| value.as_primitive().as_::<i32>());
846            assert_eq!(actual, Some(expected), "row {idx}");
847        }
848        Ok(())
849    }
850}
vortex_array/scalar_fn/fns/variant_get/mod.rs

vortex_array/scalar_fn/fns/variant_get/
mod.rs
