diff --git a/Makefile b/Makefile index 93e0953..5b065e2 100644 --- a/Makefile +++ b/Makefile @@ -18,7 +18,14 @@ build: sync ## Build sdist and wheel packages .PHONY: test test: sync ## Run tests - DAFT_ANALYTICS_ENABLED=0 uv run pytest tests/ -v $(EXTRA_ARGS) + # RAY_ENABLE_UV_RUN_RUNTIME_ENV=0: Ray walks the driver's process tree, sees + # the `uv run` below, and then replicates that environment to every worker -- + # packaging the whole working directory and rebuilding the venv per worker. + # Under DAFT_RUNNER=ray that makes worker startup fail once the suite has + # created enough actors: 277 tests died with ActorUnschedulableError, and the + # same suite passes in 46s with it off. Workers inherit this venv's + # interpreter from the raylet regardless. + DAFT_ANALYTICS_ENABLED=0 RAY_ENABLE_UV_RUN_RUNTIME_ENV=0 uv run pytest tests/ -v $(EXTRA_ARGS) .PHONY: format format: sync ## Format Python code diff --git a/README.md b/README.md index a5ff276..44428bf 100644 --- a/README.md +++ b/README.md @@ -414,6 +414,108 @@ Native namespace support in `daft.read_lance` / `DataFrame.write_lance` is track [Eventual-Inc/Daft#7282](https://github.com/Eventual-Inc/Daft/issues/7282); until that lands, use the `daft_lance` entry points shown above for namespace-addressed tables. +## Environment Variables + +| Variable | Default | Effect | +| --- | --- | --- | +| `DAFT_LANCE_NAMESPACE_CACHE_SIZE` | `16` | Size of the per-(implementation, properties) namespace client cache. Read once at import time. | +| `DAFT_LANCE_SHARE_SESSION` | `1` | Every dataset this extension opens in a process shares one lance `Session`, so a table's manifest and lance's caches are held once instead of once per scan task, maintenance UDF and writer. Set to `0` to give each open its own session (lance's own default); on a many-fragment table that costs one copy of the fragment list per concurrently open dataset. | + +### What sharing a session actually saves, per path + +The performance numbers below are from the original PR measurements; the +implementation update did not rerun the large-table benchmark. + +Measured, because the three paths differ and only one of them multiplies by +concurrency. The numbers come from a `read_lance` -> `write_parquet` pipeline in +a Ray worker, which is the shape this matters most in: a task holds its parquet +writer's buffers at the same time as its manifest, so what sharing frees becomes +headroom the writer can use. + +| Path | Opens per worker process | What sharing changes | +| --- | --- | --- | +| Scans, `read_blobs`, writes | `scantask_max_parallel` concurrent opens in one process | Collapses them to one manifest. Measured: 49 MB less resident on a 4000-fragment table at concurrency 24, and 479 MB on a 100k-fragment one at concurrency 8 -- where the pipeline also finished **9.8x sooner** (610s to 62s), because each unshared task re-parsed the same 8 MB manifest. At that scale the cost of not sharing is mostly time, not memory. | +| Indexed (`nearest`, point lookups) | One index-driven task, whatever the concurrency | The index is loaded once per process instead of once per query. Measured at 200k rows: repeated queries 1.7-2.0x faster. | +| Maintenance (`compact_files`, `update_columns`, `merge_columns`) | One, always | Nothing to collapse: Daft gives each `daft.cls` instance a process of its own, so a maintenance worker opens the table exactly once however high the concurrency. | + +An index does not change the maintenance row. A maintenance task's open brings +in the index's metadata entry and nothing else; the index data is read once, on +the driver, when the commit remaps every entry onto the new fragment ids. The +tests that pin both halves down are +`tests/io/lance/test_ray_session_memory.py::test_maintenance_instances_each_get_their_own_process` +and +`tests/io/lance/test_shared_session.py::test_compaction_loads_the_index_in_its_commit_not_in_its_tasks`. + +Two knobs keep a large table affordable without weakening any of that, because +a manifest is the table's whole fragment list and an open costs the same +whatever its task goes on to read: `read_lance(..., fragment_group_size=N)` +puts several fragments in one scan task, and reading only a prefix of the table +prunes whole fragments. At 100k fragments, reading 1% through groups of ten +turns a billion-row scan into ten million rows and 100k tasks into a hundred, +while each of the concurrent opens still carries the same 47 MB fragment list. +What has to survive is the overlap: enough tasks to keep the concurrency in +flight. + +One caveat if you measure this yourself: a `Session` reports both `size_bytes()` +and `index_cache_size_bytes()`, and they are not the same accounting. Only +`size_bytes()` tracks resident memory — on a 200k-row IVF_PQ index it moves by +about 5 MB per loaded copy, matching the process's own RSS, where +`index_cache_size_bytes()` reads 37 MB. + + +### Shared Lance caches + +Daft-Lance lazily creates one automatically managed Lance Session per process, +shared by all tables for dataset opens and `write_fragments` calls. Tasks carry +serializable open arguments; the native Session is created in the worker. +Dataset and network operations run outside the Session initialization lock. +Set `DAFT_LANCE_SHARE_SESSION=0` (also `false` or `no`) to disable sharing +across all opens and fragment writes; each call then receives a fresh Session. + +Configure cache capacities in bytes in the driver and worker environments: + +```bash +export DAFT_LANCE_METADATA_CACHE_SIZE_BYTES=1073741824 +export DAFT_LANCE_INDEX_CACHE_SIZE_BYTES=268435456 +``` + +The first Session initialization fixes the process-wide capacities: explicit +read sizes take precedence over environment defaults, then Lance defaults. +Later calls without explicit sizes use these fixed capacities; conflicting +explicit sizes fail clearly instead of being silently ignored. Set the same +capacity environment variables in driver and workers before first use. Capacity +environment changes require restarting the process or calling the internal +`clear_caches()` helper. Legacy `index_cache_size` uses pylance's conversion of +20 MiB per entry. Explicit externally supplied Sessions are preserved. + +Capacities are shared across tables within a process, not multiplied per table. +Each worker has its own budget. All users of the global Session must agree that +a URI identifies the same dataset; identical URIs addressing different data on +different endpoints require separate processes or explicitly isolated Sessions. +Storage options and namespace credentials continue to be passed to each open +and write operation. + +Each operation opens its own Dataset while reusing a Session. There is no +Dataset cache or special handling for V1 manifest naming. Manifest and index +cache admission, identity, and eviction are managed by Lance itself; the intended +metadata-reuse path uses V2 manifest naming. + +`DAFT_LANCE_CACHE_TTL_SECONDS` controls idle Session retention (default 300 +seconds) and requires a positive integer. Byte capacities require nonnegative +integers. Invalid settings report the variable name and value. + +A process-local daemon sweeps the global Session every 30 seconds, including +when the worker receives no further requests. An idle Session can remain for +up to TTL plus one sweep interval. Cleanup drops the global reference without +closing handles held by active tasks. A later open recreates the Session with +the same fixed capacities; explicit `clear_caches()` also resets configuration. +Fork resets the Session, configuration, lock, and cleanup thread. Native +allocators may retain memory after references are released, so TTL is not an +RSS limit. There is no Session LRU pool or per-key Future. + +Worker reads retain their pinned version, while operations opening latest +still resolve the current version. + ## Migration The migration only requires replacing `daft.io.lance` with `daft_lance`. diff --git a/daft_lance/_blob.py b/daft_lance/_blob.py index 3ce44b8..560fc2f 100644 --- a/daft_lance/_blob.py +++ b/daft_lance/_blob.py @@ -12,6 +12,7 @@ from daft.dependencies import pa from daft.series import Series from daft_lance._metadata import LANCE_BLOB_DESCRIPTOR_TYPE +from daft_lance.namespace import open_dataset_from_open_kwargs if TYPE_CHECKING: from collections.abc import Iterable @@ -340,13 +341,26 @@ def read_blobs( if batch_size < 1: raise ValueError(f"batch_size must be at least 1, got {batch_size}") - # 2. Create a UDF closure over the lance dataset so we can read the blobs. The dataset - # pickles with its version and manifest, so workers read the same snapshot. + # 2. Create a UDF closure that reopens the pinned snapshot on the worker. + # Closing over ``ds`` would pickle it, and ``LanceDataset.__reduce__`` + # carries the whole re-encoded manifest -- one copy of the table's entire + # fragment list per task, which also bypasses the shared session's + # manifest cache. The uri, version and storage options are enough. + uri = ds.uri + open_kwargs: dict[str, Any] = {"version": ds.version, "storage_options": ds._storage_options} + worker_ds: list[lance.LanceDataset] = [] + + def _worker_dataset() -> lance.LanceDataset: + # Opened once per worker process, then reused for every batch. + if not worker_ds: + worker_ds.append(open_dataset_from_open_kwargs(uri, open_kwargs)) + return worker_ds[0] + @daft.func.batch(return_dtype=DataType.binary(), batch_size=batch_size) def read_blobs_udf(row_ids: Series) -> Series: ids = row_ids.to_pylist() valid = [i for i, row_id in enumerate(ids) if row_id is not None] - blobs = ds.read_blobs( + blobs = _worker_dataset().read_blobs( column, ids=[ids[i] for i in valid], io_buffer_size=io_buffer_size, diff --git a/daft_lance/_runtime.py b/daft_lance/_runtime.py new file mode 100644 index 0000000..4bc323e --- /dev/null +++ b/daft_lance/_runtime.py @@ -0,0 +1,189 @@ +from __future__ import annotations + +import os +from dataclasses import dataclass +from datetime import datetime +from pathlib import Path +from threading import Lock, Thread +from time import monotonic +from typing import TYPE_CHECKING, Any + +import lance + +# pylance exposes this public alias at runtime but omits its stub export. +from lance import Session as Session # type: ignore[attr-defined] +from lance_namespace import LanceNamespace + +if TYPE_CHECKING: + from lance.commit import CommitLock + +DatasetUri = str | Path | None +_session: Session | None = None +_configuration: tuple[int | None, int | None] | None = None +_deadline = 0.0 +_session_lock = Lock() +_reaper: Thread | None = None + + +class CacheConfigurationError(ValueError): + """Invalid runtime configuration; never a missing-table signal.""" + + +def _env_integer(name: str, default: int | None = None, *, minimum: int = 0) -> int | None: + raw = os.environ.get(name) + if raw is None: + return default + try: + value = int(raw) + if value < minimum: + raise ValueError + except ValueError: + raise CacheConfigurationError(f"{name}={raw!r}: expected an integer >= {minimum}") from None + return value + + +def _expire_session() -> None: + global _session + released = None + with _session_lock: + if _session is not None and _deadline <= monotonic(): + released, _session = _session, None + del released + + +def _sweep() -> None: + from time import sleep + + while True: + sleep(30) + _expire_session() + + +def clear_caches() -> None: + """Release the global Session and configuration; active operations retain it.""" + global _session, _configuration + with _session_lock: + released, _session = _session, None + _configuration = None + del released + + +def _reset_after_fork() -> None: + global _session, _configuration, _session_lock, _reaper, _deadline + _session = None + _configuration = None + _session_lock = Lock() + _reaper = None + _deadline = 0.0 + + +if hasattr(os, "register_at_fork"): + os.register_at_fork(after_in_child=_reset_after_fork) + + +def shared_session( + *, + index_cache_size: int | None = None, + index_cache_size_bytes: int | None = None, + metadata_cache_size_bytes: int | None = None, +) -> Session: + """One automatically managed Session per process, shared by all tables. + + The first call establishes process-wide capacities. Later unspecified sizes + reuse them; conflicting explicit sizes raise instead of being + silently ignored. Idle expiration keeps the configuration for the next open. + """ + global _session, _configuration, _deadline, _reaper + index_bytes = index_cache_size_bytes + if index_bytes is None and index_cache_size is not None: + index_bytes = index_cache_size * 20 * 1024 * 1024 + metadata_bytes = metadata_cache_size_bytes + for name, value in (("index_cache_size_bytes", index_bytes), ("metadata_cache_size_bytes", metadata_bytes)): + if value is not None and value < 0: + raise CacheConfigurationError(f"{name}={value!r}: expected a nonnegative integer") + if os.environ.get("DAFT_LANCE_SHARE_SESSION", "1").strip().lower() in {"0", "false", "no"}: + # Preserve the original PR's worker-side A/B baseline for every path. + if index_bytes is None: + index_bytes = _env_integer("DAFT_LANCE_INDEX_CACHE_SIZE_BYTES") + if metadata_bytes is None: + metadata_bytes = _env_integer("DAFT_LANCE_METADATA_CACHE_SIZE_BYTES") + return Session(index_cache_size_bytes=index_bytes, metadata_cache_size_bytes=metadata_bytes) + ttl = _env_integer("DAFT_LANCE_CACHE_TTL_SECONDS", 300, minimum=1) + assert ttl is not None + with _session_lock: + if _configuration is not None: + for name, requested, configured in zip( + ("index_cache_size_bytes", "metadata_cache_size_bytes"), + (index_bytes, metadata_bytes), + _configuration, + ): + if requested is not None and requested != configured: + raise CacheConfigurationError( + f"{name}={requested}: global Session already configured with {configured!r} " + "(None means Lance default); use consistent process-wide cache sizes" + ) + index_bytes, metadata_bytes = _configuration + if _configuration is None: + if index_bytes is None: + index_bytes = _env_integer("DAFT_LANCE_INDEX_CACHE_SIZE_BYTES") + if metadata_bytes is None: + metadata_bytes = _env_integer("DAFT_LANCE_METADATA_CACHE_SIZE_BYTES") + if _session is None: + # Only Session creation is locked; dataset/network operations run outside. + _session = Session(index_cache_size_bytes=index_bytes, metadata_cache_size_bytes=metadata_bytes) + _configuration = (index_bytes, metadata_bytes) + _deadline = monotonic() + ttl + if _reaper is None: + _reaper = Thread(target=_sweep, name="daft-lance-cache-cleanup", daemon=True) + _reaper.start() + return _session + + +@dataclass(frozen=True) +class DatasetOpenOptions: + """Typed pylance open parameters; live handles stay in the current process.""" + + version: int | str | None = None + asof: str | datetime | None = None + block_size: int | None = None + commit_lock: CommitLock | None = None + index_cache_size: int | None = None + storage_options: dict[str, str] | None = None + default_scan_options: dict[str, Any] | None = None + metadata_cache_size_bytes: int | None = None + index_cache_size_bytes: int | None = None + read_params: dict[str, Any] | None = None + session: Session | None = None + namespace_client: LanceNamespace | None = None + table_id: list[str] | None = None + base_store_params: dict[str, dict[str, str]] | None = None + + def resolve_session(self) -> Session: + if self.session is not None: + return self.session + return shared_session( + index_cache_size=self.index_cache_size, + index_cache_size_bytes=self.index_cache_size_bytes, + metadata_cache_size_bytes=self.metadata_cache_size_bytes, + ) + + +def open_dataset(uri: DatasetUri, options: DatasetOpenOptions | None = None) -> lance.LanceDataset: + options = options if options is not None else DatasetOpenOptions() + return lance.dataset( + uri, + version=options.version, + asof=options.asof, + block_size=options.block_size, + commit_lock=options.commit_lock, + index_cache_size=options.index_cache_size, + storage_options=options.storage_options, + default_scan_options=options.default_scan_options, + metadata_cache_size_bytes=options.metadata_cache_size_bytes, + index_cache_size_bytes=options.index_cache_size_bytes, + read_params=options.read_params, + session=options.resolve_session(), + namespace_client=options.namespace_client, + table_id=options.table_id, + base_store_params=options.base_store_params, + ) diff --git a/daft_lance/lance_data_sink.py b/daft_lance/lance_data_sink.py index c67cbc8..9c0567a 100644 --- a/daft_lance/lance_data_sink.py +++ b/daft_lance/lance_data_sink.py @@ -26,6 +26,7 @@ detect_blob_v2_columns, resolve_storage_version, ) +from daft_lance._runtime import CacheConfigurationError, DatasetOpenOptions, open_dataset, shared_session from daft_lance.namespace import ( ResolvedNamespaceTable, get_namespace_commit_kwargs, @@ -346,12 +347,13 @@ def _absorb_existing_dataset(self) -> lance.LanceDataset | None: """ dataset: lance.LanceDataset | None try: - dataset = lance.dataset( - self._dataset_uri_arg, storage_options=self._storage_options, **self._namespace_kwargs + dataset = open_dataset( + self._dataset_uri_arg, + DatasetOpenOptions(storage_options=self._storage_options, **self._namespace_kwargs), ) except (ValueError, FileNotFoundError, OSError) as e: # Pinned to the Rust message format; lance has no typed exception. See test_lance_message_format_unchanged. - if "was not found" in str(e): + if not isinstance(e, CacheConfigurationError) and "was not found" in str(e): dataset = None else: # Re-raise other errors (permissions, network, etc.) @@ -429,6 +431,7 @@ def _write_arrow_table(self, table: pa.Table) -> WriteResult[list[FragmentMetada wrapped = self._blob.wrap_table(table) fragments = lance.fragment.write_fragments( wrapped, + session=shared_session(), dataset_uri=self._table_uri, mode=self._write_mode, storage_options=self._storage_options, @@ -450,8 +453,13 @@ def _write_arrow_table(self, table: pa.Table) -> WriteResult[list[FragmentMetada def _ensure_mem_wal_dataset(self) -> lance.LanceDataset: assert self._effective_pyarrow_schema is not None, "LanceDataSink.start() must run before writes" try: - ds = lance.dataset(self._dataset_uri_arg, storage_options=self._storage_options, **self._namespace_kwargs) - except (ValueError, FileNotFoundError, OSError): + ds = open_dataset( + self._dataset_uri_arg, + DatasetOpenOptions(storage_options=self._storage_options, **self._namespace_kwargs), + ) + except (ValueError, FileNotFoundError, OSError) as error: + if isinstance(error, CacheConfigurationError) or "was not found" not in str(error): + raise ds = None if ds is None: @@ -631,7 +639,9 @@ def checkpoint_commit_exists(self) -> bool: if self._checkpoint is None: raise RuntimeError("checking a Lance checkpoint commit requires the driver-side CheckpointStore") self._ensure_started() - dataset = lance.dataset(self._dataset_uri_arg, storage_options=self._storage_options, **self._namespace_kwargs) + dataset = open_dataset( + self._dataset_uri_arg, DatasetOpenOptions(storage_options=self._storage_options, **self._namespace_kwargs) + ) return self._idempotence_key_exists(dataset) def _finalize_uncheckpointed(self, write_results: list[WriteResult[list[FragmentMetadata]]]) -> MicroPartition: @@ -654,7 +664,7 @@ def _finalize_cow(self, write_results: list[WriteResult[list[FragmentMetadata]]] operation = lance.LanceOperation.Append(fragments) assert self._table_uri is not None, "LanceDataSink.start() must run before finalize" - # Unlike lance.dataset(), commit() requires a str base_uri even when a + # Unlike open_dataset(), commit() requires a str base_uri even when a # namespace client is passed; the namespace kwargs only register the commit. dataset = lance.LanceDataset.commit( self._table_uri, @@ -675,11 +685,9 @@ def _finalize_insert_overwrite(self, fragments: list[FragmentMetadata]) -> Micro # Pinned to the version start() read: the deletions describe that snapshot, # and the commit declares it as its read version. - pinned = lance.dataset( + pinned = open_dataset( self._dataset_uri_arg, - version=self._version, - storage_options=self._storage_options, - **self._namespace_kwargs, + DatasetOpenOptions(version=self._version, storage_options=self._storage_options, **self._namespace_kwargs), ) open_context = DatasetOpenContext.from_dataset( pinned, @@ -700,13 +708,16 @@ def _finalize_insert_overwrite(self, fragments: list[FragmentMetadata]) -> Micro "insert_overwrite matched no rows and wrote no data for predicate %r; no version created", self._overwrite_where, ) - dataset = lance.dataset( - self._dataset_uri_arg, storage_options=self._storage_options, **self._namespace_kwargs + dataset = open_dataset( + self._dataset_uri_arg, + DatasetOpenOptions(storage_options=self._storage_options, **self._namespace_kwargs), ) return _dataset_stats(dataset) def _finalize_mem_wal(self, write_results: list[WriteResult[list[FragmentMetadata]]]) -> MicroPartition: - dataset = lance.dataset(self._dataset_uri_arg, storage_options=self._storage_options, **self._namespace_kwargs) + dataset = open_dataset( + self._dataset_uri_arg, DatasetOpenOptions(storage_options=self._storage_options, **self._namespace_kwargs) + ) if self._compact_after_write: logger.info( @@ -727,8 +738,9 @@ def _finalize_mem_wal(self, write_results: list[WriteResult[list[FragmentMetadat storage_options=self._storage_options, ), ) - dataset = lance.dataset( - self._dataset_uri_arg, storage_options=self._storage_options, **self._namespace_kwargs + dataset = open_dataset( + self._dataset_uri_arg, + DatasetOpenOptions(storage_options=self._storage_options, **self._namespace_kwargs), ) return _dataset_stats(dataset) @@ -743,7 +755,9 @@ def _finalize_checkpointed(self, write_results: list[WriteResult[list[FragmentMe assert self._checkpoint is not None self._ensure_started() - dataset = lance.dataset(self._dataset_uri_arg, storage_options=self._storage_options, **self._namespace_kwargs) + dataset = open_dataset( + self._dataset_uri_arg, DatasetOpenOptions(storage_options=self._storage_options, **self._namespace_kwargs) + ) pending_ids = self._pending_checkpoint_ids() # Check the Lance history again at finalize time. The pre-check in diff --git a/daft_lance/namespace.py b/daft_lance/namespace.py index ea7c8a4..331fd7f 100644 --- a/daft_lance/namespace.py +++ b/daft_lance/namespace.py @@ -36,9 +36,21 @@ import lance from lance_namespace import DeclareTableResponse, DescribeTableResponse, LanceNamespace +from ._runtime import DatasetOpenOptions, Session, clear_caches, open_dataset, shared_session + _NAMESPACE_CACHE_SIZE = int(os.environ.get("DAFT_LANCE_NAMESPACE_CACHE_SIZE", "16")) +def get_shared_session() -> Session: + """Return the runtime's process Session, honoring DAFT_LANCE_SHARE_SESSION.""" + return shared_session() + + +def reset_shared_session() -> None: + """Release the runtime Session and its configuration (primarily for tests).""" + clear_caches() + + def has_namespace_params(namespace_impl: str | None, table_id: list[str] | None) -> bool: """Whether namespace addressing is in effect (``namespace_properties`` stays optional).""" return namespace_impl is not None and table_id is not None @@ -353,6 +365,10 @@ def _open(self, version: int | None) -> lance.LanceDataset: """ return lance.LanceDataset( self.uri, + session=shared_session( + index_cache_size=self.index_cache_size, + metadata_cache_size_bytes=self.metadata_cache_size_bytes, + ), version=version, storage_options=self.storage_options, block_size=self.block_size, @@ -385,8 +401,7 @@ def open_dataset_from_open_kwargs(ds_uri: str | None, open_kwargs: dict[str, Any """ open_kwargs = dict(open_kwargs or {}) namespace_impl, namespace_properties, table_id = pop_namespace_params(open_kwargs) - return lance.dataset( + return open_dataset( None if has_namespace_params(namespace_impl, table_id) else ds_uri, - **get_namespace_kwargs(namespace_impl, namespace_properties, table_id), - **open_kwargs, + DatasetOpenOptions(**get_namespace_kwargs(namespace_impl, namespace_properties, table_id), **open_kwargs), ) diff --git a/daft_lance/utils.py b/daft_lance/utils.py index c3cf79d..4bfbb9e 100644 --- a/daft_lance/utils.py +++ b/daft_lance/utils.py @@ -9,6 +9,7 @@ from daft.dependencies import pa from daft.io.object_store_options import io_config_to_storage_options from daft.logical.schema import Schema as DaftSchema +from daft_lance._runtime import DatasetOpenOptions, open_dataset from daft_lance.namespace import ( DatasetOpenContext, get_namespace_commit_kwargs, @@ -210,12 +211,14 @@ def construct_lance_dataset_handle( kwargs["default_scan_options"] = original_default_scan_options dataset_uri = None if has_namespace_params(namespace_impl, table_id) else resolved_uri - dataset = lance.dataset( + dataset = open_dataset( dataset_uri, - storage_options=merged_storage_options, - version=version, - **get_namespace_kwargs(namespace_impl, namespace_properties, table_id), - **kwargs, + DatasetOpenOptions( + storage_options=merged_storage_options, + version=version, + **get_namespace_kwargs(namespace_impl, namespace_properties, table_id), + **kwargs, + ), ) effective_kwargs = { diff --git a/tests/io/lance/test_lance_runtime.py b/tests/io/lance/test_lance_runtime.py new file mode 100644 index 0000000..2086733 --- /dev/null +++ b/tests/io/lance/test_lance_runtime.py @@ -0,0 +1,108 @@ +from __future__ import annotations + +import pickle +from concurrent.futures import ThreadPoolExecutor + +import lance +import pyarrow as pa +import pytest + +from daft_lance._runtime import CacheConfigurationError, DatasetOpenOptions, clear_caches, open_dataset, shared_session +from daft_lance.lance_data_sink import LanceDataSink +from daft_lance.namespace import open_dataset_from_open_kwargs +from daft_lance.utils import construct_lance_dataset_handle + + +@pytest.fixture(autouse=True) +def reset_global_session(): + clear_caches() + yield + clear_caches() + + +def test_concurrent_opens_share_session_and_keep_snapshot(tmp_path): + uri = str(tmp_path / "table") + lance.write_dataset(pa.table({"a": [1, 2]}), uri) + handle = construct_lance_dataset_handle(uri, metadata_cache_size_bytes=32 * 1024 * 1024) + with ThreadPoolExecutor(max_workers=4) as executor: + datasets = list(executor.map(lambda _: open_dataset_from_open_kwargs(uri, handle.open_kwargs), range(8))) + assert len({id(ds) for ds in datasets}) == len(datasets) + assert all(ds.session().is_same_as(datasets[0].session()) for ds in datasets) + assert datasets[0].session().is_same_as(handle.dataset.session()) + + lance.write_dataset(pa.table({"a": [99]}), uri, mode="overwrite") + pinned = open_dataset_from_open_kwargs(uri, handle.open_kwargs) + assert pinned.version == handle.dataset.version + assert pinned.to_table().to_pydict() == {"a": [1, 2]} + assert open_dataset(uri).to_table().to_pydict() == {"a": [99]} + + +def test_global_session_capacity(tmp_path, monkeypatch): + monkeypatch.setenv("DAFT_LANCE_INDEX_CACHE_SIZE_BYTES", "1048576") + first = shared_session() + assert first.is_same_as(shared_session(index_cache_size_bytes=1048576)) + with pytest.raises(CacheConfigurationError, match="global Session already configured"): + shared_session(index_cache_size_bytes=2097152) + clear_caches() + assert shared_session(index_cache_size=1).is_same_as(shared_session(index_cache_size_bytes=20 * 1024 * 1024)) + clear_caches() + datasets = [] + for i in range(2): + uri = str(tmp_path / f"table-{i}") + lance.write_dataset(pa.table({"a": [i]}), uri) + datasets.append(open_dataset(uri)) + assert datasets[0].session().is_same_as(datasets[1].session()) + assert [ds.to_table()["a"].to_pylist() for ds in datasets] == [[0], [1]] + + +def test_fragment_writes_share_worker_session_without_serializing_it(tmp_path, monkeypatch): + uri = str(tmp_path / "writes") + schema = pa.schema([("a", pa.int64())]) + lance.write_dataset(pa.table({"a": [0]}), uri) + sink = LanceDataSink(uri=uri, schema=schema, mode="append") + sink.start() + worker = pickle.loads(pickle.dumps(sink)) + sessions = [] + original_write = lance.fragment.write_fragments + + def write(*args, **kwargs): + sessions.append(kwargs["session"]) + return original_write(*args, **kwargs) + + monkeypatch.setattr(lance.fragment, "write_fragments", write) + results = [worker._write_arrow_table(pa.table({"a": [value]})) for value in [1, 2]] + assert sessions[0].is_same_as(sessions[1]) + assert sessions[0].is_same_as(open_dataset(uri).session()) + fragments = [fragment for result in results for fragment in result.result] + committed = lance.LanceDataset.commit(uri, lance.LanceOperation.Append(fragments), read_version=1) + assert sorted(committed.to_table()["a"].to_pylist()) == [0, 1, 2] + + +def test_explicit_session_is_preserved(tmp_path): + uri = str(tmp_path / "explicit") + lance.write_dataset(pa.table({"a": [1]}), uri) + session = lance.Session(metadata_cache_size_bytes=1024) + assert open_dataset(uri, DatasetOpenOptions(session=session)).session().is_same_as(session) + + +def test_maintenance_context_reuses_session_and_latest_still_advances(tmp_path): + uri = str(tmp_path / "maintenance") + lance.write_dataset(pa.table({"a": [1]}), uri) + handle = construct_lance_dataset_handle(uri) + context = pickle.loads(pickle.dumps(handle.worker_open_context())) + first = context.open_pinned() + second = context.open_pinned() + assert first.session().is_same_as(second.session()) + assert first.session().is_same_as(handle.dataset.session()) + lance.write_dataset(pa.table({"a": [2]}), uri, mode="append") + assert context.open_pinned().to_table().to_pydict() == {"a": [1]} + assert context.open_latest().to_table().to_pydict() == {"a": [1, 2]} + + +def test_namespace_read_and_maintenance_share_session(tmp_path): + from tests.io.lance.test_namespace import _ns_handle + + handle = _ns_handle(tmp_path) + scan = open_dataset_from_open_kwargs(handle.uri, handle.open_kwargs) + maintenance = handle.worker_open_context().open_pinned() + assert scan.session().is_same_as(maintenance.session()) diff --git a/tests/io/lance/test_ray_session_memory.py b/tests/io/lance/test_ray_session_memory.py new file mode 100644 index 0000000..2278ebe --- /dev/null +++ b/tests/io/lance/test_ray_session_memory.py @@ -0,0 +1,357 @@ +"""Under Ray, the dataset opens happen in a worker process, not in the driver. + +``daft_lance`` hands one process-level ``lance.Session`` to every open, so a +table's fragment list is held once per process instead of once per scan task, +blob reader and writer. None of that is observable from the driver: a +monkeypatch or a spy only ever patches the driver, which is why the spy tests in +``test_shared_session.py`` are native-only. + +So these ask the worker itself. A plain batch UDF is fused into the scan's own +``SwordfishTask``, so a UDF that reports ``Session.size_bytes()`` reports it from +the very process that opened the datasets. ``size_bytes()`` is used rather than +RSS on purpose: it is exact, it needs no large table, and it is zero for any open +that did not go through the process session -- which is exactly what the +pre-sharing behavior looks like. These tests validate cache reuse rather than +RSS or timing thresholds. + +The last two tests cover the maintenance path, which behaves differently enough +to be worth stating: its opens do not multiply by concurrency, because Daft runs +each ``daft.cls`` instance in a process of its own. They assert that, so the day +it stops being true is the day the tests say so. +""" + +from __future__ import annotations + +import json +import os +from typing import Any + +import lance +import numpy as np +import pyarrow as pa +import pytest +from pytest import TempPathFactory + +import daft +import daft_lance +from daft_lance._blob import read_blobs +from daft_lance.utils import construct_lance_dataset_handle + +ray = pytest.importorskip("ray") + +pytestmark = pytest.mark.skipif( + os.environ.get("DAFT_RUNNER") != "ray", + reason="a worker process's lance session is only observable under the Ray runner", +) + +NUM_FRAGMENTS = 32 +DRIVER_FLAG = os.environ.get("DAFT_LANCE_SHARE_SESSION") + + +def report_from_worker() -> dict[int, dict[str, Any]]: + """Ask every worker that runs a task for its lance session state. + + The probe is defined here, not at module scope, so cloudpickle ships it by + value and the worker does not have to import this test module. It must stay + a plain batch UDF: ``concurrency=`` or ``use_process=True`` would move it + into a process of its own, and it would then describe the wrong process. + """ + + @daft.func.batch(return_dtype=daft.DataType.string()) + def probe(column: daft.Series) -> list[str]: + from daft_lance.namespace import get_shared_session + + session = get_shared_session() + payload = json.dumps( + { + "pid": os.getpid(), + "session_size": session.size_bytes(), + "session_obj_id": id(session), + # ``is_same_as`` exists at runtime; pylance 12's stub omits it. + "same_session": bool(session.is_same_as(get_shared_session())), # type: ignore[attr-defined] + "flag": os.environ.get("DAFT_LANCE_SHARE_SESSION"), + } + ) + return [payload] * len(column) + + df = daft.from_pydict({"x": [0]}).with_column("p", probe(daft.col("x"))) + rows = [json.loads(value) for value in df.to_pydict()["p"]] + return {row["pid"]: row for row in rows} + + +def one_worker(report: dict[int, dict[str, Any]]) -> dict[str, Any]: + assert len(report) == 1, f"expected a single worker process, got pids {sorted(report)}" + return next(iter(report.values())) + + +def session_growth(before: dict[int, dict[str, Any]], after: dict[int, dict[str, Any]]) -> int: + """How much the worker's process-wide session grew between two reports.""" + assert sorted(before) == sorted(after), ( + f"the reports ran on different workers ({sorted(before)} then {sorted(after)}), so nothing can be compared" + ) + pid = next(iter(before)) + return int(after[pid]["session_size"]) - int(before[pid]["session_size"]) + + +def write_table(uri: str, fragments: int = NUM_FRAGMENTS) -> lance.LanceDataset: + """One row per fragment, so the manifest is mostly its fragment list.""" + table = pa.table({"i": pa.array(range(fragments), type=pa.int64())}) + return lance.write_dataset(table, uri, mode="create", max_rows_per_file=1) + + +def test_worker_sees_the_drivers_share_session_flag() -> None: + """Every other test here, and the A/B benchmark, rests on this. + + Ray actors do not inherit the driver's environment by contract, so if the + kill switch stopped reaching the worker, the benchmark would quietly compare + two identical runs and report a flattering number. + """ + worker = one_worker(report_from_worker()) + + assert worker["flag"] == DRIVER_FLAG + + +def test_scan_tasks_open_through_the_process_session(tmp_path_factory: TempPathFactory) -> None: + uri = str(tmp_path_factory.mktemp("read") / "table.lance") + dataset = write_table(uri) + + before = report_from_worker() + # Never count_rows(): count pushdown collapses the plan to a single task + # that opens no fragment at all. + rows = daft_lance.read_lance(uri).to_arrow().num_rows + after = report_from_worker() + + assert rows == NUM_FRAGMENTS + growth = session_growth(before, after) + assert growth >= len(dataset._ds.serialized_manifest()), ( + f"the worker's session grew by only {growth} bytes, which cannot hold this table's manifest: " + "the scan tasks opened their datasets outside the process session" + ) + assert one_worker(after)["same_session"] + + +def test_scan_tasks_reuse_one_session_per_process(tmp_path_factory: TempPathFactory) -> None: + """Two queries in a row must land on the same session object, not two.""" + uri = str(tmp_path_factory.mktemp("reuse") / "table.lance") + write_table(uri) + + daft_lance.read_lance(uri).to_arrow() + first = one_worker(report_from_worker()) + daft_lance.read_lance(uri).to_arrow() + second = one_worker(report_from_worker()) + + assert first["same_session"] and second["same_session"], ( + "the worker handed out a fresh session per call, so nothing is shared between its tasks" + ) + assert first["session_obj_id"] == second["session_obj_id"], ( + "the two queries saw different session objects, so the session is per-query, not per-process" + ) + # The second read of the same pinned snapshot is served from the manifest + # cache, so it adds nothing. + assert second["session_size"] == first["session_size"], ( + "re-reading the same pinned snapshot grew the session again: the manifest cache was not hit" + ) + + +def test_writes_open_through_the_process_session(tmp_path_factory: TempPathFactory) -> None: + """``write_fragments`` reloads the manifest per flush unless it gets a session.""" + uri = str(tmp_path_factory.mktemp("write") / "table.lance") + write_table(uri) + + before = report_from_worker() + daft.from_pydict({"i": list(range(8))}).write_lance(uri, mode="append", max_rows_per_file=1).collect() + after = report_from_worker() + + assert session_growth(before, after) > 0, ( + "the worker's session did not grow across the write: the writer opened outside it" + ) + assert lance.dataset(uri).count_rows() == NUM_FRAGMENTS + 8 + + +def test_indexed_queries_load_the_index_once_per_process(tmp_path_factory: TempPathFactory) -> None: + """A vector index is the other half of what a session holds, and the costly half. + + A session carries an index cache as well as a metadata cache, and index size + scales with rows where a manifest scales with fragments: measured on a + 200k-row table, one IVF_PQ index costs about 5 MB per session, against 9 KB + of manifest. A ``nearest`` read plans a single index-driven scan task, so + nothing here multiplies by concurrency -- what sharing changes is that the + index is loaded once per worker process instead of once per query, which is + a storage round trip and a decode every query without it. + """ + uri = str(tmp_path_factory.mktemp("indexed") / "table.lance") + rows, dim = 512, 8 + rng = np.random.default_rng(0) + vectors = rng.random((rows, dim), dtype=np.float32) + table = pa.table( + { + "id": pa.array(range(rows), type=pa.int64()), + "v": pa.FixedSizeListArray.from_arrays(pa.array(vectors.reshape(-1), type=pa.float32()), dim), + } + ) + dataset = lance.write_dataset(table, uri, max_rows_per_file=128) + dataset.create_index("v", index_type="IVF_PQ", num_partitions=4, num_sub_vectors=2) + + # The same vector every time, on purpose: a different one probes different + # IVF partitions and legitimately loads more of the index, which would make + # "was it reloaded?" unanswerable on a table this small. + query = rng.random(dim, dtype=np.float32) + + def search() -> int: + nearest = {"column": "v", "q": query, "k": 4} + return daft_lance.read_lance(uri, default_scan_options={"nearest": nearest}).to_arrow().num_rows + + before = report_from_worker() + assert search() == 4 + after_first = report_from_worker() + assert search() == 4 + after_second = report_from_worker() + + first_growth = session_growth(before, after_first) + assert first_growth > 0, ( + "the worker's session did not grow across an indexed query: the index was loaded into a session of its own" + ) + # The second query must be served from the index cache the first one filled. + assert session_growth(after_first, after_second) == 0, ( + "the second indexed query grew the session again, so it reloaded the index instead of reusing it" + ) + + +def test_read_blobs_opens_through_the_process_session(tmp_path_factory: TempPathFactory) -> None: + """The blob reader reopens the pinned snapshot on the worker. + + That the reopen uses a uri plus open kwargs -- rather than closing over the + driver's ``LanceDataset``, whose ``__reduce__`` ships the whole re-encoded + manifest per task -- is asserted by the spy test in ``test_shared_session``, + which can only run natively. What is checked here is the other half: under + Ray the payload really does travel and the open really does land in the + worker's process-wide session. + """ + uri = str(tmp_path_factory.mktemp("blobs") / "table.lance") + rows = 4 + table = pa.table( + { + "id": pa.array(range(rows), type=pa.int64()), + "blob": lance.blob_array([b"payload-%d" % i for i in range(rows)]), + } + ) + dataset = lance.write_dataset(table, uri, data_storage_version="2.2", max_rows_per_file=1) + + before = report_from_worker() + df = daft.read_lance(uri, default_scan_options={"with_row_id": True}) + blobs = read_blobs(df, dataset, "blob").sort("id").to_pydict()["blob"] + after = report_from_worker() + + assert blobs == [b"payload-%d" % i for i in range(rows)] + assert session_growth(before, after) > 0, ( + "the worker's session did not grow across the blob read: it opened outside the process session" + ) + + +MAINT_CONCURRENCY = 2 + + +def maintenance_census(open_context: Any, concurrency: int) -> list[dict[str, Any]]: + """Open the table the way a maintenance worker does, and report who did it. + + Daft's maintenance operators are ``daft.cls`` instances -- one open per + instance, cached on ``self`` -- so what their memory costs turns entirely on + where Daft puts those instances. The real operators cannot be asked (they + return rewrite results, not pids), so this stands in for them, decorated the + same way: ``max_concurrency`` set, ``use_process`` left alone. + + Defined inside the function so cloudpickle ships it by value, exactly as in + :func:`report_from_worker`. + """ + + class OpenCensus: + def __init__(self, context: Any) -> None: + self.context = context + self._lance_ds: lance.LanceDataset | None = None + + @daft.method.batch(return_dtype=daft.DataType.string()) + def __call__(self, column: daft.Series) -> list[str]: + import json as json_ + import os as os_ + + from daft_lance.namespace import get_shared_session + + if self._lance_ds is None: + self._lance_ds = self.context.open_pinned() + session = get_shared_session() + payload = json_.dumps( + { + "pid": os_.getpid(), + "instance": id(self), + "session_obj_id": id(session), + "session_size": session.size_bytes(), + "fragments": len(self._lance_ds.get_fragments()), + } + ) + return [payload] * len(column) + + census = daft.cls(OpenCensus, max_concurrency=concurrency) + # One partition per instance several times over, so every instance in the + # pool is given work and reports back. + partitions = concurrency * 8 + df = daft.from_pydict({"x": list(range(partitions))}).into_partitions(partitions) + rows = [json.loads(value) for value in df.select(census(open_context)(daft.col("x")).alias("p")).to_pydict()["p"]] + instances = {(row["pid"], row["instance"]): row for row in rows} + assert len(instances) >= 2, ( + f"only {len(instances)} instance ran at concurrency {concurrency}, so this measured nothing about concurrency" + ) + return list(instances.values()) + + +def test_maintenance_instances_each_get_their_own_process(tmp_path_factory: TempPathFactory) -> None: + """Why concurrency does not multiply a maintenance worker's manifests. + + A scan task shares its process with ``scantask_max_parallel``-1 others, so + without a shared session their manifests add up and sharing collapses them. + A maintenance instance does not: Daft runs each one in a process of its own, + so it opens the table exactly once however high the concurrency, and there + is nothing for a session to collapse. This is asserted rather than assumed + because it is the whole reason the maintenance path needs no saving -- if + Daft ever fuses these instances into one process, the manifests start + multiplying and both this file and the A/B benchmark need a stage for it. + """ + uri = str(tmp_path_factory.mktemp("maint_procs") / "table.lance") + write_table(uri) + open_context = construct_lance_dataset_handle(uri).worker_open_context() + + instances = maintenance_census(open_context, MAINT_CONCURRENCY) + + pids = {row["pid"] for row in instances} + assert len(pids) == len(instances), ( + f"{len(instances)} maintenance instances shared {len(pids)} processes. Daft used to give each its own; " + "now that they overlap, their opens multiply within a process and that saving should be measured" + ) + assert {row["fragments"] for row in instances} == {NUM_FRAGMENTS}, ( + "the instances disagreed about the table's fragment count, so they did not all open the pinned snapshot" + ) + + +def test_maintenance_workers_open_through_the_process_session(tmp_path_factory: TempPathFactory) -> None: + """The open itself still has to go through the session, per process. + + One open per process is the floor, not zero: the worker pays one manifest + read either way. What this checks is that it lands in the process-wide + session rather than beside it, which is what keeps the floor at one -- a + maintenance process that also scans, or that Daft later gives a second + instance, then shares the one copy. + """ + uri = str(tmp_path_factory.mktemp("maint_session") / "table.lance") + dataset = write_table(uri) + manifest_bytes = len(dataset._ds.serialized_manifest()) + open_context = construct_lance_dataset_handle(uri).worker_open_context() + + instances = maintenance_census(open_context, MAINT_CONCURRENCY) + + for row in instances: + assert row["session_size"] >= manifest_bytes, ( + f"a maintenance instance's process session holds {row['session_size']} bytes, too little for this " + f"table's {manifest_bytes}-byte manifest: it opened outside the process session" + ) + assert len({row["session_obj_id"] for row in instances}) == len(instances), ( + "two instances in different processes reported the same session object id, which cannot be" + ) diff --git a/tests/io/lance/test_shared_session.py b/tests/io/lance/test_shared_session.py new file mode 100644 index 0000000..30c6d7f --- /dev/null +++ b/tests/io/lance/test_shared_session.py @@ -0,0 +1,208 @@ +"""Every dataset a process opens must share one lance Session. + +Without it each ``lance.dataset()`` builds its own ``Arc`` holding +the table's whole fragment list, plus its own index/metadata cache budgets and +object-store client. Daft opens one dataset per scan task, per maintenance UDF +instance and per writer, so on a many-fragment table that cost is paid again +for every concurrently open dataset. + +How much that saves depends on how many of those opens share a process, which +differs by path and is measured rather than assumed: see +``test_ray_session_memory.py`` for the scan and maintenance paths, and +``test_compaction_loads_the_index_in_its_commit_not_in_its_tasks`` below for +where an indexed table's maintenance cost actually sits. +""" + +from __future__ import annotations + +import os +import pickle + +import lance +import numpy as np +import pyarrow as pa +import pytest +from pytest import TempPathFactory + +import daft +from daft_lance._blob import read_blobs +from daft_lance.namespace import ( + DatasetOpenContext, + get_shared_session, + open_dataset_from_open_kwargs, + reset_shared_session, +) +from daft_lance.utils import construct_lance_dataset_handle + +NUM_FRAGMENTS = 32 + +ON_RAY = os.environ.get("DAFT_RUNNER") == "ray" + + +def same_session(left: object, right: object) -> bool: + """``Session.is_same_as`` exists at runtime; pylance 12's stub omits it.""" + return bool(left.is_same_as(right)) # type: ignore[attr-defined] + + +@pytest.fixture +def many_fragment_uri(tmp_path_factory: TempPathFactory) -> str: + """A table with one fragment per row, so the manifest is dominated by fragments.""" + uri = str(tmp_path_factory.mktemp("many_fragments") / "table.lance") + table = pa.table({"i": pa.array(range(NUM_FRAGMENTS), type=pa.int64())}) + lance.write_dataset(table, uri, mode="create", max_rows_per_file=1) + return uri + + +@pytest.fixture(autouse=True) +def fresh_session() -> None: + reset_shared_session() + + +def test_get_shared_session_returns_one_session_per_process() -> None: + first = get_shared_session() + assert get_shared_session() is first + assert same_session(first, get_shared_session()) + + +def test_reset_shared_session_builds_a_new_one() -> None: + first = get_shared_session() + reset_shared_session() + assert not same_session(get_shared_session(), first) + + +def test_kill_switch_gives_every_open_its_own_session(monkeypatch: pytest.MonkeyPatch) -> None: + """``DAFT_LANCE_SHARE_SESSION=0`` restores lance's per-open sessions. + + This is the only way to get pre-sharing behavior inside a Ray worker, so + the A/B benchmark depends on it: a driver-side patch never reaches there. + """ + monkeypatch.setenv("DAFT_LANCE_SHARE_SESSION", "0") + unshared = get_shared_session() + assert not same_session(get_shared_session(), unshared) + + monkeypatch.delenv("DAFT_LANCE_SHARE_SESSION") + shared = get_shared_session() + assert same_session(get_shared_session(), shared) + assert not same_session(shared, unshared) + + +def test_scan_task_opens_share_the_session(many_fragment_uri: str) -> None: + """``open_dataset_from_open_kwargs`` is what every scan task calls.""" + shared = get_shared_session() + open_kwargs = {"version": 1} + datasets = [open_dataset_from_open_kwargs(many_fragment_uri, dict(open_kwargs)) for _ in range(NUM_FRAGMENTS)] + + assert all(same_session(ds.session(), shared) for ds in datasets) + assert all(len(ds.get_fragments()) == NUM_FRAGMENTS for ds in datasets) + + +def test_driver_and_worker_opens_share_the_session(many_fragment_uri: str) -> None: + """The native runner plans and executes in one process, so these must match.""" + handle = construct_lance_dataset_handle(many_fragment_uri) + worker_ds = handle.worker_open_context().open_pinned() + + assert same_session(handle.dataset.session(), worker_ds.session()) + assert worker_ds.version == handle.dataset.version + + +def test_open_context_still_pickles(many_fragment_uri: str) -> None: + """The session must never become a field: it does not pickle, the context must.""" + context = DatasetOpenContext.from_dataset(lance.dataset(many_fragment_uri), many_fragment_uri) + + restored = pickle.loads(pickle.dumps(context)) + + assert restored == context + assert same_session(restored.open_pinned().session(), get_shared_session()) + + +def test_lance_session_itself_does_not_pickle() -> None: + with pytest.raises(TypeError): + pickle.dumps(get_shared_session()) + + +@pytest.mark.skipif(ON_RAY, reason="the spy patches the driver process, which is not where a Ray worker runs") +def test_read_blobs_reopens_on_the_worker(tmp_path_factory: TempPathFactory, monkeypatch: pytest.MonkeyPatch) -> None: + """The UDF must reopen the snapshot, not close over the driver's dataset. + + Closing over it pickles ``LanceDataset``, whose ``__reduce__`` carries the + whole re-encoded manifest -- one copy per task -- and a dataset rebuilt + from a serialized manifest also bypasses the session's manifest cache. + """ + uri = str(tmp_path_factory.mktemp("blobs")) + table = pa.table({"id": pa.array([1, 2], pa.int64()), "blob": lance.blob_array([b"first", b"second"])}) + ds = lance.write_dataset(table, uri, data_storage_version="2.2", max_rows_per_file=1) + + calls: list[dict[str, object]] = [] + + def spy(ds_uri: str | None, open_kwargs: dict[str, object] | None) -> lance.LanceDataset: + calls.append({"uri": ds_uri, **(open_kwargs or {})}) + return open_dataset_from_open_kwargs(ds_uri, open_kwargs) + + monkeypatch.setattr("daft_lance._blob.open_dataset_from_open_kwargs", spy) + + df = daft.read_lance(ds.uri, default_scan_options={"with_row_id": True}) + blobs = read_blobs(df, ds, "blob").sort("id").to_pydict()["blob"] + + assert blobs == [b"first", b"second"] + assert calls, "read_blobs did not reopen the dataset on the worker" + assert all(call["uri"] == ds.uri and call["version"] == ds.version for call in calls) + + +def test_compaction_loads_the_index_in_its_commit_not_in_its_tasks(tmp_path_factory: TempPathFactory) -> None: + """On an indexed table, compaction's index cost lands on the driver. + + The worry this settles: a maintenance task opens the table itself, so on an + indexed table every concurrent task might load the index and the cost would + multiply by the concurrency. It does not. A task's open brings in the + index's *metadata* entry and nothing more; the index data is read when the + rewritten fragments are committed, because that is when every index entry + has to be remapped onto the new fragment ids -- and the commit runs once, + on the driver, on ``lance_ds`` rather than on the workers' reopened copy. + + Run in-process on purpose: it needs the tasks and the commit to be + separable, which ``compact_files`` does not expose. ``CompactionTaskUDF`` is + exactly what Daft instantiates on a worker, so calling it directly measures + the same opens. + + The numbers come from ``index_cache_size_bytes()``, which is the accounting + lance caps its index budget with rather than resident bytes (it reads about + seven times what ``size_bytes()`` and the process's own RSS move by). Only + growth within that one accounting is compared here, never against a size in + bytes. + """ + from lance.optimize import Compaction, CompactionOptions + + from daft_lance.lance_compaction import CompactionTaskUDF + + uri = str(tmp_path_factory.mktemp("compact_index") / "table.lance") + rows, dim = 512, 8 + rng = np.random.default_rng(0) + vectors = rng.random((rows, dim), dtype=np.float32) + table = pa.table( + { + "i": pa.array(range(rows), type=pa.int64()), + "v": pa.FixedSizeListArray.from_arrays(pa.array(vectors.reshape(-1), type=pa.float32()), dim), + } + ) + dataset = lance.write_dataset(table, uri, max_rows_per_file=rows // 16) + dataset.create_index("v", index_type="IVF_PQ", num_partitions=4, num_sub_vectors=2) + + session = get_shared_session() + handle = construct_lance_dataset_handle(uri) + plan = Compaction.plan(handle.dataset, CompactionOptions(target_rows_per_fragment=rows // 4)) + assert plan.num_tasks() > 1, "nothing to compact, so this measured no task opens at all" + + before_tasks = session.index_cache_size_bytes() + worker = CompactionTaskUDF(handle.worker_open_context()) + rewrites = [worker(task) for task in plan.tasks] + after_tasks = session.index_cache_size_bytes() + Compaction.commit(handle.dataset, rewrites) + after_commit = session.index_cache_size_bytes() + + task_growth = after_tasks - before_tasks + commit_growth = after_commit - after_tasks + assert commit_growth > 3 * task_growth, ( + f"the {plan.num_tasks()} compaction tasks grew the index cache by {task_growth} bytes and the commit by " + f"{commit_growth}; the index is no longer remapped in the commit, so a maintenance task's index cost may " + "now multiply by the concurrency and the memory assumptions need revisiting" + )