Merge remote-tracking branch 'origin/main' into gatekeeper/fix-2923-1

python/src/table.rs

@@ -28,11 +28,72 @@ use pyo3::{
     Bound, FromPyObject, Py, PyAny, PyRef, PyResult, Python,
     exceptions::{PyRuntimeError, PyValueError},
     pyclass, pyfunction, pymethods,
-    types::{IntoPyDict, PyAnyMethods, PyBytes, PyDict, PyDictMethods},
+    types::{IntoPyDict, PyAnyMethods, PyBytes, PyDict, PyDictMethods, PyList, PyListMethods},
 };
 
 mod scannable;
 
+/// Convert `LsmStats` to a Python dict, preserving the per-bucket list.
+///
+/// Deliberately not flattened to a table-level summary: a table is N
+/// buckets on one node, and the per-bucket detail is the reason the
+/// endpoint exists — flattening hides the single hot bucket someone opened
+/// it to find.
+fn lsm_stats_to_py(py: Python<'_>, stats: &lancedb::table::LsmStats) -> PyResult<Py<PyDict>> {
+    let out = PyDict::new(py);
+    let buckets = PyList::empty(py);
+    for b in &stats.buckets {
+        let e = PyDict::new(py);
+        e.set_item("shard_id", &b.shard_id)?;
+        e.set_item("status", &b.status)?;
+        e.set_item("writer_epoch", b.writer_epoch)?;
+        e.set_item("manifest_version", b.manifest_version)?;
+        e.set_item("current_generation", b.current_generation)?;
+        e.set_item(
+            "replay_after_wal_entry_position",
+            b.replay_after_wal_entry_position,
+        )?;
+        e.set_item(
+            "wal_entry_position_last_seen",
+            b.wal_entry_position_last_seen,
+        )?;
+
+        let generations = PyList::empty(py);
+        for g in &b.generations {
+            let ge = PyDict::new(py);
+            ge.set_item("generation", g.generation)?;
+            ge.set_item("bytes", g.bytes)?;
+            ge.set_item("rows", g.rows)?;
+            generations.append(ge)?;
+        }
+        e.set_item("generations", generations)?;
+        e.set_item("compacting", b.compacting)?;
+
+        e.set_item(
+            "memtables",
+            b.memtables
+                .as_ref()
+                .map(|ms| {
+                    let l = PyList::empty(py);
+                    for m in ms {
+                        let d = PyDict::new(py);
+                        d.set_item("generation", m.generation)?;
+                        d.set_item("rows", m.rows)?;
+                        d.set_item("bytes", m.bytes)?;
+                        d.set_item("batches", m.batches)?;
+                        d.set_item("indexes", m.indexes.clone())?;
+                        l.append(d)?;
+                    }
+                    PyResult::Ok(l.unbind())
+                })
+                .transpose()?,
+        )?;
+        buckets.append(e)?;
+    }
+    out.set_item("buckets", buckets)?;
+    Ok(out.unbind())
+}
+
 #[derive(FromPyObject)]
 enum PredicateArg {
     Expr(PyExpr),
@@ -185,12 +246,22 @@ impl From<lancedb::table::MergeResult> for MergeResult {
     }
 }
 
+/// Render for `__repr__`, so the default reads as Python's `None` rather than
+/// Rust's `Some([..])`.
+fn fmt_maintained(maintained: &Option<Vec<String>>) -> String {
+    match maintained {
+        Some(names) => format!("{:?}", names),
+        None => "None".to_string(),
+    }
+}
+
 /// Specification selecting Lance's MemWAL LSM-style write path for
 /// `merge_insert`.
 ///
 /// Constructed via the `bucket(...)`, `identity(...)`, or `unsharded()`
 /// classmethods, then optionally chain `with_maintained_indexes(...)` and
-/// `with_writer_config_defaults(...)`.
+/// `with_writer_config_defaults(...)`. A fresh spec maintains every index the
+/// MemWAL supports, resolved on install.
 #[pyclass(from_py_object)]
 #[derive(Clone, Debug)]
 pub struct LsmWriteSpec {
@@ -230,11 +301,11 @@ impl LsmWriteSpec {
         }
     }
 
-    /// Replace the list of indexes the MemWAL should keep up to date as
-    /// rows are appended. Each name must reference an index that
-    /// already exists on the table at the time `set_lsm_write_spec`
-    /// is called.
-    pub fn with_maintained_indexes(&self, indexes: Vec<String>) -> Self {
+    /// Set which indexes the MemWAL maintains. `None` (the default)
+    /// resolves every supported index on install; a list is verbatim,
+    /// and an empty list maintains nothing.
+    #[pyo3(signature = (indexes))]
+    pub fn with_maintained_indexes(&self, indexes: Option<Vec<String>>) -> Self {
         Self {
             inner: self.inner.clone().with_maintained_indexes(indexes),
         }
@@ -256,23 +327,29 @@ impl LsmWriteSpec {
                 maintained_indexes,
                 writer_config_defaults,
             } => format!(
-                "LsmWriteSpec.bucket(column={:?}, num_buckets={}, maintained_indexes={:?}, writer_config_defaults={:?})",
-                column, num_buckets, maintained_indexes, writer_config_defaults,
+                "LsmWriteSpec.bucket(column={:?}, num_buckets={}, maintained_indexes={}, writer_config_defaults={:?})",
+                column,
+                num_buckets,
+                fmt_maintained(maintained_indexes),
+                writer_config_defaults,
             ),
             lancedb::table::LsmWriteSpec::Identity {
                 column,
                 maintained_indexes,
                 writer_config_defaults,
             } => format!(
-                "LsmWriteSpec.identity(column={:?}, maintained_indexes={:?}, writer_config_defaults={:?})",
-                column, maintained_indexes, writer_config_defaults,
+                "LsmWriteSpec.identity(column={:?}, maintained_indexes={}, writer_config_defaults={:?})",
+                column,
+                fmt_maintained(maintained_indexes),
+                writer_config_defaults,
             ),
             lancedb::table::LsmWriteSpec::Unsharded {
                 maintained_indexes,
                 writer_config_defaults,
             } => format!(
-                "LsmWriteSpec.unsharded(maintained_indexes={:?}, writer_config_defaults={:?})",
-                maintained_indexes, writer_config_defaults,
+                "LsmWriteSpec.unsharded(maintained_indexes={}, writer_config_defaults={:?})",
+                fmt_maintained(maintained_indexes),
+                writer_config_defaults,
             ),
         }
     }
@@ -307,10 +384,10 @@ impl LsmWriteSpec {
         }
     }
 
-    /// Names of indexes the MemWAL should keep up to date during writes.
+    /// Indexes the MemWAL keeps up to date, or `None` for every supported one.
     #[getter]
-    pub fn maintained_indexes(&self) -> Vec<String> {
-        self.inner.maintained_indexes().to_vec()
+    pub fn maintained_indexes(&self) -> Option<Vec<String>> {
+        self.inner.maintained_indexes().map(<[String]>::to_vec)
     }
 
     /// Default `ShardWriter` configuration recorded by this spec.
@@ -745,6 +822,9 @@ impl Table {
 
     #[allow(private_interfaces)]
     pub fn delete(self_: PyRef<'_, Self>, condition: PredicateArg) -> PyResult<Bound<'_, PyAny>> {
+        // Do not hold the Python borrow across the await.  The cloned Rust table
+        // handle is thread-safe and allows deletes on the same Python table to
+        // run concurrently without PyO3 reporting "Already borrowed".
         let inner = self_.inner_ref()?.clone();
         future_into_py(self_.py(), async move {
             let result = match &condition {
@@ -1336,6 +1416,51 @@ impl Table {
         })
     }
 
+    /// Converge the table's LSM write path into its base table.
+    ///
+    /// Best-effort: with writes flowing, new rows may land after the last
+    /// pass. Errors if the table stops making progress.
+    pub fn checkpoint_lsm(self_: PyRef<'_, Self>) -> PyResult<Bound<'_, PyAny>> {
+        let inner = self_.inner_ref()?.clone();
+        future_into_py(self_.py(), async move {
+            inner.checkpoint_lsm().await.infer_error()
+        })
+    }
+
+    /// Seal every bucket's active memtable into L0.
+    pub fn flush_lsm(self_: PyRef<'_, Self>) -> PyResult<Bound<'_, PyAny>> {
+        let inner = self_.inner_ref()?.clone();
+        future_into_py(
+            self_.py(),
+            async move { inner.flush_lsm().await.infer_error() },
+        )
+    }
+
+    /// Trigger a background L0 → base pass per bucket. Returns once the
+    /// passes are dispatched, not once they finish — watch `get_lsm_stats`.
+    pub fn compact_lsm(self_: PyRef<'_, Self>) -> PyResult<Bound<'_, PyAny>> {
+        let inner = self_.inner_ref()?.clone();
+        future_into_py(self_.py(), async move {
+            inner.compact_lsm().await.infer_error()
+        })
+    }
+
+    /// Live LSM state, or `None` when the LSM write path is not enabled.
+    #[pyo3(signature = (include_generation_rows=false))]
+    pub fn get_lsm_stats(
+        self_: PyRef<'_, Self>,
+        include_generation_rows: bool,
+    ) -> PyResult<Bound<'_, PyAny>> {
+        let inner = self_.inner_ref()?.clone();
+        future_into_py(self_.py(), async move {
+            let stats = inner
+                .get_lsm_stats(include_generation_rows)
+                .await
+                .infer_error()?;
+            Python::attach(|py| stats.map(|s| lsm_stats_to_py(py, &s)).transpose())
+        })
+    }
+
     pub fn close_lsm_writers(self_: PyRef<'_, Self>) -> PyResult<Bound<'_, PyAny>> {
         let inner = self_.inner_ref()?.clone();
         future_into_py(self_.py(), async move {