feat: distributed cleanup_old_versions

docs/README.md

@@ -44,6 +44,7 @@ The built documentation will be in the `site/` directory.
   - `write.md` - Write operations guide  
   - `distributed-indexing.md` - Distributed indexing guide
   - `compaction.md` - Distributed compaction guide
+  - `cleanup.md` - Distributed cleanup guide
   - `data-evolution.md` - Data evolution guide
   - `examples.md` - Usage examples
   - `.pages` - Navigation configuration

docs/src/.pages

@@ -5,5 +5,6 @@ nav:
   - Write: write.md
   - Index: distributed-indexing.md
   - Compaction: compaction.md
+  - Cleanup: cleanup.md
   - Data Evolution: data-evolution.md
   - Examples: examples.md

docs/src/cleanup.md

@@ -0,0 +1,87 @@
+# Distributed Cleanup
+
+Overwrites, compactions and index rebuilds leave old versions behind, together with the data, deletion and index files only those versions reference. `cleanup_old_versions` removes them. On a large table most of its time goes into listing the dataset's directories, which a single process does one page at a time, and deleting the files it finds one request at a time.
+
+Lance-Ray splits that work across Ray workers. Lance plans the cleanup on the driver: it reads every manifest, removes the old ones, and splits the listing of each directory into key ranges. Each range is swept by a worker, which lists it and removes the files there that no retained version references.
+
+Which files are removed is decided by Lance with the same rules as `lance.LanceDataset.cleanup_old_versions`. Lance-Ray only schedules the ranges.
+
+Requires a pylance release with `LanceDataset.plan_cleanup_old_versions`. With older releases the cleanup runs on the driver, and a warning is logged.
+
+## `cleanup_old_versions`
+
+```python
+cleanup_old_versions(
+    uri=None,
+    *,
+    table_id=None,
+    older_than=None,
+    retain_versions=None,
+    versions=None,
+    delete_unverified=False,
+    error_if_tagged_old_versions=True,
+    delete_rate_limit=None,
+    delete_concurrency=None,
+    num_workers=4,
+    max_tasks=None,
+    max_task_retries=2,
+    storage_options=None,
+    namespace_impl=None,
+    namespace_properties=None,
+    ray_remote_args=None,
+)
+```
+
+**Parameters:**
+
+- `uri`: Dataset URI (either `uri` OR `namespace_impl` + `table_id` required)
+- `table_id`: Table identifier as a list of strings (requires `namespace_impl`)
+- `older_than`, `retain_versions`, `versions`, `delete_unverified`, `error_if_tagged_old_versions`: As for `lance.LanceDataset.cleanup_old_versions`. Without `older_than`, `retain_versions` or `versions`, versions older than two weeks are removed.
+- `delete_rate_limit`: Approximate maximum delete requests per second for the whole cleanup
+- `delete_concurrency`: Maximum in-flight deletes per worker
+- `num_workers`: Number of Ray workers (default: 4)
+- `max_tasks`: Maximum number of key ranges (default: four per worker). Every listed directory gets at least one.
+- `max_task_retries`: Retries per failed range (default: 2)
+- `storage_options`: Optional storage configuration dictionary
+- `namespace_impl`: Namespace implementation type (e.g., `"rest"`, `"dir"`)
+- `namespace_properties`: Properties for connecting to the namespace
+- `ray_remote_args`: Optional kwargs for the Ray workers
+
+**Returns:** `CleanupStats` for the whole cleanup.
+
+The driver logs a line as each range finishes, with the running number of files and bytes removed.
+
+## `explain_cleanup_old_versions`
+
+Takes the same arguments, without the delete limits, plus `max_files` (default: 1000), and returns the `CleanupExplanation` of what `cleanup_old_versions` would remove. Nothing is deleted.
+
+## Safety
+
+- A file is removed only if no retained version references it. Before listing its range, each worker reads the versions committed since the plan and keeps their files.
+- A file that no version references at all, for example one left by a failed write, is removed only once it is seven days old, counted from when the plan was created. `delete_unverified=True` removes such files regardless of age. Only set it when nothing else writes to the dataset until the cleanup ends, or the files of an in-progress write can be removed.
+- The old versions are removed while planning, before any file is swept, so none of them can be restored, tagged or branched from while its files are being removed. If a tag or branch is created on one of them while the plan is being made, the cleanup fails before removing anything.
+- If a range still fails after its retries, a `RuntimeError` is raised. The old versions are already gone, and no version references the files the range left, so later cleanups only remove them as unverified files.
+- Every worker must run the same pylance version as the driver.
+- Cleaning referenced branches is not supported.
+
+## Limits
+
+- `delete_rate_limit` is divided evenly between the workers, at least one request per second each, so the total is approximate.
+- On a store that does not list keys in lexical order (for example S3 Express One Zone), a range would have to list its whole directory, so each directory gets a single range.
+- Reading the manifests and removing the old ones still happens on the driver.
+
+## Example
+
+```python
+from datetime import timedelta
+
+import lance_ray as lr
+
+stats = lr.cleanup_old_versions(
+    "s3://bucket/path/to/table.lance",
+    older_than=timedelta(days=7),
+    num_workers=16,
+    delete_rate_limit=2000,
+)
+print(stats.old_versions, stats.bytes_removed)
+```

lance_ray/__init__.py

@@ -8,6 +8,7 @@
 __version__ = "0.5.0"
 __author__ = "LanceDB Devs"
 __email__ = "dev@lancedb.com"
+from .cleanup import cleanup_old_versions, explain_cleanup_old_versions
 from .compaction import compact_database, compact_files
 
 # Main imports
@@ -42,6 +43,8 @@
     "optimize_indices",
     "compact_files",
     "compact_database",
+    "cleanup_old_versions",
+    "explain_cleanup_old_versions",
     "LanceFragmentWriter",
     "LanceFragmentCommitter",
 ]

lance_ray/_dispatch.py

@@ -0,0 +1,122 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright The Lance Authors
+
+"""Run independent, idempotent work items on a set of Ray actors."""
+
+from __future__ import annotations
+
+import logging
+import time
+from collections import deque
+from collections.abc import Callable
+from typing import Any, Optional
+
+import ray
+from ray.exceptions import ActorUnavailableError
+
+logger = logging.getLogger(__name__)
+
+Handler = Callable[[Any], dict[str, Any]]
+
+
+class _BatchWorker:
+    """Ray actor that runs one handler call at a time."""
+
+    def run(self, handler: Handler, item: Any) -> dict[str, Any]:
+        return handler(item)
+
+
+def _map_async_with_pool(
+    create_handler: Callable[[], Handler],
+    items: list[Any],
+    *,
+    num_workers: int,
+    ray_remote_args: Optional[dict[str, Any]],
+    error_prefix: str,
+    max_task_retries: int = 0,
+    on_result: Optional[Callable[[int, dict[str, Any]], None]] = None,
+) -> list[dict[str, Any]]:
+    """Run items on a pool of Ray actors, one item per idle worker.
+
+    Each item is handed to whichever worker is idle at that moment, so a slow
+    node only delays the items it is actually running. (Ray's
+    ``multiprocessing.Pool`` assigns items round-robin at submission time,
+    which makes a map of N items wait for the slowest worker's fixed share
+    of N.) Items whose handler returned ``{"status": "error"}`` or whose
+    Ray task/actor failed are resubmitted immediately, up to
+    ``max_task_retries`` times each, rather than after every other item has
+    finished; handlers must therefore be idempotent.
+
+    ``on_result(index, result)`` is called once per item, with its final result.
+
+    Workers are actors created with ``ray_remote_args`` (``num_cpus=0`` unless
+    given, like Ray's Pool). A worker whose process dies is dropped from the
+    pool unless ``{"max_restarts": ...}`` is set, in which case Ray restarts
+    it and it keeps receiving items.
+    """
+    remote_args = {"num_cpus": 0, **(ray_remote_args or {})}
+    worker_cls: Any = ray.remote(**remote_args)(_BatchWorker)
+    workers: list[Any] = [worker_cls.remote() for _ in range(num_workers)]
+    try:
+        handler = create_handler()
+        results: dict[int, dict[str, Any]] = {}
+        pending: deque[tuple[int, int]] = deque(
+            (index, 0) for index in range(len(items))
+        )
+        idle: list[Any] = []
+        # A worker joins the idle list once Ray has actually started it, so an
+        # item never queues on an actor still waiting for cluster resources.
+        starting: dict[Any, Any] = {w.__ray_ready__.remote(): w for w in workers}
+        in_flight: dict[Any, tuple[Any, int, int]] = {}
+        while pending or in_flight:
+            while pending and idle:
+                index, attempt = pending.popleft()
+                worker = idle.pop()
+                ref = worker.run.remote(handler, items[index])
+                in_flight[ref] = (worker, index, attempt)
+            if not in_flight and not starting:
+                raise RuntimeError("every worker has died")
+            [ref], _ = ray.wait(list(in_flight) + list(starting), num_returns=1)
+            if ref in starting:
+                worker = starting.pop(ref)
+                try:
+                    ray.get(ref)
+                    idle.append(worker)
+                except ActorUnavailableError:
+                    # Restarting under max_restarts: probe again shortly.
+                    time.sleep(0.5)
+                    starting[worker.__ray_ready__.remote()] = worker
+                except Exception as exc:
+                    logger.warning(
+                        "Worker unavailable (%s); dropping it. Pass max_restarts "
+                        "in ray_remote_args to restart dead workers",
+                        exc,
+                    )
+                continue
+            worker, index, attempt = in_flight.pop(ref)
+            try:
+                result = ray.get(ref)
+                idle.append(worker)
+            except Exception as exc:
+                result = {"status": "error", "error": f"{type(exc).__name__}: {exc}"}
+                # A live worker is ready again at once; a dead one is dropped
+                # (or, with max_restarts, restarted) by the start path above.
+                starting[worker.__ray_ready__.remote()] = worker
+            results[index] = result
+            if result.get("status") != "success" and attempt < max_task_retries:
+                logger.warning(
+                    "Batch %d failed (attempt %d of %d): %s; resubmitting",
+                    index,
+                    attempt + 1,
+                    max_task_retries + 1,
+                    result.get("error"),
+                )
+                pending.append((index, attempt + 1))
+            elif on_result is not None:
+                on_result(index, result)
+        return [results[i] for i in range(len(items))]
+    except Exception as exc:
+        raise RuntimeError(f"{error_prefix}: {exc}") from exc
+    finally:
+        for worker in workers:
+            ray.kill(worker)