/// A scalar function applied to child expressions.

Scalar {

/// The scalar fn for this node.

scalar_fn: ScalarFnRef,

/// Any children of this expression.

children: Arc<Vec<Expression>>,

},

/// The full scope of the expression evaluation.

Root,

/// Create a new expression node from a scalar_fn expression and its children.

pub fn try_new(

scalar_fn: ScalarFnRef,

children: impl IntoIterator<Item = Expression>,

) -> 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()

);

Ok(Self::Scalar {

scalar_fn,

children: children.into(),

})

}

/// Whether this expression is the scope root.

pub fn is_root(&self) -> bool {

matches!(self, Self::Root)

}

/// Returns the scalar fn for this expression, or `None` if it is not a scalar node.

pub fn as_scalar(&self) -> Option<&ScalarFnRef> {

match self {

Self::Scalar { scalar_fn, .. } => Some(scalar_fn),

Self::Root => None,

}

}

/// Whether this expression's scalar fn is of the given vtable type.

pub fn is<V: ScalarFnVTable>(&self) -> bool {

self.as_scalar().is_some_and(|sf| sf.is::<V>())

}

/// The typed options for this expression if its scalar fn matches the given vtable type.

pub fn as_opt<V: ScalarFnVTable>(&self) -> Option<&V::Options> {

self.as_scalar().and_then(|sf| sf.as_opt::<V>())

}

/// The typed options for this expression.

///

/// # Panics

///

/// Panics if the vtable type does not match.

pub fn as_<V: ScalarFnVTable>(&self) -> &V::Options {

self.as_opt::<V>()

.vortex_expect("Expression options type mismatch")

}

/// Returns the children of this expression.

pub fn children(&self) -> &[Expression] {

match self {

Self::Scalar { children, .. } => children.as_slice(),

Self::Root => NO_CHILDREN,

}

}

/// Returns the n'th child of this expression.

pub fn child(&self, n: usize) -> &Expression {

&self.children()[n]

}

/// Replace the children of this expression with the provided new children.

pub fn with_children(

self,

children: impl IntoIterator<Item = Expression>,

) -> VortexResult<Self> {

let children = Vec::from_iter(children);

match &self {

Self::Root => {

vortex_ensure!(

children.is_empty(),

"Expression arity mismatch: root expects 0 children but got {}",

children.len()

);

Ok(Self::Root)

}

Self::Scalar { scalar_fn, .. } => {

vortex_ensure!(

scalar_fn.signature().arity().matches(children.len()),

"Expression arity mismatch: expected {} children but got {}",

scalar_fn.signature().arity(),

children.len()

);

Ok(Self::Scalar {

scalar_fn: scalar_fn.clone(),

children: children.into(),

})

}

}

}

/// Computes the return dtype of this expression given the input dtype.

pub fn return_dtype(&self, scope: &DType) -> VortexResult<DType> {

match self {

Self::Root => Ok(scope.clone()),

Self::Scalar {

scalar_fn,

children,

} => {

let dtypes: Vec<_> = children

.iter()

.map(|c| c.return_dtype(scope))

.try_collect()?;

scalar_fn.return_dtype(&dtypes)

}

}

}

/// Returns a new expression representing the validity mask output of this expression.

///

/// The returned expression evaluates to a non-nullable boolean array.

pub fn validity(&self) -> VortexResult<Expression> {

match self {

// The scope is exactly as valid as itself.

Self::Root => Ok(Self::Root),

Self::Scalar { scalar_fn, .. } => scalar_fn.validity(self),

}

}

/// Format the expression as a compact string.

///

/// Since this is a recursive formatter, it is exposed on the public Expression type.

/// See fmt_data that is only implemented on the vtable trait.

pub fn fmt_sql(&self, f: &mut Formatter<'_>) -> fmt::Result {

match self {

Self::Root => write!(f, "$"),

Self::Scalar { scalar_fn, .. } => scalar_fn.fmt_sql(self, f),

}

}

/// Display the expression as a formatted tree structure.

///

/// This provides a hierarchical view of the expression that shows the relationships

/// between parent and child expressions, making complex nested expressions easier

/// to understand and debug.

///

/// # Example

///

/// ```rust

/// # use vortex_array::dtype::{DType, Nullability, PType};

/// # use vortex_array::scalar_fn::fns::like::{Like, LikeOptions};

/// # use vortex_array::scalar_fn::ScalarFnVTableExt;

/// # use vortex_array::expr::{and, cast, eq, get_item, gt, lit, not, root, select};

/// // Build a complex nested expression

/// let complex_expr = select(

/// ["result"],

/// and(

/// not(eq(get_item("status", root()), lit("inactive"))),

/// and(

/// Like.new_expr(LikeOptions::default(), [get_item("name", root()), lit("%admin%")]),

/// gt(

/// cast(get_item("score", root()), DType::Primitive(PType::F64, Nullability::NonNullable)),

/// lit(75.0)

/// )

/// )

/// )

/// );

///

/// println!("{}", complex_expr.display_tree());

/// ```

///

/// This produces output like:

///

/// ```text

/// Select(include): {result}

/// └── Binary(and)

/// ├── lhs: Not

/// │ └── Binary(=)

/// │ ├── lhs: GetItem(status)

/// │ │ └── Root

/// │ └── rhs: Literal(value: "inactive", dtype: utf8)

/// └── rhs: Binary(and)

/// ├── lhs: Like

/// │ ├── child: GetItem(name)

/// │ │ └── Root

/// │ └── pattern: Literal(value: "%admin%", dtype: utf8)

/// └── rhs: Binary(>)

/// ├── lhs: Cast(target: f64)

/// │ └── GetItem(score)

/// │ └── Root

/// └── rhs: Literal(value: 75f64, dtype: f64)

/// ```

pub fn display_tree(&self) -> impl Display {

DisplayTreeExpr(self)

}

/// Returns true if this expression contains expression E inside.

///

/// # Example

///

/// ```rust

/// # use vortex_array::scalar_fn::fns::literal::Literal;

/// # use vortex_array::expr::{eq, lit, root};

/// let expression = &eq(root(), lit(3u64));

/// assert!(expression.contains::<Literal>().unwrap());

/// let expression = root();

/// assert!(!expression.contains::<Literal>().unwrap());

/// ```

pub fn contains<E: ScalarFnVTable>(&self) -> VortexResult<bool> {

let mut contains = false;

pre_order_visit_down(self, |node| {

if node.is::<E>() {

contains = true;

return Ok(TraversalOrder::Stop);

}

Ok(TraversalOrder::Continue)

})?;

Ok(contains)

}

fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {

self.fmt_sql(f)

}

fn drop(&mut self) {

let Self::Scalar { children, .. } = self else {

return;

};

let Some(children) = Arc::get_mut(children) else {

return;

};

let mut children_to_drop = std::mem::take(children);

while let Some(mut child) = children_to_drop.pop() {

if let Self::Scalar { children, .. } = &mut child

&& let Some(expr_children) = Arc::get_mut(children)

{

children_to_drop.append(expr_children);

}

}

}