Enabling symbolic execution support for filters, joins and unions

arrow-array/src/array/byte_array.rs

@@ -530,7 +530,6 @@ impl<T: ByteArrayType> Array for GenericByteArray<T> {
     }
 
     fn to_symbolic_data(&self) -> super::SymbolicArrayData {
-        dbg!(&self.data_type, &self, &self.symbolic_data);
         self.symbolic_data
             .clone()
             .unwrap_or(self.make_symbolic_data())

arrow-array/src/array/mod.rs

@@ -80,6 +80,9 @@ use crate::iterator::ArrayIter;
 
 mod symbolic_expr;
 pub use symbolic_expr::make_colref_symbolic_expr_array;
+pub use symbolic_expr::arrow_type_to_var_type;
+pub use symbolic_expr::add_constraint_to_row;
+pub use symbolic_expr::VarType as SymbolicVarType;
 pub use symbolic_expr::Expr as SymbolicExpr;
 pub use symbolic_expr::Operator as SymbolicOperator;
 pub use symbolic_expr::ScalarValue as SymbolicScalarValue;
@@ -116,8 +119,6 @@ pub trait Array: std::fmt::Debug + Send + Sync {
     /// Returns the underlying data of this array
     fn to_data(&self) -> ArrayData;
 
-    // fn to_symbolic_data(&self) -> SymbolicArrayData;
-
     /// Returns the underlying data of this array
     ///
     /// Unlike [`Array::to_data`] this consumes self, allowing it avoid unnecessary clones
@@ -368,10 +369,6 @@ impl Array for ArrayRef {
         self.as_ref().to_data()
     }
 
-    // fn to_symbolic_data(&self) -> SymbolicArrayData {
-    //     self.as_ref().to_symbolic_data()
-    // }
-
     fn into_data(self) -> ArrayData {
         self.to_data()
     }
@@ -459,10 +456,6 @@ impl<T: Array> Array for &T {
         T::to_data(self)
     }
 
-    // fn to_symbolic_data(&self) -> SymbolicArrayData {
-    //     T::to_symbolic_data(self)
-    // }
-
     fn into_data(self) -> ArrayData {
         self.to_data()
     }

arrow-array/src/array/symbolic_expr.rs

@@ -1,3 +1,5 @@
+use arrow_schema::DataType;
+
 /// Literal values used in symbolic expressions
 #[derive(Debug, Clone, PartialEq)]
 pub enum ScalarValue {
@@ -31,6 +33,34 @@ pub enum ScalarValue {
     Utf8(Option<String>),
 }
 
+#[derive(Debug, Copy, Clone, PartialEq)]
+pub enum VarType {
+    /// A boolean variable
+    Boolean,
+    /// A string variable
+    String,
+    /// A int variable
+    Int,
+}
+
+pub fn arrow_type_to_var_type(arrow_type: DataType) -> VarType {
+    match arrow_type {
+        DataType::Boolean => VarType::Boolean,
+        DataType::Utf8 | DataType::LargeUtf8 => VarType::String,
+        DataType::Int8
+        | DataType::Int16
+        | DataType::Int32
+        | DataType::Int64
+        | DataType::UInt8
+        | DataType::UInt16
+        | DataType::UInt32
+        | DataType::UInt64
+        | DataType::Decimal128(_, _)
+        | DataType::Decimal256(_, _) => VarType::Int,
+        _ => VarType::String, // Default to String for other types
+    }
+}
+
 /// Operators applied to symbolic expressions
 #[derive(Debug, Copy, Clone, PartialEq, Eq, PartialOrd, Hash)]
 pub enum Operator {
@@ -60,6 +90,10 @@ pub enum Operator {
     And,
     /// Logical OR, like `||`
     Or,
+    /// Logical NOT, like `!`
+    Not,
+    /// In, like `in`
+    In,
 }
 
 /// A symbolic expression for a columnar array.
@@ -79,6 +113,29 @@ pub enum Expr {
         column: usize,
         /// The position of the row
         row: usize,
+        /// The variable type
+        var_type: VarType,
+    },
+
+    /// A symbolic variable representing a row in a table
+    Row {
+        /// The name of the table
+        table: String,
+        /// The position of the row
+        row_id: usize,
+        /// The row constraint
+        constraint: Option<Vec<Expr>>,
+    },
+
+    /// A list of symbolic expressions
+    List(Vec<Expr>),
+
+    /// A count symbolic expression
+    Count {
+        /// The column to count
+        expr: Box<Expr>,
+        /// Whether the count is distinct
+        is_distinct: bool,
     },
 
     /// A binary symbolic expression
@@ -108,8 +165,39 @@ impl Expr {
     }
 
     /// Create a symbolic variable
-    pub fn variable(table: String, column: usize, row: usize) -> Self {
-        Expr::Variable { table, column, row }
+    pub fn variable(table: String, column: usize, row: usize, var_type: VarType) -> Self {
+        Expr::Variable {
+            table,
+            column,
+            row,
+            var_type,
+        }
+    }
+
+    /// Create a symbolic row
+    pub fn row(table: String, row_id: usize) -> Self {
+        Expr::Row {
+            table,
+            row_id,
+            constraint: None,
+        }
+    }
+
+    /// Create a symbolic row with a constraint
+    pub fn row_with_constraint(table: String, row_id: usize, constraint: Vec<Expr>) -> Self {
+        Expr::Row {
+            table,
+            row_id,
+            constraint: Some(constraint),
+        }
+    }
+
+    /// Create a count symbolic expression
+    pub fn count(expr: Expr, is_distinct: bool) -> Self {
+        Expr::Count {
+            expr: Box::new(expr),
+            is_distinct,
+        }
     }
 
     /// Create a binary symbolic expression
@@ -138,14 +226,44 @@ impl Expr {
 }
 
 /// Create a symbolic expression array for a column reference using table name and column position
-pub fn make_colref_symbolic_expr_array(table: String, column: usize, len: usize) -> Vec<Expr> {
+pub fn make_colref_symbolic_expr_array(
+    table: String,
+    column: usize,
+    len: usize,
+    row_offset: usize,
+    var_type: VarType,
+) -> Vec<Expr> {
     let mut res = Vec::with_capacity(len);
     for i in 0..len {
         res.push(Expr::Variable {
             table: table.clone(),
             column,
-            row: i,
+            row: i + row_offset,
+            var_type,
         });
     }
     res
 }
+
+pub fn add_constraint_to_row(row: &mut Expr, constraint: Expr) {
+    if let Expr::Row {
+        constraint: Some(constraints),
+        ..
+    } = row
+    {
+        constraints.push(constraint);
+    } else if let Expr::Row {
+        table,
+        row_id,
+        constraint: None,
+    } = row
+    {
+        *row = Expr::Row {
+            table: table.clone(),
+            row_id: *row_id,
+            constraint: Some(vec![constraint]),
+        };
+    } else {
+        panic!("Expr is not a symbolic row");
+    }
+}

arrow-array/src/record_batch.rs

@@ -205,6 +205,7 @@ pub struct RecordBatch {
     ///
     /// This is stored separately from the columns to handle the case of no columns
     row_count: usize,
+    row_symbolic_data: Option<Vec<crate::SymbolicExpr>>,
     constraints: Option<Vec<crate::SymbolicExpr>>,
 }
 
@@ -238,7 +239,7 @@ impl RecordBatch {
     /// ```
     pub fn try_new(schema: SchemaRef, columns: Vec<ArrayRef>) -> Result<Self, ArrowError> {
         let options = RecordBatchOptions::new();
-        Self::try_new_impl(schema, columns, &options, None)
+        Self::try_new_impl(schema, columns, &options, None, None)
     }
 
     /// Creates a `RecordBatch` from a schema and columns, with additional options,
@@ -250,7 +251,7 @@ impl RecordBatch {
         columns: Vec<ArrayRef>,
         options: &RecordBatchOptions,
     ) -> Result<Self, ArrowError> {
-        Self::try_new_impl(schema, columns, options, None)
+        Self::try_new_impl(schema, columns, options, None, None)
     }
 
     /// Creates a `RecordBatch` from a schema and columns, with additional options and symbolic constraints.
@@ -260,7 +261,48 @@ impl RecordBatch {
         options: &RecordBatchOptions,
         constraints: Option<Vec<crate::SymbolicExpr>>,
     ) -> Result<Self, ArrowError> {
-        Self::try_new_impl(schema, columns, options, constraints)
+        Self::try_new_impl(schema, columns, options, None, constraints)
+    }
+
+    /// Creates a `RecordBatch` from a schema and columns, with symbolic data.
+    pub fn try_new_with_symbolic_data(
+        schema: SchemaRef,
+        columns: Vec<ArrayRef>,
+        row_symbolic_data: Option<Vec<crate::SymbolicExpr>>,
+    ) -> Result<Self, ArrowError> {
+        Self::try_new_impl(
+            schema,
+            columns,
+            &RecordBatchOptions::new(),
+            row_symbolic_data,
+            None,
+        )
+    }
+
+    /// Creates a `RecordBatch` from a schema and columns, with constraints.
+    pub fn try_new_with_constraints(
+        schema: SchemaRef,
+        columns: Vec<ArrayRef>,
+        constraints: Option<Vec<crate::SymbolicExpr>>,
+    ) -> Result<Self, ArrowError> {
+        Self::try_new_impl(
+            schema,
+            columns,
+            &RecordBatchOptions::new(),
+            None,
+            constraints,
+        )
+    }
+
+    /// Creates a `RecordBatch` from a schema and columns, with additional options and symbolic constraints.
+    pub fn try_new_with_options_and_constraints_and_symbolic_data(
+        schema: SchemaRef,
+        columns: Vec<ArrayRef>,
+        options: &RecordBatchOptions,
+        row_symbolic_data: Option<Vec<crate::SymbolicExpr>>,
+        constraints: Option<Vec<crate::SymbolicExpr>>,
+    ) -> Result<Self, ArrowError> {
+        Self::try_new_impl(schema, columns, options, row_symbolic_data, constraints)
     }
 
     /// Creates a new empty [`RecordBatch`].
@@ -275,6 +317,7 @@ impl RecordBatch {
             schema,
             columns,
             row_count: 0,
+            row_symbolic_data: None,
             constraints: None,
         }
     }
@@ -285,9 +328,9 @@ impl RecordBatch {
         schema: SchemaRef,
         columns: Vec<ArrayRef>,
         options: &RecordBatchOptions,
+        row_symbolic_data: Option<Vec<crate::SymbolicExpr>>,
         constraints: Option<Vec<crate::SymbolicExpr>>,
     ) -> Result<Self, ArrowError> {
-        // dbg!(&columns);
         // check that number of fields in schema match column length
         if schema.fields().len() != columns.len() {
             return Err(ArrowError::InvalidArgumentError(format!(
@@ -351,6 +394,7 @@ impl RecordBatch {
             schema,
             columns,
             row_count,
+            row_symbolic_data,
             constraints,
         })
     }
@@ -371,6 +415,7 @@ impl RecordBatch {
             schema,
             columns: self.columns,
             row_count: self.row_count,
+            row_symbolic_data: self.row_symbolic_data,
             constraints: self.constraints,
         })
     }
@@ -390,6 +435,11 @@ impl RecordBatch {
         self.constraints.as_ref().map(|c| &c[..])
     }
 
+    /// Returns the symbolic data of the record batch.
+    pub fn row_symbolic_data(&self) -> Option<&[crate::SymbolicExpr]> {
+        self.row_symbolic_data.as_ref().map(|c| &c[..])
+    }
+
     /// Projects the schema onto the specified columns
     pub fn project(&self, indices: &[usize]) -> Result<RecordBatch, ArrowError> {
         let projected_schema = self.schema.project(indices)?;
@@ -632,15 +682,13 @@ impl RecordBatch {
             .map(|column| column.slice(offset, length))
             .collect();
 
-        let constraints = self
-            .constraints
-            .as_ref()
-            .map(|c| c[offset..offset + length].to_vec());
+        let constraints = self.constraints.clone();
 
         Self {
             schema: self.schema.clone(),
             columns,
             row_count: length,
+            row_symbolic_data: self.row_symbolic_data.clone(),
             constraints,
         }
     }
@@ -801,6 +849,7 @@ impl From<StructArray> for RecordBatch {
             schema: Arc::new(Schema::new(fields)),
             row_count,
             columns,
+            row_symbolic_data: None,
             constraints: None,
         }
     }

arrow-ord/src/cmp.rs

@@ -73,13 +73,15 @@ macro_rules! sym {
 
         let mut sym_res = vec![];
         for i in 0..syms1.len() {
-            let l = &syms1[i];
-            let r = &syms2[i];
-            sym_res.push(SymbolicExpr::binary(
-                l.clone(),
-                SymbolicOperator::$op,
-                r.clone(),
-            ));
+            for j in 0..syms2.len() {
+                let l = &syms1[i];
+                let r = &syms2[j];
+                sym_res.push(SymbolicExpr::binary(
+                    l.clone(),
+                    SymbolicOperator::$op,
+                    r.clone(),
+                ));
+            }
         }
         Ok(expr.with_symbolic_data(&sym_res))
     }};