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

vortex_array/scalar_fn/fns/variant_get/
mod.rs
