feat(csr): fold the delta into a new base with optimize_csr()

crates/lance-csr-python/src/graph.rs

@@ -1117,6 +1117,11 @@ impl CypherEngine {
             .map_err(graph_error_to_pyerr)
     }
 
+    fn optimize_csr(&self, py: Python) -> PyResult<()> {
+        RT.block_on(Some(py), self.inner.optimize_csr())?
+            .map_err(graph_error_to_pyerr)
+    }
+
     fn node_index_stats(&self, py: Python) -> PyResult<Py<PyAny>> {
         let stats = PyDict::new(py);
         for (label, entry) in self.inner.node_index_stats() {

crates/lance-csr/src/engine.rs

@@ -171,8 +171,8 @@ pub struct CsrStats {
     /// Live edges outside the base CSR: appended since its last full build and not
     /// since removed. Lance's `num_unindexed_rows`, counted in edges.
     pub unindexed_edges: u64,
-    /// How the last refresh produced it: `build` (first), `rebuild` (full), `delta`
-    /// (changes applied) or `reused` (no edge change).
+    /// How the last refresh or `optimize_csr()` produced it: `build` (first), `rebuild`
+    /// (full), `delta` (changes applied) or `reused` (no edge change).
     pub update: &'static str,
     /// Relationship dataset version it reflects (`None` for Arrow input).
     pub version: Option<u64>,
@@ -273,7 +273,7 @@ struct CsrEntry {
     base_version: Option<u64>,
     /// Wall time of the last build or delta application.
     build_ms: f64,
-    /// How the last refresh produced this entry: build, rebuild, delta or reused.
+    /// How the last refresh or optimize produced this entry: build, rebuild, delta or reused.
     update: &'static str,
 }
 
@@ -460,12 +460,14 @@ fn validate_sources(config: &GraphConfig, schemas: &HashMap<String, SchemaRef>)
 }
 
 /// Validate sources and register them in a fresh catalog and session context.
+/// `fold_csr_deltas` rebuilds every CSR of `previous` that carries a delta.
 async fn build_engine_state(
     config: &GraphConfig,
     sources: HashMap<String, EngineSource>,
     use_csr: bool,
     target_partitions: Option<usize>,
     previous: Option<&EngineState>,
+    fold_csr_deltas: bool,
 ) -> Result<EngineState> {
     let mut session_config = SessionConfig::new();
     if let Some(partitions) = target_partitions {
@@ -508,7 +510,7 @@ async fn build_engine_state(
     validate_sources(config, &schemas)?;
     let (csr, node_index) = if use_csr {
         (
-            build_csr_entries(config, &sources, &schemas, previous).await?,
+            build_csr_entries(config, &sources, &schemas, previous, fold_csr_deltas).await?,
             build_node_index(config, &sources, previous).await?,
         )
     } else {
@@ -752,6 +754,7 @@ async fn build_csr_entries(
     sources: &HashMap<String, EngineSource>,
     schemas: &HashMap<String, SchemaRef>,
     previous: Option<&EngineState>,
+    fold_deltas: bool,
 ) -> Result<HashMap<String, CsrEntry>> {
     let mut entries = HashMap::new();
     for (_, mapping) in config.iter_relationship_mappings() {
@@ -763,7 +766,8 @@ async fn build_csr_entries(
             continue;
         };
         if let Some(old) = previous.and_then(|p| p.csr.get(rel_type)) {
-            if old.source.same_data(source) {
+            let delta = old.outgoing.delta_edges() + old.outgoing.tombstones() > 0;
+            if old.source.same_data(source) && !(fold_deltas && delta) {
                 entries.insert(rel_type.clone(), old.clone());
                 continue;
             }
@@ -786,7 +790,7 @@ async fn build_csr_entries(
         }
 
         let previous_entry = previous.and_then(|p| p.csr.get(rel_type));
-        if let Some(old) = previous_entry {
+        if let Some(old) = previous_entry.filter(|_| !fold_deltas) {
             if let Some(entry) = try_apply_csr_delta(old, source, src.name(), dst.name()).await {
                 entries.insert(rel_type.clone(), entry);
                 continue;
@@ -887,7 +891,8 @@ impl CypherEngine {
                 )));
             }
         }
-        let state = build_engine_state(&config, sources, use_csr, target_partitions, None).await?;
+        let state =
+            build_engine_state(&config, sources, use_csr, target_partitions, None, false).await?;
         Ok(Self {
             config,
             state: RwLock::new(Arc::new(state)),
@@ -1048,6 +1053,31 @@ impl CypherEngine {
             self.use_csr,
             self.target_partitions,
             Some(&current),
+            false,
+        )
+        .await?;
+        *self.state.write().expect("engine state lock poisoned") = Arc::new(state);
+        Ok(())
+    }
+
+    /// Fold each CSR's delta into a new base, at the snapshot the engine is pinned to: the
+    /// CSR's counterpart of Lance's `optimize_indices()`.
+    ///
+    /// `refresh()` applies changes as a delta over the base CSR and rebuilds on its own
+    /// only once they pass 25% of the base. This rebuilds every CSR that carries a delta
+    /// now, so lookups stop carrying it; one without a delta, and everything else, is
+    /// reused. Results never change. The replacement is validated in full before it is
+    /// swapped in, as for `refresh()`, and the two run one at a time.
+    pub async fn optimize_csr(&self) -> Result<()> {
+        let _refreshing = self.refresh_lock.lock().await;
+        let current = self.snapshot();
+        let state = build_engine_state(
+            &self.config,
+            current.sources.clone(),
+            self.use_csr,
+            self.target_partitions,
+            Some(&current),
+            true,
         )
         .await?;
         *self.state.write().expect("engine state lock poisoned") = Arc::new(state);

crates/lance-csr/tests/test_engine.rs

@@ -161,6 +161,47 @@ async fn lance_input_is_pinned_until_refresh() {
     assert_eq!(engine.dataset_versions()["Person"], 1);
 }
 
+#[tokio::test]
+async fn optimize_csr_folds_the_delta_into_the_base() {
+    let dir = tempfile::tempdir().unwrap();
+    let person_uri = dir.path().join("person.lance");
+    let knows_uri = dir.path().join("knows.lance");
+    let (person_uri, knows_uri) = (person_uri.to_str().unwrap(), knows_uri.to_str().unwrap());
+    write(people(), person_uri, WriteMode::Create).await;
+    // A base large enough that one appended edge stays under the 25% rebuild threshold.
+    let base: Vec<_> = [(1, 2), (2, 3)].into_iter().chain([(3, 4); 8]).collect();
+    write(knows(&base), knows_uri, WriteMode::Create).await;
+    let mut sources = HashMap::new();
+    for (name, uri) in [("Person", person_uri), ("KNOWS", knows_uri)] {
+        let source = EngineSource::open_lance(uri, None, HashMap::new());
+        sources.insert(name.to_string(), source.await.unwrap());
+    }
+    let engine = CypherEngine::new(config(), sources, true).await.unwrap();
+
+    write(knows(&[(2, 4)]), knows_uri, WriteMode::Append).await;
+    engine.refresh().await.unwrap();
+    let csr = engine.csr_stats()["KNOWS"].clone();
+    assert_eq!(
+        (csr.update, csr.delta_edges, csr.unindexed_edges),
+        ("delta", 1, 1)
+    );
+    assert_eq!((csr.version, csr.base_version), (Some(2), Some(1)));
+
+    engine.optimize_csr().await.unwrap();
+    let csr = engine.csr_stats()["KNOWS"].clone();
+    assert_eq!((csr.update, csr.edges, csr.delta_edges), ("rebuild", 11, 0));
+    assert_eq!((csr.tombstones, csr.unindexed_edges), (0, 0));
+    assert_eq!((csr.version, csr.base_version), (Some(2), Some(2)));
+    assert_eq!(
+        names(&engine.execute(TWO_HOPS, HashMap::new()).await.unwrap()),
+        ["Carol", "David"]
+    );
+
+    // Nothing left to fold: the CSR is reused as it is.
+    engine.optimize_csr().await.unwrap();
+    assert_eq!(engine.csr_stats()["KNOWS"].build_ms, csr.build_ms);
+}
+
 #[tokio::test]
 async fn construction_errors_are_config_errors() {
     let empty = CypherEngine::new(config(), HashMap::new(), false).await;

docs/csr.md

@@ -245,7 +245,7 @@ print(engine.csr_stats()["KNOWS"])
 | `delta_edges` | Appended edges held outside the base CSR since its last full build. |
 | `tombstones` | Removed edges masked out since the last full build. |
 | `unindexed_edges` | Live edges outside the base CSR: appended since its last full build and not since removed. Lance index statistics call this `num_unindexed_rows`. |
-| `update` | How the last build or `refresh()` produced it: `build` (first), `rebuild` (full), `delta` (changes applied) or `reused` (no edge change). |
+| `update` | How the last build, `refresh()` or [`optimize_csr()`](#folding-the-delta-in) produced it: `build` (first), `rebuild` (full), `delta` (changes applied) or `reused` (no edge change). |
 | `version` | The relationship dataset's Lance version, or `None` for PyArrow input. |
 | `base_version` | The version of the last full build, or `None` for PyArrow input. It trails `version` while deltas apply. |
 | `build_ms` | Wall time of that build or delta, both directions. |
@@ -287,7 +287,8 @@ flowchart TD
   previous snapshot stays usable. Queries already running finish on the snapshot they started with.
 
 The 25% threshold is fixed: lookups pay a little for every pending delta edge, and a rebuild clears them.
-`python/benchmarks/csr_delta_accumulated.py` measures that trade-off.
+`python/benchmarks/csr_delta_accumulated.py` measures that trade-off. To fold the delta in sooner, call
+[`optimize_csr()`](#folding-the-delta-in).
 
 ### How to keep the CSR fresh
 
@@ -300,6 +301,8 @@ The engine doesn't watch your datasets for changes. Keeping the CSR in step with
 3. **Check what it did with [`csr_stats()`](#csr_stats)**: `update` says whether the CSR was `reused`, updated
    with a `delta` or fully rebuilt (`rebuild`). `base_version` and `unindexed_edges` say how far the base
    CSR trails the data. Answers are exact either way: delta edges are part of every lookup.
+4. **Optionally, fold the delta in with [`optimize_csr()`](#folding-the-delta-in)** at a quiet moment, so
+   lookups stop carrying it.
 
 === "Python"
 
@@ -337,6 +340,43 @@ than you clean up old versions to keep updates incremental.
 
 In-memory PyArrow tables never change. To query new in-memory data, build a new engine.
 
+### Folding the delta in
+
+The CSR follows Lance's own index design. A Lance index covers the fragments it was built from, and data
+written after that stays unindexed: searches combine indexed results with a scan of the unindexed rows, so
+answers stay correct while latency grows, until you call `optimize_indices()` (see Lance's
+[index maintenance guide](https://lance.org/guide/indexing/#maintain-indexes-after-appends) and the
+[index format](https://lance.org/format/index/#basic-concepts)). The CSR maps onto that one to one:
+
+| Lance index | CSR |
+| --- | --- |
+| The fragments an index was built from | The base CSR, at `base_version` |
+| Unindexed rows, scanned alongside the index | The delta (appended edges and tombstones), read by every lookup |
+| `num_unindexed_rows` in index statistics | `unindexed_edges` in `csr_stats()` |
+| `optimize_indices()` | [`optimize_csr()`][lance_csr.CypherEngine.optimize_csr] |
+
+`optimize_csr()` rebuilds every CSR that carries a delta, at the version the engine is pinned to, so the base
+covers all the data again: `unindexed_edges`, `delta_edges` and `tombstones` drop to 0, `base_version` catches
+up with `version` and `update` says `rebuild`. A CSR without a delta is kept as it is, and answers never change.
+
+```python
+engine.refresh()  # applies the new edges as a delta
+engine.optimize_csr()  # folds it into a new base
+stats = engine.csr_stats()["KNOWS"]
+print(stats["update"], stats["unindexed_edges"])  # rebuild 0
+```
+
+!!! tip "`refresh()` or `optimize_csr()`?"
+
+    Call `refresh()` whenever your queries should see new writes. It's cheap and incremental, and it already
+    rebuilds a CSR on its own once the delta passes 25% of the base. Call `optimize_csr()` only when
+    `csr_stats()` shows a delta built up below that threshold and you want lookups to stop carrying it: after a
+    bulk load, before a latency-sensitive run of queries, or on a schedule at a quiet time. It scans every edge
+    of the relationship dataset again, like the first build, so don't call it after every refresh.
+
+`optimize_csr()` doesn't move to newer data, so call `refresh()` first. Like `refresh()`, it builds the
+replacement in full before swapping it in, and the two run one at a time.
+
 ## Should I turn it on?
 
 | Situation | Recommendation |

docs/rust.md

@@ -148,6 +148,7 @@ The deadline applies only to work inside the call; it isn't stored in the engine
 | `engine.execute(query, params)` → `pyarrow.Table` | `engine.execute(query, params).await` → `RecordBatch` |
 | `engine.execute_batches(...)` → `RecordBatchReader` | `engine.execute_stream(...).await` → stream |
 | `engine.refresh()` | `engine.refresh().await` |
+| `engine.optimize_csr()` | `engine.optimize_csr().await` |
 | `engine.execute(query, timeout=2.0)` | `timeout_at(deadline, with_execution_deadline(Some(deadline), engine.execute(...)))` |
 | `CypherEngine(..., target_partitions=4)` | `CypherEngine::with_target_partitions(config, sources, use_csr, Some(4)).await` |
 | `engine.csr_stats()` → `dict` | `engine.csr_stats()` → `HashMap<String, CsrStats>` |

python/python/lance_csr/_internal/graph.pyi

@@ -133,6 +133,7 @@ class CypherEngine:
         self, query: str, params: dict[str, _ParameterT] | None = None
     ) -> str: ...
     def refresh(self) -> None: ...
+    def optimize_csr(self) -> None: ...
     def node_index_stats(self) -> dict[str, _NodeIndexStats]: ...
     def csr_stats(self) -> dict[str, _CsrStats]: ...
     def dataset_versions(self) -> dict[str, int]: ...

python/python/lance_csr/cypher.py

@@ -258,6 +258,26 @@ def refresh(self) -> None:
         """
         self._inner.refresh()
 
+    def optimize_csr(self) -> None:
+        """Fold each CSR's delta into a new base (experimental; needs `use_csr=True`).
+
+        The CSR's counterpart of Lance's `optimize_indices()`.
+        [`refresh`][lance_csr.CypherEngine.refresh] applies Lance writes as a delta
+        over the base CSR, and rebuilds on its own only once pending changes pass 25%
+        of the base. This rebuilds every CSR that carries a delta now, at the version
+        the engine is pinned to, so lookups stop carrying it. It scans every edge
+        again, so call it after a bulk load or at a quiet time, not after every
+        refresh. It doesn't move to
+        newer versions: call `refresh` first for that. CSRs without a delta, and
+        PyArrow inputs, are left as they are, and query results never change.
+        Afterwards [`csr_stats`][lance_csr.CypherEngine.csr_stats] reports `update`
+        `"rebuild"`, `unindexed_edges` 0 and `base_version` equal to `version`. As for
+        `refresh`, the replacement is validated in full before it's swapped in, and the
+        two run one at a time. See
+        [Folding the delta in](csr.md#folding-the-delta-in).
+        """
+        self._inner.optimize_csr()
+
     def node_index_stats(self) -> dict[str, NodeIndexStats]:
         """Statistics for each node index (experimental; needs `use_csr=True`).
 
@@ -289,8 +309,10 @@ def csr_stats(self) -> dict[str, CsrStats]:
               full build, held outside the base CSR.
             - `unindexed_edges`: live edges outside the base CSR, appended since its
               last full build and not since removed.
-            - `update`: how the last refresh produced it: `"build"` (first), `"rebuild"`
-              (full), `"delta"` (changes applied) or `"reused"` (no edge change).
+            - `update`: how the last refresh or
+              [`optimize_csr`][lance_csr.CypherEngine.optimize_csr] produced it:
+              `"build"` (first), `"rebuild"` (full), `"delta"` (changes applied) or
+              `"reused"` (no edge change).
             - `version`: the relationship dataset version it reflects (`None` for
               PyArrow).
             - `base_version`: the relationship dataset version of the last full

python/python/tests/test_engine_csr_delta.py

@@ -166,6 +166,45 @@ def test_sequential_deltas_accumulate(graph):
     check(engine, uri)
 
 
+def test_optimize_folds_delta_into_base(graph):
+    engine, uri = graph()
+    lance.write_dataset(edges([(7, 8), (7, 9)], 2000), uri, mode="append")
+    lance.dataset(uri).delete("since = 2000 OR since = 0")
+    engine.refresh()
+    before = stats(engine)
+    assert before["update"] == "delta" and before["unindexed_edges"] == 1
+
+    engine.optimize_csr()
+    s = stats(engine)
+    assert (s["update"], s["delta_edges"], s["tombstones"]) == ("rebuild", 0, 0)
+    assert s["unindexed_edges"] == 0 and s["edges"] == before["edges"]
+    assert s["base_version"] == s["version"] == before["version"]
+    check(engine, uri)
+
+    # Nothing left to fold, so the CSR is kept as it is.
+    engine.optimize_csr()
+    assert stats(engine) == s
+
+
+def test_optimize_stays_on_the_pinned_version(graph):
+    engine, uri = graph()
+    before = rows(engine.execute(QUERIES[0]))
+    lance.write_dataset(edges([(7, 8)], 2000), uri, mode="append")
+    engine.optimize_csr()
+    s = stats(engine)
+    # A fresh build has no delta: nothing is rebuilt and the append stays unseen.
+    assert s["update"] == "build" and s["version"] == 1
+    assert rows(engine.execute(QUERIES[0])) == before
+
+
+def test_optimize_without_csr_does_nothing():
+    datasets = {"Person": PEOPLE, "KNOWS": edges(BASE)}
+    engine = CypherEngine(CONFIG, datasets)
+    engine.optimize_csr()
+    assert engine.csr_stats() == {}
+    assert len(engine.execute(QUERIES[0])) == len(BASE)
+
+
 def test_compaction_keeps_results(graph):
     engine, uri = graph()
     ds = lance.dataset(uri)