Skip to content

Commit c26c658

Browse files
CentrilCentril
authored
WASM ABI: add datastore_table_row_count (#1636)
1 parent 2b69583 commit c26c658

File tree

14 files changed

+133
-37
lines changed

14 files changed

+133
-37
lines changed

‎crates/bindings-csharp/Runtime/Internal/FFI.cs‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -120,6 +120,12 @@ public static partial CheckedStatus _table_id_from_name(
120120
out TableId out_
121121
);
122122

123+
[LibraryImport(StdbNamespace)]
124+
public static partial CheckedStatus _datastore_table_row_count(
125+
TableId table_id,
126+
out ulong out_
127+
);
128+
123129
[LibraryImport(StdbNamespace)]
124130
public static partial CheckedStatus _iter_by_col_eq(
125131
TableId table_id,

‎crates/bindings-csharp/Runtime/bindings.c‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -45,6 +45,9 @@ IMPORT(void, _console_log,
4545
IMPORT(Status, _table_id_from_name,
4646
(const uint8_t* name, uint32_t name_len, TableId* id),
4747
(name, name_len, id));
48+
IMPORT(Status, _datastore_table_row_count,
49+
(TableId table_id, uint64_t* count),
50+
(table_id, count));
4851
IMPORT(Status, _iter_by_col_eq,
4952
(TableId table_id, ColId col_id, const uint8_t* value,
5053
uint32_t value_len, RowIter* iter),

‎crates/bindings-sys/src/lib.rs‎

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -40,6 +40,21 @@ pub mod raw {
4040
/// - `NO_SUCH_TABLE`, when `name` is not the name of a table.
4141
pub fn _table_id_from_name(name: *const u8, name_len: usize, out: *mut TableId) -> u16;
4242

43+
/// Writes the number of rows currently in table identified by `table_id` to `out`.
44+
///
45+
/// # Traps
46+
///
47+
/// Traps if:
48+
/// - `out` is NULL or `out[..size_of::<u64>()]` is not in bounds of WASM memory.
49+
///
50+
/// # Errors
51+
///
52+
/// Returns an error:
53+
///
54+
/// - `NOT_IN_TRANSACTION`, when called outside of a transaction.
55+
/// - `NO_SUCH_TABLE`, when `table_id` is not a known ID of a table.
56+
pub fn _datastore_table_row_count(table_id: TableId, out: *mut u64) -> u16;
57+
4358
/// Finds all rows in the table identified by `table_id`,
4459
/// where the row has a column, identified by `col_id`,
4560
/// with data matching the byte string, in WASM memory, pointed to at by `val`.
@@ -523,6 +538,14 @@ pub fn table_id_from_name(name: &str) -> Result<TableId, Errno> {
523538
unsafe { call(|out| raw::_table_id_from_name(name.as_ptr(), name.len(), out)) }
524539
}
525540

541+
/// Queries and returns the number of rows in the table identified by `table_id`.
542+
///
543+
/// Returns an error if the table does not exist or if not in a transaction.
544+
#[inline]
545+
pub fn datastore_table_row_count(table_id: TableId) -> Result<u64, Errno> {
546+
unsafe { call(|out| raw::_datastore_table_row_count(table_id, out)) }
547+
}
548+
526549
/// Finds all rows in the table identified by `table_id`,
527550
/// where the row has a column, identified by `col_id`,
528551
/// with data matching the byte string `val`.

‎crates/core/src/db/datastore/locking_tx_datastore/committed_state.rs‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -57,6 +57,11 @@ impl StateView for CommittedState {
5757
fn get_schema(&self, table_id: TableId) -> Option<&Arc<TableSchema>> {
5858
self.tables.get(&table_id).map(|table| table.get_schema())
5959
}
60+
61+
fn table_row_count(&self, table_id: TableId) -> Option<u64> {
62+
self.get_table(table_id).map(|table| table.row_count)
63+
}
64+
6065
fn iter<'a>(&'a self, ctx: &'a ExecutionContext, table_id: TableId) -> Result<Iter<'a>> {
6166
if let Some(table_name) = self.table_name(table_id) {
6267
return Ok(Iter::new(ctx, table_id, table_name, None, self));

‎crates/core/src/db/datastore/locking_tx_datastore/mut_tx.rs‎

Lines changed: 34 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -7,16 +7,13 @@ use super::{
77
tx_state::TxState,
88
SharedMutexGuard, SharedWriteGuard,
99
};
10-
use crate::db::{
11-
datastore::{
12-
system_tables::{
13-
table_name_is_system, StColumnFields, StColumnRow, StConstraintFields, StConstraintRow, StFields as _,
14-
StIndexFields, StIndexRow, StScheduledRow, StSequenceFields, StSequenceRow, StTableFields, StTableRow,
15-
SystemTable, ST_COLUMN_ID, ST_CONSTRAINT_ID, ST_INDEX_ID, ST_SCHEDULED_ID, ST_SEQUENCE_ID, ST_TABLE_ID,
16-
},
17-
traits::{RowTypeForTable, TxData},
10+
use crate::db::datastore::{
11+
system_tables::{
12+
table_name_is_system, StColumnFields, StColumnRow, StConstraintFields, StConstraintRow, StFields as _,
13+
StIndexFields, StIndexRow, StScheduledRow, StSequenceFields, StSequenceRow, StTableFields, StTableRow,
14+
SystemTable, ST_COLUMN_ID, ST_CONSTRAINT_ID, ST_INDEX_ID, ST_SCHEDULED_ID, ST_SEQUENCE_ID, ST_TABLE_ID,
1815
},
19-
db_metrics::table_num_rows,
16+
traits::{RowTypeForTable, TxData},
2017
};
2118
use crate::{
2219
error::{DBError, IndexError, SequenceError, TableError},
@@ -1082,6 +1079,18 @@ impl StateView for MutTxId {
10821079
.map(|table| table.get_schema())
10831080
}
10841081

1082+
fn table_row_count(&self, table_id: TableId) -> Option<u64> {
1083+
let commit_count = self.committed_state_write_lock.table_row_count(table_id);
1084+
let (tx_ins_count, tx_del_count) = self.tx_state.table_row_count(table_id);
1085+
let commit_count = commit_count.map(|cc| cc - tx_del_count);
1086+
// Keep track of whether `table_id` exists.
1087+
match (commit_count, tx_ins_count) {
1088+
(Some(cc), Some(ic)) => Some(cc + ic),
1089+
(Some(c), None) | (None, Some(c)) => Some(c),
1090+
(None, None) => None,
1091+
}
1092+
}
1093+
10851094
fn iter<'a>(&'a self, ctx: &'a ExecutionContext, table_id: TableId) -> Result<Iter<'a>> {
10861095
if let Some(table_name) = self.table_name(table_id) {
10871096
return Ok(Iter::new(
@@ -1137,29 +1146,31 @@ impl StateView for MutTxId {
11371146
))),
11381147
None => {
11391148
#[cfg(feature = "unindexed_iter_by_col_range_warn")]
1140-
match self.schema_for_table(ctx, table_id) {
1149+
match self.table_row_count(table_id) {
11411150
// TODO(ux): log these warnings to the module logs rather than host logs.
1142-
Err(e) => log::error!(
1143-
"iter_by_col_range on unindexed column, but got error from `schema_for_table` during diagnostics: {e:?}",
1151+
None => log::error!(
1152+
"iter_by_col_range on unindexed column, but couldn't fetch table `{table_id}`s row count",
11441153
),
1145-
Ok(schema) => {
1154+
Some(num_rows) => {
11461155
const TOO_MANY_ROWS_FOR_SCAN: u64 = 1000;
1147-
1148-
let table_name = &schema.table_name;
1149-
let num_rows = table_num_rows(ctx.database(), table_id, table_name);
1150-
11511156
if num_rows >= TOO_MANY_ROWS_FOR_SCAN {
1152-
let col_names = cols.iter()
1153-
.map(|col_id| schema.columns()
1154-
.get(col_id.idx())
1155-
.map(|col| &col.col_name[..])
1156-
.unwrap_or("[unknown column]"))
1157+
let schema = self.schema_for_table(ctx, table_id).unwrap();
1158+
let table_name = &schema.table_name;
1159+
let col_names = cols
1160+
.iter()
1161+
.map(|col_id| {
1162+
schema
1163+
.columns()
1164+
.get(col_id.idx())
1165+
.map(|col| &col.col_name[..])
1166+
.unwrap_or("[unknown column]")
1167+
})
11571168
.collect::<Vec<_>>();
11581169
log::warn!(
11591170
"iter_by_col_range without index: table {table_name} has {num_rows} rows; scanning columns {col_names:?}",
11601171
);
11611172
}
1162-
},
1173+
}
11631174
}
11641175

11651176
Ok(IterByColRange::Scan(ScanIterByColRange::new(

‎crates/core/src/db/datastore/locking_tx_datastore/state_view.rs‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -32,6 +32,9 @@ pub trait StateView {
3232
Ok(row.map(|row| row.read_col(StTableFields::TableId).unwrap()))
3333
}
3434

35+
/// Returns the number of rows in the table identified by `table_id`.
36+
fn table_row_count(&self, table_id: TableId) -> Option<u64>;
37+
3538
fn iter<'a>(&'a self, ctx: &'a ExecutionContext, table_id: TableId) -> Result<Iter<'a>>;
3639

3740
fn table_name(&self, table_id: TableId) -> Option<&str> {

‎crates/core/src/db/datastore/locking_tx_datastore/tx.rs‎

Lines changed: 4 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,10 @@ impl StateView for TxId {
2626
self.committed_state_shared_lock.get_schema(table_id)
2727
}
2828

29+
fn table_row_count(&self, table_id: TableId) -> Option<u64> {
30+
self.committed_state_shared_lock.table_row_count(table_id)
31+
}
32+
2933
fn iter<'a>(&'a self, ctx: &'a ExecutionContext, table_id: TableId) -> Result<Iter<'a>> {
3034
self.committed_state_shared_lock.iter(ctx, table_id)
3135
}
@@ -60,12 +64,6 @@ impl TxId {
6064
record_metrics(ctx, self.timer, self.lock_wait_time, true, None, None);
6165
}
6266

63-
pub(crate) fn get_row_count(&self, table_id: TableId) -> Option<u64> {
64-
self.committed_state_shared_lock
65-
.get_table(table_id)
66-
.map(|table| table.row_count)
67-
}
68-
6967
/// The Number of Distinct Values (NDV) for a column or list of columns,
7068
/// if there's an index available on `cols`.
7169
pub(crate) fn num_distinct_values(&self, table_id: TableId, cols: &ColList) -> Option<u64> {

‎crates/core/src/db/datastore/locking_tx_datastore/tx_state.rs‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -64,6 +64,14 @@ pub(super) struct TxState {
6464
}
6565

6666
impl TxState {
67+
/// Returns the row count in insert tables
68+
/// and the number of rows deleted from committed state.
69+
pub(super) fn table_row_count(&self, table_id: TableId) -> (Option<u64>, u64) {
70+
let del_count = self.delete_tables.get(&table_id).map(|dt| dt.len() as u64).unwrap_or(0);
71+
let ins_count = self.insert_tables.get(&table_id).map(|it| it.row_count);
72+
(ins_count, del_count)
73+
}
74+
6775
/// When there's an index on `cols`,
6876
/// returns an iterator over the [BTreeIndex] that yields all the `RowId`s
6977
/// that match the specified `value` in the indexed column.

‎crates/core/src/db/relational_db.rs‎

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -937,6 +937,12 @@ impl RelationalDB {
937937
self.inner.table_name_from_id_mut_tx(ctx, tx, table_id)
938938
}
939939

940+
pub fn table_row_count_mut(&self, tx: &MutTx, table_id: TableId) -> Option<u64> {
941+
// TODO(Centril): Go via MutTxDatastore trait instead.
942+
// Doing this for now to ship this quicker.
943+
tx.table_row_count(table_id)
944+
}
945+
940946
pub fn column_constraints(
941947
&self,
942948
tx: &mut MutTx,
@@ -1821,7 +1827,7 @@ mod tests {
18211827

18221828
let stdb = stdb.reopen()?;
18231829
let tx = stdb.begin_tx();
1824-
assert_eq!(tx.get_row_count(table_id).unwrap(), 2);
1830+
assert_eq!(tx.table_row_count(table_id).unwrap(), 2);
18251831
Ok(())
18261832
}
18271833

‎crates/core/src/estimation.rs‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
use crate::db::relational_db::Tx;
1+
use crate::db::{datastore::locking_tx_datastore::state_view::StateView as _, relational_db::Tx};
22
use spacetimedb_primitives::{ColList, TableId};
33
use spacetimedb_vm::expr::{Query, QueryExpr, SourceExpr};
44

@@ -11,7 +11,7 @@ pub fn num_rows(tx: &Tx, expr: &QueryExpr) -> u64 {
1111
fn row_est(tx: &Tx, src: &SourceExpr, ops: &[Query]) -> u64 {
1212
match ops {
1313
// The base case is the table row count.
14-
[] => src.table_id().and_then(|id| tx.get_row_count(id)).unwrap_or(0),
14+
[] => src.table_id().and_then(|id| tx.table_row_count(id)).unwrap_or(0),
1515
// Walk in reverse from the end (`op`) to the beginning.
1616
[input @ .., op] => match op {
1717
// How selective is an index lookup?
@@ -61,7 +61,7 @@ fn row_est(tx: &Tx, src: &SourceExpr, ops: &[Query]) -> u64 {
6161
/// Note this method is not applicable to range scans.
6262
fn index_row_est(tx: &Tx, table_id: TableId, cols: &ColList) -> u64 {
6363
tx.num_distinct_values(table_id, cols)
64-
.map_or(0, |ndv| tx.get_row_count(table_id).unwrap_or(0) / ndv)
64+
.map_or(0, |ndv| tx.table_row_count(table_id).unwrap_or(0) / ndv)
6565
}
6666

6767
#[cfg(test)]

0 commit comments

Comments
 (0)