From 0c9b5c8244dfda63a1aa3ad9f08d4fcd8daa3dde Mon Sep 17 00:00:00 2001 From: ChaosJu Date: Mon, 9 Nov 2026 21:10:30 +0800 Subject: [PATCH] perf: reuse a pinned dataset as the write_fragments target Passing the table URI makes Lance load the latest manifest on every write_fragments call, which costs O(fragments) per batch on appends to large tables. Workers now open the dataset once per process at the version start() pinned and pass it as the write destination; create and namespace writes keep the URI. --- daft_lance/lance_data_sink.py | 30 ++++- .../io/lance/test_lance_write_target_cache.py | 114 ++++++++++++++++++ 2 files changed, 142 insertions(+), 2 deletions(-) create mode 100644 tests/io/lance/test_lance_write_target_cache.py diff --git a/daft_lance/lance_data_sink.py b/daft_lance/lance_data_sink.py index c67cbc8..07897cb 100644 --- a/daft_lance/lance_data_sink.py +++ b/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, diff --git a/tests/io/lance/test_lance_write_target_cache.py b/tests/io/lance/test_lance_write_target_cache.py new file mode 100644 index 0000000..67e4b3e --- /dev/null +++ b/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]