wip

Author: syvb

pgvectorscale/src/access_method/build.rs

@@ -487,6 +487,7 @@ unsafe fn insert_storage<S: Storage>(
         spare_vector,
         storage,
         stats,
+        false,
     );
 }
 
@@ -662,12 +663,9 @@ fn do_heap_scan(
                     );
                 }
 
-                finalize_index_build(
-                    &mut plain,
-                    bs.into_build_state(),
-                    index_relation,
-                    write_stats,
-                )
+                // In parallel mode, nodes are finalized during insertion via streaming
+                // Just need to handle any remaining cached nodes and update meta page
+                finalize_remaining_parallel_nodes(&mut plain, bs, index_relation, write_stats)
             }
             StorageType::SbqCompression => {
                 let mut bq = unsafe {
@@ -701,7 +699,9 @@ fn do_heap_scan(
                     );
                 }
 
-                finalize_index_build(&mut bq, bs.into_build_state(), index_relation, write_stats)
+                // In parallel mode, nodes are finalized during insertion via streaming
+                // Just need to handle any remaining cached nodes and update meta page
+                finalize_remaining_parallel_nodes(&mut bq, bs, index_relation, write_stats)
             }
         }
     } else {
@@ -774,6 +774,17 @@ fn do_heap_scan(
     }
 }
 
+fn finalize_remaining_parallel_nodes<S: Storage>(
+    storage: &mut S,
+    state: BuildStateParallel,
+    index_relation: &PgRelation,
+    write_stats: WriteStats,
+) -> usize {
+    // Convert parallel state to regular build state for final processing
+    let build_state = state.into_build_state();
+    finalize_index_build(storage, build_state, index_relation, write_stats)
+}
+
 fn finalize_index_build<S: Storage>(
     storage: &mut S,
     state: BuildState,
@@ -1001,6 +1012,7 @@ fn build_callback_internal<S: Storage>(
         spare_vector,
         storage,
         &mut state.stats,
+        false,
     );
 }
 
@@ -1060,14 +1072,16 @@ fn build_callback_parallel_internal<S: Storage>(
         &mut state.local_stats,
     );
 
-    // Insert node into graph - PostgreSQL page locking handles concurrency
+    // Insert node into graph with parallel build mode enabled
+    // PostgreSQL page locking handles concurrency when finalizing nodes
     state.graph.insert(
         index,
         index_pointer,
         vector,
         spare_vector,
         storage,
         &mut state.local_stats,
+        true,
     );
 }
 

pgvectorscale/src/access_method/graph/mod.rs

@@ -642,6 +642,7 @@ digraph G {
         spare_vec: LabeledVector,
         storage: &S,
         stats: &mut InsertStats,
+        parallel_build: bool,
     ) {
         self.update_start_nodes(
             index,
@@ -658,6 +659,13 @@ digraph G {
 
         // Insert starting from default start node and avoid label filtering
         self.insert_internal(index_pointer, vec, true, storage, stats);
+
+        // In parallel build mode, periodically flush cached nodes to disk to avoid
+        // accumulating large amounts of data in memory that would need to be
+        // processed sequentially at the end
+        if parallel_build {
+            self.maybe_flush_neighbor_cache(storage, stats);
+        }
     }
 
     fn insert_internal<S: Storage>(
@@ -736,4 +744,18 @@ digraph G {
         let (_pruned, n) = self.add_neighbors(storage, from, from_labels, new.clone(), prune_stats);
         n.contains(&new[0])
     }
+
+    /// In parallel builds, periodically flush cached neighbor data to disk
+    /// to prevent accumulating too much data in memory
+    fn maybe_flush_neighbor_cache<S: Storage>(
+        &mut self,
+        storage: &S,
+        stats: &mut InsertStats,
+    ) {
+        if let GraphNeighborStore::Builder(ref cache) = self.neighbor_store {
+            // Flush cache when it's getting full to avoid memory pressure
+            // and ensure nodes get written to disk during parallel processing
+            cache.flush_if_above_threshold(storage, &mut stats.prune_neighbor_stats, 0.8);
+        }
+    }
 }

pgvectorscale/src/access_method/graph/neighbor_store.rs

@@ -136,6 +136,45 @@ impl BuilderNeighborCache {
             storage.set_neighbors_on_disk(key, new_neighbors.as_slice(), stats);
         }
     }
+
+    /// Flush cached entries to disk if cache usage is above the given threshold.
+    /// This helps prevent memory buildup during parallel builds.
+    pub fn flush_if_above_threshold<S: Storage>(
+        &self,
+        storage: &S,
+        stats: &mut PruneNeighborStats,
+        threshold: f64,
+    ) {
+        let mut cache = self.neighbor_map.borrow_mut();
+        let current_size = cache.len();
+        let capacity_val = cache.cap().get();
+        
+        if current_size as f64 / capacity_val as f64 > threshold {
+            let target_size = (capacity_val as f64 * 0.5) as usize;
+            
+            // Flush least recently used entries to disk
+            while cache.len() > target_size {
+                if let Some((neighbors_of, entry)) = cache.pop_lru() {
+                    let pruned_neighbors = if entry.neighbors.len() > self.num_neighbors {
+                        Graph::prune_neighbors(
+                            self.max_alpha,
+                            self.num_neighbors,
+                            entry.labels.as_ref(),
+                            entry.neighbors,
+                            storage,
+                            stats,
+                        )
+                    } else {
+                        entry.neighbors
+                    };
+                    
+                    storage.set_neighbors_on_disk(neighbors_of, &pruned_neighbors, stats);
+                } else {
+                    break;
+                }
+            }
+        }
+    }
 }
 
 pub enum GraphNeighborStore {

pgvectorscale/src/util/lru.rs

@@ -110,4 +110,14 @@ impl<K: Hash + Eq + Clone, V> LruCacheWithStats<K, V> {
     pub fn into_parts(self) -> (LruCache<K, V>, CacheStats) {
         (self.cache, self.stats)
     }
+
+    /// Remove and return the least recently used key-value pair
+    pub fn pop_lru(&mut self) -> Option<(K, V)> {
+        if let Some((key, value)) = self.cache.pop_lru() {
+            self.stats.evictions += 1;
+            Some((key, value))
+        } else {
+            None
+        }
+    }
 }