Skip to content

perf: reuse a pinned dataset as the write_fragments target - #91

Open
ChaosJu wants to merge 1 commit into
daft-engine:mainfrom
ChaosJu:perf/cache-write-target-dataset
Open

ChaosJu wants to merge 1 commit into
daft-engine:mainfrom
ChaosJu:perf/cache-write-target-dataset

Conversation

@ChaosJu

@ChaosJu ChaosJu commented Nov 9, 2026

Copy link
Copy Markdown

Changes Made

LanceDataSink._write_arrow_table passes the table URI to lance.fragment.write_fragments. For a URI destination Lance's InsertBuilder::resolve_context calls DatasetBuilder::load() on every call, so each write loads and parses the latest manifest. On an append to a large table this costs O(fragments) per batch, and since Daft calls the sink once per micropartition, a job with small batches reloads the manifest thousands of times.

We hit this in production: a table with ~4.56M fragments has a 635 MB manifest that takes ~10 s to open, and appends ran at ~4 MB/s regardless of worker count (20 and 200 workers gave the same throughput), while the same writes to a fresh path ran at GB/s.

This change:

  • Adds a process-level cache (_WRITE_TARGETS, _cached_write_target) that opens the dataset once per worker process at the version start() pinned on the driver, and passes the LanceDataset as the write destination. With a dataset destination Lance skips the load.
  • Keys the entry by (version, storage_options) and keeps one entry per URI, so a long-lived worker replaces the old snapshot instead of accumulating manifests.
  • Keeps the URI for create (Lance's lookup is a cheap NotFound there) and for namespace-addressed tables (the namespace client stays in charge of access).

Writing against the pinned version also matches the read_version the driver commits with, instead of whatever is latest at each flush.

Relation to #86

#86 passes a shared Session to write_fragments, which turns the repeated load into a metadata-cache hit. The two are complementary, but the session only helps while the parsed manifest fits the session's metadata cache. Measured locally (pylance 12.0.0, 10k one-row fragments, 6 appends of 3 rows each):

Destination first call later calls (avg)
URI (current) 59.5 ms 30.4 ms
URI + shared session, default caches 25.5 ms 6.1 ms
URI + shared session, 1 MiB metadata cache 27.7 ms 28.3 ms
Pinned LanceDataset (this PR) 4.3 ms 4.5 ms

The 10k-fragment manifest took 7.8 MB of session cache (~780 B per fragment). Extrapolating linearly, the 4.56M-fragment table above needs ~3.5 GB, beyond the 1 GiB default metadata cache, so a URI destination would still reload it on every flush. A URI destination also still resolves the latest version on each call, which is an object-store list on S3. Holding the dataset object avoids both and does not depend on cache sizing.

Memory note: each worker process keeps one parsed manifest per target URI alive for the duration of the process, which is what a cache hit holds in #86 as well.

Related Issues

None.

Checklist

  • Unit tests added in tests/io/lance/test_lance_write_target_cache.py: multi-batch appends reuse one pinned dataset opened once, sequential appends refresh the pinned version, create keeps the URI, overwrite uses the pinned dataset. The spies patch the driver, so they are native-only (same as the spy tests in perf: share one lance Session per process #86).
  • Under Ray, a probe UDF in the worker process confirmed the cache holds a LanceDataset pinned to the driver's version after a 10-batch append, and the data round-trips.
  • DAFT_RUNNER=native pytest tests/: 519 passed; the 37 failures are the same set as on main (Windows path / namespace issues in my local environment, unrelated to this change).
  • ruff format --check, ruff check, mypy --strict --ignore-missing-imports daft_lance: clean.
  • No README change: no user-facing API or configuration changes.

Passing the table URI makes Lance load the latest manifest on every write_fragments call, which costs O(fragments) per batch on appends to large tables. Workers now open the dataset once per process at the version start() pinned and pass it as the write destination; create and namespace writes keep the URI.
@ChaosJu
ChaosJu force-pushed the perf/cache-write-target-dataset branch from f6a0726 to 0c9b5c8 Compare November 9, 2026 13:10
@FANNG1

FANNG1 commented Nov 11, 2026

Copy link
Copy Markdown
Contributor

Thanks for investigating this and sharing the measurements. I think sharing a Lance Session is the more general approach to repeated manifest loading: it lets multiple Dataset opens reuse Lance's metadata and index caches, covering reads as well as writes, Blob reads, and maintenance operations.

Could you also review #86? It introduces a process-wide shared Session across these paths. It would be useful to evaluate this optimization on top of that common foundation.

As noted in this PR, Session reuse still depends on the parsed manifest fitting the configured metadata cache, and it does not eliminate latest-version resolution. Passing a pinned Dataset can therefore still be useful for the large-manifest case; that optimization would need its own carefully bounded lifetime and snapshot identity.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants