Use BTreeMap instead of OrdMap in Funrun (#24677)

Use BTreeMap instead of OrdMap in Funrun:

1. BTreeMaps are generally much more efficient memory structures. The only reason to use OrdMap is because they have cheap copy-on-write. However, we use them as read-only in FunrunInMemory indexes so there is no need for the copy-on-write.

2. The implementation of WithHeapSize for OrdMap is very inaccurate since it doesn't take into account the tree structure overhead like WithHeapSize for BTreeMap does. We could fix it, but implementing WithHeapSize for OrdMap is close to impossible due to the copy-on-write and efficient shared pointers. You can have 10000 OrdMaps that consume the same memory as 1 depending on which nodes are shared. Remove the implementation to avoid future confusion here.

GitOrigin-RevId: fe2173bddd79714e8c258e4b31a7124c504dc229

Author: Preslav Le

Cargo.lock

@@ -19,7 +19,6 @@ database = { path = "../database" }
 errors = { path = "../errors" }
 file_storage = { path = "../file_storage" }
 futures = { workspace = true }
-imbl = { workspace = true }
 indexing = { path = "../indexing" }
 isolate = { path = "../isolate" }
 keybroker = { path = "../keybroker" }

crates/function_runner/Cargo.toml

@@ -54,7 +54,6 @@ use futures::{
     FutureExt,
     TryStreamExt,
 };
-use imbl::OrdMap;
 use indexing::{
     backend_in_memory_indexes::{
         DatabaseIndexSnapshot,
@@ -158,7 +157,7 @@ struct CacheKey {
 /// The cache value is the same as [DatabaseIndexMap] apart from keeping track
 /// of last modified timestamps. The [OrdMap] keys are the index keys.
 #[derive(Clone)]
-struct CacheValue(WithHeapSize<OrdMap<Vec<u8>, (Timestamp, PackedDocument)>>);
+struct CacheValue(WithHeapSize<BTreeMap<Vec<u8>, (Timestamp, PackedDocument)>>);
 
 impl SizedValue for CacheValue {
     fn size(&self) -> u64 {
@@ -181,12 +180,11 @@ async fn load_index(
     table_name: String,
 ) -> anyhow::Result<CacheValue> {
     let _timer = load_index_timer(&table_name, &instance_name);
-    let index_map: OrdMap<Vec<u8>, (Timestamp, PackedDocument)> = persistence_snapshot
+    let index_map: BTreeMap<Vec<u8>, (Timestamp, PackedDocument)> = persistence_snapshot
         .index_scan(index_id, table_id, &Interval::all(), Order::Asc, usize::MAX)
         .map_ok(|(key, ts, doc)| (key.0, (ts, PackedDocument::pack(doc))))
-        .try_collect::<Vec<_>>()
-        .await?
-        .into();
+        .try_collect()
+        .await?;
     log_funrun_index_load_rows(index_map.len() as u64, &table_name, &instance_name);
     Ok(CacheValue(index_map.into()))
 }

crates/function_runner/src/in_memory_indexes.rs

@@ -26,7 +26,6 @@ use std::{
 
 use ::bytes;
 use futures::channel::oneshot;
-use imbl::OrdMap;
 #[cfg(any(test, feature = "testing"))]
 use proptest::{
     prelude::Arbitrary,
@@ -709,25 +708,6 @@ impl<T: HeapSize> From<WithHeapSize<BTreeSet<T>>> for BTreeSet<T> {
     }
 }
 
-impl<K: HeapSize + Ord, V: HeapSize> From<OrdMap<K, V>> for WithHeapSize<OrdMap<K, V>> {
-    fn from(value: OrdMap<K, V>) -> Self {
-        let elements_heap_size = value
-            .iter()
-            .map(|(k, v)| k.heap_size() + v.heap_size())
-            .sum();
-        Self {
-            inner: value,
-            elements_heap_size,
-        }
-    }
-}
-
-impl<K: HeapSize + Ord, V: HeapSize> HeapSize for WithHeapSize<OrdMap<K, V>> {
-    fn heap_size(&self) -> usize {
-        self.elements_heap_size
-    }
-}
-
 impl HeapSize for serde_json::Map<String, JsonValue> {
     fn heap_size(&self) -> usize {
         // It's actually an IndexMap since we enable the preserve_order feature.