Move rules onto array vtable so we can optimize during construction without a session

gatesn · gatesn · commit 157579efd978 · 2025-12-12T15:42:51.000Z
Signed-off-by: Nicholas Gates &lt;nick@nickgates.com&gt;
diff --git a/vortex-array/src/arrays/scalar_fn/mod.rs b/vortex-array/src/arrays/scalar_fn/mod.rs
@@ -4,6 +4,7 @@
 mod array;
 mod kernel;
 mod metadata;
+mod rules;
 mod vtable;
 
 pub use array::*;
diff --git a/vortex-array/src/arrays/scalar_fn/rules.rs b/vortex-array/src/arrays/scalar_fn/rules.rs
@@ -0,0 +1,144 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+use std::any::Any;
+
+use vortex_dtype::DType;
+use vortex_error::VortexExpect;
+use vortex_error::VortexResult;
+use vortex_error::vortex_ensure;
+use vortex_scalar::Scalar;
+use vortex_vector::Datum;
+use vortex_vector::VectorOps;
+use vortex_vector::datum_matches_dtype;
+
+use crate::ArrayRef;
+use crate::ArrayVisitor;
+use crate::IntoArray;
+use crate::arrays::ConstantArray;
+use crate::arrays::ConstantVTable;
+use crate::arrays::ScalarFnArray;
+use crate::arrays::ScalarFnVTable;
+use crate::expr::ExecutionArgs;
+use crate::expr::ReduceCtx;
+use crate::expr::ReduceNode;
+use crate::expr::ScalarFn;
+use crate::optimizer::rules::ArrayReduceRule;
+use crate::optimizer::rules::ReduceRuleSet;
+
+pub(super) const RULES: ReduceRuleSet<ScalarFnVTable> =
+    ReduceRuleSet::new(&[&ScalarFnConstantRule, &ScalarFnAbstractReduceRule]);
+
+#[derive(Debug)]
+struct ScalarFnConstantRule;
+impl ArrayReduceRule<ScalarFnVTable> for ScalarFnConstantRule {
+    fn reduce(&self, array: &ScalarFnArray) -> VortexResult<Option<ArrayRef>> {
+        if !array.children.iter().all(|c| c.is::<ConstantVTable>()) {
+            return Ok(None);
+        }
+
+        let input_datums: Vec<_> = array
+            .children
+            .iter()
+            .map(|c| c.as_::<ConstantVTable>().scalar().to_vector_scalar())
+            .map(Datum::Scalar)
+            .collect();
+        let input_dtypes = array.children.iter().map(|c| c.dtype().clone()).collect();
+
+        let result = array.scalar_fn.execute(ExecutionArgs {
+            datums: input_datums,
+            dtypes: input_dtypes,
+            row_count: array.len,
+            return_dtype: array.dtype.clone(),
+        })?;
+        vortex_ensure!(
+            datum_matches_dtype(&result, &array.dtype),
+            "Scalar function {} result does not match expected dtype",
+            array.scalar_fn
+        );
+
+        let result = match result {
+            Datum::Scalar(s) => s,
+            Datum::Vector(v) => {
+                tracing::warn!(
+                    "Scalar function {} returned vector from execution over all scalar inputs",
+                    array.scalar_fn
+                );
+                v.scalar_at(0)
+            }
+        };
+
+        Ok(Some(
+            ConstantArray::new(Scalar::from_vector_scalar(result, &array.dtype)?, array.len)
+                .into_array(),
+        ))
+    }
+}
+
+#[derive(Debug)]
+struct ScalarFnAbstractReduceRule;
+impl ArrayReduceRule<ScalarFnVTable> for ScalarFnAbstractReduceRule {
+    fn reduce(&self, array: &ScalarFnArray) -> VortexResult<Option<ArrayRef>> {
+        if let Some(reduced) = array.scalar_fn.reduce(
+            // Blergh, re-boxing
+            &array.to_array(),
+            &ArrayReduceCtx { len: array.len },
+        )? {
+            return Ok(Some(
+                reduced
+                    .as_any()
+                    .downcast_ref::<ArrayRef>()
+                    .vortex_expect("ReduceNode is not an ArrayRef")
+                    .clone(),
+            ));
+        }
+        Ok(None)
+    }
+}
+
+impl ReduceNode for ArrayRef {
+    fn as_any(&self) -> &dyn Any {
+        self
+    }
+
+    fn node_dtype(&self) -> VortexResult<DType> {
+        Ok(self.as_ref().dtype().clone())
+    }
+
+    fn scalar_fn(&self) -> Option<&ScalarFn> {
+        self.as_opt::<ScalarFnVTable>().map(|a| a.scalar_fn())
+    }
+
+    fn child(&self, idx: usize) -> Box<dyn ReduceNode> {
+        Box::new(self.children()[idx].clone())
+    }
+}
+
+struct ArrayReduceCtx {
+    // The length of the array being reduced
+    len: usize,
+}
+impl ReduceCtx for ArrayReduceCtx {
+    fn create_node(
+        &self,
+        scalar_fn: ScalarFn,
+        children: &[Box<dyn ReduceNode>],
+    ) -> VortexResult<Box<dyn ReduceNode>> {
+        Ok(Box::new(
+            ScalarFnArray::try_new(
+                scalar_fn,
+                children
+                    .iter()
+                    .map(|c| {
+                        c.as_any()
+                            .downcast_ref::<ArrayRef>()
+                            .vortex_expect("ReduceNode is not an ArrayRef")
+                            .clone()
+                    })
+                    .collect(),
+                self.len,
+            )?
+            .into_array(),
+        ))
+    }
+}
diff --git a/vortex-array/src/arrays/scalar_fn/vtable/mod.rs b/vortex-array/src/arrays/scalar_fn/vtable/mod.rs
@@ -16,22 +16,17 @@ use vortex_dtype::DType;
 use vortex_error::VortexResult;
 use vortex_error::vortex_bail;
 use vortex_error::vortex_ensure;
-use vortex_scalar::Scalar;
-use vortex_vector::Datum;
-use vortex_vector::VectorOps;
-use vortex_vector::datum_matches_dtype;
 
 use crate::Array;
 use crate::ArrayRef;
 use crate::IntoArray;
-use crate::arrays::ConstantArray;
 use crate::arrays::ConstantVTable;
 use crate::arrays::scalar_fn::array::ScalarFnArray;
 use crate::arrays::scalar_fn::kernel::KernelInput;
 use crate::arrays::scalar_fn::kernel::ScalarFnKernel;
 use crate::arrays::scalar_fn::metadata::ScalarFnMetadata;
+use crate::arrays::scalar_fn::rules::RULES;
 use crate::expr;
-use crate::expr::ExecutionArgs;
 use crate::expr::ExprVTable;
 use crate::expr::ScalarFn;
 use crate::kernel::BindCtx;
@@ -170,45 +165,7 @@ impl VTable for ScalarFnVTable {
     }
 
     fn reduce(array: &Self::Array) -> VortexResult<Option<ArrayRef>> {
-        if !array.children.iter().all(|c| c.is::<ConstantVTable>()) {
-            return Ok(None);
-        }
-
-        let input_datums: Vec<_> = array
-            .children
-            .iter()
-            .map(|c| c.as_::<ConstantVTable>().scalar().to_vector_scalar())
-            .map(Datum::Scalar)
-            .collect();
-        let input_dtypes = array.children.iter().map(|c| c.dtype().clone()).collect();
-
-        let result = array.scalar_fn.execute(ExecutionArgs {
-            datums: input_datums,
-            dtypes: input_dtypes,
-            row_count: array.len,
-            return_dtype: array.dtype.clone(),
-        })?;
-        vortex_ensure!(
-            datum_matches_dtype(&result, &array.dtype),
-            "Scalar function {} result does not match expected dtype",
-            array.scalar_fn
-        );
-
-        let result = match result {
-            Datum::Scalar(s) => s,
-            Datum::Vector(v) => {
-                tracing::warn!(
-                    "Scalar function {} returned vector from execution over all scalar inputs",
-                    array.scalar_fn
-                );
-                v.scalar_at(0)
-            }
-        };
-
-        Ok(Some(
-            ConstantArray::new(Scalar::from_vector_scalar(result, &array.dtype)?, array.len)
-                .into_array(),
-        ))
+        RULES.evaluate(array)
     }
 }
 
diff --git a/vortex-array/src/expr/exprs/get_item.rs b/vortex-array/src/expr/exprs/get_item.rs
@@ -23,10 +23,15 @@ use crate::builtins::ExprBuiltins;
 use crate::compute::mask;
 use crate::expr::Arity;
 use crate::expr::ChildName;
+use crate::expr::EmptyOptions;
 use crate::expr::ExecutionArgs;
 use crate::expr::ExprId;
 use crate::expr::Expression;
+use crate::expr::Literal;
+use crate::expr::Mask;
 use crate::expr::Pack;
+use crate::expr::ReduceCtx;
+use crate::expr::ReduceNode;
 use crate::expr::StatsCatalog;
 use crate::expr::VTable;
 use crate::expr::VTableExt;
@@ -135,6 +140,33 @@ impl VTable for GetItem {
         }
     }
 
+    fn reduce(
+        &self,
+        field_name: &FieldName,
+        node: &dyn ReduceNode,
+        ctx: &dyn ReduceCtx,
+    ) -> VortexResult<Option<Box<dyn ReduceNode>>> {
+        let child = node.child(0);
+        if let Some(child_fn) = child.scalar_fn()
+            && let Some(pack) = child_fn.as_opt::<Pack>()
+            && let Some(idx) = pack.names.find(field_name)
+        {
+            let mut field = child.child(idx);
+
+            // Possibly mask the field if the pack is nullable
+            if pack.nullability.is_nullable() {
+                field = ctx.create_node(
+                    Mask.bind(EmptyOptions),
+                    &[field, ctx.create_node(Literal.bind(true.into()), &[])?],
+                )?;
+            }
+
+            return Ok(Some(field));
+        }
+
+        Ok(None)
+    }
+
     fn simplify_untyped(
         &self,
         field_name: &FieldName,
diff --git a/vortex-array/src/expr/exprs/select.rs b/vortex-array/src/expr/exprs/select.rs
@@ -136,41 +136,6 @@ impl VTable for Select {
         Ok(DType::Struct(projected, child_dtype.nullability()))
     }
 
-    fn simplify(
-        &self,
-        options: &Self::Options,
-        expr: &Expression,
-        ctx: &dyn SimplifyCtx,
-    ) -> VortexResult<Option<Expression>> {
-        let child = expr.child(0);
-        let child_dtype = ctx.return_dtype(child)?;
-        let child_nullability = child_dtype.nullability();
-
-        let child_dtype = child_dtype.as_struct_fields_opt().ok_or_else(|| {
-            vortex_err!(
-                "Select child must return a struct dtype, however it was a {}",
-                child_dtype
-            )
-        })?;
-
-        let expr = pack(
-            options
-                .as_include_names(child_dtype.names())
-                .map_err(|e| {
-                    e.with_context(format!(
-                        "Select fields {:?} must be a subset of child fields {:?}",
-                        options,
-                        child_dtype.names()
-                    ))
-                })?
-                .iter()
-                .map(|name| (name.clone(), get_item(name.clone(), child.clone()))),
-            child_nullability,
-        );
-
-        Ok(Some(expr))
-    }
-
     fn evaluate(
         &self,
         selection: &FieldSelection,
@@ -238,6 +203,41 @@ impl VTable for Select {
         .into())
     }
 
+    fn simplify(
+        &self,
+        options: &Self::Options,
+        expr: &Expression,
+        ctx: &dyn SimplifyCtx,
+    ) -> VortexResult<Option<Expression>> {
+        let child = expr.child(0);
+        let child_dtype = ctx.return_dtype(child)?;
+        let child_nullability = child_dtype.nullability();
+
+        let child_dtype = child_dtype.as_struct_fields_opt().ok_or_else(|| {
+            vortex_err!(
+                "Select child must return a struct dtype, however it was a {}",
+                child_dtype
+            )
+        })?;
+
+        let expr = pack(
+            options
+                .as_include_names(child_dtype.names())
+                .map_err(|e| {
+                    e.with_context(format!(
+                        "Select fields {:?} must be a subset of child fields {:?}",
+                        options,
+                        child_dtype.names()
+                    ))
+                })?
+                .iter()
+                .map(|name| (name.clone(), get_item(name.clone(), child.clone()))),
+            child_nullability,
+        );
+
+        Ok(Some(expr))
+    }
+
     fn is_null_sensitive(&self, _instance: &Self::Options) -> bool {
         true
     }
diff --git a/vortex-array/src/expr/scalar_fn.rs b/vortex-array/src/expr/scalar_fn.rs
@@ -19,6 +19,8 @@ use crate::expr::ExecutionArgs;
 use crate::expr::ExprId;
 use crate::expr::ExprVTable;
 use crate::expr::Expression;
+use crate::expr::ReduceCtx;
+use crate::expr::ReduceNode;
 use crate::expr::VTable;
 use crate::expr::options::ExpressionOptions;
 use crate::expr::signature::ExpressionSignature;
@@ -74,6 +76,16 @@ impl ScalarFn {
         }
     }
 
+    /// Returns whether the scalar function is of the given vtable type.
+    pub fn is<V: VTable>(&self) -> bool {
+        self.vtable.is::<V>()
+    }
+
+    /// Returns the typed options for this expression if it matches the given vtable type.
+    pub fn as_opt<V: VTable>(&self) -> Option<&V::Options> {
+        self.options().as_any().downcast_ref::<V::Options>()
+    }
+
     /// Signature information for this expression.
     pub fn signature(&self) -> ExpressionSignature<'_> {
         ExpressionSignature {
@@ -101,6 +113,15 @@ impl ScalarFn {
     pub fn execute(&self, ctx: ExecutionArgs) -> VortexResult<Datum> {
         self.vtable.as_dyn().execute(self.options.deref(), ctx)
     }
+
+    /// Perform abstract reduction on this scalar function node.
+    pub fn reduce(
+        &self,
+        node: &dyn ReduceNode,
+        ctx: &dyn ReduceCtx,
+    ) -> VortexResult<Option<Box<dyn ReduceNode>>> {
+        self.vtable.as_dyn().reduce(self.options.deref(), node, ctx)
+    }
 }
 
 impl Clone for ScalarFn {
diff --git a/vortex-array/src/expr/vtable.rs b/vortex-array/src/expr/vtable.rs
diff --git a/vortex-array/src/optimizer/rules.rs b/vortex-array/src/optimizer/rules.rs
diff --git a/vortex-dtype/src/field_names.rs b/vortex-dtype/src/field_names.rs
