fix(python): added cleanup_frag_reuse_index

python/python/lance/dataset.py

@@ -6283,10 +6283,13 @@ def batch_readahead(self, nbatches: Optional[int] = None) -> ScannerBuilder:
         """
         Set the maximum number of batches to decode concurrently.
 
-        This parameter must be greater than zero.
+        This parameter must be a positive integer.
         """
-        if nbatches is not None and int(nbatches) <= 0:
-            raise ValueError("batch_readahead must be greater than 0")
+        if nbatches is not None:
+            if not isinstance(nbatches, int) or isinstance(nbatches, bool):
+                raise ValueError("batch_readahead must be an integer")
+            if nbatches <= 0:
+                raise ValueError("batch_readahead must be greater than 0")
         self._batch_readahead = nbatches
         return self
 
@@ -7062,6 +7065,28 @@ def optimize_indices(self, **kwargs):
         """
         self._dataset._ds.optimize_indices(**kwargs)
 
+    def cleanup_frag_reuse_index(self) -> None:
+        """Prune obsolete generations from the ``__lance_frag_reuse`` system index.
+
+        The fragment-reuse index is created by :meth:`compact_files` when
+        ``defer_index_remap=True``. It retains one generation per deferred
+        compaction so indices can be remapped lazily. Once every index has caught
+        up to a generation (for example after rebuilding it or calling
+        :meth:`optimize_indices`), that generation is obsolete and can be pruned.
+
+        This is safe to call at any time: a generation that an index has not yet
+        caught up to is always retained.
+
+        Examples
+        --------
+        >>> import lance
+        >>> import pyarrow as pa
+        >>> data = pa.table({"id": [1, 2, 3]})
+        >>> dataset = lance.write_dataset(data, "memory://frag_reuse_example")
+        >>> dataset.optimize.cleanup_frag_reuse_index()
+        """
+        self._dataset._ds.cleanup_frag_reuse_index()
+
     def enable_auto_cleanup(self, auto_cleanup_config: AutoCleanupConfig, **kwargs):
         """Enable autocleaning for an existing dataset.
 

python/python/lance/lance/__init__.pyi

@@ -555,6 +555,7 @@ class _Dataset:
     ) -> Tuple[_Dataset, Transaction]: ...
     def validate(self): ...
     def migrate_manifest_paths_v2(self): ...
+    def cleanup_frag_reuse_index(self) -> None: ...
     def drop_columns(self, columns: List[str]): ...
     def add_columns_from_reader(
         self, reader: pa.RecordBatchReader, batch_size: Optional[int] = None

python/python/tests/test_dataset.py

@@ -1166,6 +1166,16 @@ def test_to_batches(tmp_path: Path):
     assert sorted_batches == table
 
 
+@pytest.mark.parametrize("batch_readahead", [1.5, "1", True])
+def test_batch_readahead_requires_integer(tmp_path: Path, batch_readahead):
+    table = pa.Table.from_pydict({"a": range(10)})
+    base_dir = tmp_path / "test"
+    lance.write_dataset(table, base_dir)
+
+    with pytest.raises(ValueError, match="batch_readahead must be an integer"):
+        lance.dataset(base_dir).scanner(batch_readahead=batch_readahead)
+
+
 def test_list_from_parquet(tmp_path: Path):
     # This is a regression for GH-1482, the parquet reader creates
     # list fields with the name 'element' instead of 'item'.  We should

python/python/tests/test_optimize.py

@@ -324,6 +324,62 @@ def test_defer_index_remap(tmp_path: Path):
     assert any(idx.name == "__lance_frag_reuse" for idx in indices)
 
 
+def test_cleanup_frag_reuse_index(tmp_path: Path):
+    """cleanup_frag_reuse_index prunes only the reuse generations that every user
+    index has caught up to, and never the ones still in use.
+
+    Setup: 6 fragments, one BTREE scalar index.  Compact with
+    defer_index_remap=True so the frag-reuse index is populated.  While the
+    scalar index has not caught up, cleanup must retain the generation.  After
+    rebuilding the scalar index (create_scalar_index with replace=True) it is
+    newer than the reuse generation, so cleanup may prune it and num_versions
+    drops to 0.
+    """
+    base_dir = tmp_path / "dataset"
+    data = pa.table({"i": range(6_000), "val": range(6_000)})
+    dataset = lance.write_dataset(data, base_dir, max_rows_per_file=1_000)
+    dataset.create_scalar_index("i", "BTREE")
+
+    dataset.delete("i < 500")
+    dataset.optimize.compact_files(
+        target_rows_per_fragment=2_000, defer_index_remap=True, num_threads=1
+    )
+
+    dataset = lance.dataset(base_dir)
+    assert any(
+        idx.name == "__lance_frag_reuse" for idx in dataset.describe_indices()
+    ), "precondition: defer_index_remap must have created the frag-reuse index"
+
+    before_stats = dataset.stats.index_stats("__lance_frag_reuse")
+    versions_before = before_stats["num_versions"]
+    assert versions_before >= 1, (
+        "precondition: frag-reuse index must have at least one version before cleanup"
+    )
+
+    # Negative case: index not yet remapped, so cleanup must retain the generation.
+    dataset.optimize.cleanup_frag_reuse_index()
+    dataset = lance.dataset(base_dir)
+    retained_stats = dataset.stats.index_stats("__lance_frag_reuse")
+    assert retained_stats["num_versions"] == versions_before, (
+        f"cleanup_frag_reuse_index must retain generations that an index has not "
+        f"caught up to, but num_versions went from {versions_before} to "
+        f"{retained_stats['num_versions']}"
+    )
+
+    # Positive case: rebuilding catches the index up, so cleanup can now prune all.
+    dataset.create_scalar_index("i", "BTREE", replace=True)
+    dataset = lance.dataset(base_dir)
+
+    dataset.optimize.cleanup_frag_reuse_index()
+
+    dataset = lance.dataset(base_dir)
+    after_stats = dataset.stats.index_stats("__lance_frag_reuse")
+    assert after_stats["num_versions"] == 0, (
+        f"cleanup_frag_reuse_index should have pruned all reuse generations "
+        f"but num_versions={after_stats['num_versions']}"
+    )
+
+
 @pytest.mark.parametrize("use_commit_options", [True, False])
 def test_defer_index_remap_via_commit_options(tmp_path: Path, use_commit_options: bool):
     """Compaction.commit respects defer_index_remap passed in options.

python/src/dataset.rs

@@ -2995,6 +2995,17 @@ impl Dataset {
         Ok(())
     }
 
+    fn cleanup_frag_reuse_index(&mut self) -> PyResult<()> {
+        let mut new_self = self.ds.as_ref().clone();
+        rt().block_on(
+            None,
+            lance::dataset::index::frag_reuse::cleanup_frag_reuse_index(&mut new_self),
+        )?
+        .infer_error()?;
+        self.ds = Arc::new(new_self);
+        Ok(())
+    }
+
     fn drop_columns(&mut self, columns: Vec<String>) -> PyResult<()> {
         let mut new_self = self.ds.as_ref().clone();
         let columns: Vec<_> = columns.iter().map(|s| s.as_str()).collect();