perf: share one lance Session per process

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

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`.

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,

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,
+    )