perf: read DataKey instead of recomputing it in subscriptions (#276)

Fixes #259 

(1) Updates MemTable to use RelValue instead of ProductValue
(2) Adds a DataKey member to DataRef and RelValue
(3) Subscriptions compute DataKey only when not present on row

Author: joshua-spacetime

crates/client-api/src/routes/database.rs

@@ -557,7 +557,7 @@ pub async fn sql(
         .into_iter()
         .map(|result| StmtResultJson {
             schema: result.head.ty(),
-            rows: result.data.into_iter().map(|x| x.elements).collect::<Vec<_>>(),
+            rows: result.data.into_iter().map(|x| x.data.elements).collect::<Vec<_>>(),
         })
         .collect::<Vec<_>>();
 

crates/core/src/db/datastore/locking_tx_datastore/mod.rs

@@ -92,19 +92,29 @@ impl traits::Data for Data {
     }
 }
 
+/// A `DataRef` represents a row stored in a table.
+///
+/// A table row always has a [`DataKey`] associated with it.
+/// This is in contrast to rows that are materialized during query execution
+/// which may or may not have an associated `DataKey`.
 #[derive(Clone)]
 pub struct DataRef {
+    id: DataKey,
     data: ProductValue,
 }
 
 impl DataRef {
-    fn new(data: ProductValue) -> Self {
-        Self { data }
+    fn new(id: DataKey, data: ProductValue) -> Self {
+        Self { id, data }
     }
 
     pub fn view(&self) -> &ProductValue {
         &self.data
     }
+
+    pub fn id(&self) -> &DataKey {
+        &self.id
+    }
 }
 
 pub struct MutTxId {
@@ -1371,7 +1381,7 @@ impl Inner {
         match self.tx_state.as_ref().unwrap().get_row_op(table_id, row_id) {
             RowState::Committed(_) => unreachable!("a row cannot be committed in a tx state"),
             RowState::Insert(row) => {
-                return Ok(Some(DataRef::new(row)));
+                return Ok(Some(DataRef::new(row_id.0, row)));
             }
             RowState::Delete => {
                 return Ok(None);
@@ -1383,7 +1393,7 @@ impl Inner {
             .tables
             .get(table_id)
             .and_then(|table| table.get_row(row_id))
-            .map(|row| DataRef::new(row.clone())))
+            .map(|row| DataRef::new(row_id.0, row.clone())))
     }
 
     fn get_row_type(&self, table_id: &TableId) -> Option<&ProductType> {
@@ -1741,10 +1751,10 @@ impl Iterator for Iter<'_> {
                             Some(RowState::Insert(_)) => (), // Do nothing, we'll get it in the next stage
                             Some(RowState::Delete) => (),    // Skip it, it's been deleted
                             Some(RowState::Absent) => {
-                                return Some(DataRef::new(row.clone()));
+                                return Some(DataRef::new(row_id.0, row.clone()));
                             }
                             None => {
-                                return Some(DataRef::new(row.clone()));
+                                return Some(DataRef::new(row_id.0, row.clone()));
                             }
                         }
                     }
@@ -1762,8 +1772,8 @@ impl Iterator for Iter<'_> {
                     }
                 }
                 ScanStage::CurrentTx { iter } => {
-                    if let Some((_, row)) = iter.next() {
-                        return Some(DataRef::new(row.clone()));
+                    if let Some((id, row)) = iter.next() {
+                        return Some(DataRef::new(id.0, row.clone()));
                     }
                     break;
                 }
@@ -1854,6 +1864,7 @@ impl Iterator for IndexSeekIterInner<'_> {
     fn next(&mut self) -> Option<Self::Item> {
         if let Some(row_id) = self.inserted_rows.next() {
             return Some(DataRef::new(
+                row_id.0,
                 self.tx_state.get_row(&self.table_id, &row_id).unwrap().clone(),
             ));
         }
@@ -1902,7 +1913,10 @@ impl Iterator for CommittedIndexIterByColEq<'_> {
 /// Retrieve a commited row. Panics if `table_id` and `row_id` do not identify an actually
 /// present row.
 fn get_committed_row(state: &CommittedState, table_id: &TableId, row_id: &RowId) -> DataRef {
-    DataRef::new(state.tables.get(table_id).unwrap().get_row(row_id).unwrap().clone())
+    DataRef::new(
+        row_id.0,
+        state.tables.get(table_id).unwrap().get_row(row_id).unwrap().clone(),
+    )
 }
 
 pub enum IterByColRange<'a, R: RangeBounds<AlgebraicValue>> {

crates/core/src/host/instance_env.rs

@@ -427,7 +427,7 @@ impl InstanceEnv {
             _ => unreachable!("query should always return a table"),
         };
         Ok(std::iter::once(bsatn::to_vec(&row_type))
-            .chain(results.data.into_iter().map(|row| bsatn::to_vec(&row)))
+            .chain(results.data.into_iter().map(|row| bsatn::to_vec(&row.data)))
             .map(|bytes| bytes.expect("encoding algebraic values should never fail")))
     }
 }

crates/core/src/sql/execute.rs

@@ -103,9 +103,10 @@ pub(crate) mod tests {
     use crate::db::relational_db::tests_utils::make_test_db;
     use crate::db::relational_db::{ST_TABLES_ID, ST_TABLES_NAME};
     use crate::vm::tests::create_table_with_rows;
+    use itertools::Itertools;
     use spacetimedb_lib::auth::{StAccess, StTableType};
     use spacetimedb_lib::error::ResultTest;
-    use spacetimedb_lib::relation::Header;
+    use spacetimedb_lib::relation::{Header, RelValue};
     use spacetimedb_sats::{product, AlgebraicType, BuiltinType, ProductType};
     use spacetimedb_vm::dsl::{mem_table, scalar};
     use spacetimedb_vm::eval::create_game_data;
@@ -357,9 +358,27 @@ pub(crate) mod tests {
         let (db, _tmp_dir) = make_test_db()?;
 
         let mut tx = db.begin_tx();
-        create_table_with_rows(&db, &mut tx, "Inventory", data.inv.head.into(), &data.inv.data)?;
-        create_table_with_rows(&db, &mut tx, "Player", data.player.head.into(), &data.player.data)?;
-        create_table_with_rows(&db, &mut tx, "Location", data.location.head.into(), &data.location.data)?;
+        create_table_with_rows(
+            &db,
+            &mut tx,
+            "Inventory",
+            data.inv.head.into(),
+            &data.inv.data.iter().map(|row| row.data.clone()).collect_vec(),
+        )?;
+        create_table_with_rows(
+            &db,
+            &mut tx,
+            "Player",
+            data.player.head.into(),
+            &data.player.data.iter().map(|row| row.data.clone()).collect_vec(),
+        )?;
+        create_table_with_rows(
+            &db,
+            &mut tx,
+            "Location",
+            data.location.head.into(),
+            &data.location.data.iter().map(|row| row.data.clone()).collect_vec(),
+        )?;
 
         let result = &run_for_testing(
             &db,
@@ -427,7 +446,7 @@ pub(crate) mod tests {
         let mut result = result.first().unwrap().clone();
 
         let row = product!(scalar(2u64), scalar("test"));
-        input.data.push(row);
+        input.data.push(RelValue::new(&input.head, &row, None));
         input.data.sort();
         result.data.sort();
 
@@ -509,7 +528,7 @@ pub(crate) mod tests {
 
         let mut change = input;
         change.data.clear();
-        change.data.push(row);
+        change.data.push(RelValue::new(&change.head, &row, None));
 
         assert_eq!(
             change.as_without_table_name(),
@@ -526,7 +545,7 @@ pub(crate) mod tests {
             .map(|x| {
                 x.data
                     .into_iter()
-                    .map(|x| x.field_as_str(1, None).unwrap().to_string())
+                    .map(|x| x.data.field_as_str(1, None).unwrap().to_string())
                     .collect::<Vec<_>>()
             })
             .collect();

crates/core/src/subscription/query.rs

@@ -6,7 +6,8 @@ use crate::sql::compiler::compile_sql;
 use crate::sql::execute::execute_single_sql;
 use crate::subscription::subscription::QuerySet;
 use spacetimedb_lib::identity::AuthCtx;
-use spacetimedb_lib::relation::{Column, FieldName, MemTable};
+use spacetimedb_lib::relation::{Column, FieldName, MemTable, RelValue};
+use spacetimedb_lib::DataKey;
 use spacetimedb_sats::AlgebraicType;
 use spacetimedb_vm::expr::{Crud, CrudExpr, DbType, QueryExpr, SourceExpr};
 
@@ -53,7 +54,8 @@ pub fn to_mem_table(of: QueryExpr, data: &DatabaseTableUpdate) -> QueryExpr {
         t.data.extend(data.ops.iter().map(|row| {
             let mut new = row.row.clone();
             new.elements[pos] = row.op_type.into();
-            new
+            let mut bytes: &[u8] = row.row_pk.as_ref();
+            RelValue::new(&t.head, &new, Some(DataKey::decode(&mut bytes).unwrap()))
         }));
     } else {
         t.head.fields.push(Column::new(
@@ -63,7 +65,9 @@ pub fn to_mem_table(of: QueryExpr, data: &DatabaseTableUpdate) -> QueryExpr {
         for row in &data.ops {
             let mut new = row.row.clone();
             new.elements.push(row.op_type.into());
-            t.data.push(new);
+            let mut bytes: &[u8] = row.row_pk.as_ref();
+            t.data
+                .push(RelValue::new(&t.head, &new, Some(DataKey::decode(&mut bytes).unwrap())));
         }
     }
 
@@ -305,6 +309,46 @@ mod tests {
         Ok(())
     }
 
+    #[test]
+    fn test_eval_incr_maintains_row_ids() -> ResultTest<()> {
+        let (db, _) = make_test_db()?;
+        let mut tx = db.begin_tx();
+
+        let schema = ProductType::from_iter([("u8", BuiltinType::U8)]);
+        let row = product!(1u8);
+
+        // generate row id from row
+        let id1 = &row.to_data_key().to_bytes();
+
+        // create table empty table "test"
+        let table_id = create_table_with_rows(&db, &mut tx, "test", schema.clone(), &[])?;
+
+        // select * from test
+        let query = QueryExpr::new(db_table(schema.clone(), "test", table_id));
+        let query = QuerySet(vec![Query { queries: vec![query] }]);
+
+        let op = TableOp {
+            op_type: 0,
+            row_pk: id1.clone(),
+            row: row.clone(),
+        };
+
+        let update = DatabaseTableUpdate {
+            table_id,
+            table_name: "test".into(),
+            ops: vec![op],
+        };
+
+        let update = DatabaseUpdate { tables: vec![update] };
+
+        let result = query.eval_incr(&db, &mut tx, &update, AuthCtx::for_testing())?;
+        let id2 = &result.tables[0].ops[0].row_pk;
+
+        // check that both row ids are the same
+        assert_eq!(id1, id2);
+        Ok(())
+    }
+
     #[test]
     fn test_subscribe() -> ResultTest<()> {
         let (db, _tmp_dir) = make_test_db()?;

crates/core/src/subscription/subscription.rs

@@ -1,5 +1,7 @@
 use spacetimedb_lib::auth::{StAccess, StTableType};
 use spacetimedb_lib::identity::AuthCtx;
+use spacetimedb_lib::relation::RelValue;
+use spacetimedb_lib::PrimaryKey;
 use spacetimedb_sats::{AlgebraicValue, BuiltinValue};
 use spacetimedb_vm::expr::QueryExpr;
 use std::collections::HashSet;
@@ -50,6 +52,15 @@ impl Subscription {
     }
 }
 
+// If a RelValue has an id (DataKey) return it directly, otherwise we must construct it from the
+// row itself which can be an expensive operation.
+fn pk_for_row(row: &RelValue) -> PrimaryKey {
+    match row.id {
+        Some(data_key) => PrimaryKey { data_key },
+        None => RelationalDB::pk_for_row(&row.data),
+    }
+}
+
 impl QuerySet {
     /// Queries all the [`StTableType::User`] tables *right now*
     /// and turns them into [`QueryExpr`],
@@ -103,14 +114,14 @@ impl QuerySet {
                             //Hack: remove the hidden field OP_TYPE_FIELD_NAME. see `to_mem_table`
                             // Needs to be done before calculating the PK.
                             let op_type = if let AlgebraicValue::Builtin(BuiltinValue::U8(op)) =
-                                row.elements.remove(pos_op_type)
+                                row.data.elements.remove(pos_op_type)
                             {
                                 op
                             } else {
                                 panic!("Fail to extract `{OP_TYPE_FIELD_NAME}` on `{}`", result.head.table_name)
                             };
 
-                            let row_pk = RelationalDB::pk_for_row(&row);
+                            let row_pk = pk_for_row(&row);
 
                             //Skip rows that are already resolved in a previous subscription...
                             if seen.contains(&(table.table_id, row_pk)) {
@@ -120,6 +131,7 @@ impl QuerySet {
                             seen.insert((table.table_id, row_pk));
 
                             let row_pk = row_pk.to_bytes();
+                            let row = row.data;
                             table_row_operations.ops.push(TableOp { op_type, row_pk, row });
                         }
                         output.tables.push(table_row_operations);
@@ -156,7 +168,7 @@ impl QuerySet {
                             let mut table_row_operations = Vec::new();
 
                             for row in table.data {
-                                let row_pk = RelationalDB::pk_for_row(&row);
+                                let row_pk = pk_for_row(&row);
 
                                 //Skip rows that are already resolved in a previous subscription...
                                 if seen.contains(&(t.table_id, row_pk)) {
@@ -165,6 +177,7 @@ impl QuerySet {
                                 seen.insert((t.table_id, row_pk));
 
                                 let row_pk = row_pk.to_bytes();
+                                let row = row.data;
                                 table_row_operations.push(TableOp {
                                     op_type: 1, // Insert
                                     row_pk,

crates/core/src/vm.rs

@@ -4,6 +4,7 @@ use crate::db::datastore::locking_tx_datastore::MutTxId;
 use crate::db::datastore::traits::{ColumnDef, IndexDef, IndexId, SequenceId, TableDef};
 use crate::db::relational_db::RelationalDB;
 use crate::error::DBError;
+use itertools::Itertools;
 use spacetimedb_lib::auth::{StAccess, StTableType};
 use spacetimedb_lib::identity::AuthCtx;
 use spacetimedb_lib::relation::{FieldExpr, Relation};
@@ -177,15 +178,19 @@ impl<'db, 'tx> DbProgram<'db, 'tx> {
         let result = self._eval_query(query)?;
 
         match result {
-            Code::Table(result) => self._execute_delete(&table, result.data),
+            Code::Table(result) => {
+                self._execute_delete(&table, result.data.into_iter().map(|row| row.data).collect_vec())
+            }
             _ => Ok(result),
         }
     }
 
     fn insert_query(&mut self, table: &Table, query: QueryCode) -> Result<Code, ErrorVm> {
         let result = self._eval_query(query)?;
         match result {
-            Code::Table(result) => self._execute_insert(table, result.data),
+            Code::Table(result) => {
+                self._execute_insert(table, result.data.into_iter().map(|row| row.data).collect_vec())
+            }
             _ => Ok(result),
         }
     }
@@ -282,9 +287,12 @@ impl ProgramVm for DbProgram<'_, '_> {
                     Code::Table(result) => result,
                     _ => return Ok(result),
                 };
-                self._execute_delete(&table, deleted.data.clone())?;
+                self._execute_delete(
+                    &table,
+                    deleted.data.clone().into_iter().map(|row| row.data).collect_vec(),
+                )?;
 
-                let to_insert = mem_table(table.head(), deleted.data);
+                let to_insert = mem_table(table.head(), deleted.data.into_iter().map(|row| row.data));
                 insert.table = Table::MemTable(to_insert);
 
                 let result = self.insert_query(&table, insert)?;
@@ -335,7 +343,7 @@ impl RelOps for TableCursor<'_> {
     #[tracing::instrument(skip_all)]
     fn next(&mut self) -> Result<Option<RelValue>, ErrorVm> {
         if let Some(row) = self.iter.next() {
-            return Ok(Some(RelValue::new(self.head(), row.view())));
+            return Ok(Some(RelValue::new(self.head(), row.view(), Some(*row.id()))));
         };
         Ok(None)
     }
@@ -362,7 +370,7 @@ where
     #[tracing::instrument(skip_all)]
     fn next(&mut self) -> Result<Option<RelValue>, ErrorVm> {
         if let Some(row) = self.iter.next() {
-            return Ok(Some(RelValue::new(self.head(), &row)));
+            return Ok(Some(RelValue::new(self.head(), &row, None)));
         };
         Ok(None)
     }

crates/lib/src/primary_key.rs

@@ -3,6 +3,10 @@ use std::fmt;
 use crate::buffer::{BufReader, BufWriter, DecodeError};
 use crate::DataKey;
 
+// TODO(280): Remove PrimaryKey.
+// PrimaryKey is a wrapper for DataKey which identifies each row in the database.
+// It is not a column and therefore does correspond to a traditional primary key
+// in a relational database.
 #[derive(Copy, Clone, PartialEq, Eq, Hash)]
 pub struct PrimaryKey {
     pub data_key: DataKey,