Nuke to_mem_table_with_op_type (#990)

* Nuke `to_mem_table_with_op_type`

Rather than annotating rows with `__op_type` during `eval_incr` of selects,
partition the rows before evaluation, then merge after.

* Add historical comment.

Co-authored-by: Mazdak Farrokhzad <[email protected]>
Signed-off-by: Phoebe Goldman <[email protected]>

* Remove `_replaced_source_id`

---------

Signed-off-by: Phoebe Goldman <[email protected]>
Co-authored-by: Mazdak Farrokhzad <[email protected]>

Author: gefjon

crates/core/src/subscription/execution_unit.rs

@@ -1,4 +1,4 @@
-use super::query::{self, find_op_type_col_pos, Supported, OP_TYPE_FIELD_NAME};
+use super::query::{self, Supported};
 use super::subscription::{IncrementalJoin, SupportedQuery};
 use crate::db::relational_db::{RelationalDB, Tx};
 use crate::error::DBError;
@@ -9,11 +9,11 @@ use crate::vm::{build_query, TxMode};
 use spacetimedb_client_api_messages::client_api::{TableRowOperation, TableUpdate};
 use spacetimedb_lib::bsatn::to_writer;
 use spacetimedb_primitives::TableId;
-use spacetimedb_sats::relation::{DbTable, Header};
+use spacetimedb_sats::relation::DbTable;
 use spacetimedb_vm::eval::IterRows;
 use spacetimedb_vm::expr::{Query, QueryExpr, SourceExpr, SourceSet};
 use spacetimedb_vm::rel_ops::RelOps;
-use spacetimedb_vm::relation::RelValue;
+use spacetimedb_vm::relation::{MemTable, RelValue};
 use std::hash::Hash;
 
 /// A hash for uniquely identifying query execution units,
@@ -58,7 +58,7 @@ enum EvalIncrPlan {
     Semijoin(IncrementalJoin),
 
     /// For single-table selects, store only one version of the plan,
-    /// which has a single source, a [`MemTable`] produced by [`query::to_mem_table_with_op_type`].
+    /// which has a single source, a [`MemTable`], produced by [`query::query_to_mem_table`].
     Select(QueryExpr),
 }
 
@@ -100,7 +100,6 @@ impl From<SupportedQuery> for ExecutionUnit {
 
 impl ExecutionUnit {
     /// Pre-compute a plan for `eval_incr` which reads from a `MemTable`
-    /// whose rows are augmented with an `__op_type` column,
     /// rather than re-planning on every incremental update.
     fn compile_select_eval_incr(expr: &QueryExpr) -> QueryExpr {
         let source = expr
@@ -126,7 +125,7 @@ impl ExecutionUnit {
         // Some day down the line, when we have a real query planner,
         // we may need to provide a row count estimation that is, if not accurate,
         // at least less specifically inaccurate.
-        let (eval_incr_plan, _source_set) = query::to_mem_table(expr.clone(), &table_update);
+        let (eval_incr_plan, _source_set) = query::query_to_mem_table(expr.clone(), &table_update);
         debug_assert_eq!(_source_set.len(), 1);
 
         eval_incr_plan
@@ -278,6 +277,21 @@ impl ExecutionUnit {
         }))
     }
 
+    fn eval_query_expr_against_memtable<'a>(
+        ctx: &'a ExecutionContext,
+        db: &'a RelationalDB,
+        tx: &'a TxMode,
+        mem_table: MemTable,
+        eval_incr_plan: &'a QueryExpr,
+    ) -> Result<Box<IterRows<'a>>, DBError> {
+        // Build a `SourceSet` containing the updates from `table`.
+        let mut sources = SourceSet::default();
+        sources.add_mem_table(mem_table);
+        // Evaluate the saved plan against the new `SourceSet`,
+        // returning an iterator over the selected rows.
+        build_query(ctx, db, tx, eval_incr_plan, &mut sources).map_err(Into::into)
+    }
+
     fn eval_incr_query_expr<'a>(
         db: &RelationalDB,
         tx: &Tx,
@@ -289,51 +303,64 @@ impl ExecutionUnit {
         let tx: TxMode = tx.into();
 
         let SourceExpr::MemTable {
-            source_id: _source_id,
             ref header,
             table_access,
             ..
         } = eval_incr_plan.source
         else {
             panic!("Expected MemTable in `eval_incr_plan`, but found `DbTable`");
         };
+
+        // Partition the `update` into two `MemTable`s, `(inserts, deletes)`,
+        // so that we can remember which are which without adding a column to each row.
+        // Previously, we used to add such a column `"__op_type: AlgebraicType::U8"`.
+        let partition_updates = |update: &DatabaseTableUpdate| -> (Option<MemTable>, Option<MemTable>) {
+            // Pre-allocate with capacity given by an upper bound,
+            // because realloc is worse than over-allocing.
+            let mut inserts = Vec::with_capacity(update.ops.len());
+            let mut deletes = Vec::with_capacity(update.ops.len());
+            for op in update.ops.iter() {
+                // 0 = delete, 1 = insert
+                if op.op_type == 0 { &mut deletes } else { &mut inserts }.push(op.row.clone());
+            }
+            (
+                (!inserts.is_empty()).then(|| MemTable::new(header.clone(), table_access, inserts)),
+                (!deletes.is_empty()).then(|| MemTable::new(header.clone(), table_access, deletes)),
+            )
+        };
+
         let mut ops = Vec::new();
 
         for table in tables.filter(|table| table.table_id == return_table) {
-            // Build a `SourceSet` containing the updates from `table`.
-            let mem_table = query::to_mem_table_with_op_type(header.clone(), table_access, table);
-            let mut sources = SourceSet::default();
-            let _source_expr = sources.add_mem_table(mem_table);
-            debug_assert_eq!(_source_expr.source_id(), Some(_source_id));
-            // Evaluate the saved plan against the new `SourceSet`
-            // and capture the new row operations.
-            let query = build_query(&ctx, db, &tx, eval_incr_plan, &mut sources)?;
-            Self::collect_rows_remove_table_ops(&mut ops, query, header)?;
+            // Evaluate the query separately against inserts and deletes,
+            // so that we can pass each row to the query engine unaltered,
+            // without forgetting which are inserts and which are deletes.
+            // Then, collect the rows into the single `ops` vec,
+            // restoring the appropriate `op_type`.
+            let (inserts, deletes) = partition_updates(table);
+            if let Some(inserts) = inserts {
+                let query = Self::eval_query_expr_against_memtable(&ctx, db, &tx, inserts, eval_incr_plan)?;
+                // op_type 1: insert
+                Self::collect_rows_with_table_op(&mut ops, query, 1)?;
+            }
+            if let Some(deletes) = deletes {
+                let query = Self::eval_query_expr_against_memtable(&ctx, db, &tx, deletes, eval_incr_plan)?;
+                // op_type 0: delete
+                Self::collect_rows_with_table_op(&mut ops, query, 0)?;
+            }
         }
         Ok(ops)
     }
 
-    /// Convert a set of rows annotated with the `__op_type` fields into a set of [`TableOp`]s,
-    /// and collect them into a vec `into`.
-    fn collect_rows_remove_table_ops(
+    /// Collect the results of `query` into a vec `into`,
+    /// annotating each as a `TableOp` with the `op_type`.
+    fn collect_rows_with_table_op(
         into: &mut Vec<TableOp>,
         mut query: Box<IterRows<'_>>,
-        header: &Header,
+        op_type: u8,
     ) -> Result<(), DBError> {
-        let pos_op_type = find_op_type_col_pos(header).unwrap_or_else(|| {
-            panic!(
-                "Failed to locate `{OP_TYPE_FIELD_NAME}` in `{}`, fields: {:?}",
-                header.table_name,
-                header.fields.iter().map(|x| &x.field).collect::<Vec<_>>()
-            )
-        });
-        let pos_op_type = pos_op_type.idx();
         while let Some(row_ref) = query.next()? {
-            let mut row = row_ref.into_product_value();
-            let op_type =
-                row.elements.remove(pos_op_type).into_u8().unwrap_or_else(|_| {
-                    panic!("Failed to extract `{OP_TYPE_FIELD_NAME}` from `{}`", header.table_name)
-                });
+            let row = row_ref.into_product_value();
             into.push(TableOp::new(op_type, row));
         }
         Ok(())

crates/core/src/subscription/query.rs

@@ -9,13 +9,8 @@ use once_cell::sync::Lazy;
 use regex::Regex;
 use spacetimedb_lib::identity::AuthCtx;
 use spacetimedb_lib::Address;
-use spacetimedb_primitives::ColId;
-use spacetimedb_sats::db::auth::StAccess;
-use spacetimedb_sats::relation::{Column, FieldName, Header};
-use spacetimedb_sats::AlgebraicType;
 use spacetimedb_vm::expr::{self, Crud, CrudExpr, DbType, QueryExpr, SourceSet};
 use spacetimedb_vm::relation::MemTable;
-use std::sync::Arc;
 use std::time::Instant;
 
 use super::subscription::get_all;
@@ -28,90 +23,14 @@ pub enum QueryDef {
     Sql(String),
 }
 
-pub const OP_TYPE_FIELD_NAME: &str = "__op_type";
-
-/// Locate the `__op_type` column in the table described by `header`,
-/// if it exists.
-///
-/// The current version of this function depends on the fact that
-/// the `__op_type` column is always the final column in the schema.
-/// This is true because the `__op_type` column is added by [`to_mem_table_with_op_type`] at the end,
-/// and never originates anywhere else.
-///
-/// If we ever change to having the `__op_type` column in any other position,
-/// e.g. by projecting together two `MemTables` from [`to_mem_table_with_op_type`],
-/// this function may need to change, possibly to:
-/// ```ignore
-/// header.find_pos_by_name(OP_TYPE_FIELD_NAME)
-/// ```
-pub fn find_op_type_col_pos(header: &Header) -> Option<ColId> {
-    if let Some(last_col) = header.fields.last() {
-        if last_col.field.field_name() == Some(OP_TYPE_FIELD_NAME) {
-            return Some(ColId((header.fields.len() - 1) as u32));
-        }
-    }
-    None
-}
-
-/// Create a virtual table from a sequence of table updates.
-/// Add a special column __op_type to distinguish inserts and deletes.
-pub fn to_mem_table_with_op_type(head: Arc<Header>, table_access: StAccess, data: &DatabaseTableUpdate) -> MemTable {
-    let mut t = MemTable::new(head, table_access, Vec::with_capacity(data.ops.len()));
-
-    if let Some(pos) = find_op_type_col_pos(&t.head) {
-        for op in data.ops.iter().map(|row| {
-            let mut new = row.row.clone();
-
-            match new.elements.len().cmp(&pos.idx()) {
-                std::cmp::Ordering::Equal => {
-                    // When we enter through `ExecutionUnit::eval_incr`,
-                    // we will have a `head` computed by a previous call to `to_mem_table`,
-                    // and therefore will have an op_type column,
-                    // but the `data` will be fresh for a newly committed transaction,
-                    // and therefore the rows will not include the op_type column.
-                    // In that case, push the op_type onto the end of each row.
-                    new.elements.push(row.op_type.into());
-                }
-                std::cmp::Ordering::Greater => {
-                    new.elements[pos.idx()] = row.op_type.into();
-                }
-                std::cmp::Ordering::Less => {
-                    panic!(
-                        "Expected {} either in-bounds or as the last column, but found at position {} in {:?}",
-                        OP_TYPE_FIELD_NAME, pos, t.head,
-                    );
-                }
-            }
-
-            new
-        }) {
-            t.data.push(op);
-        }
-    } else {
-        // TODO(perf): Eliminate this `clone_for_error` call, as we're not in an error path.
-        let mut head = t.head.clone_for_error();
-        head.fields.push(Column::new(
-            FieldName::named(&t.head.table_name, OP_TYPE_FIELD_NAME),
-            AlgebraicType::U8,
-            t.head.fields.len().into(),
-        ));
-        t.head = Arc::new(head);
-        for row in &data.ops {
-            let mut new = row.row.clone();
-            new.elements.push(row.op_type.into());
-            t.data.push(new);
-        }
-    }
-    t
-}
-
 /// Replace the primary (ie. `source`) table of the given [`QueryExpr`] with
 /// a virtual [`MemTable`] consisting of the rows in [`DatabaseTableUpdate`].
-///
-/// To be able to reify the `op_type` of the individual operations in the update,
-/// each virtual row is extended with a column [`OP_TYPE_FIELD_NAME`].
-pub fn to_mem_table(mut of: QueryExpr, data: &DatabaseTableUpdate) -> (QueryExpr, SourceSet) {
-    let mem_table = to_mem_table_with_op_type(of.source.head().clone(), of.source.table_access(), data);
+pub fn query_to_mem_table(mut of: QueryExpr, data: &DatabaseTableUpdate) -> (QueryExpr, SourceSet) {
+    let mem_table = MemTable::new(
+        of.source.head().clone(),
+        of.source.table_access(),
+        data.ops.iter().map(|op| op.row.clone()).collect(),
+    );
     let mut sources = SourceSet::default();
     let source_expr = sources.add_mem_table(mem_table);
     of.source = source_expr;
@@ -247,7 +166,7 @@ mod tests {
     use spacetimedb_sats::db::auth::{StAccess, StTableType};
     use spacetimedb_sats::db::def::*;
     use spacetimedb_sats::relation::FieldName;
-    use spacetimedb_sats::{product, ProductType, ProductValue};
+    use spacetimedb_sats::{product, AlgebraicType, ProductType, ProductValue};
     use spacetimedb_vm::dsl::{mem_table, scalar};
     use spacetimedb_vm::operator::OpCmp;
 
@@ -366,7 +285,7 @@ mod tests {
         q: &QueryExpr,
         data: &DatabaseTableUpdate,
     ) -> ResultTest<()> {
-        let (q, sources) = to_mem_table(q.clone(), data);
+        let (q, sources) = query_to_mem_table(q.clone(), data);
         let result = run_query(
             &ExecutionContext::default(),
             db,
@@ -556,7 +475,7 @@ mod tests {
 
         let q = QueryExpr::new(&schema);
 
-        let (q, sources) = to_mem_table(q, &data);
+        let (q, sources) = query_to_mem_table(q, &data);
         //Try access the private table
         match run_query(
             &ExecutionContext::default(),