Add GetItemFn (#5609)

Signed-off-by: Nicholas Gates <[email protected]>

Author: gatesn

vortex-array/src/scalar_fns/get_item/mod.rs

@@ -0,0 +1,115 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+use prost::Message;
+use vortex_dtype::DType;
+use vortex_dtype::FieldName;
+use vortex_dtype::FieldPath;
+use vortex_dtype::Nullability;
+use vortex_error::VortexResult;
+use vortex_error::vortex_err;
+use vortex_proto::expr as pb;
+use vortex_vector::Datum;
+use vortex_vector::ScalarOps;
+use vortex_vector::VectorOps;
+
+use crate::expr::Expression;
+use crate::expr::StatsCatalog;
+use crate::expr::functions::ArgName;
+use crate::expr::functions::Arity;
+use crate::expr::functions::ExecutionArgs;
+use crate::expr::functions::FunctionId;
+use crate::expr::functions::VTable;
+use crate::expr::stats::Stat;
+
+pub struct GetItemFn;
+impl VTable for GetItemFn {
+    type Options = FieldName;
+
+    fn id(&self) -> FunctionId {
+        FunctionId::from("vortex.get_item")
+    }
+
+    fn serialize(&self, field_name: &FieldName) -> VortexResult<Option<Vec<u8>>> {
+        Ok(Some(
+            pb::GetItemOpts {
+                path: field_name.to_string(),
+            }
+            .encode_to_vec(),
+        ))
+    }
+
+    fn deserialize(&self, bytes: &[u8]) -> VortexResult<Self::Options> {
+        let opts = pb::GetItemOpts::decode(bytes)?;
+        Ok(FieldName::from(opts.path))
+    }
+
+    fn arity(&self, _field_name: &FieldName) -> Arity {
+        Arity::Exact(1)
+    }
+
+    fn arg_name(&self, _field_name: &FieldName, _arg_idx: usize) -> ArgName {
+        ArgName::from("input")
+    }
+
+    fn stat_expression(
+        &self,
+        field_name: &FieldName,
+        _expr: &Expression,
+        stat: Stat,
+        catalog: &dyn StatsCatalog,
+    ) -> Option<Expression> {
+        // TODO(ngates): I think we can do better here and support stats over nested fields.
+        //  It would be nice if delegating to our child would return a struct of statistics
+        //  matching the nested DType such that we can write:
+        //    `get_item(expr.child(0).stat_expression(...), expr.data().field_name())`
+
+        // TODO(ngates): this is a bug whereby we may return stats for a nested field of the same
+        //  name as a field in the root struct. This should be resolved with upcoming change to
+        //  falsify expressions, but for now I'm preserving the existing buggy behavior.
+        catalog.stats_ref(&FieldPath::from_name(field_name.clone()), stat)
+    }
+
+    fn return_dtype(&self, field_name: &FieldName, arg_types: &[DType]) -> VortexResult<DType> {
+        let struct_dtype = &arg_types[0];
+        let field_dtype = struct_dtype
+            .as_struct_fields_opt()
+            .and_then(|st| st.field(field_name))
+            .ok_or_else(|| {
+                vortex_err!("Couldn't find the {} field in the input scope", field_name)
+            })?;
+
+        // Match here to avoid cloning the dtype if nullability doesn't need to change
+        if matches!(
+            (struct_dtype.nullability(), field_dtype.nullability()),
+            (Nullability::Nullable, Nullability::NonNullable)
+        ) {
+            return Ok(field_dtype.with_nullability(Nullability::Nullable));
+        }
+
+        Ok(field_dtype)
+    }
+
+    fn execute(&self, field_name: &FieldName, args: &ExecutionArgs) -> VortexResult<Datum> {
+        let struct_dtype = args
+            .input_type(0)
+            .as_struct_fields_opt()
+            .ok_or_else(|| vortex_err!("Expected struct dtype for child of GetItem expression"))?;
+        let field_idx = struct_dtype
+            .find(field_name)
+            .ok_or_else(|| vortex_err!("Field {} not found in struct dtype", field_name))?;
+
+        match args.input_datums(0) {
+            Datum::Scalar(s) => {
+                let mut field = s.as_struct().field(field_idx);
+                field.mask_validity(s.is_valid());
+                Ok(Datum::Scalar(field))
+            }
+            Datum::Vector(v) => {
+                let mut field = v.as_struct().fields()[field_idx].clone();
+                field.mask_validity(v.validity());
+                Ok(Datum::Vector(field))
+            }
+        }
+    }
+}

vortex-array/src/scalar_fns/mod.rs

@@ -10,6 +10,7 @@
 //! the equivalent Arrow compute function.
 
 use vortex_dtype::DType;
+use vortex_dtype::FieldName;
 use vortex_error::VortexResult;
 
 use crate::Array;
@@ -20,6 +21,7 @@ use crate::expr::ScalarFnExprExt;
 use crate::expr::functions::EmptyOptions;
 
 pub mod cast;
+pub mod get_item;
 pub mod is_null;
 pub mod mask;
 pub mod not;
@@ -29,66 +31,80 @@ pub trait ExprBuiltins: Sized {
     /// Cast to the given data type.
     fn cast(&self, dtype: DType) -> VortexResult<Expression>;
 
+    /// Get item by field name (for struct types).
+    fn get_item(&self, field_name: impl Into<FieldName>) -> VortexResult<Expression>;
+
     /// Is null check.
     fn is_null(&self) -> VortexResult<Expression>;
 
-    /// Boolean negation.
-    fn not(&self) -> VortexResult<Expression>;
-
     /// Mask the expression using the given boolean mask.
     /// The resulting expression's validity is the intersection of the original expression's
     /// validity.
     fn mask(&self, mask: Expression) -> VortexResult<Expression>;
+
+    /// Boolean negation.
+    fn not(&self) -> VortexResult<Expression>;
 }
 
 impl ExprBuiltins for Expression {
     fn cast(&self, dtype: DType) -> VortexResult<Expression> {
         cast::CastFn.try_new_expr(dtype, [self.clone()])
     }
 
-    fn is_null(&self) -> VortexResult<Expression> {
-        is_null::IsNullFn.try_new_expr(EmptyOptions, [self.clone()])
+    fn get_item(&self, field_name: impl Into<FieldName>) -> VortexResult<Expression> {
+        get_item::GetItemFn.try_new_expr(field_name.into(), [self.clone()])
     }
 
-    fn not(&self) -> VortexResult<Expression> {
-        not::NotFn.try_new_expr(EmptyOptions, [self.clone()])
+    fn is_null(&self) -> VortexResult<Expression> {
+        is_null::IsNullFn.try_new_expr(EmptyOptions, [self.clone()])
     }
 
     fn mask(&self, mask: Expression) -> VortexResult<Expression> {
         mask::MaskFn.try_new_expr(EmptyOptions, [self.clone(), mask])
     }
+
+    fn not(&self) -> VortexResult<Expression> {
+        not::NotFn.try_new_expr(EmptyOptions, [self.clone()])
+    }
 }
 
 pub trait ArrayBuiltins: Sized {
     /// Cast to the given data type.
     fn cast(&self, dtype: DType) -> VortexResult<ArrayRef>;
 
+    /// Get item by field name (for struct types).
+    fn get_item(&self, field_name: impl Into<FieldName>) -> VortexResult<ArrayRef>;
+
     /// Is null check.
     fn is_null(&self) -> VortexResult<ArrayRef>;
 
-    /// Boolean negation.
-    fn not(&self) -> VortexResult<ArrayRef>;
-
     /// Mask the array using the given boolean mask.
     /// The resulting array's validity is the intersection of the original array's validity
     /// and the mask's validity.
     fn mask(&self, mask: &ArrayRef) -> VortexResult<ArrayRef>;
+
+    /// Boolean negation.
+    fn not(&self) -> VortexResult<ArrayRef>;
 }
 
 impl ArrayBuiltins for ArrayRef {
     fn cast(&self, dtype: DType) -> VortexResult<ArrayRef> {
         cast::CastFn.try_new_array(self.len(), dtype, [self.clone()])
     }
 
-    fn is_null(&self) -> VortexResult<ArrayRef> {
-        is_null::IsNullFn.try_new_array(self.len(), EmptyOptions, [self.clone()])
+    fn get_item(&self, field_name: impl Into<FieldName>) -> VortexResult<ArrayRef> {
+        get_item::GetItemFn.try_new_array(self.len(), field_name.into(), [self.clone()])
     }
 
-    fn not(&self) -> VortexResult<ArrayRef> {
-        not::NotFn.try_new_array(self.len(), EmptyOptions, [self.clone()])
+    fn is_null(&self) -> VortexResult<ArrayRef> {
+        is_null::IsNullFn.try_new_array(self.len(), EmptyOptions, [self.clone()])
     }
 
     fn mask(&self, mask: &ArrayRef) -> VortexResult<ArrayRef> {
         mask::MaskFn.try_new_array(self.len(), EmptyOptions, [self.clone(), mask.clone()])
     }
+
+    fn not(&self) -> VortexResult<ArrayRef> {
+        not::NotFn.try_new_array(self.len(), EmptyOptions, [self.clone()])
+    }
 }

vortex-vector/src/scalar.rs

@@ -130,3 +130,77 @@ impl Scalar {
         vortex_panic!("Cannot convert non-struct Scalar into StructScalar");
     }
 }
+
+impl Scalar {
+    /// Converts the `Scalar` into a `NullScalar`.
+    pub fn as_null(&self) -> &NullScalar {
+        if let Scalar::Null(scalar) = self {
+            return scalar;
+        }
+        vortex_panic!("Cannot convert non-null Scalar into NullScalar");
+    }
+
+    /// Converts the `Scalar` into a `BoolScalar`.
+    pub fn as_bool(&self) -> &BoolScalar {
+        if let Scalar::Bool(scalar) = self {
+            return scalar;
+        }
+        vortex_panic!("Cannot convert non-bool Scalar into BoolScalar");
+    }
+
+    /// Converts the `Scalar` into a `DecimalScalar`.
+    pub fn as_decimal(&self) -> &DecimalScalar {
+        if let Scalar::Decimal(scalar) = self {
+            return scalar;
+        }
+        vortex_panic!("Cannot convert non-decimal Scalar into DecimalScalar");
+    }
+
+    /// Converts the `Scalar` into a `PrimitiveScalar`.
+    pub fn as_primitive(&self) -> &PrimitiveScalar {
+        if let Scalar::Primitive(scalar) = self {
+            return scalar;
+        }
+        vortex_panic!("Cannot convert non-primitive Scalar into PrimitiveScalar");
+    }
+
+    /// Converts the `Scalar` into a `StringScalar`.
+    pub fn as_string(&self) -> &StringScalar {
+        if let Scalar::String(scalar) = self {
+            return scalar;
+        }
+        vortex_panic!("Cannot convert non-string Scalar into StringScalar");
+    }
+
+    /// Converts the `Scalar` into a `BinaryScalar`.
+    pub fn as_binary(&self) -> &BinaryScalar {
+        if let Scalar::Binary(scalar) = self {
+            return scalar;
+        }
+        vortex_panic!("Cannot convert non-binary Scalar into BinaryScalar");
+    }
+
+    /// Converts the `Scalar` into a `ListViewScalar`.
+    pub fn as_list(&self) -> &ListViewScalar {
+        if let Scalar::List(scalar) = self {
+            return scalar;
+        }
+        vortex_panic!("Cannot convert non-list Scalar into ListViewScalar");
+    }
+
+    /// Converts the `Scalar` into a `FixedSizeListScalar`.
+    pub fn as_fixed_size_list(&self) -> &FixedSizeListScalar {
+        if let Scalar::FixedSizeList(scalar) = self {
+            return scalar;
+        }
+        vortex_panic!("Cannot convert non-fixed-size-list Scalar into FixedSizeListScalar");
+    }
+
+    /// Converts the `Scalar` into a `StructScalar`.
+    pub fn as_struct(&self) -> &StructScalar {
+        if let Scalar::Struct(scalar) = self {
+            return scalar;
+        }
+        vortex_panic!("Cannot convert non-struct Scalar into StructScalar");
+    }
+}

vortex-vector/src/struct_/scalar.rs

@@ -29,6 +29,11 @@ impl StructScalar {
     pub fn value(&self) -> &StructVector {
         &self.0
     }
+
+    /// Returns the nth field scalar of the struct.
+    pub fn field(&self, field_idx: usize) -> Scalar {
+        self.0.fields()[field_idx].scalar_at(0)
+    }
 }
 
 impl ScalarOps for StructScalar {