Re-land mult-col index selection for queries (#918)

* Revert "Revert "Adding an index selector that take in account multi-column indexes (and improve the `query!` macro) (#694)" (#914)"

This reverts commit 8e5ce79.

* drive-by: refactor impl From<IndexScan> for ColumnOp

* reactor IndexScan bounds structure

* remove temp allocation in extract_fields

* skip index scan for NotEq

* drive-by: simplify Select

* clarify unreachable!(...) for NotEq

Signed-off-by: Mazdak Farrokhzad <[email protected]>

* address Joshua's review + refactor compiler tests

---------

Signed-off-by: Mazdak Farrokhzad <[email protected]>

Author: Centril

Cargo.lock

@@ -210,6 +210,7 @@ tracing-core = "0.1.31"
 tracing-flame = "0.2.0"
 tracing-log = "0.1.3"
 tracing-subscriber = { version = "0.3.17", features = ["env-filter"] }
+typed-arena = "2.0"
 url = "2.3.1"
 urlencoding = "2.1.2"
 uuid = { version = "1.2.1", features = ["v4"] }

Cargo.toml

@@ -6,7 +6,7 @@ use spacetimedb::host::module_host::{DatabaseTableUpdate, DatabaseUpdate, TableO
 use spacetimedb::subscription::query::compile_read_only_query;
 use spacetimedb::subscription::subscription::ExecutionSet;
 use spacetimedb_lib::identity::AuthCtx;
-use spacetimedb_primitives::TableId;
+use spacetimedb_primitives::{col_list, TableId};
 use spacetimedb_sats::{product, AlgebraicType, AlgebraicValue, ProductValue};
 use tempdir::TempDir;
 
@@ -18,8 +18,7 @@ fn create_table_location(db: &RelationalDB) -> Result<TableId, DBError> {
         ("z", AlgebraicType::I32),
         ("dimension", AlgebraicType::U32),
     ];
-    let indexes = &[(0.into(), "entity_id"), (1.into(), "chunk_index"), (2.into(), "x")];
-    db.create_table_for_test("location", schema, indexes)
+    db.create_table_for_test_multi_column("location", schema, col_list![2, 3, 4])
 }
 
 fn create_table_footprint(db: &RelationalDB) -> Result<TableId, DBError> {
@@ -70,7 +69,7 @@ fn eval(c: &mut Criterion) {
             for i in 0u64..1200 {
                 let entity_id = chunk_index * 1200 + i;
                 let x = 0i32;
-                let z = 0i32;
+                let z = entity_id as i32;
                 let dimension = 0u32;
                 let row = product!(entity_id, chunk_index, x, z, dimension);
                 let _ = db.insert(tx, rhs, row)?;
@@ -98,43 +97,34 @@ fn eval(c: &mut Criterion) {
         ],
     };
 
+    let bench_eval = |c: &mut Criterion, name, sql| {
+        c.bench_function(name, |b| {
+            let auth = AuthCtx::for_testing();
+            let tx = db.begin_tx();
+            let query = compile_read_only_query(&db, &tx, &auth, sql).unwrap();
+            let query: ExecutionSet = query.into();
+
+            b.iter(|| drop(black_box(query.eval(&db, &tx, auth).unwrap())))
+        });
+    };
+
     // To profile this benchmark for 30s
     // samply record -r 10000000 cargo bench --bench=subscription --profile=profiling -- full-scan --exact --profile-time=30
-    c.bench_function("full-scan", |b| {
-        // Iterate 1M rows.
-        let scan = "select * from footprint";
-        let auth = AuthCtx::for_testing();
-        let tx = db.begin_tx();
-        let query = compile_read_only_query(&db, &tx, &auth, scan).unwrap();
-        let query: ExecutionSet = query.into();
-
-        b.iter(|| {
-            let out = query.eval(&db, &tx, auth).unwrap();
-            black_box(out);
-        })
-    });
+    // Iterate 1M rows.
+    bench_eval(c, "full-scan", "select * from footprint");
 
     // To profile this benchmark for 30s
     // samply record -r 10000000 cargo bench --bench=subscription --profile=profiling -- full-join --exact --profile-time=30
-    c.bench_function("full-join", |b| {
-        // Join 1M rows on the left with 12K rows on the right.
-        // Note, this should use an index join so as not to read the entire lhs table.
-        let join = format!(
-            "\
-            select footprint.* \
-            from footprint join location on footprint.entity_id = location.entity_id \
-            where location.chunk_index = {chunk_index}"
-        );
-        let auth = AuthCtx::for_testing();
-        let tx = db.begin_tx();
-        let query = compile_read_only_query(&db, &tx, &auth, &join).unwrap();
-        let query: ExecutionSet = query.into();
-
-        b.iter(|| {
-            let out = query.eval(&db, &tx, AuthCtx::for_testing()).unwrap();
-            black_box(out);
-        })
-    });
+    // Join 1M rows on the left with 12K rows on the right.
+    // Note, this should use an index join so as not to read the entire lhs table.
+    let name = format!(
+        r#"
+        select footprint.*
+        from footprint join location on footprint.entity_id = location.entity_id
+        where location.chunk_index = {chunk_index}
+        "#
+    );
+    bench_eval(c, "full-join", &name);
 
     // To profile this benchmark for 30s
     // samply record -r 10000000 cargo bench --bench=subscription --profile=profiling -- incr-select --exact --profile-time=30
@@ -174,6 +164,15 @@ fn eval(c: &mut Criterion) {
             black_box(out);
         })
     });
+
+    // To profile this benchmark for 30s
+    // samply record -r 10000000 cargo bench --bench=subscription --profile=profiling -- query-indexes-multi --exact --profile-time=30
+    // Iterate 1M rows.
+    bench_eval(
+        c,
+        "query-indexes-multi",
+        "select * from location WHERE x = 0 AND z = 10000 AND dimension = 0",
+    );
 }
 
 criterion_group!(benches, eval);

crates/bench/benches/subscription.rs

@@ -18,7 +18,6 @@ use crate::error::{DBError, DatabaseError, TableError};
 use crate::execution_context::ExecutionContext;
 use crate::hash::Hash;
 use fs2::FileExt;
-use itertools::Itertools;
 use spacetimedb_primitives::*;
 use spacetimedb_sats::db::auth::{StAccess, StTableType};
 use spacetimedb_sats::db::def::{ColumnDef, IndexDef, SequenceDef, TableDef, TableSchema};
@@ -412,33 +411,47 @@ impl RelationalDB {
         self.inner.create_table_mut_tx(tx, schema.into())
     }
 
+    fn col_def_for_test(schema: &[(&str, AlgebraicType)]) -> Vec<ColumnDef> {
+        schema
+            .iter()
+            .cloned()
+            .map(|(col_name, col_type)| ColumnDef {
+                col_name: col_name.into(),
+                col_type,
+            })
+            .collect()
+    }
+
     pub fn create_table_for_test(
         &self,
         name: &str,
         schema: &[(&str, AlgebraicType)],
         indexes: &[(ColId, &str)],
     ) -> Result<TableId, DBError> {
-        let table_name = name.to_string();
-        let table_type = StTableType::User;
-        let table_access = StAccess::Public;
-
-        let columns = schema
-            .iter()
-            .map(|(col_name, col_type)| ColumnDef {
-                col_name: col_name.to_string(),
-                col_type: col_type.clone(),
-            })
-            .collect_vec();
-
         let indexes = indexes
             .iter()
-            .map(|(col_id, index_name)| IndexDef::btree(index_name.to_string(), *col_id, false))
-            .collect_vec();
+            .copied()
+            .map(|(col_id, index_name)| IndexDef::btree(index_name.into(), col_id, false))
+            .collect();
 
-        let schema = TableDef::new(table_name, columns)
+        let schema = TableDef::new(name.into(), Self::col_def_for_test(schema))
             .with_indexes(indexes)
-            .with_type(table_type)
-            .with_access(table_access);
+            .with_type(StTableType::User)
+            .with_access(StAccess::Public);
+
+        self.with_auto_commit(&ExecutionContext::default(), |tx| self.create_table(tx, schema))
+    }
+
+    pub fn create_table_for_test_multi_column(
+        &self,
+        name: &str,
+        schema: &[(&str, AlgebraicType)],
+        idx_cols: ColList,
+    ) -> Result<TableId, DBError> {
+        let schema = TableDef::new(name.into(), Self::col_def_for_test(schema))
+            .with_column_index(idx_cols, false)
+            .with_type(StTableType::User)
+            .with_access(StAccess::Public);
 
         self.with_auto_commit(&ExecutionContext::default(), |tx| self.create_table(tx, schema))
     }
@@ -548,7 +561,7 @@ impl RelationalDB {
     /// Returns the `index_id`
     ///
     /// NOTE: It loads the data from the table into it before returning
-    #[tracing::instrument(skip(self, tx, index), fields(index=index.index_name))]
+    #[tracing::instrument(skip(self, tx, index), fields(index = index.index_name))]
     pub fn create_index(&self, tx: &mut MutTx, table_id: TableId, index: IndexDef) -> Result<IndexId, DBError> {
         self.inner.create_index_mut_tx(tx, table_id, index)
     }
@@ -691,7 +704,7 @@ impl RelationalDB {
     }
 
     /// Add a [Sequence] into the database instance, generates a stable [SequenceId] for it that will persist on restart.
-    #[tracing::instrument(skip(self, tx, seq), fields(seq=seq.sequence_name))]
+    #[tracing::instrument(skip(self, tx, seq), fields(seq = seq.sequence_name))]
     pub fn create_sequence(
         &mut self,
         tx: &mut MutTx,

crates/core/src/db/relational_db.rs

@@ -1,3 +1,4 @@
+use itertools::Itertools;
 use std::borrow::Cow;
 use std::collections::HashMap;
 
@@ -32,6 +33,7 @@ impl Unsupported for bool {
         *self
     }
 }
+
 impl<T> Unsupported for Option<T> {
     fn unsupported(&self) -> bool {
         self.is_some()
@@ -59,7 +61,7 @@ impl Unsupported for sqlparser::ast::GroupByExpr {
     }
 }
 
-macro_rules! unsupported{
+macro_rules! unsupported {
     ($name:literal,$a:expr)=>{{
         let name = stringify!($name);
         let it = stringify!($a);
@@ -163,6 +165,13 @@ impl From {
         }))
     }
 
+    /// Returns all the columns from [Self::iter_tables], removing duplicates by `col_name`
+    pub fn iter_columns_dedup(&self) -> impl Iterator<Item = (&TableSchema, &ColumnSchema)> {
+        self.iter_tables()
+            .flat_map(|t| t.columns().iter().map(move |column| (t, column)))
+            .dedup_by(|(_, x), (_, y)| x.col_name == y.col_name)
+    }
+
     /// Returns all the table names as a `Vec<String>`, including the ones inside the joins.
     pub fn table_names(&self) -> Vec<String> {
         self.iter_tables().map(|x| x.table_name.clone()).collect()
@@ -317,7 +326,7 @@ fn compile_expr_value(table: &From, field: Option<&ProductTypeElement>, of: SqlE
             x => {
                 return Err(PlanError::Unsupported {
                     feature: format!("Unsupported value: {x}."),
-                })
+                });
             }
         }),
         SqlExpr::BinaryOp { left, op, right } => {
@@ -331,7 +340,7 @@ fn compile_expr_value(table: &From, field: Option<&ProductTypeElement>, of: SqlE
         x => {
             return Err(PlanError::Unsupported {
                 feature: format!("Unsupported expression: {x}"),
-            })
+            });
         }
     }))
 }
@@ -385,7 +394,7 @@ fn compile_bin_op(
         x => {
             return Err(PlanError::Unsupported {
                 feature: format!("BinaryOperator not supported in WHERE: {x}."),
-            })
+            });
         }
     };
 
@@ -775,7 +784,7 @@ fn column_def_type(named: &String, is_null: bool, data_type: &DataType) -> Resul
         x => {
             return Err(PlanError::Unsupported {
                 feature: format!("Column {} of type {}", named, x),
-            })
+            });
         }
     };
 
@@ -816,15 +825,15 @@ fn compile_column_option(col: &SqlColumnDef) -> Result<(bool, Constraints), Plan
                     x => {
                         return Err(PlanError::Unsupported {
                             feature: format!("IDENTITY option {x:?}"),
-                        })
+                        });
                     }
                 }
             }
             ColumnOption::Comment(_) => {}
             x => {
                 return Err(PlanError::Unsupported {
                     feature: format!("Column option {x}"),
-                })
+                });
             }
         }
     }
@@ -875,7 +884,7 @@ fn compile_drop(name: &ObjectName, kind: ObjectType) -> Result<SqlAst, PlanError
         x => {
             return Err(PlanError::Unsupported {
                 feature: format!("DROP {x}"),
-            })
+            });
         }
     };
 
@@ -920,7 +929,7 @@ fn compile_statement<T: TableSchemaView>(db: &RelationalDB, tx: &T, statement: S
                     _ => {
                         return Err(PlanError::Unsupported {
                             feature: "Insert WITHOUT values".into(),
-                        })
+                        });
                     }
                 };