perf(623): Use row counts in query optimization

Closes #623.

Before this patch query optimization was entirely syntax driven.
Now that we keep table size metrics we can be somewhat data driven.

This patch improves index joins,
by using row counts to determine the index side and the probe side.

Author: joshua-spacetime

crates/core/src/db/relational_db.rs

@@ -36,6 +36,7 @@ pub struct RelationalDB {
     commit_log: Option<CommitLogMut>,
     _lock: Arc<File>,
     address: Address,
+    row_count_fn: Arc<dyn Fn(TableId) -> i64 + Send + Sync>,
 }
 
 impl DataRow for RelationalDB {
@@ -139,12 +140,29 @@ impl RelationalDB {
             commit_log,
             _lock: Arc::new(lock),
             address: db_address,
+            row_count_fn: Arc::new(move |table| {
+                METRICS
+                    .rdb_num_table_rows
+                    .with_label_values(&db_address, &table.into())
+                    .get()
+            }),
         };
 
         log::info!("[{}] DATABASE: OPENED", address);
         Ok(db)
     }
 
+    /// Returns an approximate row count for a particular table.
+    pub fn row_count(&self, table: TableId) -> i64 {
+        (self.row_count_fn)(table)
+    }
+
+    /// Update this `RelationalDB` with an approximate row count function.
+    pub fn with_row_count(mut self, row_count: Arc<dyn Fn(TableId) -> i64 + Send + Sync>) -> Self {
+        self.row_count_fn = row_count;
+        self
+    }
+
     /// Returns the address for this database
     pub fn address(&self) -> Address {
         self.address
@@ -654,7 +672,7 @@ pub(crate) mod tests_utils {
         let tmp_dir = TempDir::with_prefix("stdb_test")?;
         let in_memory = false;
         let fsync = false;
-        let stdb = open_db(&tmp_dir, in_memory, fsync)?;
+        let stdb = open_db(&tmp_dir, in_memory, fsync)?.with_row_count(Arc::new(|_| i64::MAX));
         Ok((stdb, tmp_dir))
     }
 }

crates/core/src/sql/compiler.rs

@@ -270,7 +270,7 @@ fn compile_statement(db: &RelationalDB, statement: SqlAst) -> Result<CrudExpr, P
         } => compile_drop(name, kind, table_access)?,
     };
 
-    Ok(q.optimize(Some(db.address())))
+    Ok(q.optimize(&|table| db.row_count(table)))
 }
 
 #[cfg(test)]

crates/core/src/subscription/subscription.rs

@@ -657,7 +657,7 @@ impl<'a> IncrementalJoin<'a> {
             // Replan query after replacing the indexed table with a virtual table,
             // since join order may need to be reversed.
             let join_a = with_delta_table(self.join.clone(), true, self.index_side.inserts());
-            let join_a = QueryExpr::from(join_a).optimize(Some(db.address()));
+            let join_a = QueryExpr::from(join_a).optimize(&|table| db.row_count(table));
 
             // No need to replan after replacing the probe side with a virtual table,
             // since no new constraints have been added.
@@ -685,7 +685,7 @@ impl<'a> IncrementalJoin<'a> {
             // Replan query after replacing the indexed table with a virtual table,
             // since join order may need to be reversed.
             let join_a = with_delta_table(self.join.clone(), true, self.index_side.deletes());
-            let join_a = QueryExpr::from(join_a).optimize(Some(db.address()));
+            let join_a = QueryExpr::from(join_a).optimize(&|table| db.row_count(table));
 
             // No need to replan after replacing the probe side with a virtual table,
             // since no new constraints have been added.
@@ -880,7 +880,7 @@ mod tests {
 
         // Optimize the query plan for the incremental update.
         let expr: QueryExpr = with_delta_table(join, true, delta).into();
-        let mut expr = expr.optimize(Some(db.address()));
+        let mut expr = expr.optimize(&|_| i64::MAX);
 
         assert_eq!(expr.source.table_name(), "lhs");
         assert_eq!(expr.query.len(), 1);
@@ -975,7 +975,7 @@ mod tests {
 
         // Optimize the query plan for the incremental update.
         let expr: QueryExpr = with_delta_table(join, false, delta).into();
-        let mut expr = expr.optimize(Some(db.address()));
+        let mut expr = expr.optimize(&|_| i64::MAX);
 
         assert_eq!(expr.source.table_name(), "lhs");
         assert_eq!(expr.query.len(), 1);

crates/vm/src/eval.rs

@@ -77,7 +77,7 @@ fn build_typed<P: ProgramVm>(p: &mut P, node: Expr) -> ExprOpt {
             }
         }
         Expr::Crud(q) => {
-            let q = q.optimize(p.address());
+            let q = q.optimize(&|_| i64::MAX);
             match q {
                 CrudExpr::Query(q) => {
                     let source = build_query_opt(q);

crates/vm/src/expr.rs

@@ -9,7 +9,7 @@ use crate::errors::{ErrorKind, ErrorLang, ErrorType, ErrorVm};
 use crate::functions::{FunDef, Param};
 use crate::operator::{Op, OpCmp, OpLogic, OpQuery};
 use crate::types::Ty;
-use spacetimedb_lib::{Address, Identity};
+use spacetimedb_lib::Identity;
 use spacetimedb_primitives::*;
 use spacetimedb_sats::algebraic_type::AlgebraicType;
 use spacetimedb_sats::algebraic_value::AlgebraicValue;
@@ -426,7 +426,7 @@ impl IndexJoin {
     // Reorder the index and probe sides of an index join.
     // This is necessary if the indexed table has been replaced by a delta table.
     // A delta table is a virtual table consisting of changes or updates to a physical table.
-    pub fn reorder(self) -> Self {
+    pub fn reorder(self, row_count: impl Fn(TableId) -> i64) -> Self {
         // The probe table must be a physical table.
         if matches!(self.probe_side.source, SourceExpr::MemTable(_)) {
             return self;
@@ -449,26 +449,25 @@ impl IndexJoin {
         {
             return self;
         }
-        // Returns the `FieldName` for this `ColId`.
-        // Note there is an unwrap which the compiler will ensure is safe.
-        // The existence of `col` will already have been verified,
-        // during construction of the index join.
-        fn field(table: &MemTable, col: &ColId) -> FieldName {
-            table
-                .head()
-                .fields
-                .iter()
-                .find(|Column { col_id, .. }| col_id == col)
-                .unwrap()
-                .field
-                .clone()
-        }
         // For the same reason the compiler also ensures this unwrap is safe.
         let probe_column = self.probe_side.source.head().column(&self.probe_field).unwrap().col_id;
         match self.index_side {
-            Table::MemTable(delta) => {
-                let index_field = field(&delta, &self.index_col);
-
+            // If the size of the indexed table is sufficiently large, do not reorder.
+            Table::DbTable(DbTable { table_id, .. }) if row_count(table_id) > 1000 => self,
+            // If this is a delta table, we must reorder.
+            // If this is a sufficiently small physical table, we should reorder.
+            table => {
+                // The compiler ensures the following unwrap is safe.
+                // The existence of the index `ColId` as already been verified,
+                // during construction of the index join.
+                let index_field = table
+                    .head()
+                    .fields
+                    .iter()
+                    .find(|col| col.col_id == self.index_col)
+                    .unwrap()
+                    .field
+                    .clone();
                 // Merge all selections from the original probe side into a single predicate.
                 // This includes an index scan if present.
                 let predicate = self.probe_side.query.into_iter().fold(None, |acc, op| {
@@ -483,11 +482,11 @@ impl IndexJoin {
                 // Push any selections on the index side to the probe side.
                 let probe_side = if let Some(predicate) = self.index_select {
                     QueryExpr {
-                        source: delta.into(),
+                        source: table.into(),
                         query: vec![predicate.into()],
                     }
                 } else {
-                    delta.into()
+                    table.into()
                 };
                 IndexJoin {
                     // The new probe side consists of the updated rows.
@@ -506,7 +505,6 @@ impl IndexJoin {
                     return_index_rows: !self.return_index_rows,
                 }
             }
-            Table::DbTable(_) => self,
         }
     }
 
@@ -624,9 +622,9 @@ pub enum CrudExpr {
 }
 
 impl CrudExpr {
-    pub fn optimize(self, db: Option<Address>) -> Self {
+    pub fn optimize(self, row_count: &impl Fn(TableId) -> i64) -> Self {
         match self {
-            CrudExpr::Query(x) => CrudExpr::Query(x.optimize(db)),
+            CrudExpr::Query(x) => CrudExpr::Query(x.optimize(row_count)),
             _ => self,
         }
     }
@@ -1286,7 +1284,8 @@ impl QueryExpr {
     //
     // Ex. SELECT Left.* FROM Left JOIN Right ON Left.id = Right.id ...
     // where `Left` has an index defined on `id`.
-    fn try_index_join(mut query: QueryExpr, _db: Option<Address>) -> QueryExpr {
+    fn try_index_join(self) -> QueryExpr {
+        let mut query = self;
         // We expect 2 and only 2 operations - a join followed by a wildcard projection.
         if query.query.len() != 2 {
             return query;
@@ -1391,7 +1390,7 @@ impl QueryExpr {
         q
     }
 
-    pub fn optimize(mut self, db: Option<Address>) -> Self {
+    pub fn optimize(mut self, row_count: &impl Fn(TableId) -> i64) -> Self {
         let mut q = Self {
             source: self.source.clone(),
             query: Vec::with_capacity(self.query.len()),
@@ -1405,7 +1404,7 @@ impl QueryExpr {
 
         if self.query.len() == 1 && matches!(self.query[0], Query::IndexJoin(_)) {
             if let Some(Query::IndexJoin(join)) = self.query.pop() {
-                q.query.push(Query::IndexJoin(join.reorder()));
+                q.query.push(Query::IndexJoin(join.reorder(row_count)));
                 return q;
             }
         }
@@ -1415,11 +1414,18 @@ impl QueryExpr {
                 Query::Select(op) => {
                     q = Self::optimize_select(q, op, &tables);
                 }
-                Query::JoinInner(join) => q = q.with_join_inner(join.rhs.optimize(db), join.col_lhs, join.col_rhs),
+                Query::JoinInner(join) => {
+                    q = q.with_join_inner(join.rhs.optimize(row_count), join.col_lhs, join.col_rhs)
+                }
                 _ => q.query.push(query),
             };
         }
-        Self::try_index_join(q, db)
+
+        let q = q.try_index_join();
+        if q.query.len() == 1 && matches!(q.query[0], Query::IndexJoin(_)) {
+            return q.optimize(row_count);
+        }
+        q
     }
 }