perf: reuse a pinned dataset as the write_fragments target

daft_lance/lance_data_sink.py

@@ -2,6 +2,7 @@
 
 import logging
 import pathlib
+import threading
 import uuid
 import warnings
 from itertools import chain
@@ -54,6 +55,24 @@
 # storage-version compatibility check entirely).
 LancePhysicalWriteMode = Literal["create", "append", "overwrite"]
 
+# Lance resolves a URI write destination by loading the latest manifest on every
+# write_fragments call, which costs O(fragments) per batch. Workers instead reuse
+# one dataset handle per process, pinned to the version start() read on the driver.
+# One entry per URI keeps a long-lived worker from holding every old manifest.
+_WRITE_TARGETS: dict[str, tuple[tuple[Any, ...], lance.LanceDataset]] = {}
+_WRITE_TARGETS_LOCK = threading.Lock()
+
+
+def _cached_write_target(uri: str, version: int, storage_options: dict[str, str] | None) -> lance.LanceDataset:
+    key = (version, tuple(sorted((storage_options or {}).items())))
+    with _WRITE_TARGETS_LOCK:
+        cached = _WRITE_TARGETS.get(uri)
+        if cached is not None and cached[0] == key:
+            return cached[1]
+        dataset = lance.dataset(uri, version=version, storage_options=storage_options)
+        _WRITE_TARGETS[uri] = (key, dataset)
+        return dataset
+
 
 def _dataset_stats(dataset: lance.LanceDataset) -> MicroPartition:
     """The single-row write result: dataset stats plus the version just produced."""
@@ -424,12 +443,19 @@ def _prepare_arrow_table(self, input_table: pa.Table) -> pa.Table:
             return input_table.cast(target_schema)
         return pa.Table.from_batches(input_table.to_batches(), target_schema)
 
-    def _write_arrow_table(self, table: pa.Table) -> WriteResult[list[FragmentMetadata]]:
+    def _write_target(self) -> str | lance.LanceDataset:
         assert self._table_uri is not None, "LanceDataSink.start() must run before writes"
+        # create targets a missing dataset, where Lance's per-call lookup is a cheap NotFound.
+        # Namespace tables keep the URI path so the namespace client stays in charge of access.
+        if self._table_schema is None or self._dataset_uri_arg is None:
+            return self._table_uri
+        return _cached_write_target(self._table_uri, self._version, self._storage_options)
+
+    def _write_arrow_table(self, table: pa.Table) -> WriteResult[list[FragmentMetadata]]:
         wrapped = self._blob.wrap_table(table)
         fragments = lance.fragment.write_fragments(
             wrapped,
-            dataset_uri=self._table_uri,
+            dataset_uri=self._write_target(),
             mode=self._write_mode,
             storage_options=self._storage_options,
             max_rows_per_file=self._max_rows_per_file,

tests/io/lance/test_lance_write_target_cache.py

@@ -0,0 +1,114 @@
+from __future__ import annotations
+
+import os
+from typing import TYPE_CHECKING, Any
+
+import lance
+import pytest
+
+import daft
+from daft_lance import lance_data_sink
+
+if TYPE_CHECKING:
+    import pathlib
+    from collections.abc import Iterator
+
+pytestmark = pytest.mark.skipif(
+    os.environ.get("DAFT_RUNNER") == "ray",
+    reason="the spies patch the driver process, while Ray writes from worker processes",
+)
+
+
+@pytest.fixture(autouse=True)
+def clear_write_targets() -> Iterator[None]:
+    lance_data_sink._WRITE_TARGETS.clear()
+    yield
+    lance_data_sink._WRITE_TARGETS.clear()
+
+
+@pytest.fixture
+def write_targets(monkeypatch: pytest.MonkeyPatch) -> list[Any]:
+    """Record the destination of every lance.fragment.write_fragments call."""
+    targets: list[Any] = []
+    original = lance.fragment.write_fragments
+
+    def spy(data: Any, dataset_uri: Any, *args: Any, **kwargs: Any) -> Any:
+        targets.append(dataset_uri)
+        return original(data, dataset_uri, *args, **kwargs)
+
+    monkeypatch.setattr(lance.fragment, "write_fragments", spy)
+    return targets
+
+
+@pytest.fixture
+def pinned_opens(monkeypatch: pytest.MonkeyPatch) -> list[int]:
+    """Count lance.dataset calls that pin a version, i.e. write-target opens."""
+    opens: list[int] = []
+    original = lance.dataset
+
+    def spy(*args: Any, **kwargs: Any) -> Any:
+        if kwargs.get("version") is not None:
+            opens.append(kwargs["version"])
+        return original(*args, **kwargs)
+
+    monkeypatch.setattr(lance, "dataset", spy)
+    return opens
+
+
+def _ids(path: str) -> list[int]:
+    return sorted(lance.dataset(path).to_table(columns=["id"]).to_pydict()["id"])
+
+
+def test_append_reuses_one_pinned_dataset_across_writes(
+    tmp_path: pathlib.Path, write_targets: list[Any], pinned_opens: list[int]
+) -> None:
+    path = str(tmp_path / "t.lance")
+    daft.from_pydict({"id": [0, 1]}).write_lance(path, mode="create")
+    version_before = lance.dataset(path).version
+    write_targets.clear()
+
+    df = daft.from_pydict({"id": list(range(2, 12))}).into_batches(2)
+    df.write_lance(path, mode="append", max_rows_per_file=2)
+
+    assert len(write_targets) >= 2
+    assert all(isinstance(t, lance.LanceDataset) for t in write_targets)
+    assert {t.version for t in write_targets} == {version_before}
+    assert pinned_opens == [version_before]
+    assert _ids(path) == list(range(12))
+
+
+def test_sequential_appends_refresh_pinned_version(tmp_path: pathlib.Path, write_targets: list[Any]) -> None:
+    path = str(tmp_path / "t.lance")
+    daft.from_pydict({"id": [0]}).write_lance(path, mode="create")
+    v1 = lance.dataset(path).version
+
+    daft.from_pydict({"id": [1]}).write_lance(path, mode="append")
+    v2 = lance.dataset(path).version
+    daft.from_pydict({"id": [2]}).write_lance(path, mode="append")
+
+    assert [t.version for t in write_targets[1:]] == [v1, v2]
+    assert len(lance_data_sink._WRITE_TARGETS) == 1
+    assert _ids(path) == [0, 1, 2]
+
+
+def test_create_writes_to_uri(tmp_path: pathlib.Path, write_targets: list[Any], pinned_opens: list[int]) -> None:
+    path = str(tmp_path / "t.lance")
+    daft.from_pydict({"id": list(range(6))}).into_batches(2).write_lance(path, mode="create", max_rows_per_file=2)
+
+    assert write_targets
+    assert all(isinstance(t, str) for t in write_targets)
+    assert pinned_opens == []
+    assert _ids(path) == list(range(6))
+
+
+def test_overwrite_existing_uses_pinned_dataset(tmp_path: pathlib.Path, write_targets: list[Any]) -> None:
+    path = str(tmp_path / "t.lance")
+    daft.from_pydict({"id": [0, 1, 2]}).write_lance(path, mode="create")
+    version_before = lance.dataset(path).version
+    write_targets.clear()
+
+    daft.from_pydict({"id": [7, 8]}).write_lance(path, mode="overwrite")
+
+    assert write_targets
+    assert all(isinstance(t, lance.LanceDataset) and t.version == version_before for t in write_targets)
+    assert _ids(path) == [7, 8]