feat(expr): add bound expressions

vortex-array/src/expr/bound_expression.rs

@@ -0,0 +1,390 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+use std::fmt;
+use std::fmt::Display;
+use std::fmt::Formatter;
+use std::hash::Hash;
+use std::hash::Hasher;
+use std::sync::Arc;
+
+use itertools::Itertools;
+use vortex_error::VortexExpect;
+use vortex_error::VortexResult;
+use vortex_error::vortex_ensure;
+
+use crate::dtype::DType;
+use crate::expr::Expression;
+use crate::expr::display::DisplayTreeExpr;
+use crate::expr::scope::Scope;
+use crate::scalar_fn::ScalarFnRef;
+use crate::scalar_fn::fns::root::Root;
+
+/// An [`Expression`] that has been type-checked against a [`Scope`].
+///
+/// Every node carries its own dtype, so reading one is a field access rather than a walk of the
+/// subtree. Holding a `BoundExpression` is proof that the whole tree type-checked.
+///
+/// Binding is purely logical: it deals only in [`DType`]s and never sees an array, a length, or an
+/// encoding.
+#[derive(Clone, Debug, PartialEq, Eq, Hash)]
+pub struct BoundExpression {
+    kind: BoundKind,
+    dtype: DType,
+}
+
+/// The per-variant contents of a [`BoundExpression`], mirroring the logical variants of
+/// [`Expression`].
+#[derive(Clone, Debug, PartialEq, Eq, Hash)]
+pub enum BoundKind {
+    /// A scalar function applied to bound children.
+    Scalar {
+        /// The scalar function for this node.
+        scalar_fn: ScalarFnRef,
+        /// The bound children, in argument order.
+        ///
+        /// Sharing keeps clones cheap even though the iterative [`Drop`] implementation prevents
+        /// consumers from destructuring a `BoundExpression` by value.
+        children: Arc<Vec<BoundExpression>>,
+    },
+    /// The scope itself. Its dtype is the scope's root dtype.
+    Root,
+}
+
+/// A bound-expression wrapper that compares shared tree identity instead of structure.
+#[derive(Clone, Debug)]
+pub struct ExactBoundExpr(pub BoundExpression);
+
+impl PartialEq for ExactBoundExpr {
+    fn eq(&self, other: &Self) -> bool {
+        if self.0.dtype != other.0.dtype {
+            return false;
+        }
+
+        match (&self.0.kind, &other.0.kind) {
+            (BoundKind::Root, BoundKind::Root) => true,
+            (
+                BoundKind::Scalar {
+                    scalar_fn: lhs_fn,
+                    children: lhs_children,
+                },
+                BoundKind::Scalar {
+                    scalar_fn: rhs_fn,
+                    children: rhs_children,
+                },
+            ) => lhs_fn == rhs_fn && Arc::ptr_eq(lhs_children, rhs_children),
+            _ => false,
+        }
+    }
+}
+
+impl Eq for ExactBoundExpr {}
+
+impl Hash for ExactBoundExpr {
+    fn hash<H: Hasher>(&self, state: &mut H) {
+        self.0.dtype.hash(state);
+        match &self.0.kind {
+            BoundKind::Root => state.write_u8(0),
+            BoundKind::Scalar {
+                scalar_fn,
+                children,
+            } => {
+                state.write_u8(1);
+                scalar_fn.hash(state);
+                Arc::as_ptr(children).hash(state);
+            }
+        }
+    }
+}
+
+impl BoundExpression {
+    /// Create a bound root expression with the given dtype.
+    pub fn new_root(dtype: DType) -> Self {
+        Self {
+            kind: BoundKind::Root,
+            dtype,
+        }
+    }
+
+    /// Create a bound scalar node from a scalar function and already-bound children.
+    pub fn try_new(
+        scalar_fn: ScalarFnRef,
+        children: impl IntoIterator<Item = BoundExpression>,
+    ) -> VortexResult<Self> {
+        let children = Vec::from_iter(children);
+        vortex_ensure!(
+            scalar_fn.signature().arity().matches(children.len()),
+            "Expression arity mismatch: expected {} children but got {}",
+            scalar_fn.signature().arity(),
+            children.len()
+        );
+
+        let arg_dtypes = children
+            .iter()
+            .map(|child| child.dtype().clone())
+            .collect_vec();
+        let dtype = scalar_fn.return_dtype(&arg_dtypes)?;
+
+        Ok(Self {
+            kind: BoundKind::Scalar {
+                scalar_fn,
+                children: children.into(),
+            },
+            dtype,
+        })
+    }
+
+    /// Rebuild this node with new bound children, recomputing its dtype.
+    pub fn with_children(
+        self,
+        children: impl IntoIterator<Item = BoundExpression>,
+    ) -> VortexResult<Self> {
+        let children = Vec::from_iter(children);
+        let BoundKind::Scalar { scalar_fn, .. } = &self.kind else {
+            vortex_ensure!(
+                children.is_empty(),
+                "Root expression cannot have {} children",
+                children.len()
+            );
+            return Ok(self);
+        };
+
+        Self::try_new(scalar_fn.clone(), children)
+    }
+
+    /// The dtype this expression evaluates to.
+    pub fn dtype(&self) -> &DType {
+        &self.dtype
+    }
+
+    /// The per-variant contents of this node.
+    pub fn kind(&self) -> &BoundKind {
+        &self.kind
+    }
+
+    /// The bound children of this node, in argument order. Empty for [`BoundKind::Root`].
+    pub fn children(&self) -> &[BoundExpression] {
+        match &self.kind {
+            BoundKind::Scalar { children, .. } => children.as_slice(),
+            BoundKind::Root => &[],
+        }
+    }
+
+    /// The scalar function for this node, or `None` if it is the scope root.
+    pub fn as_scalar(&self) -> Option<&ScalarFnRef> {
+        match &self.kind {
+            BoundKind::Scalar { scalar_fn, .. } => Some(scalar_fn),
+            BoundKind::Root => None,
+        }
+    }
+
+    /// Whether this node is the scope root.
+    pub fn is_root(&self) -> bool {
+        matches!(self.kind, BoundKind::Root)
+    }
+
+    /// Display the bound expression as a formatted tree structure.
+    pub fn display_tree(&self) -> impl Display {
+        DisplayTreeExpr(self)
+    }
+
+    /// Convert this bound tree back into its unbound logical representation.
+    ///
+    /// This rebuilds the expression iteratively; the bound representation does not retain a
+    /// second expression tree.
+    // TODO: This is temporary artifact of the migration from using `Expression`s to
+    // `BoundExpression`s
+    pub fn unbind(&self) -> Expression {
+        let mut pending = vec![(self, false)];
+        let mut expressions = Vec::new();
+
+        while let Some((node, visited)) = pending.pop() {
+            match node.kind() {
+                BoundKind::Root => expressions.push(crate::expr::root()),
+                BoundKind::Scalar {
+                    scalar_fn,
+                    children,
+                } if visited => {
+                    let child_start = expressions.len() - children.len();
+                    let child_expressions = expressions.split_off(child_start);
+                    expressions.push(
+                        Expression::try_new(scalar_fn.clone(), child_expressions)
+                            .vortex_expect("a bound expression always has valid arity"),
+                    );
+                }
+                BoundKind::Scalar { children, .. } => {
+                    pending.push((node, true));
+                    pending.extend(children.iter().rev().map(|child| (child, false)));
+                }
+            }
+        }
+
+        expressions
+            .pop()
+            .vortex_expect("binding always produces one expression root")
+    }
+}
+
+impl Display for BoundExpression {
+    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
+        Display::fmt(&self.unbind(), f)
+    }
+}
+
+impl Expression {
+    /// Bind this expression against a root dtype, type-checking every node in a single walk.
+    ///
+    /// The returned tree carries a dtype on each node, so callers needing types at more than one
+    /// node should bind once and read fields rather than calling
+    /// [`return_dtype`](Expression::return_dtype) repeatedly.
+    pub fn bind(&self, dtype: &DType) -> VortexResult<BoundExpression> {
+        self.bind_scope(&Scope::new(dtype.clone()))
+    }
+
+    /// Bind this expression against an explicit [`Scope`].
+    pub fn bind_scope(&self, scope: &Scope) -> VortexResult<BoundExpression> {
+        if self.is::<Root>() {
+            return Ok(BoundExpression::new_root(scope.root().clone()));
+        }
+
+        let children: Vec<_> = self
+            .children()
+            .iter()
+            .map(|child| child.bind_scope(scope))
+            .try_collect()?;
+        BoundExpression::try_new(self.scalar_fn().clone(), children)
+    }
+}
+
+/// Iterative drop to avoid stack overflows on deep trees.
+impl Drop for BoundExpression {
+    fn drop(&mut self) {
+        let BoundKind::Scalar { children, .. } = &mut self.kind else {
+            return;
+        };
+        let Some(children) = Arc::get_mut(children) else {
+            return;
+        };
+
+        let mut to_drop = std::mem::take(children);
+        while let Some(mut child) = to_drop.pop() {
+            if let BoundKind::Scalar { children, .. } = &mut child.kind
+                && let Some(grandchildren) = Arc::get_mut(children)
+            {
+                to_drop.append(grandchildren);
+            }
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use vortex_error::VortexResult;
+
+    use super::*;
+    use crate::dtype::Nullability;
+    use crate::dtype::PType;
+    use crate::expr::col;
+    use crate::expr::eq;
+    use crate::expr::lit;
+    use crate::expr::root;
+    use crate::expr::test_harness::struct_dtype;
+
+    fn scope() -> Scope {
+        Scope::new(struct_dtype())
+    }
+
+    #[test]
+    fn root_binds_to_the_scope() -> VortexResult<()> {
+        let bound = root().bind_scope(&scope())?;
+        assert!(bound.is_root());
+        assert_eq!(bound.dtype(), &struct_dtype());
+        assert_eq!(bound.unbind(), root());
+        Ok(())
+    }
+
+    #[test]
+    fn every_node_carries_its_dtype() -> VortexResult<()> {
+        let expr = eq(col("a"), lit(1_i32));
+        let bound = expr.bind_scope(&scope())?;
+
+        assert_eq!(bound.dtype(), &DType::Bool(Nullability::NonNullable));
+
+        let lhs = &bound.children()[0];
+        assert_eq!(
+            lhs.dtype(),
+            &DType::Primitive(PType::I32, Nullability::NonNullable)
+        );
+        assert_eq!(lhs.children()[0].dtype(), &struct_dtype());
+        Ok(())
+    }
+
+    #[test]
+    fn bind_agrees_with_return_dtype() -> VortexResult<()> {
+        for expr in [root(), col("a"), eq(col("a"), lit(1_i32)), lit(true)] {
+            assert_eq!(
+                expr.bind(&struct_dtype())?.dtype(),
+                &expr.return_dtype(&struct_dtype())?,
+                "disagreement for {expr}"
+            );
+        }
+        Ok(())
+    }
+
+    #[test]
+    fn bound_tree_display_matches_unbound() -> VortexResult<()> {
+        for expr in [root(), col("a"), eq(col("a"), lit(1_i32)), lit(true)] {
+            let bound = expr.bind_scope(&scope())?;
+            assert_eq!(
+                bound.display_tree().to_string(),
+                expr.display_tree().to_string()
+            );
+        }
+        Ok(())
+    }
+
+    #[test]
+    fn clone_shares_children() -> VortexResult<()> {
+        let bound = eq(col("a"), lit(1_i32)).bind_scope(&scope())?;
+        let cloned = bound.clone();
+
+        let (BoundKind::Scalar { children: a, .. }, BoundKind::Scalar { children: b, .. }) =
+            (bound.kind(), cloned.kind())
+        else {
+            unreachable!("eq is a scalar node")
+        };
+        assert!(Arc::ptr_eq(a, b));
+        Ok(())
+    }
+
+    #[test]
+    fn repeated_subtree_is_bound_per_occurrence() -> VortexResult<()> {
+        let shared = col("a");
+        let bound = eq(shared.clone(), shared).bind_scope(&scope())?;
+        let children = bound.children();
+        assert_eq!(children[0].dtype(), children[1].dtype());
+        Ok(())
+    }
+
+    #[test]
+    fn structural_and_exact_equality_are_distinct() -> VortexResult<()> {
+        let expr = eq(col("a"), lit(1_i32));
+        let bound = expr.bind_scope(&scope())?;
+        let independently_bound = expr.bind_scope(&scope())?;
+
+        assert_eq!(bound, independently_bound);
+        assert_eq!(ExactBoundExpr(bound.clone()), ExactBoundExpr(bound.clone()));
+        assert_ne!(
+            ExactBoundExpr(bound.clone()),
+            ExactBoundExpr(independently_bound)
+        );
+        assert_eq!(bound.unbind(), expr);
+        Ok(())
+    }
+
+    #[test]
+    fn binding_reports_a_type_error() {
+        let expr = eq(col("a"), lit("nope"));
+        assert!(expr.bind_scope(&scope()).is_err());
+    }
+}