Minimal expressions API for vortex

Cargo.toml

@@ -11,6 +11,7 @@ members = [
     "vortex-datetime-parts",
     "vortex-dict",
     "vortex-error",
+    "vortex-expr",
     "vortex-fastlanes",
     "vortex-flatbuffers",
     "vortex-ipc",

vortex-dtype/src/dtype.rs

@@ -8,7 +8,8 @@ use DType::*;
 use crate::nullability::Nullability;
 use crate::{ExtDType, PType};
 
-pub type FieldNames = Arc<[Arc<str>]>;
+pub type FieldName = Arc<str>;
+pub type FieldNames = Arc<[FieldName]>;
 
 pub type Metadata = Vec<u8>;
 

vortex-dtype/src/ptype.rs

@@ -90,6 +90,26 @@ macro_rules! match_each_native_ptype {
     })
 }
 
+#[macro_export]
+macro_rules! match_each_potentially_negative_ptype {
+    ($self:expr, | $_:tt $enc:ident | $($body:tt)*) => ({
+        macro_rules! __with__ {( $_ $enc:ident ) => ( $($body)* )}
+        use $crate::PType;
+        use $crate::half::f16;
+        match $self {
+            PType::I8 => __with__! { i8 },
+            PType::I16 => __with__! { i16 },
+            PType::I32 => __with__! { i32 },
+            PType::I64 => __with__! { i64 },
+            PType::F16 => __with__! { f16 },
+            PType::F32 => __with__! { f32 },
+            PType::F64 => __with__! { f64 },
+            _ => panic!("Unsupported ptype {}", $self),
+
+        }
+    })
+}
+
 #[macro_export]
 macro_rules! match_each_integer_ptype {
     ($self:expr, | $_:tt $enc:ident | $($body:tt)*) => ({

vortex-expr/Cargo.toml

@@ -0,0 +1,23 @@
+[package]
+name = "vortex-expr"
+version = { workspace = true }
+description = "Vortex Expressions"
+homepage = { workspace = true }
+repository = { workspace = true }
+authors = { workspace = true }
+license = { workspace = true }
+keywords = { workspace = true }
+include = { workspace = true }
+edition = { workspace = true }
+rust-version = { workspace = true }
+
+[lints]
+workspace = true
+
+[dependencies]
+vortex-dtype = { path = "../vortex-dtype" }
+vortex-error = { path = "../vortex-error" }
+vortex-scalar = { path = "../vortex-scalar" }
+
+
+[dev-dependencies]

vortex-expr/README.md

@@ -0,0 +1,7 @@
+# Vortex Expressions
+
+Expressions for querying vortex arrays, designed to express a minimal
+superset of the row predicates that can be pushed down to vortex metadata.
+
+Takes inspiration from postgres https://www.postgresql.org/docs/current/sql-expressions.html
+and datafusion https://github.com/apache/datafusion/tree/5fac581efbaffd0e6a9edf931182517524526afd/datafusion/expr

vortex-expr/src/display.rs

@@ -0,0 +1,141 @@
+use core::fmt;
+use std::fmt::{Display, Formatter};
+
+use vortex_dtype::{match_each_native_ptype, DType};
+use vortex_scalar::{BoolScalar, PrimitiveScalar};
+
+use crate::expressions::{BinaryExpr, Expr};
+use crate::operators::{Associativity, Operator};
+use crate::scalar::ScalarDisplayWrapper;
+
+impl Display for Expr {
+    fn fmt(&self, f: &mut Formatter) -> fmt::Result {
+        match self {
+            Expr::Binary(expr) => write!(f, "{expr}"),
+            Expr::Field(d) => write!(f, "{d}"),
+            Expr::Literal(v) => {
+                let wrapped = ScalarDisplayWrapper(v);
+                write!(f, "{wrapped}")
+            }
+            Expr::Not(expr) => write!(f, "NOT {expr}"),
+            Expr::IsNull(expr) => write!(f, "{expr} IS NULL"),
+        }
+    }
+}
+
+enum Side {
+    Left,
+    Right,
+}
+
+impl Display for BinaryExpr {
+    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
+        fn write_inner(
+            f: &mut Formatter<'_>,
+            outer: &Expr,
+            outer_op: Operator,
+            side: Side,
+        ) -> fmt::Result {
+            if let Expr::Binary(inner) = outer {
+                let inner_op_precedence = inner.op.precedence();
+                // if the inner operator has higher precedence than the outer expression,
+                // wrap it in parentheses to prevent inversion of priority
+                if inner_op_precedence > outer_op.precedence() ||
+                    // if the inner and outer operators have the same precedence, we need to
+                    // account for operator associativity when determining grouping
+                    (inner_op_precedence == outer_op.precedence() &&
+                        match side {
+                            Side::Left => {
+                                outer_op.associativity() == Associativity::Left
+                            }
+                            Side::Right => {
+                                outer_op.associativity() == Associativity::Right
+                            }
+                        })
+                {
+                    write!(f, "({inner})")?;
+                } else {
+                    write!(f, "{inner}")?;
+                }
+            } else if let Expr::Literal(scalar) = outer {
+                // use alternative formatting for scalars
+                let wrapped = ScalarDisplayWrapper(scalar);
+                write!(f, "{wrapped}")?;
+            } else {
+                write!(f, "{outer}")?;
+            }
+            Ok(())
+        }
+
+        write_inner(f, self.left.as_ref(), self.op, Side::Left)?;
+        write!(f, " {} ", self.op)?;
+        write_inner(f, self.right.as_ref(), self.op, Side::Right)
+    }
+}
+
+/// Alternative display for scalars
+impl Display for ScalarDisplayWrapper<'_> {
+    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
+        match self.0.dtype() {
+            DType::Null => write!(f, "null"),
+            DType::Bool(_) => match BoolScalar::try_from(self.0).expect("bool").value() {
+                None => write!(f, "null"),
+                Some(b) => write!(f, "{}", b),
+            },
+            DType::Primitive(ptype, _) => match_each_native_ptype!(ptype, |$T| {
+                match PrimitiveScalar::try_from(self.0).expect("primitive").typed_value::<$T>() {
+                    None => write!(f, "null"),
+                    Some(v) => write!(f, "{}{}", v,  std::any::type_name::<$T>()),
+                }
+            }),
+            DType::Utf8(_) => todo!(),
+            DType::Binary(_) => todo!(),
+            DType::Struct(..) => todo!(),
+            DType::List(..) => todo!(),
+            DType::Extension(..) => todo!(),
+        }
+    }
+}
+
+impl Display for Operator {
+    fn fmt(&self, f: &mut Formatter) -> fmt::Result {
+        let display = match &self {
+            Operator::And => "AND",
+            Operator::Or => "OR",
+            Operator::EqualTo => "=",
+            Operator::NotEqualTo => "!=",
+            Operator::GreaterThan => ">",
+            Operator::GreaterThanOrEqualTo => ">=",
+            Operator::LessThan => "<",
+            Operator::LessThanOrEqualTo => "<=",
+        };
+        write!(f, "{display}")
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use crate::expression_fns::{equals, field};
+    use crate::literal::lit;
+
+    #[test]
+    fn test_formatting() {
+        // And
+        let string = format!("{}", lit(true).and(lit(false)));
+        assert_eq!(string, "true AND false");
+        // Or
+        let string = format!("{}", lit(true).or(lit(false)));
+        assert_eq!(string, "true OR false");
+        // Not
+        let string = format!("{}", !lit(1u32));
+        assert_eq!(string, "NOT 1u32");
+    }
+
+    #[test]
+    fn test_format_respects_operator_associativity() {
+        let left = field("id").eq(lit(1));
+        let right = field("id2").eq(lit(2));
+        let s = format!("{}", equals(left, right));
+        assert_eq!(s, "id = 1i32 = (id2 = 2i32)")
+    }
+}

vortex-expr/src/expression_fns.rs

@@ -0,0 +1,40 @@
+#![allow(dead_code)]
+
+use vortex_dtype::FieldName;
+
+use crate::expressions::BinaryExpr;
+use crate::expressions::Expr;
+use crate::operators::Operator;
+
+pub fn binary_expr(left: Expr, op: Operator, right: Expr) -> Expr {
+    Expr::Binary(BinaryExpr::new(Box::new(left), op, Box::new(right)))
+}
+
+/// Create a field expression based on a qualified field name.
+pub fn field(field: impl Into<FieldName>) -> Expr {
+    Expr::Field(field.into())
+}
+
+pub fn equals(left: Expr, right: Expr) -> Expr {
+    Expr::Binary(BinaryExpr::new(
+        Box::new(left),
+        Operator::EqualTo,
+        Box::new(right),
+    ))
+}
+
+pub fn and(left: Expr, right: Expr) -> Expr {
+    Expr::Binary(BinaryExpr::new(
+        Box::new(left),
+        Operator::And,
+        Box::new(right),
+    ))
+}
+
+pub fn or(left: Expr, right: Expr) -> Expr {
+    Expr::Binary(BinaryExpr::new(
+        Box::new(left),
+        Operator::Or,
+        Box::new(right),
+    ))
+}

vortex-expr/src/expressions.rs

@@ -0,0 +1,78 @@
+use vortex_dtype::FieldName;
+use vortex_scalar::Scalar;
+
+use crate::expression_fns::binary_expr;
+use crate::operators::Operator;
+
+#[derive(Clone, Debug, PartialEq)]
+pub enum Expr {
+    /// A binary expression such as "duration_seconds == 100"
+    Binary(BinaryExpr),
+
+    /// A named reference to a qualified field in a dtype.
+    Field(FieldName),
+
+    /// True if argument is NULL, false otherwise. This expression itself is never NULL.
+    IsNull(Box<Expr>),
+
+    /// A constant scalar value.
+    Literal(Scalar),
+
+    /// Negation of an expression. The expression's type must be a boolean.
+    Not(Box<Expr>),
+}
+
+impl Expr {
+    // binary logic
+
+    pub fn and(self, other: Expr) -> Expr {
+        binary_expr(self, Operator::And, other)
+    }
+
+    pub fn or(self, other: Expr) -> Expr {
+        binary_expr(self, Operator::Or, other)
+    }
+
+    // comparisons
+
+    pub fn eq(self, other: Expr) -> Expr {
+        binary_expr(self, Operator::EqualTo, other)
+    }
+
+    pub fn not_eq(self, other: Expr) -> Expr {
+        binary_expr(self, Operator::NotEqualTo, other)
+    }
+
+    pub fn gt(self, other: Expr) -> Expr {
+        binary_expr(self, Operator::GreaterThan, other)
+    }
+
+    pub fn gte(self, other: Expr) -> Expr {
+        binary_expr(self, Operator::GreaterThanOrEqualTo, other)
+    }
+
+    pub fn lt(self, other: Expr) -> Expr {
+        binary_expr(self, Operator::LessThan, other)
+    }
+
+    pub fn lte(self, other: Expr) -> Expr {
+        binary_expr(self, Operator::LessThanOrEqualTo, other)
+    }
+
+    pub fn is_null(self) -> Expr {
+        Expr::IsNull(Box::new(self))
+    }
+}
+
+#[derive(Clone, Debug, PartialEq)]
+pub struct BinaryExpr {
+    pub left: Box<Expr>,
+    pub op: Operator,
+    pub right: Box<Expr>,
+}
+
+impl BinaryExpr {
+    pub fn new(left: Box<Expr>, op: Operator, right: Box<Expr>) -> Self {
+        Self { left, op, right }
+    }
+}

vortex-expr/src/lib.rs

@@ -0,0 +1,8 @@
+extern crate core;
+
+mod display;
+mod expression_fns;
+pub mod expressions;
+mod literal;
+pub mod operators;
+pub mod scalar;

vortex-expr/src/literal.rs

@@ -0,0 +1,26 @@
+use vortex_scalar::Scalar;
+
+use crate::expressions::Expr;
+
+pub trait Literal {
+    fn lit(&self) -> Expr;
+}
+
+#[allow(dead_code)]
+pub fn lit<T: Into<Scalar>>(n: T) -> Expr {
+    n.into().lit()
+}
+
+#[cfg(test)]
+mod test {
+    use super::*;
+    use crate::expression_fns::field;
+
+    #[test]
+    fn test_lit() {
+        let scalar: Scalar = 1.into();
+        let rhs: Expr = lit(scalar);
+        let expr = field("id").eq(rhs);
+        assert_eq!(format!("{}", expr), "id = 1i32");
+    }
+}