Fixes #3240, perf of indexing many rows with same key (#3971)

# Description of Changes

Fixes #3240.

Non-unique indices are now backed by a type `SameKeyEntry` which holds
the `RowPointer`s for the same key.
When these `RowPointer`s exceed 4KiB (512 entries), the data structure
switches from using an array list to a hash set.

# API and ABI breaking changes

None

# Expected complexity level and risk

2?

# Testing

Covered by existing tests, though more test will come in future PRs.

Author: Centril

crates/bindings-typescript/src/server/index.ts

@@ -7,6 +7,8 @@ export { type Reducer, type ReducerCtx } from '../lib/reducers';
 export { type DbView } from './db_view';
 export { and, or, not } from './query';
 export type { ProcedureCtx, TransactionCtx } from '../lib/procedures';
+export { toCamelCase } from '../lib/util';
+export { type Uuid } from '../lib/uuid';
 
 import './polyfills'; // Ensure polyfills are loaded
 import './register_hooks'; // Ensure module hooks are registered

crates/table/src/table_index/mod.rs

@@ -24,6 +24,7 @@
 /// We also represent unique indices more compactly than non-unique ones, avoiding the multi-map.
 /// Additionally, beyond our btree indices,
 /// we support direct unique indices, where key are indices into `Vec`s.
+use self::same_key_entry::SameKeyEntryIter;
 use super::indexes::RowPointer;
 use super::table::RowRef;
 use crate::{read_column::ReadColumn, static_assert_size};
@@ -37,6 +38,7 @@ use spacetimedb_sats::{
 
 mod key_size;
 mod multimap;
+mod same_key_entry;
 pub mod unique_direct_fixed_cap_index;
 pub mod unique_direct_index;
 pub mod uniquemap;
@@ -47,7 +49,7 @@ use unique_direct_fixed_cap_index::{UniqueDirectFixedCapIndex, UniqueDirectFixed
 use unique_direct_index::{UniqueDirectIndex, UniqueDirectIndexPointIter, UniqueDirectIndexRangeIter};
 
 type BtreeIndex<K> = multimap::MultiMap<K, RowPointer>;
-type BtreeIndexPointIter<'a> = multimap::MultiMapPointIter<'a, RowPointer>;
+type BtreeIndexPointIter<'a> = SameKeyEntryIter<'a, RowPointer>;
 type BtreeIndexRangeIter<'a, K> = multimap::MultiMapRangeIter<'a, K, RowPointer>;
 type BtreeUniqueIndex<K> = uniquemap::UniqueMap<K, RowPointer>;
 type BtreeUniqueIndexPointIter<'a> = uniquemap::UniqueMapPointIter<'a, RowPointer>;

crates/table/src/table_index/multimap.rs

@@ -1,38 +1,39 @@
-use core::ops::RangeBounds;
-use core::slice;
-use smallvec::SmallVec;
+use super::same_key_entry::{same_key_iter, SameKeyEntry, SameKeyEntryIter};
+use core::{hash::Hash, ops::RangeBounds};
 use spacetimedb_sats::memory_usage::MemoryUsage;
 use std::collections::btree_map::{BTreeMap, Range};
 
 /// A multi map that relates a `K` to a *set* of `V`s.
 #[derive(Debug, PartialEq, Eq)]
-pub struct MultiMap<K, V> {
+pub struct MultiMap<K, V: Eq + Hash> {
     /// The map is backed by a `BTreeMap` for relating keys to values.
     ///
     /// A value set is stored as a `SmallVec`.
     /// This is an optimization over a `Vec<_>`
     /// as we allow a single element to be stored inline
     /// to improve performance for the common case of one element.
-    map: BTreeMap<K, SmallVec<[V; 1]>>,
+    map: BTreeMap<K, SameKeyEntry<V>>,
 }
 
-impl<K, V> Default for MultiMap<K, V> {
+impl<K, V: Eq + Hash> Default for MultiMap<K, V> {
     fn default() -> Self {
         Self { map: BTreeMap::new() }
     }
 }
 
-impl<K: MemoryUsage, V: MemoryUsage> MemoryUsage for MultiMap<K, V> {
+impl<K: MemoryUsage, V: MemoryUsage + Eq + Hash> MemoryUsage for MultiMap<K, V> {
     fn heap_usage(&self) -> usize {
         let Self { map } = self;
         map.heap_usage()
     }
 }
 
-impl<K: Ord, V: Ord> MultiMap<K, V> {
+impl<K: Ord, V: Ord + Hash> MultiMap<K, V> {
     /// Inserts the relation `key -> val` to this multimap.
     ///
     /// The map does not check whether `key -> val` was already in the map.
+    /// It's assumed that the same `val` is never added twice,
+    /// and multimaps do not bind one `key` to the same `val`.
     pub fn insert(&mut self, key: K, val: V) {
         self.map.entry(key).or_default().push(val);
     }
@@ -41,30 +42,31 @@ impl<K: Ord, V: Ord> MultiMap<K, V> {
     ///
     /// Returns whether `key -> val` was present.
     pub fn delete(&mut self, key: &K, val: &V) -> bool {
-        if let Some(vset) = self.map.get_mut(key) {
-            // The `vset` is not sorted, so we have to do a linear scan first.
-            if let Some(idx) = vset.iter().position(|v| v == val) {
-                vset.swap_remove(idx);
-                return true;
-            }
+        let Some(vset) = self.map.get_mut(key) else {
+            return false;
+        };
+
+        let (deleted, is_empty) = vset.delete(val);
+
+        if is_empty {
+            self.map.remove(key);
         }
-        false
+
+        deleted
     }
 
     /// Returns an iterator over the multimap that yields all the `V`s
     /// of the `K`s that fall within the specified `range`.
     pub fn values_in_range(&self, range: &impl RangeBounds<K>) -> MultiMapRangeIter<'_, K, V> {
         MultiMapRangeIter {
             outer: self.map.range((range.start_bound(), range.end_bound())),
-            inner: None,
+            inner: SameKeyEntry::empty_iter(),
         }
     }
 
     /// Returns an iterator over the multimap that yields all the `V`s of the `key: &K`.
-    pub fn values_in_point(&self, key: &K) -> MultiMapPointIter<'_, V> {
-        let vals = self.map.get(key).map(|vs| &**vs).unwrap_or_default();
-        let iter = vals.iter();
-        MultiMapPointIter { iter }
+    pub fn values_in_point(&self, key: &K) -> SameKeyEntryIter<'_, V> {
+        same_key_iter(self.map.get(key))
     }
 
     /// Returns the number of unique keys in the multimap.
@@ -75,7 +77,7 @@ impl<K: Ord, V: Ord> MultiMap<K, V> {
     /// Returns the total number of entries in the multimap.
     #[allow(unused)] // No use for this currently.
     pub fn len(&self) -> usize {
-        self.map.values().map(|ptrs| ptrs.len()).sum()
+        self.map.values().map(|vals: &SameKeyEntry<V>| vals.len()).sum()
     }
 
     /// Returns whether there are any entries in the multimap.
@@ -91,46 +93,28 @@ impl<K: Ord, V: Ord> MultiMap<K, V> {
     }
 }
 
-/// An iterator over values in a [`MultiMap`] where the key is a point.
-pub struct MultiMapPointIter<'a, V> {
-    /// The inner iterator for the value set for a found key.
-    iter: slice::Iter<'a, V>,
-}
-
-impl<'a, V> Iterator for MultiMapPointIter<'a, V> {
-    type Item = &'a V;
-
-    fn next(&mut self) -> Option<Self::Item> {
-        self.iter.next()
-    }
-}
-
 /// An iterator over values in a [`MultiMap`] where the keys are in a certain range.
-pub struct MultiMapRangeIter<'a, K, V> {
+pub struct MultiMapRangeIter<'a, K, V: Eq + Hash> {
     /// The outer iterator seeking for matching keys in the range.
-    outer: Range<'a, K, SmallVec<[V; 1]>>,
+    outer: Range<'a, K, SameKeyEntry<V>>,
     /// The inner iterator for the value set for a found key.
-    inner: Option<slice::Iter<'a, V>>,
+    inner: SameKeyEntryIter<'a, V>,
 }
 
-impl<'a, K, V> Iterator for MultiMapRangeIter<'a, K, V> {
+impl<'a, K, V: Eq + Hash> Iterator for MultiMapRangeIter<'a, K, V> {
     type Item = &'a V;
 
     fn next(&mut self) -> Option<Self::Item> {
         loop {
-            if let Some(inner) = self.inner.as_mut() {
-                if let Some(val) = inner.next() {
-                    // While the inner iterator has elements, yield them.
-                    return Some(val);
-                }
+            // While the inner iterator has elements, yield them.
+            if let Some(val) = self.inner.next() {
+                return Some(val);
             }
 
-            // This makes the iterator fused.
-            self.inner = None;
             // Advance and get a new inner, if possible, or quit.
             // We'll come back and yield elements from it in the next iteration.
-            let (_, next) = self.outer.next()?;
-            self.inner = Some(next.iter());
+            let inner = self.outer.next().map(|(_, i)| i)?;
+            self.inner = inner.iter();
         }
     }
 }

crates/table/src/table_index/same_key_entry.rs

@@ -0,0 +1,165 @@
+use crate::{indexes::RowPointer, static_assert_size};
+use core::hash::Hash;
+use core::slice;
+use smallvec::SmallVec;
+use spacetimedb_data_structures::map::{hash_set, HashCollectionExt, HashSet};
+use spacetimedb_memory_usage::MemoryUsage;
+
+/// A supporting type for multimap implementations
+/// that handles all the values for the same key,
+/// leaving the multimap to only have to care about the keys.
+///
+/// For performance reasons,
+/// this is an enum
+/// that deals with a smaller number of values in the first variant
+/// and with a larger number in the second variant.
+#[derive(Debug, PartialEq, Eq)]
+pub(super) enum SameKeyEntry<V: Eq + Hash> {
+    /// A small number of values.
+    ///
+    /// No ordering is kept between values.
+    /// This makes insertions into amortized `O(k)`
+    /// whereas deletions become `O(|values|)` instead.
+    /// This is acceptable as `|values|` is small
+    /// and because deleting from an array list is `O(n)` either way.
+    ///
+    /// This also represents the "no values" case,
+    /// although the multimap may want to delete the key in that case.
+    ///
+    /// Up to two values are represented inline here.
+    /// It's not profitable to represent this as a separate variant
+    /// as that would increase `size_of::<SameKeyEntry>()` by 8 bytes.
+    Small(SmallVec<[V; 2]>),
+
+    /// A large number of values.
+    ///
+    /// Used when the heap size of `Small` would exceed one standard page.
+    /// See [`SameKeyEntry::LARGE_AFTER_LEN`] for details.
+    ///
+    /// Note that using a `HashSet`, with `S = RandomState`,
+    /// entails that the iteration order is not deterministic.
+    /// This is observed when doing queries against the index.
+    Large(HashSet<V>),
+}
+
+static_assert_size!(SameKeyEntry<RowPointer>, 32);
+
+impl<V: Eq + Hash> Default for SameKeyEntry<V> {
+    fn default() -> Self {
+        Self::Small(<_>::default())
+    }
+}
+
+impl<V: MemoryUsage + Eq + Hash> MemoryUsage for SameKeyEntry<V> {
+    fn heap_usage(&self) -> usize {
+        match self {
+            Self::Small(x) => x.heap_usage(),
+            Self::Large(x) => x.heap_usage(),
+        }
+    }
+}
+
+impl<V: Eq + Hash> SameKeyEntry<V> {
+    /// The number of elements
+    /// beyond which the strategy is changed from small to large storage.
+    const LARGE_AFTER_LEN: usize = 4096 / size_of::<V>();
+
+    /// Pushes `val` as an entry for the key.
+    ///
+    /// This assumes that `val` was previously not recorded.
+    /// The structure does not check whether it was previously resident.
+    /// As a consequence, the time complexity is `O(k)` amortized.
+    pub(super) fn push(&mut self, val: V) {
+        match self {
+            Self::Small(list) if list.len() <= Self::LARGE_AFTER_LEN => {
+                list.push(val);
+            }
+            Self::Small(list) => {
+                // Reconstruct into a set.
+                let mut set = HashSet::with_capacity(list.len() + 1);
+                set.extend(list.drain(..));
+
+                // Add `val`.
+                set.insert(val);
+
+                *self = Self::Large(set);
+            }
+            Self::Large(set) => {
+                set.insert(val);
+            }
+        }
+    }
+
+    /// Deletes `val` as an entry for the key.
+    ///
+    /// Returns `(was_deleted, is_empty)`.
+    pub(super) fn delete(&mut self, val: &V) -> (bool, bool) {
+        match self {
+            Self::Small(list) => {
+                // The `list` is not sorted, so we have to do a linear scan first.
+                if let Some(idx) = list.iter().position(|v| v == val) {
+                    list.swap_remove(idx);
+                    (true, list.is_empty())
+                } else {
+                    (false, false)
+                }
+            }
+            Self::Large(set) => {
+                let removed = set.remove(val);
+                let empty = set.is_empty();
+                (removed, empty)
+            }
+        }
+    }
+
+    /// Returns an iterator over all the entries for this key.
+    pub(super) fn iter(&self) -> SameKeyEntryIter<'_, V> {
+        match self {
+            Self::Small(list) => SameKeyEntryIter::Small(list.iter()),
+            Self::Large(set) => SameKeyEntryIter::Large(set.iter().into()),
+        }
+    }
+
+    /// Returns an iterator over no entries.
+    pub(super) fn empty_iter<'a>() -> SameKeyEntryIter<'a, V> {
+        SameKeyEntryIter::Small(const { &[] }.iter())
+    }
+
+    /// Returns the number of entries for the same key.
+    pub(super) fn len(&self) -> usize {
+        match self {
+            Self::Small(list) => list.len(),
+            Self::Large(set) => set.len(),
+        }
+    }
+}
+
+/// Returns an iterator for a key's entries `ske`.
+/// This efficiently handles the case where there's no key (`None`).
+pub(super) fn same_key_iter<V: Eq + Hash>(ske: Option<&SameKeyEntry<V>>) -> SameKeyEntryIter<'_, V> {
+    match ske {
+        None => SameKeyEntry::empty_iter(),
+        Some(ske) => ske.iter(),
+    }
+}
+
+/// An iterator over values in a [`SameKeyEntry`].
+pub enum SameKeyEntryIter<'a, V> {
+    Small(slice::Iter<'a, V>),
+    /// This variant doesn't occur so much
+    /// and we'd like to reduce the footprint of `SameKeyEntryIter`.
+    Large(Box<hash_set::Iter<'a, V>>),
+}
+
+static_assert_size!(SameKeyEntryIter<RowPointer>, 16);
+
+impl<'a, V> Iterator for SameKeyEntryIter<'a, V> {
+    type Item = &'a V;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        match self {
+            Self::Small(list) => list.next(),
+            Self::Large(set) => set.next(),
+        }
+    }
+}

sdks/rust/tests/test-client/src/module_bindings/mod.rs

@@ -6,6 +6,8 @@ use std::{
     time::Duration,
 };
 
+const TEST_TIMEOUT_SECS: u64 = 2 * 60;
+
 #[derive(Default)]
 struct TestCounterInner {
     /// Maps test names to their outcomes
@@ -57,7 +59,7 @@ impl TestCounter {
         let lock = self.inner.lock().expect("TestCounterInner Mutex is poisoned");
         let (lock, timeout_result) = self
             .wait_until_done
-            .wait_timeout_while(lock, Duration::from_secs(90), |inner| {
+            .wait_timeout_while(lock, Duration::from_secs(TEST_TIMEOUT_SECS), |inner| {
                 inner.outcomes.len() != inner.registered.len()
             })
             .expect("TestCounterInner Mutex is poisoned");