From c71b4c768d972f5d395f7d159cdfa5d9c39385bd Mon Sep 17 00:00:00 2001 From: fanng Date: Mon, 12 Oct 2026 17:15:52 +0900 Subject: [PATCH] perf: share one lance Session per process Every `lance.dataset()` opened without a session builds its own `Arc` holding the table's whole fragment list, plus its own index and metadata cache budgets and its own object-store client. Daft opens one per scan task, per writer flush and per maintenance instance, so a many-fragment table pays that again for every concurrently open dataset. Pass one process-level `Session` to every open instead. The manifest is a metadata-cache entry keyed by (uri, version, e_tag), so every open of the snapshot the driver pinned hands back a clone of the same `Arc`. `read_blobs` additionally stops closing over the driver's `LanceDataset`: `__reduce__` ships the whole re-encoded manifest per task, which both inflates the payload and bypasses the session's manifest cache. It now carries the uri plus the pinned open kwargs and reopens once per worker process. `DAFT_LANCE_SHARE_SESSION=0` restores lance's per-open sessions. It is an escape hatch, and the only way to get pre-sharing behavior inside a Ray worker, which is what the A/B benchmark needs: a driver-side monkeypatch never reaches the worker. Measured in a Ray worker on a KubeRay cluster, `read_lance` -> `write_parquet` (a write rather than an aggregation on purpose: a task holds its parquet buffers at the same time as its manifest, so the two compete for one memory limit): - 4000 fragments x 10k rows, concurrency 24, whole table: worker ru_maxrss grew 206.5 MB unshared against 157.4 MB shared, so 49.1 MB less resident. Daft's own sampled peak agrees (487.8 vs 442.1 MB). - 100k fragments x 10k rows, concurrency 8, reading 1% in groups of ten: 1052.2 MB against 573.3 MB, so 479.0 MB less resident -- and the pipeline finished in 61.9s instead of 609.6s. At that scale the cost of not sharing is mostly time: each unshared task re-parsed the same 8 MB manifest, about 5.5s of the 6.1s it spent. - Repeated `nearest` queries on a 200k-row IVF_PQ index run 1.7-2.0x faster, because the index is loaded once per process instead of once per query. The maintenance path is the exception and is covered so the exception stays documented: Daft runs each `daft.cls` instance in a process of its own, so a maintenance worker opens the table exactly once however high the concurrency, and nothing multiplies by it. An indexed table does not change that -- a task's open brings in the index's metadata entry only, and the index data is read once on the driver, when the compaction commit remaps it onto the new fragment ids. Tests assert the mechanism rather than the megabytes: `Session.size_bytes()` is exact, needs no large table, and is zero for any open that did not go through the process session. `benchmarks/lance_session_memory.py` (`make bench`) is the counterpart that puts numbers on it. --- Makefile | 13 +- README.md | 45 + benchmarks/lance_session_memory.py | 1298 +++++++++++++++++++++ daft_lance/_blob.py | 20 +- daft_lance/lance_data_sink.py | 43 +- daft_lance/namespace.py | 78 +- daft_lance/utils.py | 4 + tests/io/lance/test_ray_session_memory.py | 357 ++++++ tests/io/lance/test_shared_session.py | 208 ++++ 9 files changed, 2050 insertions(+), 16 deletions(-) create mode 100644 benchmarks/lance_session_memory.py create mode 100644 tests/io/lance/test_ray_session_memory.py create mode 100644 tests/io/lance/test_shared_session.py diff --git a/Makefile b/Makefile index 93e0953..2dd6b36 100644 --- a/Makefile +++ b/Makefile @@ -18,7 +18,18 @@ 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: bench +bench: sync ## Measure lance session memory in a real Ray run + DAFT_ANALYTICS_ENABLED=0 uv run --group ray python benchmarks/lance_session_memory.py $(EXTRA_ARGS) .PHONY: format format: sync ## Format Python code diff --git a/README.md b/README.md index ee6d1b5..59a9519 100644 --- a/README.md +++ b/README.md @@ -325,6 +325,51 @@ 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. `benchmarks/lance_session_memory.py` (`make bench`) measures the difference in a real Ray worker. | + +### What sharing a session actually saves, per path + +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, which `benchmarks/lance_session_memory.py` checks rather than assumes. + +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. + ## Migration The migration only requires replacing `daft.io.lance` with `daft_lance`. diff --git a/benchmarks/lance_session_memory.py b/benchmarks/lance_session_memory.py new file mode 100644 index 0000000..99146c9 --- /dev/null +++ b/benchmarks/lance_session_memory.py @@ -0,0 +1,1298 @@ +"""Measure what sharing one lance ``Session`` saves inside a real Ray worker. + +``daft_lance`` passes one process-level ``lance.Session`` to every dataset it +opens (``daft_lance.namespace.get_shared_session``), so a table's fragment list +and lance's caches are held once per process instead of once per scan task, +maintenance UDF and writer. The numbers that motivated the change were measured +in a driver process, but under Ray the memory is spent in the worker, and a +driver-side monkeypatch cannot reach it. So this script runs the real pipeline +twice in two subprocesses that differ only in ``DAFT_LANCE_SHARE_SESSION``, and +reads the worker's memory back out of the worker itself: + +* a probe UDF, fused into the scan's own ``SwordfishTask`` (no ``concurrency=``, + no ``use_process=True``, or Daft would give it a process of its own), reports + ``getrusage(RUSAGE_SELF).ru_maxrss`` -- the worker process's high-water mark -- + plus ``Session.size_bytes()`` and the flag value it actually sees; +* Daft's own sampler (``DAFT_PROCESS_MONITOR_ENABLED=1`` + ``DAFT_EVENT_LOG_DIR``) + writes ``process.memory.rss`` every 200ms from inside each worker to + ``//events.jsonl``, which is an independent second opinion. + +The main stage is the shape that matters most in practice: **read Lance, write +Parquet**. It is the one the thresholds are derived from, and it is measured as +a write rather than an aggregation on purpose -- a task holds its parquet +writer's buffers at the same time as its manifest, so the two compete for the +same memory limit, which is what a real backfill does. An aggregation keeps +almost nothing resident and would flatter the worker. + +Three more stages follow, and they make different claims. ``read_blobs`` and the +lance write share the main stage's shape: ``--parallel`` concurrent tasks in one +worker process, so their manifests multiply by concurrency and sharing collapses +them. The indexed stage multiplies by nothing -- a ``nearest`` read plans a +single index-driven task -- and what sharing buys there is loading the index +once per process instead of once per query. The maintenance stage (compaction +and ``update_columns``) is the one that turns out *not* to multiply: Daft gives +each concurrent class instance a process of its own, so a maintenance worker +opens the table exactly once however high the concurrency, and an indexed +table's real maintenance cost is the index remap in the commit, on the driver. + +Usage:: + + make bench + make bench EXTRA_ARGS="--fragments 2000 --parallel 16" + uv run --group ray python benchmarks/lance_session_memory.py --mode shared +""" + +from __future__ import annotations + +import argparse +import json +import os +import pickle +import resource +import statistics +import subprocess +import sys +import tempfile +import time +from pathlib import Path +from typing import Any + +import lance +import pyarrow as pa +import ray +from lance.lance import _Session as LanceSession + +import daft +import daft_lance +from daft import method +from daft_lance._blob import read_blobs +from daft_lance.namespace import get_shared_session +from daft_lance.utils import construct_lance_dataset_handle + +MB = 1024 * 1024 +RESULT_PREFIX = "RESULT_JSON " +# The indexed stage. A vector index costs roughly 26 bytes per row in a +# session's index cache (measured: 5.1 MB at 200k rows), which dwarfs the +# manifest on any normally shaped table -- so a benchmark that only ever reads +# an unindexed table understates what sharing a session saves. +INDEX_ROWS = 200_000 +INDEX_DIM = 128 +INDEX_PARTITIONS = 256 +INDEX_SUB_VECTORS = 16 +INDEX_QUERIES = 6 +RSS_KEY = "process.memory.rss" +WRITE_ROWS = 8 +# Rows per fragment in the table under test. Also load-bearing for the +# measurement, not just for realism -- see build_table. +ROWS_PER_FRAGMENT = 10_000 +# Worker RSS wobbles by a megabyte or so between runs (Ray's own buffers, the +# arrow allocator). Below this the manifest saving is inside the noise, and the +# table under test is simply too small to make the claim with RSS. +RSS_NOISE_FLOOR = 4 * MB +# How long one side of the A/B may take before something is clearly wrong. +CHILD_TIMEOUT = 1800 +# The maintenance stage. Fewer fragments than the read table and no vector +# column: what is being measured there is where a maintenance *worker's* memory +# goes, and a worker opens the table once whatever the fragment count. The index +# is a scalar one for the same reason -- cheap to build, and the compaction +# commit still has to remap it. +MAINT_FRAGMENTS = 200 +MAINT_CONCURRENCY = 6 +# Rows in each small fragment the maintenance stage appends to give compaction +# something to merge. Compaction's target is a few of these, so the table's own +# large fragments stay above the target and are never rewritten -- the stage +# costs seconds whatever the table's size. +MAINT_ROWS_PER_APPEND = 1_000 +MAINT_TARGET_FRAGMENT_ROWS = MAINT_ROWS_PER_APPEND * 4 +# How many extra opens calibrate_per_open_bytes may hold at once. +CALIBRATION_OPENS = 4 + + +def maxrss_bytes() -> int: + """``ru_maxrss`` normalized to bytes: macOS reports bytes, Linux kilobytes.""" + rss = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss + return int(rss) if sys.platform == "darwin" else int(rss) * 1024 + + +def worker_parallelism() -> dict[str, Any]: + """What the executing process thinks it may run in parallel. + + The scheduling ceiling is min(the Ray node's CPUs, scantask_max_parallel), + but a container's real ceiling is its cgroup quota, and that is what sizes + Daft's execution thread pool. On a laptop these agree; in a pod with + ``num-cpus: 24`` and ``limits.cpu: 2`` they do not, so both get recorded. + """ + info: dict[str, Any] = {"cpu_count": os.cpu_count()} + affinity = getattr(os, "sched_getaffinity", None) + if affinity is not None: + info["affinity"] = len(affinity(0)) + try: + quota, period = Path("/sys/fs/cgroup/cpu.max").read_text().split() + info["cgroup_cpu_max"] = None if quota == "max" else round(int(quota) / int(period), 2) + except (OSError, ValueError): + info["cgroup_cpu_max"] = None + return info + + +@daft.func.batch(return_dtype=daft.DataType.string()) +def probe(column: daft.Series) -> list[str]: + """Report the executing process's memory and session state, one row at a time. + + Deliberately a plain batch UDF: Daft fuses those into the scan's task + instead of giving them their own actor, so what this returns describes the + very process that opened the lance datasets. + """ + session = get_shared_session() + payload = json.dumps( + { + "pid": os.getpid(), + "maxrss": maxrss_bytes(), + "session_size": session.size_bytes(), + "session_obj_id": id(session), + "same_session": bool(session.is_same_as(get_shared_session())), + "flag": os.environ.get("DAFT_LANCE_SHARE_SESSION"), + "parallelism": worker_parallelism(), + } + ) + return [payload] * len(column) + + +def build_table(uri: str, fragments: int, rows_per_fragment: int = ROWS_PER_FRAGMENT) -> None: + """``fragments`` files of ``rows_per_fragment`` rows each. + + A real lance table holds tens of thousands of rows per fragment, and the row + count matters to this benchmark beyond realism: with one row per fragment a + scan task opens its dataset, reads a row and drops it, so the opens barely + overlap and far fewer than ``--parallel`` manifests are ever resident at + once. Filling the fragments keeps each task busy long enough for the + concurrency to actually materialize. + """ + # Written in chunks of about ten million rows so the whole table is never in + # memory at once, and each chunk is a whole number of fragments so no chunk + # boundary leaves a half-filled one. Chunks are not just a memory knob: each + # one is an append, and an append rewrites the whole manifest, so chunking a + # 100k-fragment table a million rows at a time spends most of its time + # rewriting an ever-growing fragment list. Ten million keeps the arrow + # buffer under a hundred megabytes and cuts the commits tenfold. + rows_per_chunk = max(rows_per_fragment, (10_000_000 // max(rows_per_fragment, 1)) * rows_per_fragment) + total = fragments * rows_per_fragment + start = 0 + while start < total: + count = min(rows_per_chunk, total - start) + table = pa.table({"i": pa.array(range(start, start + count), type=pa.int64())}) + lance.write_dataset( + table, + uri, + mode="create" if start == 0 else "append", + max_rows_per_file=rows_per_fragment, + ) + start += count + + +def build_blob_table(uri: str, fragments: int) -> None: + rows = max(fragments, 2) + table = pa.table( + { + "id": pa.array(range(rows), type=pa.int64()), + "blob": lance.blob_array([b"payload-%d" % i for i in range(rows)]), + } + ) + lance.write_dataset(table, uri, mode="create", data_storage_version="2.2", max_rows_per_file=1) + + +def build_indexed_table(uri: str, rows: int = INDEX_ROWS, dim: int = INDEX_DIM) -> None: + """A normally shaped table -- few fragments, many rows -- carrying a vector index. + + The point of this one is the index, not the fragment list: index size scales + with rows while a manifest scales with fragments, so this is the mirror image + of the one-row-per-fragment table above. + """ + import numpy as np + + rng = np.random.default_rng(0) + chunk = 50_000 + for start in range(0, rows, chunk): + count = min(chunk, rows - start) + vectors = rng.random((count, dim), dtype=np.float32) + table = pa.table( + { + "id": pa.array(range(start, start + count), type=pa.int64()), + "v": pa.FixedSizeListArray.from_arrays(pa.array(vectors.reshape(-1), type=pa.float32()), dim), + } + ) + lance.write_dataset(table, uri, mode="create" if start == 0 else "append", max_rows_per_file=25_000) + dataset = lance.dataset(uri) + dataset.create_index( + "v", + index_type="IVF_PQ", + num_partitions=INDEX_PARTITIONS, + num_sub_vectors=INDEX_SUB_VECTORS, + ) + + +def indexed_query_stage(uri: str) -> dict[str, Any]: + """Run the same vector search repeatedly and time each one. + + ``read_lance`` with ``nearest`` plans a *single* index-driven scan task, so + this is not about concurrency: it is about whether the index is loaded once + per process or once per query. Shared, the first query pays for the load and + the rest are served from the session's index cache; unshared, every query + re-reads and re-decodes the index. + """ + import numpy as np + + rng = np.random.default_rng(7) + durations: list[float] = [] + sessions: list[int] = [] + for _ in range(INDEX_QUERIES): + query = rng.random(INDEX_DIM, dtype=np.float32) + df = daft_lance.read_lance(uri, default_scan_options={"nearest": {"column": "v", "q": query, "k": 10}}) + started = time.monotonic() + rows = df.to_arrow().num_rows + durations.append((time.monotonic() - started) * 1000) + if rows != 10: + raise SystemExit(f"the indexed query returned {rows} rows, expected 10") + report = report_from_worker() + sessions.append(max(int(row["session_size"]) for row in report.values())) + return { + "durations_ms": durations, + # Steady state, i.e. every query but the first: the first one is where a + # shared session does its one load, and including it would hide the + # difference the rest of them show. + "steady_median_ms": statistics.median(durations[1:]), + "first_ms": durations[0], + "worker_session_sizes": sessions, + } + + +def build_maintenance_table(uri: str, fragments: int, rows_per_fragment: int) -> None: + """A many-fragment table with a scalar index, for the maintenance stage. + + Compaction rewrites every row it touches, so this one carries no vector + column: the index is a BTREE over a string key, which still has to be + remapped at commit but costs seconds to build rather than gigabytes to + rewrite. Written with pyarrow rather than through Daft because Daft emits + ``large_string``, which lance will not index. + """ + total = fragments * rows_per_fragment + rows_per_chunk = max(rows_per_fragment, (1_000_000 // max(rows_per_fragment, 1)) * rows_per_fragment) + start = 0 + while start < total: + count = min(rows_per_chunk, total - start) + table = pa.table( + { + "i": pa.array(range(start, start + count), type=pa.int64()), + "k": pa.array([f"k{value:09d}" for value in range(start, start + count)], type=pa.string()), + } + ) + lance.write_dataset( + table, + uri, + mode="create" if start == 0 else "append", + max_rows_per_file=rows_per_fragment, + ) + start += count + lance.dataset(uri).create_scalar_index("k", index_type="BTREE") + + +def append_compaction_work(uri: str, groups: int, rows_per_append: int = MAINT_ROWS_PER_APPEND) -> int: + """Append small fragments so there is always something to compact. + + The data volume is deliberately trivial. Compaction consumes its own input, + so a stage that compacts whatever it finds works exactly once per table and + reports "nothing to do" on every later run -- which is how a reused table + turned a passing A/B into a failure. Appending is also the honest shape of + the problem: small fragments arriving from streaming writes are what + compaction exists for. + + None of this touches what the stage measures. A maintenance instance's cost + is one open of the table, and an open carries the whole fragment list + whatever the task then rewrites. + """ + dataset = lance.dataset(uri) + start = dataset.count_rows() + # Enough fragments that every instance in the pool gets work, with slack so + # the plan does not degenerate to one task. + count = max(groups, 2) * 4 * rows_per_append + table = pa.table( + { + "i": pa.array(range(start, start + count), type=pa.int64()), + "k": pa.array([f"k{value:09d}" for value in range(start, start + count)], type=pa.string()), + } + ) + lance.write_dataset(table, uri, mode="append", max_rows_per_file=rows_per_append) + return count // rows_per_append + + +class OpenCensus: + """Open the table the way a maintenance worker does, and report who did it. + + Decorated at the call site exactly like ``CompactionTaskUDF`` and the + ``update_columns`` handlers -- ``max_concurrency`` set, ``use_process`` + left alone -- because the question this answers is about placement: if + Daft puts the N instances in one process, the N opens collapse into one + shared session and sharing saves N-1 manifests; if it gives each its own + process, there is exactly one open per process and nothing to collapse. + The real maintenance UDFs cannot be asked this (they return rewrite + results, not pids), so this stands in for them. + """ + + def __init__(self, open_context: Any) -> None: + self.open_context = open_context + self._lance_ds: lance.LanceDataset | None = None + + @method.batch(return_dtype=daft.DataType.string()) + def __call__(self, column: daft.Series) -> list[str]: + if self._lance_ds is None: + # Once per instance, which is what the maintenance UDFs do. + self._lance_ds = self.open_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()), + "flag": os.environ.get("DAFT_LANCE_SHARE_SESSION"), + } + ) + return [payload] * len(column) + + +def maintenance_census(open_context: Any, concurrency: int) -> dict[str, Any]: + """Run ``concurrency`` maintenance-shaped instances and count processes.""" + census = daft.cls(OpenCensus, max_concurrency=concurrency) + 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} + return { + "requested": concurrency, + "instances": len(instances), + "processes": len({pid for pid, _ in instances}), + "sessions": len({row["session_obj_id"] for row in instances.values()}), + "flags": sorted({row["flag"] for row in instances.values()}, key=str), + "fragments": sorted({row["fragments"] for row in instances.values()}), + # Per process, so it is comparable between the sides regardless of how + # many processes each got. + "session_size_max": max(int(row["session_size"]) for row in instances.values()), + } + + +def maintenance_stage(data_dir: str, side: str, fragments: int, concurrency: int) -> dict[str, Any]: + """Compaction and ``update_columns`` on an indexed table, with the driver watched too. + + Two separate things are measured here, and they answer different halves of + "does an index make maintenance expensive": + + * the worker side, via :func:`maintenance_census`: how many processes the + concurrent instances land in, and what one of them holds; + * the driver side, directly, because under Ray the driver *is* this process: + what the compaction leaves in it. Compaction has to remap every index + entry onto the new fragment ids, and that happens in the commit, on the + driver, not in the tasks -- so this is where an indexed table's + maintenance cost shows up at all. The figure is manifest and index + together; ``test_shared_session.py`` is where the two are separated, + which needs the tasks and the commit pulled apart. + + Each side of the A/B gets its own copy of the table: compaction consumes its + own input, so whoever ran second would otherwise find nothing to compact. + The stage also makes its own work rather than compacting whatever happens to + be there, because the tables outlive a run -- see :func:`append_compaction_work`. + """ + uri = str(Path(data_dir) / f"maint-{side}-{fragments}.lance") + appended = append_compaction_work(uri, concurrency) + session = get_shared_session() + + def snapshot() -> dict[str, int]: + # ``size_bytes()`` only. A session also reports + # ``index_cache_size_bytes()``, but that is the accounting lance + # enforces its index budget with, not resident bytes: measured on a + # 200k-row IVF_PQ index it reads 37 MB where both ``size_bytes()`` and + # the process's own RSS move by about 5 MB per loaded copy. Printing it + # as memory would overstate an index's cost sevenfold. + return {"total": session.size_bytes()} + + driver: dict[str, dict[str, int]] = {"start": snapshot()} + handle = construct_lance_dataset_handle(uri) + driver["after_open"] = snapshot() + result: dict[str, Any] = { + "uri": uri, + "fragments_appended": appended, + "fragments_before": len(handle.dataset.get_fragments()), + "rows": handle.dataset.count_rows(), + "indices": [index["name"] for index in handle.dataset.list_indices()], + "concurrency": concurrency, + "census": maintenance_census(handle.worker_open_context(), concurrency), + } + driver["after_census"] = snapshot() + + started = time.monotonic() + metrics = daft_lance.compact_files( + uri, + # Keyed to the appended fragments, not to the table's own: the large + # fragments already sit above this and are not candidates, so only the + # freshly appended ones are rewritten. + compaction_options={"target_rows_per_fragment": MAINT_TARGET_FRAGMENT_ROWS}, + partition_num=concurrency, + concurrency=concurrency, + ) + result["compact_seconds"] = time.monotonic() - started + driver["after_compact"] = snapshot() + if metrics is None: + raise SystemExit( + f"compaction found nothing to do on {uri} even after appending {appended} small fragments; " + "check that target_rows_per_fragment is above the appended fragments' size" + ) + result["compaction"] = { + "fragments_removed": metrics.fragments_removed, + "fragments_added": metrics.fragments_added, + } + + started = time.monotonic() + updated = daft_lance.update_columns( + uri, + transform={"i": "i + 0"}, + columns=["i"], + max_concurrency=concurrency, + ) + result["update_seconds"] = time.monotonic() - started + driver["after_update"] = snapshot() + if updated.rows_updated != result["rows"]: + raise SystemExit(f"update_columns touched {updated.rows_updated} rows, expected {result['rows']}") + result["rows_updated"] = updated.rows_updated + result["driver_session"] = driver + return result + + +def report_from_worker() -> dict[int, dict[str, Any]]: + """Run a one-row query so the worker reports its current state back.""" + 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 sampled_rss_peaks(event_log_dir: str) -> dict[int, dict[str, Any]]: + """Peak ``process.memory.rss`` per pid, from Daft's own 200ms sampler.""" + peaks: dict[int, dict[str, Any]] = {} + for path in sorted(Path(event_log_dir).glob("*/events.jsonl")): + for line in path.read_text().splitlines(): + try: + record = json.loads(line) + except ValueError: + continue + if record.get("event") != "process_stats": + continue + rss = record.get("metrics", {}).get(RSS_KEY) + pid = record.get("pid") + if rss is None or pid is None: + continue + entry = peaks.setdefault(int(pid), {"component": record.get("component"), "peak": 0}) + entry["peak"] = max(entry["peak"], int(rss)) + return peaks + + +def payload_sizes(uri: str) -> dict[str, int]: + """Bytes a scan task used to carry versus what it carries now. + + ``LanceDataset.__reduce__`` ships the re-encoded manifest, so closing over a + dataset -- which ``read_blobs`` used to do -- cost that much per task. The + replacement is a uri plus the pinned open kwargs. + """ + dataset = lance.dataset(uri) + open_kwargs = {"version": dataset.version, "storage_options": dataset._storage_options} + return { + "pickled_dataset": len(pickle.dumps(dataset)), + "pickled_open_kwargs": len(pickle.dumps((uri, open_kwargs))), + } + + +def run_child(args: argparse.Namespace) -> int: + """One half of the A/B: a full read + read_blobs + write under Ray. + + Either starts a cluster of its own (the local ``--mode both`` path, which + spawns this twice) or attaches to one that is already there, which is how it + runs as a RayJob entrypoint: there the two sides are separate jobs and the + flag arrives through the pod spec instead of process inheritance. + """ + address = os.environ.get("RAY_ADDRESS") + if address: + # An attached run must not ask for CPUs: the cluster already decided, + # and Ray ignores num_cpus with a warning when it is only connecting. + # Concurrency comes from the worker's ``ray start --num-cpus``, which is + # a logical count and need not match the pod's cpu limit. + ray.init() + else: + ray.init(num_cpus=args.parallel, include_dashboard=False) + daft.set_runner_ray(noop_if_initialized=True) + daft.set_execution_config(scantask_max_parallel=args.parallel) + + uri = str(Path(args.data_dir) / f"table-{args.fragments}.lance") + blob_uri = str(Path(args.data_dir) / "blobs.lance") + dataset = lance.dataset(uri) + result: dict[str, Any] = { + "mode": args.mode, + "parallel": args.parallel, + "flag": os.environ.get("DAFT_LANCE_SHARE_SESSION"), + "ray_address": address, + "rows_per_fragment": args.rows_per_fragment, + "cluster_cpus": ray.cluster_resources().get("CPU"), + "fragments_at_read": len(dataset.get_fragments()), + "manifest_bytes": len(dataset._ds.serialized_manifest()), + } + + durations: dict[str, float] = {} + before = report_from_worker() + + # The pipeline under test: read Lance, write Parquet. One task per fragment + # group, `parallel` of them open at once, and each task holds a parquet + # writer's buffers at the same time as its manifest -- which is the point of + # measuring this shape rather than an aggregation. An aggregation keeps + # almost nothing resident, so it flatters the worker; a write makes the + # manifests compete with the writer for the same memory limit, which is what + # a real backfill does. + # + # Deliberately not count_rows(): count pushdown would collapse the whole + # plan into a single task that opens no fragment at all. + # + # count_rows() here is the dataset's own metadata, read on the driver before + # the scan: cheap, and it says how much the scan below had to touch. The two + # sides see slightly different totals because each appends in its lance + # write stage, which is why this is recorded rather than asserted against a + # formula. + rows_total = dataset.count_rows() + result["rows_total"] = rows_total + # Two knobs for keeping a large table affordable, and neither weakens the + # measurement: a manifest is the table's whole fragment list, so an open + # costs the same whether its task reads one fragment or fifty, one row or + # ten thousand. --read-fraction prunes whole fragments (``i`` is written + # strictly increasing, so a cutoff is a fragment prefix) and + # --fragment-group-size puts several fragments in one task. + # + # What does have to survive is the overlap: enough tasks to keep + # --parallel of them in flight, each long enough that their opens coexist. + # Both counts are reported so that is checkable rather than hoped for. + result["read_fraction"] = args.read_fraction + result["fragment_group_size"] = args.fragment_group_size + reader = daft_lance.read_lance(uri, fragment_group_size=args.fragment_group_size) + if args.read_fraction < 1.0: + cutoff = int(rows_total * args.read_fraction) + result["read_cutoff"] = cutoff + # Arithmetic, not a second pass over 100k fragments: build_table fills + # fragments in order, so a cutoff of N rows is the first N/rows + # fragments. + touched = -(-cutoff // args.rows_per_fragment) + result["fragments_touched"] = touched + result["tasks_estimated"] = -(-touched // max(args.fragment_group_size or 1, 1)) + reader = reader.where(daft.col("i") < cutoff) + else: + result["fragments_touched"] = result["fragments_at_read"] + result["tasks_estimated"] = -(-result["fragments_at_read"] // max(args.fragment_group_size or 1, 1)) + if result["tasks_estimated"] < 2 * args.parallel: + raise SystemExit( + f"lance -> parquet would plan about {result['tasks_estimated']} tasks at concurrency " + f"{args.parallel}: too few for the opens to overlap, so raise --read-fraction or lower " + "--fragment-group-size" + ) + # One output directory per side. Overwrite rather than append so a rerun + # does not measure the previous run's files, and so the two sides never read + # each other's output back as input to anything. + parquet_dir = Path(args.data_dir) / f"out-parquet-{args.mode}" + started = time.monotonic() + written = reader.write_parquet(str(parquet_dir), write_mode="overwrite").to_pydict() + durations["read_to_parquet"] = time.monotonic() - started + files = written.get("path") or [] + if not files: + raise SystemExit("the pipeline wrote no parquet files, so it never reached the data") + result["parquet_files"] = len(files) + result["parquet_bytes"] = sum(path.stat().st_size for path in parquet_dir.rglob("*.parquet")) + after_read = report_from_worker() + + started = time.monotonic() + blob_ds = lance.dataset(blob_uri) + blob_df = daft.read_lance(blob_uri, default_scan_options={"with_row_id": True}) + result["blobs_read"] = len(read_blobs(blob_df, blob_ds, "blob").to_pydict()["blob"]) + durations["read_blobs"] = time.monotonic() - started + after_blobs = report_from_worker() + + # A small append with tiny files, so the writer flushes `write_fragments` + # several times; each flush is a manifest load when no session is passed. + started = time.monotonic() + daft.from_pydict({"i": list(range(WRITE_ROWS))}).write_lance(uri, mode="append", max_rows_per_file=1).collect() + durations["write"] = time.monotonic() - started + after_write = report_from_worker() + + # Last, because it loads an index into whatever session it gets and would + # otherwise contaminate the manifest measurements above. + indexed_uri = str(Path(args.data_dir) / f"indexed-{INDEX_ROWS}.lance") + started = time.monotonic() + result["indexed"] = indexed_query_stage(indexed_uri) + durations["indexed_queries"] = time.monotonic() - started + + # After everything above, and on a table of its own, because compaction + # rewrites its input: it must not disturb the read table the stages above + # measured, and its own driver-side numbers are read as deltas so the + # earlier stages' caches do not matter. + started = time.monotonic() + result["maintenance"] = maintenance_stage( + args.data_dir, + args.mode, + args.maint_fragments, + args.maint_concurrency, + ) + durations["maintenance"] = time.monotonic() - started + result["durations"] = durations + + # Calibrated here rather than by whoever orchestrates the A/B: the threshold + # then comes from the same architecture, the same lance build and the same + # allocator as the measurement it gates. Last, not first: it holds several + # unshared copies of the manifest, which at 100k fragments is a few hundred + # megabytes, and the driver needs that headroom while it is planning one + # scan task per fragment. Nothing above depends on the number. + result["per_open_bytes"] = calibrate_per_open_bytes(uri, args.parallel) + + workers: dict[str, Any] = {} + for pid, row in after_write.items(): + workers[str(pid)] = { + "flag": row["flag"], + "same_session": row["same_session"], + "parallelism": row.get("parallelism"), + "session_obj_ids": sorted( + { + stage[pid]["session_obj_id"] + for stage in (before, after_read, after_blobs, after_write) + if pid in stage + } + ), + "maxrss": { + "before": before.get(pid, {}).get("maxrss"), + "after_read": after_read.get(pid, {}).get("maxrss"), + "after_blobs": after_blobs.get(pid, {}).get("maxrss"), + "after_write": after_write.get(pid, {}).get("maxrss"), + }, + "session_size": { + "before": before.get(pid, {}).get("session_size"), + "after_read": after_read.get(pid, {}).get("session_size"), + "after_blobs": after_blobs.get(pid, {}).get("session_size"), + "after_write": after_write.get(pid, {}).get("session_size"), + }, + } + result["workers"] = workers + result["worker_pids_per_stage"] = { + "before": sorted(before), + "after_read": sorted(after_read), + "after_blobs": sorted(after_blobs), + "after_write": sorted(after_write), + } + event_log_dir = os.environ.get("DAFT_EVENT_LOG_DIR") + result["sampled_rss_peak"] = ( + {str(pid): entry for pid, entry in sampled_rss_peaks(event_log_dir).items()} if event_log_dir else {} + ) + print(RESULT_PREFIX + json.dumps(result)) + return 0 + + +def calibrate_per_open_bytes(uri: str, opens: int) -> float: + """What one *additional* unshared open of this table costs, measured here and now. + + Thresholds are derived from this rather than hardcoded, so the script stays + honest on a different table or a different lance version. + + Marginal rather than average, and capped at a handful of opens, for two + reasons. The cap: at 100k fragments one open is around 64 MB, so holding + ``--parallel`` of them puts half a gigabyte in the driver -- a calibration + must not be able to push over the process it is calibrating, and at 100k + fragments on a 3Gi head it did. Marginal: the first open also pays one-time + costs (the object-store client, lance's own statics) that do not repeat per + open, so averaging them in overstates the per-open cost and with it the + threshold. + """ + first = [lance.dataset(uri, session=LanceSession())] + extra = max(min(opens - 1, CALIBRATION_OPENS), 1) + before = maxrss_bytes() + held = [lance.dataset(uri, session=LanceSession()) for _ in range(extra)] + after = maxrss_bytes() + assert len(first) == 1 and len(held) == extra + return (after - before) / extra + + +def spawn_child(mode: str, args: argparse.Namespace, event_log_dir: Path) -> dict[str, Any]: + env = dict(os.environ) + env.update( + { + "DAFT_LANCE_SHARE_SESSION": "1" if mode == "shared" else "0", + "DAFT_RUNNER": "ray", + "DAFT_PROGRESS_BAR": "0", + "DAFT_ANALYTICS_ENABLED": "0", + "DAFT_PROCESS_MONITOR_ENABLED": "1", + "DAFT_EVENT_LOG_DIR": str(event_log_dir), + # Ray notices when the driver's process tree was started by `uv run` + # -- which it is, via `make bench` -- and then replicates that + # environment to the workers: it packages the whole working + # directory and reinstalls it in a fresh venv per worker. That costs + # minutes here and has hung outright. The workers already inherit + # this venv's interpreter from the raylet, so turn it off. + "RAY_ENABLE_UV_RUN_RUNTIME_ENV": "0", + } + ) + command = [ + sys.executable, + str(Path(__file__).resolve()), + "--mode", + mode, + "--fragments", + str(args.fragments), + "--parallel", + str(args.parallel), + "--data-dir", + str(args.data_dir), + "--rows-per-fragment", + str(args.rows_per_fragment), + "--maint-fragments", + str(args.maint_fragments), + "--maint-concurrency", + str(args.maint_concurrency), + "--read-fraction", + str(args.read_fraction), + ] + if args.fragment_group_size: + command += ["--fragment-group-size", str(args.fragment_group_size)] + print(f"\n=== running {mode} (DAFT_LANCE_SHARE_SESSION={env['DAFT_LANCE_SHARE_SESSION']}) ===", flush=True) + try: + process = subprocess.run(command, env=env, capture_output=True, text=True, check=False, timeout=CHILD_TIMEOUT) + except subprocess.TimeoutExpired as expired: + sys.stderr.write(expired.stderr or "") + raise SystemExit( + f"{mode} child did not finish within {CHILD_TIMEOUT}s. The read is one task per fragment, so a large " + "--fragments takes a long time; lower it, or raise CHILD_TIMEOUT." + ) from expired + payload = None + for line in process.stdout.splitlines(): + if line.startswith(RESULT_PREFIX): + payload = json.loads(line[len(RESULT_PREFIX) :]) + if payload is None: + sys.stdout.write(process.stdout) + sys.stderr.write(process.stderr) + raise SystemExit(f"{mode} child produced no result (exit code {process.returncode})") + return payload + + +def worker_row(result: dict[str, Any]) -> tuple[str, dict[str, Any]]: + workers = result["workers"] + if len(workers) != 1: + raise SystemExit(f"expected one worker process, got {sorted(workers)}; run this on a single-node cluster") + return next(iter(workers.items())) + + +def compare(baseline: dict[str, Any], shared: dict[str, Any], per_open_bytes: float) -> list[str]: + """Print the comparison and return the failures, if any.""" + failures: list[str] = [] + baseline_pid, baseline_worker = worker_row(baseline) + shared_pid, shared_worker = worker_row(shared) + + # Before anything else: the A/B only means something if the flag actually + # reached the worker. Ray actors do not inherit the driver's environment by + # contract, so this must be observed, never assumed. + for name, result, worker, expected in ( + ("baseline", baseline, baseline_worker, "0"), + ("shared", shared, shared_worker, "1"), + ): + if worker["flag"] != expected: + failures.append( + f"{name}: worker saw DAFT_LANCE_SHARE_SESSION={worker['flag']!r}, driver set {expected!r} -- " + "the environment did not reach the worker, so these numbers compare nothing" + ) + if result["flag"] != expected: + failures.append(f"{name}: driver saw flag {result['flag']!r}, expected {expected!r}") + + # Concurrency decides how many manifests can be resident at once, so the two + # sides have to have had the same ceiling, and the reader has to be able to + # see what it was. + for name, result, worker in ( + ("baseline", baseline, baseline_worker), + ("shared", shared, shared_worker), + ): + limits = worker.get("parallelism") or {} + print( + f"{name:10} rows/fragment {result.get('rows_per_fragment')}, " + f"scantask_max_parallel {result['parallel']}, ray cluster CPUs {result.get('cluster_cpus')}, " + f"worker sees cpu_count={limits.get('cpu_count')} affinity={limits.get('affinity')} " + f"cgroup_cpu_max={limits.get('cgroup_cpu_max')}" + ) + for key in ("parallel", "rows_per_fragment", "read_fraction", "fragment_group_size"): + if baseline.get(key) != shared.get(key): + failures.append( + f"the two sides ran with different {key} ({baseline.get(key)!r} vs {shared.get(key)!r}), " + "so their memory is not comparable" + ) + + def delta(worker: dict[str, Any], field: str, stage: str) -> int: + return int(worker[field][stage]) - int(worker[field]["before"]) + + rows = [] + for name, result, worker, pid in ( + ("baseline", baseline, baseline_worker, baseline_pid), + ("shared", shared, shared_worker, shared_pid), + ): + sampled = result["sampled_rss_peak"].get(pid, {}) + rows.append( + { + "name": name, + "pid": pid, + "fragments": result["fragments_at_read"], + "rss_delta_read": delta(worker, "maxrss", "after_read"), + "rss_delta_total": delta(worker, "maxrss", "after_write"), + "sampled_peak": sampled.get("peak"), + "sampled_component": sampled.get("component"), + "session_delta_read": delta(worker, "session_size", "after_read"), + "session_delta_blobs": delta(worker, "session_size", "after_blobs"), + "session_delta_total": delta(worker, "session_size", "after_write"), + "same_session": worker["same_session"], + "sessions_seen": len(worker["session_obj_ids"]), + } + ) + + print(f"\nper-open cost of this table, measured here: {per_open_bytes / MB:.2f} MB") + for name, result in (("baseline", baseline), ("shared", shared)): + stages = ", ".join(f"{stage} {seconds:.1f}s" for stage, seconds in result["durations"].items()) + print(f"{name:10} {stages}") + print( + f"manifest: {baseline['manifest_bytes'] / 1024:.1f} KiB, table rows: {baseline.get('rows_total')}, " + f"parquet written: {baseline.get('parquet_files')} files / " + f"{(baseline.get('parquet_bytes') or 0) / MB:.1f} MB" + ) + print( + f"lance -> parquet: {baseline.get('read_fraction')} of the table, fragment_group_size " + f"{baseline.get('fragment_group_size')}, so about {baseline.get('fragments_touched')} fragments in " + f"{baseline.get('tasks_estimated')} tasks at concurrency {baseline['parallel']} " + "(a manifest costs the same per open whatever the task reads)" + ) + print( + f"{'':10} {'pid':>7} {'frags':>7} {'rssDpq':>9} {'rssDall':>9} {'rssPeak':>9} " + f"{'sessDpq':>10} {'sessDblob':>10} {'sessDall':>10}" + ) + for row in rows: + sampled = "-" if row["sampled_peak"] is None else f"{row['sampled_peak'] / MB:.1f}" + print( + f"{row['name']:10} {row['pid']:>7} {row['fragments']:>7} " + f"{row['rss_delta_read'] / MB:>9.1f} {row['rss_delta_total'] / MB:>9.1f} {sampled:>9} " + f"{row['session_delta_read'] / MB:>10.3f} " + f"{(row['session_delta_blobs'] - row['session_delta_read']) / MB:>10.3f} " + f"{row['session_delta_total'] / MB:>10.3f}" + ) + print( + "(MB. rssDpq: the worker's ru_maxrss growth across the lance-to-parquet stage. rssPeak: Daft's\n" + " own sampled rss peak for that pid, an independent reading of the same process. sessD: growth of\n" + " the process-wide lance session, which is zero whenever an open did not go through it.)" + ) + + baseline_row, shared_row = rows[0], rows[1] + + # Mechanism, with no RSS noise in it: with sharing on, a scan task's open + # lands in the process-wide session, so its size grows. With sharing off + # the global session is never touched at all. + if baseline_row["session_delta_total"] != 0: + failures.append( + f"baseline: process session grew by {baseline_row['session_delta_total']} bytes; " + "with sharing off no open should touch it" + ) + if shared_row["session_delta_read"] <= 0: + failures.append( + "shared: the lance-to-parquet stage did not grow the process session, so its scan tasks bypassed it" + ) + if shared_row["session_delta_total"] <= shared_row["session_delta_read"]: + failures.append("shared: read_blobs and the write did not add to the process session") + if not shared_row["same_session"] or shared_row["sessions_seen"] != 1: + failures.append( + f"shared: worker reported {shared_row['sessions_seen']} distinct sessions, expected one per process" + ) + if baseline_row["same_session"]: + failures.append("baseline: the kill switch did not take effect in the worker") + + # The memory claim itself. Only the per-open cost beyond the first is + # saved, and never all of it: a task holds its manifest just while its + # dataset is open, so the opens overlap only partly and the peak never + # reaches the full K-way sum. Measured at --parallel 24, the saving came in + # at 33-41 MB against a 67 MB theoretical maximum, so demanding 30% of it + # keeps roughly a 2x margin over the observed value while still failing + # outright if sharing stops working, which takes the saving to zero. + # + # Measured over the lance-to-parquet stage alone, not the whole run. + # ``ru_maxrss`` is a high-water mark, and the later stages allocate about ten + # megabytes of their own: on the shared side, whose peak was lower, that + # pushes the mark back up to roughly where the baseline already was, hiding + # the very difference being measured. That stage is also the only one whose + # concurrency is pinned, and the one the pipeline this benchmark exists for + # actually runs, so it is the one the threshold is derived from. + # The indexed stage, which is a different claim from everything above: a + # `nearest` read plans one index-driven scan task, so nothing multiplies by + # concurrency. What sharing changes is that the index is loaded once per + # process instead of once per query. + baseline_index, shared_index = baseline.get("indexed"), shared.get("indexed") + if baseline_index and shared_index: + print( + f"\nindexed vector search, {INDEX_QUERIES} queries of the same table:\n" + f"{'':10} {'first':>9} {'steady':>9} {'sessionEnd':>12}" + ) + for name, side in (("baseline", baseline_index), ("shared", shared_index)): + print( + f"{name:10} {side['first_ms']:>8.1f}ms {side['steady_median_ms']:>8.1f}ms " + f"{side['worker_session_sizes'][-1] / 1024:>10.1f}KiB" + ) + print( + "(first: the query that pays for loading the index. steady: the median of the rest. sessionEnd: what\n" + " the worker's process-wide session holds at the end, so zero means every query loaded the index\n" + " into a session of its own and threw it away.)" + ) + baseline_end = baseline_index["worker_session_sizes"][-1] + shared_end = shared_index["worker_session_sizes"][-1] + if shared_end <= 0: + failures.append("shared: the indexed queries left nothing in the process session, so none of them used it") + # Compared against the shared side rather than against zero: an empty + # lance Session is not zero bytes, it is about 152 of bookkeeping, so + # "holds no index" means "holds orders of magnitude less than one". + elif baseline_end > shared_end // 10: + failures.append( + f"baseline: the process session holds {baseline_end / 1024:.1f} KiB after the indexed queries, " + f"against {shared_end / 1024:.1f} KiB shared; with sharing off no query should have filled it" + ) + # Deliberately not a hard threshold on the ratio: it depends on index + # size and on how far the storage is, and on a local disk with a warm + # page cache it is the smallest it will ever be. Measured locally at + # 200k rows: 5.6-7.0ms shared against 14.0-14.7ms unshared. + ratio = baseline_index["steady_median_ms"] / max(shared_index["steady_median_ms"], 1e-9) + print( + f"steady-state queries are {ratio:.2f}x slower without a shared session " + f"({baseline_index['steady_median_ms']:.1f}ms vs {shared_index['steady_median_ms']:.1f}ms); " + "the gap grows with index size and storage latency" + ) + if ratio < 1.0: + failures.append( + f"the unshared side was not slower on repeated indexed queries (ratio {ratio:.2f}x), so reloading " + "the index cost nothing measurable here -- check that the index exists and that the queries used it" + ) + + # The maintenance path, which is a different shape of claim again. A scan + # task shares a process with `parallel`-1 others, so its manifest is one of + # many; a maintenance instance is given a process to itself, so its open is + # the only one there. Both are measured, because only measuring says which. + baseline_maint, shared_maint = baseline.get("maintenance"), shared.get("maintenance") + if baseline_maint and shared_maint: + print( + f"\nmaintenance on an indexed table " + f"({shared_maint['fragments_before']} fragments, {shared_maint['rows']} rows, " + f"indices {shared_maint['indices']}), concurrency {shared_maint['concurrency']}:\n" + f"{'':10} {'insts':>6} {'procs':>6} {'sessPerProc':>12} {'driverCommit':>13} " + f"{'compact':>9} {'update':>9}" + ) + for name, side in (("baseline", baseline_maint), ("shared", shared_maint)): + census = side["census"] + driver = side["driver_session"] + driver_commit = driver["after_compact"]["total"] - driver["after_census"]["total"] + print( + f"{name:10} {census['instances']:>6} {census['processes']:>6} " + f"{census['session_size_max'] / 1024:>10.1f}KiB {driver_commit / 1024:>11.1f}KiB " + f"{side['compact_seconds']:>8.1f}s {side['update_seconds']:>8.1f}s" + ) + print( + "(insts/procs: concurrent maintenance instances, and how many processes they landed in. sessPerProc:\n" + " the largest process-wide session any one of them held. driverCommit: what compaction added to the\n" + " *driver's* session, manifest and remapped index together -- remapping an index onto the new fragment\n" + " ids happens in the commit, on the driver, not in the tasks.)" + ) + for name, side, expected in (("baseline", baseline_maint, "0"), ("shared", shared_maint, "1")): + census = side["census"] + if census["flags"] != [expected]: + failures.append( + f"{name}: the maintenance instances saw DAFT_LANCE_SHARE_SESSION={census['flags']}, " + f"expected [{expected!r}] -- the flag did not reach them" + ) + if census["instances"] < 2: + failures.append( + f"{name}: only {census['instances']} maintenance instance ran, so concurrency was never " + "reached and this stage measured nothing" + ) + shared_census, baseline_census = shared_maint["census"], baseline_maint["census"] + if shared_census["session_size_max"] <= baseline_census["session_size_max"]: + failures.append( + "shared: a maintenance instance's process session is no larger than the baseline's, so the " + "maintenance open did not go through it" + ) + if baseline_census["sessions"] != baseline_census["instances"]: + failures.append( + f"baseline: {baseline_census['instances']} maintenance instances reported " + f"{baseline_census['sessions']} sessions; with sharing off each open gets its own" + ) + shared_driver = shared_maint["driver_session"] + shared_driver_commit = shared_driver["after_compact"]["total"] - shared_driver["after_census"]["total"] + if shared_driver_commit <= 0: + failures.append( + "shared: compaction left nothing in the driver's session, so the driver's planning and commit " + "opened outside the process session" + ) + baseline_driver = baseline_maint["driver_session"] + if baseline_driver["after_compact"]["total"] != baseline_driver["start"]["total"]: + failures.append( + "baseline: the driver's process session grew during maintenance; with sharing off every open " + "should get a session of its own" + ) + # Not a failure, a finding: if Daft ever fuses these instances into one + # process, the maintenance path acquires a saving to measure, and this + # stage has to be rewritten to measure it rather than quietly passing. + if shared_census["processes"] < shared_census["instances"]: + print( + f"NOTE: {shared_census['instances']} instances shared {shared_census['processes']} processes. " + "Daft used to give each one its own, which is why this stage asserts nothing about concurrency " + "multiplying manifests. It now does multiply them -- this stage should be extended to measure it." + ) + else: + print( + f"Each of the {shared_census['instances']} instances got its own process, so a maintenance worker " + f"opens the table exactly once ({shared_census['session_size_max'] / 1024:.1f} KiB of session) no " + "matter the concurrency: nothing here multiplies by it, and the index cost sits in the driver's " + "commit instead." + ) + + extra_opens = max(shared["parallel"] - 1, 1) + expected_saving = 0.3 * extra_opens * per_open_bytes + observed_saving = baseline_row["rss_delta_read"] - shared_row["rss_delta_read"] + print( + f"\nsaving over lance -> parquet: {observed_saving / MB:.1f} MB observed, " + f"{expected_saving / MB:.1f} MB required " + f"(30% of {extra_opens} extra opens x {per_open_bytes / MB:.2f} MB). " + f"Over the whole run: {(baseline_row['rss_delta_total'] - shared_row['rss_delta_total']) / MB:.1f} MB, " + "which the later stages' own allocations dilute." + ) + if expected_saving < RSS_NOISE_FLOOR: + print( + f"NOT MEASURED: {extra_opens} extra opens of this table only amount to " + f"{extra_opens * per_open_bytes / MB:.1f} MB, which RSS cannot resolve. The deterministic checks above " + "still hold; raise --fragments (or --parallel) to put a number on the saving." + ) + elif observed_saving < expected_saving: + failures.append( + f"the shared run saved {observed_saving / MB:.1f} MB, expected at least {expected_saving / MB:.1f} MB" + ) + return failures + + +def run_prepare(args: argparse.Namespace) -> int: + """Write the tables and nothing else, for a one-off prep job. + + Both sides of a Kubernetes A/B read the *same* copy, so the data has to be + built once into a shared volume: lance names its data files with random + uuids, so a per-pod copy would leave the driver's manifest pointing at file + names the worker's copy does not have. + """ + data_dir = Path(args.data_dir) + data_dir.mkdir(parents=True, exist_ok=True) + uri = data_dir / f"table-{args.fragments}.lance" + blob_uri = data_dir / "blobs.lance" + if uri.exists(): + print(f"{uri} already exists, keeping it") + else: + build_table(str(uri), args.fragments, args.rows_per_fragment) + print(f"wrote {args.fragments} fragments x {args.rows_per_fragment} rows to {uri}") + if blob_uri.exists(): + print(f"{blob_uri} already exists, keeping it") + else: + build_blob_table(str(blob_uri), min(args.fragments, 64)) + print(f"wrote {blob_uri}") + indexed_uri = data_dir / f"indexed-{INDEX_ROWS}.lance" + if indexed_uri.exists(): + print(f"{indexed_uri} already exists, keeping it") + else: + build_indexed_table(str(indexed_uri)) + print(f"wrote {INDEX_ROWS} rows with an IVF_PQ index to {indexed_uri}") + # One per side: the maintenance stage compacts its table, so a shared copy + # would leave whichever side ran second with nothing to compact. They are + # built identically, which is what makes the two sides comparable. + for side in ("baseline", "shared"): + maint_uri = data_dir / f"maint-{side}-{args.maint_fragments}.lance" + if maint_uri.exists(): + print(f"{maint_uri} already exists, keeping it") + else: + build_maintenance_table(str(maint_uri), args.maint_fragments, args.rows_per_fragment) + print( + f"wrote {args.maint_fragments} fragments x {args.rows_per_fragment} rows " + f"with a BTREE index to {maint_uri}" + ) + return 0 + + +def run_compare(args: argparse.Namespace) -> int: + """Compare two result files produced by separate runs of ``--mode``.""" + baseline = json.loads(Path(args.baseline).read_text()) + shared = json.loads(Path(args.shared).read_text()) + failures = compare(baseline, shared, float(baseline["per_open_bytes"])) + if failures: + print("\nFAILED:") + for failure in failures: + print(f" - {failure}") + return 1 + print("\nOK: sharing one session cut the worker's memory as expected.") + return 0 + + +def run_parent(args: argparse.Namespace) -> int: + address = os.environ.get("RAY_ADDRESS") + if address and not address.startswith("local"): + raise SystemExit( + f"RAY_ADDRESS={address!r} points at an existing cluster. The A/B needs the worker to inherit this " + "process's environment, which only holds for a cluster started here, so refusing to report numbers." + ) + + data_dir = Path(args.data_dir) + data_dir.mkdir(parents=True, exist_ok=True) + uri = str(data_dir / f"table-{args.fragments}.lance") + blob_uri = str(data_dir / "blobs.lance") + if not Path(uri).exists(): + print(f"writing a {args.fragments}-fragment table to {uri} (reused by later runs)", flush=True) + build_table(uri, args.fragments, args.rows_per_fragment) + if not Path(blob_uri).exists(): + build_blob_table(blob_uri, min(args.fragments, 64)) + indexed_uri = str(data_dir / f"indexed-{INDEX_ROWS}.lance") + if not Path(indexed_uri).exists(): + print(f"writing a {INDEX_ROWS}-row table with an IVF_PQ index to {indexed_uri}", flush=True) + build_indexed_table(indexed_uri) + for side in ("baseline", "shared"): + maint_uri = data_dir / f"maint-{side}-{args.maint_fragments}.lance" + if not maint_uri.exists(): + print(f"writing the {side} maintenance table to {maint_uri}", flush=True) + build_maintenance_table(str(maint_uri), args.maint_fragments, args.rows_per_fragment) + + # Each child gets its own event-log directory: the sampler keys records by + # pid, and pids are not unique across runs. + with tempfile.TemporaryDirectory(prefix="lance-session-events-") as events: + baseline = spawn_child("baseline", args, Path(events) / "baseline") + shared = spawn_child("shared", args, Path(events) / "shared") + + payload = payload_sizes(uri) + small = payload_sizes(blob_uri) + print( + f"\nper-task payload: pickling the dataset costs {payload['pickled_dataset'] / 1024:.1f} KiB at " + f"{args.fragments} fragments versus {small['pickled_dataset'] / 1024:.1f} KiB at " + f"{min(args.fragments, 64)}; the uri plus pinned open kwargs that replaced it costs " + f"{payload['pickled_open_kwargs']} and {small['pickled_open_kwargs']} bytes." + ) + + per_open_bytes = float(baseline["per_open_bytes"]) + failures = compare(baseline, shared, per_open_bytes) + if args.json: + print( + json.dumps( + {"per_open_bytes": per_open_bytes, "baseline": baseline, "shared": shared, "failures": failures}, + indent=2, + ) + ) + if failures: + print("\nFAILED:") + for failure in failures: + print(f" - {failure}") + return 1 + print("\nOK: sharing one session cut the worker's memory as expected.") + return 0 + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + parser.add_argument( + "--fragments", + type=int, + default=4000, + help=( + "fragments in the table under test. The per-open cost grows with this, and so does the saving, but " + "the read is one task per fragment, so the run time grows with it too" + ), + ) + parser.add_argument( + "--rows-per-fragment", + type=int, + default=ROWS_PER_FRAGMENT, + help=( + "rows in each fragment. Below a few thousand a scan task finishes too fast for the opens to overlap, " + "so fewer manifests are resident at once than --parallel suggests" + ), + ) + parser.add_argument( + "--parallel", + type=int, + default=24, + help=( + "Ray CPUs and concurrent scan tasks. This, not --fragments, is the lever that makes the saving " + "stand out of the RSS noise: it multiplies how many manifests are held at once without adding " + "tasks, so the run takes no longer. It may exceed the machine's core count" + ), + ) + parser.add_argument( + "--read-fraction", + type=float, + default=1.0, + help=( + "fraction of the table the read stage actually reads, as a prefix of the sorted key. A manifest " + "costs the same per open whatever the task reads, so this cuts the run time without touching the " + "measurement -- it is how a 100k-fragment table stays affordable" + ), + ) + parser.add_argument( + "--fragment-group-size", + type=int, + default=None, + help=( + "fragments per scan task (read_lance's own option). Fewer, longer tasks; each still opens the " + "dataset once and still carries the whole fragment list, so the per-open cost is unchanged" + ), + ) + parser.add_argument( + "--maint-fragments", + type=int, + default=MAINT_FRAGMENTS, + help=( + "fragments in the maintenance table. Compaction rewrites all of them, so this is the stage's whole " + "cost; it does not change what the stage concludes, since a maintenance worker opens the table once" + ), + ) + parser.add_argument( + "--maint-concurrency", + type=int, + default=MAINT_CONCURRENCY, + help=( + "concurrent compaction / update_columns instances. Each one is a separate process, so this is a " + "process count, not a thread count -- keep it near the pod's real cpu limit" + ), + ) + parser.add_argument( + "--data-dir", + default=str(Path(tempfile.gettempdir()) / "daft-lance-session-bench"), + help="where the test tables live; reused across runs", + ) + parser.add_argument( + "--mode", + choices=["both", "shared", "baseline", "prepare", "compare"], + default="both", + help=( + "both: run the whole A/B locally in two subprocesses. shared/baseline: one side only, for an " + "orchestrator that runs each side separately (a RayJob per side). prepare: only write the tables. " + "compare: read two sides' result files and judge them" + ), + ) + parser.add_argument("--baseline", help="--mode compare: the baseline side's result JSON") + parser.add_argument("--shared", help="--mode compare: the shared side's result JSON") + parser.add_argument("--json", action="store_true", help="also dump the raw measurements") + args = parser.parse_args() + if args.mode == "compare": + if not args.baseline or not args.shared: + raise SystemExit("--mode compare needs both --baseline and --shared") + return run_compare(args) + if args.parallel < 2: + raise SystemExit("--parallel must be at least 2 for the comparison to mean anything") + if not 0 < args.read_fraction <= 1.0: + raise SystemExit("--read-fraction must be in (0, 1]") + if args.mode == "prepare": + return run_prepare(args) + return run_parent(args) if args.mode == "both" else run_child(args) + + +if __name__ == "__main__": + raise SystemExit(main()) 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/lance_data_sink.py b/daft_lance/lance_data_sink.py index c67cbc8..a854ea3 100644 --- a/daft_lance/lance_data_sink.py +++ b/daft_lance/lance_data_sink.py @@ -30,6 +30,7 @@ ResolvedNamespaceTable, get_namespace_commit_kwargs, get_namespace_kwargs, + get_shared_session, get_write_fragments_kwargs, merge_storage_options, resolve_namespace_table, @@ -208,6 +209,19 @@ def _validate_overwrite_where_against_table(self, dataset: lance.LanceDataset) - def _namespace_kwargs(self) -> dict[str, Any]: return get_namespace_kwargs(self._namespace_impl, self._namespace_properties, self._table_id) + @property + def _dataset_open_kwargs(self) -> dict[str, Any]: + """Namespace wiring plus this process's shared lance session. + + Every dataset this sink opens goes through here. Without the session + each open builds its own ``Arc`` holding the table's whole + fragment list, so a many-fragment table costs that again per open -- + including the one the Rust writer does inside ``write_fragments`` on + every buffer flush. Commits keep using ``_namespace_commit_kwargs``: + ``LanceDataset.commit()`` takes no session. + """ + return {"session": get_shared_session(), **self._namespace_kwargs} + @property def _namespace_commit_kwargs(self) -> dict[str, Any]: return get_namespace_commit_kwargs( @@ -347,7 +361,7 @@ 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 + self._dataset_uri_arg, storage_options=self._storage_options, **self._dataset_open_kwargs ) except (ValueError, FileNotFoundError, OSError) as e: # Pinned to the Rust message format; lance has no typed exception. See test_lance_message_format_unchanged. @@ -430,6 +444,9 @@ def _write_arrow_table(self, table: pa.Table) -> WriteResult[list[FragmentMetada fragments = lance.fragment.write_fragments( wrapped, dataset_uri=self._table_uri, + # A uri destination makes the Rust writer load the dataset on every + # flush; the shared session turns the repeats into cache hits. + session=get_shared_session(), mode=self._write_mode, storage_options=self._storage_options, max_rows_per_file=self._max_rows_per_file, @@ -450,7 +467,9 @@ 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) + ds = lance.dataset( + self._dataset_uri_arg, storage_options=self._storage_options, **self._dataset_open_kwargs + ) except (ValueError, FileNotFoundError, OSError): ds = None @@ -465,7 +484,7 @@ def _ensure_mem_wal_dataset(self) -> lance.LanceDataset: storage_options=self._storage_options, data_storage_version=self._data_storage_version, use_legacy_format=self._use_legacy_format, - **self._namespace_kwargs, + **self._dataset_open_kwargs, ) details = ds.mem_wal_index_details() @@ -631,7 +650,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 = lance.dataset( + self._dataset_uri_arg, storage_options=self._storage_options, **self._dataset_open_kwargs + ) return self._idempotence_key_exists(dataset) def _finalize_uncheckpointed(self, write_results: list[WriteResult[list[FragmentMetadata]]]) -> MicroPartition: @@ -679,7 +700,7 @@ def _finalize_insert_overwrite(self, fragments: list[FragmentMetadata]) -> Micro self._dataset_uri_arg, version=self._version, storage_options=self._storage_options, - **self._namespace_kwargs, + **self._dataset_open_kwargs, ) open_context = DatasetOpenContext.from_dataset( pinned, @@ -701,12 +722,14 @@ def _finalize_insert_overwrite(self, fragments: list[FragmentMetadata]) -> Micro self._overwrite_where, ) dataset = lance.dataset( - self._dataset_uri_arg, storage_options=self._storage_options, **self._namespace_kwargs + self._dataset_uri_arg, storage_options=self._storage_options, **self._dataset_open_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 = lance.dataset( + self._dataset_uri_arg, storage_options=self._storage_options, **self._dataset_open_kwargs + ) if self._compact_after_write: logger.info( @@ -728,7 +751,7 @@ def _finalize_mem_wal(self, write_results: list[WriteResult[list[FragmentMetadat ), ) dataset = lance.dataset( - self._dataset_uri_arg, storage_options=self._storage_options, **self._namespace_kwargs + self._dataset_uri_arg, storage_options=self._storage_options, **self._dataset_open_kwargs ) return _dataset_stats(dataset) @@ -743,7 +766,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 = lance.dataset( + self._dataset_uri_arg, storage_options=self._storage_options, **self._dataset_open_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..dcd3463 100644 --- a/daft_lance/namespace.py +++ b/daft_lance/namespace.py @@ -28,16 +28,80 @@ from __future__ import annotations import os +import threading from dataclasses import dataclass, field from functools import lru_cache from typing import Any from urllib.parse import unquote, urlparse import lance +from lance.lance import _Session as LanceSession from lance_namespace import DeclareTableResponse, DescribeTableResponse, LanceNamespace _NAMESPACE_CACHE_SIZE = int(os.environ.get("DAFT_LANCE_NAMESPACE_CACHE_SIZE", "16")) +_SHARED_SESSION: LanceSession | None = None +_SHARED_SESSION_LOCK = threading.Lock() + + +def _sharing_enabled() -> bool: + """Whether session sharing is on; ``DAFT_LANCE_SHARE_SESSION=0`` turns it off. + + An escape hatch for a process that must not pool lance's caches, and the + only way to get a pre-sharing baseline inside a Ray worker: a driver-side + monkeypatch never reaches the worker process. Read on every call rather + than once at import, which costs nothing next to a dataset open and lets a + test flip it in-process. + """ + return os.environ.get("DAFT_LANCE_SHARE_SESSION", "1").strip().lower() not in {"0", "false", "no"} + + +def get_shared_session() -> LanceSession: + """Return this process's one lance ``Session``, creating it on first use. + + Every ``lance.dataset()`` that is opened without a session builds its own: + its own ``Arc`` holding the table's whole fragment list, its own + index and metadata cache budgets (6 GiB and 1 GiB each, not a shared + process total) and its own object-store client. A table with many fragments + therefore costs that much again for every concurrently open dataset, and + Daft opens one per scan task, per maintenance UDF instance and per writer. + + Only some of those overlap inside one process, which is where the saving + comes from. Scan tasks and writers do: ``scantask_max_parallel`` of them run + in the same worker, so their manifests add up. Maintenance instances do not + -- Daft gives each ``daft.cls`` instance a process of its own, measured on + both runners -- so there the shared session keeps the cost at one manifest + per process rather than collapsing many. An indexed table does not change + that: a maintenance task's open brings in the index's metadata entry only, + and the index data is read once on the driver, when the commit remaps it + onto the new fragment ids. + + Sharing one session collapses that: the manifest is a metadata-cache entry + keyed by (uri, version, e_tag), so every open of the snapshot the driver + pinned hands back a clone of the same ``Arc``. Measured on a + 20k-fragment table, 64 opens cost 859 MB unshared and 6 MB shared. + + Created lazily, never at import: a session carries an object-store client, + which must not be inherited across a fork. + """ + if not _sharing_enabled(): + # One fresh session per open, which is byte-for-byte what lance does + # when no session is passed at all. + return LanceSession() + global _SHARED_SESSION + if _SHARED_SESSION is None: + with _SHARED_SESSION_LOCK: + if _SHARED_SESSION is None: + _SHARED_SESSION = LanceSession() + return _SHARED_SESSION + + +def reset_shared_session() -> None: + """Drop the process's shared session. For tests only.""" + global _SHARED_SESSION + with _SHARED_SESSION_LOCK: + _SHARED_SESSION = None + def has_namespace_params(namespace_impl: str | None, table_id: list[str] | None) -> bool: """Whether namespace addressing is in effect (``namespace_properties`` stays optional).""" @@ -273,14 +337,20 @@ class DatasetOpenContext: commit as if the table were an unmanaged uri table. This context is the serializable substitute. It deliberately excludes the - live dataset, the namespace client, ``serialized_manifest``, ``session``, - ``commit_lock`` and ``asof``: + live dataset, the namespace client, ``serialized_manifest``, ``session`` + (none of which pickle usefully), ``commit_lock`` and ``asof``: * the client is rebuilt per process from the triple via the ``lru_cache``; * the manifest is omitted so the payload stays independent of fragment - count -- workers pay one pinned-manifest read instead; + count -- workers pay one pinned-manifest read instead, and + :func:`get_shared_session` makes the whole process share that one read; * ``asof``/tags are already resolved by the driver into :attr:`version`, so workers cannot drift onto a different snapshot. + + Because the session is shared, :attr:`index_cache_size` and + :attr:`metadata_cache_size_bytes` no longer size a per-dataset cache: lance + ignores them once a session is supplied, and the session's own budgets + apply to every dataset in the process. """ uri: str @@ -362,6 +432,7 @@ def _open(self, version: int | None) -> lance.LanceDataset: namespace_client=get_or_create_namespace(self.namespace_impl, self.namespace_properties), table_id=self.table_id, namespace_client_managed_versioning=self.managed_versioning, + session=get_shared_session(), ) def open_pinned(self) -> lance.LanceDataset: @@ -387,6 +458,7 @@ def open_dataset_from_open_kwargs(ds_uri: str | None, open_kwargs: dict[str, Any namespace_impl, namespace_properties, table_id = pop_namespace_params(open_kwargs) return lance.dataset( None if has_namespace_params(namespace_impl, table_id) else ds_uri, + session=get_shared_session(), **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..2742bc2 100644 --- a/daft_lance/utils.py +++ b/daft_lance/utils.py @@ -13,6 +13,7 @@ DatasetOpenContext, get_namespace_commit_kwargs, get_namespace_kwargs, + get_shared_session, has_namespace_params, merge_storage_options, resolve_namespace_table, @@ -214,6 +215,9 @@ def construct_lance_dataset_handle( dataset_uri, storage_options=merged_storage_options, version=version, + # The native runner plans and executes in one process, so the driver's + # manifest is the one every scan task reuses. + session=get_shared_session(), **get_namespace_kwargs(namespace_impl, namespace_properties, table_id), **kwargs, ) 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..c05c26a --- /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. ``benchmarks/lance_session_memory.py`` is the +counterpart that puts megabytes on it. + +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..bba23e8 --- /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 benchmarks/lance_session_memory.py needs revisiting" + )