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_session::VortexSession;
13use vortex_session::registry::CachedId;
14use vortex_utils::aliases::StringEscape;
15
16use crate::ArrayRef;
17use crate::ExecutionCtx;
18use crate::IntoArray;
19use crate::arrays::ChunkedArray;
20use crate::arrays::ConstantArray;
21use crate::arrays::ScalarFnArray;
22use crate::arrays::VariantArray;
23use crate::builders::builder_with_capacity_in;
24use crate::dtype::DType;
25use crate::dtype::FieldName;
26use crate::dtype::Nullability;
27use crate::expr::display::ExprDisplay;
28use crate::proto::expr as pb;
29use crate::proto::expr::variant_path_element;
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), [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(&dtype, input.len(), ctx.allocator());
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    }
583
584    #[test]
585    fn variant_get_return_dtype_is_nullable_variant_without_requested_dtype() {
586        let expr = variant_get(root(), VariantPath::field("data"), None);
587        let dtype = expr
588            .return_dtype(&DType::Variant(Nullability::NonNullable))
589            .unwrap();
590
591        assert_eq!(dtype, DType::Variant(Nullability::Nullable));
592    }
593
594    #[test]
595    fn variant_get_return_dtype_makes_requested_dtype_nullable() {
596        let requested = DType::Primitive(PType::I64, Nullability::NonNullable);
597        let expr = variant_get(root(), VariantPath::field("data"), Some(requested));
598        let dtype = expr
599            .return_dtype(&DType::Variant(Nullability::NonNullable))
600            .unwrap();
601
602        assert_eq!(dtype, DType::Primitive(PType::I64, Nullability::Nullable));
603    }
604
605    #[test]
606    fn variant_get_rejects_non_variant_input() {
607        let expr = variant_get(root(), VariantPath::field("data"), None);
608        let err = expr
609            .return_dtype(&DType::Utf8(Nullability::NonNullable))
610            .unwrap_err();
611
612        assert!(err.to_string().contains("VariantGet input must be Variant"));
613    }
614
615    #[test]
616    fn variant_get_formats_sql() {
617        let expr = variant_get(
618            root(),
619            parse_path("$.data[1].a").unwrap(),
620            Some(DType::Utf8(Nullability::NonNullable)),
621        );
622
623        assert_eq!(expr.to_string(), "variant_get($, \"$.data[1].a\", utf8)");
624    }
625
626    #[test]
627    fn variant_get_options_roundtrip_serialization() {
628        let options = VariantGetOptions::new(
629            VariantPath::new([
630                VariantPathElement::field("data"),
631                VariantPathElement::index(1),
632                VariantPathElement::field("a"),
633            ]),
634            Some(DType::Primitive(PType::I32, Nullability::NonNullable)),
635        );
636        let metadata = VariantGet.serialize(&options).unwrap().unwrap();
637        let actual = VariantGet
638            .deserialize(&metadata, &VortexSession::empty())
639            .unwrap();
640
641        assert_eq!(actual, options);
642    }
643
644    #[test]
645    fn variant_get_expression_roundtrip_serialization() {
646        let expr: Expression = variant_get(
647            root(),
648            parse_path("$.data[1].a").unwrap(),
649            Some(DType::Primitive(PType::I32, Nullability::NonNullable)),
650        );
651        let proto = expr.serialize_proto().unwrap();
652        let actual = Expression::from_proto(&proto, &array_session()).unwrap();
653
654        assert_eq!(actual, expr);
655    }
656
657    #[test]
658    fn variant_get_generic_fallback_extracts_field_and_list_index() -> VortexResult<()> {
659        let items = Scalar::list(
660            DType::Variant(Nullability::NonNullable),
661            vec![
662                Scalar::variant(Scalar::primitive(10i32, Nullability::NonNullable)),
663                Scalar::variant(Scalar::primitive(20i32, Nullability::NonNullable)),
664            ],
665            Nullability::NonNullable,
666        );
667        let array = variant_rows([
668            Scalar::variant(variant_object([("items", items)])),
669            Scalar::variant(variant_object([(
670                "items",
671                Scalar::list_empty(
672                    DType::Variant(Nullability::NonNullable).into(),
673                    Nullability::NonNullable,
674                ),
675            )])),
676            Scalar::variant(variant_object([(
677                "items",
678                Scalar::list(
679                    DType::Variant(Nullability::NonNullable),
680                    vec![
681                        Scalar::variant(Scalar::utf8("x", Nullability::NonNullable)),
682                        Scalar::variant(Scalar::utf8("wrong", Nullability::NonNullable)),
683                    ],
684                    Nullability::NonNullable,
685                ),
686            )])),
687        ])?;
688
689        let result = execute_variant_get(
690            array,
691            "$.items[1]",
692            Some(DType::Primitive(PType::I32, Nullability::NonNullable)),
693        )?;
694        let mut ctx = array_session().create_execution_ctx();
695
696        assert_arrays_eq!(
697            result,
698            PrimitiveArray::from_option_iter([Some(20i32), None, None]),
699            &mut ctx
700        );
701        Ok(())
702    }
703
704    #[test]
705    fn variant_get_reads_chunked_variant_input() -> VortexResult<()> {
706        let array = variant_rows([
707            Scalar::variant(variant_object([(
708                "a",
709                Scalar::primitive(10i32, Nullability::NonNullable),
710            )])),
711            Scalar::variant(variant_object([(
712                "b",
713                Scalar::primitive(20i32, Nullability::NonNullable),
714            )])),
715            Scalar::variant(variant_object([(
716                "a",
717                Scalar::primitive(30i32, Nullability::NonNullable),
718            )])),
719            Scalar::null(DType::Variant(Nullability::Nullable)),
720        ])?;
721        assert!(array.is::<Chunked>());
722
723        let result = execute_variant_get(
724            array,
725            "$.a",
726            Some(DType::Primitive(PType::I32, Nullability::NonNullable)),
727        )?;
728        let mut ctx = array_session().create_execution_ctx();
729
730        assert_arrays_eq!(
731            result,
732            PrimitiveArray::from_option_iter([Some(10i32), None, Some(30), None]),
733            &mut ctx
734        );
735        Ok(())
736    }
737
738    #[test]
739    fn variant_get_fallback_typed_output_is_contiguous() -> VortexResult<()> {
740        let array = variant_rows([
741            Scalar::variant(variant_object([(
742                "a",
743                Scalar::primitive(10i32, Nullability::NonNullable),
744            )])),
745            Scalar::variant(variant_object([(
746                "a",
747                Scalar::primitive(20i32, Nullability::NonNullable),
748            )])),
749            Scalar::variant(variant_object([(
750                "b",
751                Scalar::primitive(30i32, Nullability::NonNullable),
752            )])),
753        ])?;
754
755        let result = execute_variant_get(
756            array,
757            "$.a",
758            Some(DType::Primitive(PType::I32, Nullability::NonNullable)),
759        )?;
760
761        assert!(!result.is::<Chunked>());
762        let mut ctx = array_session().create_execution_ctx();
763        assert_arrays_eq!(
764            result,
765            PrimitiveArray::from_option_iter([Some(10i32), Some(20), None]),
766            &mut ctx
767        );
768        Ok(())
769    }
770
771    #[test]
772    fn variant_get_generic_fallback_preserves_variant_null() -> VortexResult<()> {
773        let array = variant_rows([
774            Scalar::variant(variant_object([(
775                "a",
776                Scalar::utf8("ok", Nullability::NonNullable),
777            )])),
778            Scalar::null(DType::Variant(Nullability::Nullable)),
779            Scalar::variant(variant_object([("a", Scalar::null(DType::Null))])),
780            Scalar::variant(variant_object([(
781                "b",
782                Scalar::primitive(2i32, Nullability::NonNullable),
783            )])),
784        ])?;
785
786        let result = execute_variant_get(array, "$.a", None)?;
787
788        let mut ctx = array_session().create_execution_ctx();
789        let row0 = result.execute_scalar(0, &mut ctx)?;
790        assert_eq!(
791            row0.as_variant()
792                .value()
793                .and_then(|value| value.as_utf8().value())
794                .map(|value| value.as_str()),
795            Some("ok")
796        );
797        assert_nth_scalar_is_null!(result, 1, &mut ctx);
798        assert_eq!(
799            result
800                .execute_scalar(2, &mut ctx)?
801                .as_variant()
802                .is_variant_null(),
803            Some(true)
804        );
805        assert_nth_scalar_is_null!(result, 3, &mut ctx);
806        Ok(())
807    }
808
809    #[test]
810    fn variant_get_fallback_variant_output_canonicalizes() -> VortexResult<()> {
811        let array = variant_rows([
812            Scalar::variant(variant_object([(
813                "a",
814                Scalar::primitive(10i32, Nullability::NonNullable),
815            )])),
816            Scalar::variant(variant_object([(
817                "a",
818                Scalar::primitive(20i32, Nullability::NonNullable),
819            )])),
820        ])?;
821
822        let result = execute_variant_get(array, "$.a", None)?;
823        let variant = result
824            .clone()
825            .execute::<VariantArray>(&mut array_session().create_execution_ctx())?;
826        let canonical = result.execute::<Canonical>(&mut array_session().create_execution_ctx())?;
827        let Canonical::Variant(canonical_variant) = canonical else {
828            vortex_bail!("expected Variant canonical array");
829        };
830
831        assert_eq!(variant.len(), 2);
832        assert_eq!(canonical_variant.len(), 2);
833        assert_eq!(variant.core_storage().dtype(), variant.dtype());
834        assert_eq!(variant.core_storage().len(), variant.len());
835
836        let mut ctx = array_session().create_execution_ctx();
837        for array in [variant, canonical_variant] {
838            assert_eq!(array.dtype(), &DType::Variant(Nullability::Nullable));
839            for (idx, expected) in [10i32, 20].into_iter().enumerate() {
840                let scalar = array.execute_scalar(idx, &mut ctx)?;
841                let actual = scalar
842                    .as_variant()
843                    .value()
844                    .and_then(|value| value.as_primitive().as_::<i32>());
845                assert_eq!(actual, Some(expected), "row {idx}");
846            }
847        }
848        Ok(())
849    }
850}
vortex_array/scalar_fn/fns/variant_get/mod.rs

vortex_array/scalar_fn/fns/variant_get/
mod.rs
