graph: Speed up lookups in RowGroup by memoizing last entry for each id

Author: lutter

graph/src/components/store/write.rs

@@ -1,5 +1,8 @@
 //! Data structures and helpers for writing subgraph changes to the store
-use std::{collections::HashSet, sync::Arc};
+use std::{
+    collections::{HashMap, HashSet},
+    sync::Arc,
+};
 
 use crate::{
     blockchain::{block_stream::FirehoseCursor, BlockPtr, BlockTime},
@@ -8,6 +11,7 @@ use crate::{
     data::{store::Id, subgraph::schema::SubgraphError},
     data_source::CausalityRegion,
     derive::CacheWeight,
+    env::ENV_VARS,
     internal_error,
     util::cache_weight::CacheWeight,
 };
@@ -154,9 +158,10 @@ impl EntityModification {
     }
 
     pub fn creates_entity(&self) -> bool {
+        use EntityModification::*;
         match self {
-            EntityModification::Insert { .. } => true,
-            EntityModification::Overwrite { .. } | EntityModification::Remove { .. } => false,
+            Insert { .. } => true,
+            Overwrite { .. } | Remove { .. } => false,
         }
     }
 
@@ -310,6 +315,10 @@ pub struct RowGroup {
     rows: Vec<EntityModification>,
 
     immutable: bool,
+
+    /// Map the `key.entity_id` of all entries in `rows` to the index with
+    /// the most recent entry for that id to speed up lookups
+    last_mod: HashMap<Id, usize>,
 }
 
 impl RowGroup {
@@ -318,6 +327,7 @@ impl RowGroup {
             entity_type,
             rows: Vec::new(),
             immutable,
+            last_mod: HashMap::new(),
         }
     }
 
@@ -374,6 +384,21 @@ impl RowGroup {
     }
 
     pub fn last_op(&self, key: &EntityKey, at: BlockNumber) -> Option<EntityOp<'_>> {
+        if ENV_VARS.store.write_batch_memoize {
+            let idx = *self.last_mod.get(&key.entity_id)?;
+            if let Some(op) = self.rows.get(idx).and_then(|emod| {
+                if emod.block() <= at {
+                    Some(emod.as_entity_op(at))
+                } else {
+                    None
+                }
+            }) {
+                return Some(op);
+            }
+        }
+        // We are looking for the change at a block `at` that is before the
+        // change we remember in `last_mod`, and therefore have to scan
+        // through all changes
         self.rows
             .iter()
             // We are scanning backwards, i.e., in descendng order of
@@ -383,7 +408,14 @@ impl RowGroup {
             .map(|emod| emod.as_entity_op(at))
     }
 
+    /// Return an iterator over all changes that are effective at `at`. That
+    /// makes it possible to construct the state that the deployment will
+    /// have once all changes for block `at` have been written.
     pub fn effective_ops(&self, at: BlockNumber) -> impl Iterator<Item = EntityOp<'_>> {
+        // We don't use `self.last_mod` here, because we need to return
+        // operations for all entities that have pending changes at block
+        // `at`, and there is no guarantee that `self.last_mod` is visible
+        // at `at` since the change in `self.last_mod` might come after `at`
         let mut seen = HashSet::new();
         self.rows
             .iter()
@@ -400,7 +432,12 @@ impl RowGroup {
 
     /// Find the most recent entry for `id`
     fn prev_row_mut(&mut self, id: &Id) -> Option<&mut EntityModification> {
-        self.rows.iter_mut().rfind(|emod| emod.id() == id)
+        if ENV_VARS.store.write_batch_memoize {
+            let idx = *self.last_mod.get(id)?;
+            self.rows.get_mut(idx)
+        } else {
+            self.rows.iter_mut().rfind(|emod| emod.id() == id)
+        }
     }
 
     /// Append `row` to `self.rows` by combining it with a previously
@@ -433,6 +470,14 @@ impl RowGroup {
                 ));
             }
 
+            if row.id() != prev_row.id() {
+                return Err(internal_error!(
+                    "last_mod map is corrupted: got id {} looking up id {}",
+                    prev_row.id(),
+                    row.id()
+                ));
+            }
+
             // The heart of the matter: depending on what `row` is, clamp
             // `prev_row` and either ignore `row` since it is not needed, or
             // turn it into an `Insert`, which also does not require
@@ -460,25 +505,31 @@ impl RowGroup {
                     Insert { .. },
                 ) => {
                     // prev_row was deleted
-                    self.rows.push(row);
+                    self.push_row(row);
                 }
                 (
                     Insert { end: None, .. } | Overwrite { end: None, .. },
                     Overwrite { block, .. },
                 ) => {
                     prev_row.clamp(*block)?;
-                    self.rows.push(row.as_insert(&self.entity_type)?);
+                    let row = row.as_insert(&self.entity_type)?;
+                    self.push_row(row);
                 }
                 (Insert { end: None, .. } | Overwrite { end: None, .. }, Remove { block, .. }) => {
                     prev_row.clamp(*block)?;
                 }
             }
         } else {
-            self.rows.push(row);
+            self.push_row(row);
         }
         Ok(())
     }
 
+    fn push_row(&mut self, row: EntityModification) {
+        self.last_mod.insert(row.id().clone(), self.rows.len());
+        self.rows.push(row);
+    }
+
     fn append(&mut self, group: RowGroup) -> Result<(), StoreError> {
         if self.entity_type != group.entity_type {
             return Err(internal_error!(
@@ -659,7 +710,7 @@ pub struct Batch {
     pub first_block: BlockNumber,
     /// The firehose cursor corresponding to `block_ptr`
     pub firehose_cursor: FirehoseCursor,
-    mods: RowGroups,
+    pub mods: RowGroups,
     /// New data sources
     pub data_sources: DataSources,
     pub deterministic_errors: Vec<SubgraphError>,
@@ -924,6 +975,7 @@ impl<'a> Iterator for WriteChunkIter<'a> {
 
 #[cfg(test)]
 mod test {
+    use std::collections::HashMap;
     use std::sync::Arc;
 
     use crate::{
@@ -947,18 +999,27 @@ mod test {
 
         assert_eq!(values.len(), blocks.len());
 
-        let rows = values
+        let rows: Vec<_> = values
             .iter()
             .zip(blocks.iter())
             .map(|(value, block)| EntityModification::Remove {
                 key: ROW_GROUP_TYPE.key(Id::String(Word::from(value.to_string()))),
                 block: *block,
             })
             .collect();
+        let last_mod = rows
+            .iter()
+            .enumerate()
+            .fold(HashMap::new(), |mut map, (idx, emod)| {
+                map.insert(emod.id().clone(), idx);
+                map
+            });
+
         let group = RowGroup {
             entity_type: ENTRY_TYPE.clone(),
             rows,
             immutable: false,
+            last_mod,
         };
         let act = group
             .clamps_by_block()

graph/src/env/store.rs

@@ -129,6 +129,11 @@ pub struct EnvVarsStore {
     /// is 10_000 which corresponds to 10MB. Setting this to 0 disables
     /// write batching.
     pub write_batch_size: usize,
+    /// Whether to memoize the last operation for each entity in a write
+    /// batch to speed up adding more entities. Set by
+    /// `GRAPH_STORE_WRITE_BATCH_MEMOIZE`. The default is `true`.
+    /// Remove after 2025-07-01 if there have been no issues with it.
+    pub write_batch_memoize: bool,
     /// Whether to create GIN indexes for array attributes. Set by
     /// `GRAPH_STORE_CREATE_GIN_INDEXES`. The default is `false`
     pub create_gin_indexes: bool,
@@ -184,6 +189,7 @@ impl TryFrom<InnerStore> for EnvVarsStore {
             connection_min_idle: x.connection_min_idle,
             connection_idle_timeout: Duration::from_secs(x.connection_idle_timeout_in_secs),
             write_queue_size: x.write_queue_size,
+            write_batch_memoize: x.write_batch_memoize,
             batch_target_duration: Duration::from_secs(x.batch_target_duration_in_secs),
             batch_timeout: x.batch_timeout_in_secs.map(Duration::from_secs),
             batch_workers: x.batch_workers,
@@ -277,6 +283,8 @@ pub struct InnerStore {
     write_batch_duration_in_secs: u64,
     #[envconfig(from = "GRAPH_STORE_WRITE_BATCH_SIZE", default = "10000")]
     write_batch_size: usize,
+    #[envconfig(from = "GRAPH_STORE_WRITE_BATCH_MEMOIZE", default = "true")]
+    write_batch_memoize: bool,
     #[envconfig(from = "GRAPH_STORE_CREATE_GIN_INDEXES", default = "false")]
     create_gin_indexes: bool,
     #[envconfig(from = "GRAPH_STORE_USE_BRIN_FOR_ALL_QUERY_TYPES", default = "false")]