Add distributed vector search with global top-k merging

README.md

@@ -185,6 +185,38 @@ Training runs in the coordinator process. Its memory footprint grows with
 `2**num_bits * sample_rate * dimension` for PQ.
 
 
+### Distributed Vector Search
+
+```python
+from daft_lance import vector_search
+
+result = vector_search(
+    "s3://bucket/my_dataset",
+    nearest={"column": "vector", "q": [0.1, 0.2], "k": 10, "nprobes": 8},
+    columns=["id"],
+    filter="active = true",
+    num_tasks=4,
+)
+result.explain(show_all=True)
+result.collect()
+```
+
+Each scan task searches its assigned index segments at one pinned dataset
+version; Daft merges the global Top-K, ordered by `_distance` then `_rowid`.
+Equal-distance membership at the K boundary is chosen by Lance and may vary
+with segment grouping; row-ID ordering applies to the returned candidates.
+The returned lazy DataFrame also includes `_query_id`. Pass a matrix as `q` to
+search multiple vectors, each with its own global K. Filters apply before
+search. The Ray runner can execute tasks across workers; `num_tasks` controls
+segment groups balanced by fragment row counts, while the runner controls concurrency. One segment offers
+one task per query. This API requires full index coverage; after appending
+data, run `optimize_indices` or rebuild the index first.
+
+Approximate-index candidates can differ from a single scanner because each
+task refines its own candidates. Compare recall as well as latency. See the
+[benchmark guide](benchmarks/README.md) for reproducible comparisons against
+native Lance, the existing single-task path and independent-query parallelism.
+
 ### Column Merging
 
 ```python

benchmarks/README.md

@@ -0,0 +1,118 @@
+# Vector search benchmark
+
+This benchmark compares search strategies on **the same dataset snapshot and
+index files**. It does not include data generation or index training in query
+timings. It works with synthetic data or an existing table and independent query
+embeddings. Install the project's dev and Ray dependencies plus `psutil` first.
+You can run the commands below by replacing `python` with
+`uv run --group ray --with psutil python` to resolve these dependencies.
+
+## Prepare once
+
+```bash
+python benchmarks/vector_search.py prepare --root /tmp/vector-bench \
+  --rows 1000000 --dim 128 --segments 8 --queries 100 --cpu-budget 8
+```
+
+Generation uses a fixed seed and a mixture of 16 Gaussian clusters. Queries use
+an independent random stream, not vectors selected from stored rows. Each fragment
+becomes an IVF_FLAT segment; all segments share the globally trained centroids.
+The default 32 probes visit all 32 partitions, making FLAT search exact. Dataset
+metadata and query vectors are saved beside the table. Preparation refuses to
+overwrite an existing table.
+
+## Compare end-to-end query execution
+
+Run cases **serially**, each in a fresh process, on an otherwise idle machine.
+For each strategy, use the same root, query count, K, probes, filter and recall gate.
+
+```bash
+python benchmarks/vector_search.py run --root /tmp/vector-bench \
+  --backend pylance --cpu-budget 4 --queries 20 --output /tmp/results.jsonl
+python benchmarks/vector_search.py run --root /tmp/vector-bench \
+  --backend daft-single --runner ray --workers 4 --cpu-budget 4 \
+  --queries 20 --output /tmp/results.jsonl
+python benchmarks/vector_search.py run --root /tmp/vector-bench \
+  --backend segments --runner ray --workers 4 --cpu-budget 4 \
+  --queries 20 --output /tmp/results.jsonl
+python benchmarks/vector_search.py run --root /tmp/vector-bench \
+  --backend queries --runner ray --workers 4 --cpu-budget 4 \
+  --queries 20 --output /tmp/results.jsonl
+```
+
+Strategies:
+
+| Backend | Execution |
+| --- | --- |
+| `pylance` | Sequential queries on one reusable LanceDataset |
+| `daft-single` | Existing read_lance nearest path, one task per query, sequential queries |
+| `segments` | Multiple segment scan tasks per query and a global Top-K in Daft |
+| `queries` | One scan task per query over all segments, multiple queries in one Daft plan |
+
+`segments` and `queries` reuse a pinned worker dataset, preserving its index
+cache between tasks. The existing `daft-single` path opens a dataset per scan
+task. These cache lifetimes are part of the implementations being compared.
+The latter two backends time lazy planning, execution, Arrow materialization
+and grouping results by query ID. All backends project IDs and produce K IDs
+per query. Segment execution additionally carries row IDs and distances to
+perform its merge. Ray startup is excluded; first execution includes worker
+initialization and index loading. A first execution is **not a cold-storage
+measurement**: OS and object-store caches are not reset.
+
+For single-query latency use `--queries 1`. With multiple queries the recorded
+p50/p95 are **batch latency**, not individual-query percentiles. The default
+10 measured repetitions are exploratory; use `--repeats 100` for steadier tail
+estimates. Each repetition builds a fresh DataFrame, avoiding Daft result-cache
+reuse. Warmups and brute-force ground truth run outside measured repetitions.
+
+## Resource and quality controls
+
+- **Fixed CPU:** keep `--cpu-budget 4`, vary `--workers 1/2/4`. On Linux,
+  CPU affinity bounds the driver and inherited local Ray process tree to the
+  same four logical CPUs. Rayon/BLAS thread limits are set before imports.
+  `workers` configures the Ray scheduler CPU budget; it is not a promise that
+  each logical scan task owns a distinct process.
+- **Scaling:** compare `(cpu-budget, workers) = (1,1), (2,2), (4,4), (8,8)`.
+  Report the extra CPU cost with speedup. Affinity is not portable to remote
+  clusters; those require explicit worker placement and CPU limits.
+- **Quality:** every query is compared against a brute-force Lance query of
+  the same version and predicate. Below `--min-recall` (default 0.99), the run
+  fails and writes no performance record. Approximate settings or IVF_PQ need
+  a justified quality threshold; compare configurations meeting the same
+  threshold rather than claiming speed from lower recall.
+- **Filters:** `--selectivity 0.1` and `0.01` prefilter the synthetic `bucket`
+  column. `--filter` overrides this with a Lance SQL predicate. K is global
+  per query, not per segment.
+- **Data/layout:** vary rows (100K/1M/10M), dimensions (128/768), segments
+  (1/4/8/16), K (10/100), probes and query batch size independently. Start
+  with a small matrix, then expand around observed bottlenecks.
+- **Cost:** output includes sampled local process-tree CPU seconds,
+  aggregate peak RSS and physical disk read bytes. Shared pages are counted
+  in each process's RSS; remote workers are excluded. Physical reads can be
+  zero for cached data and are not object-store request/transfer statistics.
+  Use storage-side metrics for TOS/S3 requests and bytes. These are local
+  resource observations, not production billing estimates.
+
+Records append as JSONL with dependency versions, Git state, snapshot version,
+CPU affinity, parameters, raw durations, recall and execution plan. Keep the
+dataset metadata with results. Local disk results do not predict object-store
+or multi-node behavior.
+
+## Existing embeddings and object storage
+
+```bash
+python benchmarks/vector_search.py run --root /tmp/real-bench \
+  --uri s3://bucket/table.lance --query-file /path/independent-queries.npy \
+  --column embedding --id-column id --index-name embedding_idx \
+  --backend segments --runner ray --workers 4 --cpu-budget 4 \
+  --queries 20 --nprobes 32 --min-recall 0.95 --output /tmp/real-results.jsonl
+```
+
+Use environment-based object-store credentials and a completely covered index.
+Comparison runs require exactly one vector index on the selected column, so
+every backend uses the same index. `--index-name` validates that index's name,
+and `--metric` must match its metric. Synthetic preparation builds the index
+with `--metric` (default: `l2`) and records it in `dataset.json`.
+Queries must be a finite matrix with the same dimension as the vector field;
+the ID column must uniquely identify rows for recall calculation. Use an
+explicit predicate for existing tables instead of synthetic `--selectivity`.

benchmarks/results/REPORT.md

@@ -0,0 +1,98 @@
+# Exploratory local vector-search results — 2026-10-09
+
+These historical measurements precede the upstream rebase and review fixes
+for bounded local candidates and row-balanced segment grouping. The hashes
+in the raw JSONL identify the measured code; the numbers below do not measure
+the current implementation.
+
+Measured on one Intel Xeon Platinum 8457C host with local disk. All cases use
+128-dimensional clustered synthetic vectors, 8 IVF_FLAT segments, 32 partitions,
+all 32 probes, L2, K=10, no filter, complete coverage and independent query vectors.
+Linux CPU affinity confines the process tree to CPUs 0–3, except the explicitly
+labelled 8-CPU scaling case. Daft 0.7.14, pylance 12.0.0, Python 3.11.2 and
+Ray 2.54.0. Index training and Ray startup are outside the timers.
+
+Five measured repetitions follow two warmups for batch runs; single-query runs
+have ten repetitions. These are exploratory medians and percentiles, not a
+production tail-latency guarantee. Ray cases run in fresh local clusters,
+serially. First-run measurements are stored in JSONL but are not cold-storage
+measurements. Ground-truth queries and warmups precede timed repetitions.
+
+Single-task baselines receive the whole 4-thread budget. Distributed workers
+receive budget/workers threads and share the same CPU affinity budget. The
+Ray scheduler CPU budget is not an observed count of distinct Python processes.
+All records passed the brute-force recall gate at 1.00. Source hashes, where
+present, identify the measured files; earlier exploratory records identify the
+Git base and dirty tree. Dataset roots were `/tmp/daft-vector-bench-20261009`
+and `/tmp/daft-vector-bench-1m-20261009`, seed 42; queries use seed 43.
+
+## 100K rows: batches of 5 queries
+
+| Strategy | Runner | Ray CPU budget | CPU affinity budget | p50 ms | p95 ms | Recall@10 | Peak local aggregate RSS MiB |
+| --- | --- | ---: | ---: | ---: | ---: | ---: | ---: |
+| pylance | native | — | 4 | 35.73 | 39.57 | 1.00 | 400 |
+| daft-single | native | — | 4 | 243.19 | 246.95 | 1.00 | 473 |
+| segments | ray | 1 | 4 | 596.42 | 628.14 | 1.00 | 2027 |
+| segments | ray | 2 | 4 | 413.50 | 453.53 | 1.00 | 2120 |
+| segments | ray | 4 | 4 | 338.86 | 360.98 | 1.00 | 2150 |
+| queries | ray | 4 | 4 | 159.25 | 163.91 | 1.00 | 2319 |
+| daft-single | ray | 4 | 4 | 322.54 | 329.04 | 1.00 | 2054 |
+
+Raw measurements: [local-100k.jsonl](local-100k.jsonl).
+
+## 1M rows: batches of 5 queries
+
+| Strategy | Runner | Ray CPU budget | CPU affinity budget | p50 ms | p95 ms | Recall@10 | Peak local aggregate RSS MiB |
+| --- | --- | ---: | ---: | ---: | ---: | ---: | ---: |
+| pylance | native | — | 4 | 173.15 | 178.56 | 1.00 | 1026 |
+| segments | ray | 1 | 4 | 938.32 | 957.12 | 1.00 | 2773 |
+| segments | ray | 2 | 4 | 540.92 | 559.21 | 1.00 | 2897 |
+| segments | ray | 4 | 4 | 436.53 | 458.80 | 1.00 | 2981 |
+| queries | ray | 4 | 4 | 266.66 | 271.12 | 1.00 | 4371 |
+| daft-single | ray | 4 | 4 | 973.97 | 1004.66 | 1.00 | 2996 |
+
+Raw measurements: [local-1m.jsonl](local-1m.jsonl).
+
+## 1M rows: one query per run
+
+| Strategy | Runner | Ray CPU budget | CPU affinity budget | p50 ms | p95 ms | Recall@10 | Peak local aggregate RSS MiB |
+| --- | --- | ---: | ---: | ---: | ---: | ---: | ---: |
+| pylance | native | — | 4 | 34.75 | 36.33 | 1.00 | 1022 |
+| segments | ray | 1 | 4 | 180.47 | 182.43 | 1.00 | 2462 |
+| segments | ray | 2 | 4 | 117.33 | 138.01 | 1.00 | 2454 |
+| segments | ray | 4 | 4 | 91.58 | 100.47 | 1.00 | 2772 |
+| segments | ray | 8 | 8 | 68.81 | 72.84 | 1.00 | 3098 |
+| daft-single | ray | 4 | 4 | 205.47 | 223.65 | 1.00 | 2819 |
+
+Raw measurements: [local-1m-single-query.jsonl](local-1m-single-query.jsonl).
+
+## Interpretation
+
+- At 1M rows and a fixed 4-CPU budget, one query takes about **205 ms** through
+  the existing Ray single-task path and **92 ms** with segment tasks: about
+  **2.24x** faster. Four segment workers are about **1.97x** faster than the
+  new API's one-worker configuration (180 ms), which also benefits from worker
+  dataset caching. Raising the budget to 8 CPUs gives about 69 ms, at higher
+  resource cost.
+- For batches of five at 1M rows, segment tasks take about **437 ms**, versus
+  **974 ms** for the existing Ray single-task path. Independent-query tasks
+  take about **267 ms**, making them the better strategy in this measured
+  batch workload.
+- At 100K rows, segment scheduling costs dominate: **339 ms** versus
+  **323 ms** for the existing Ray single-task path. Sharding is not a general
+  latency improvement for small datasets.
+- Native pylance remains faster in every measured workload: about **35 ms**
+  per query at 1M rows. This work improves the Daft execution path and adds
+  distributed scheduling; these data do not establish an advantage over a
+  native scanner on one machine.
+- RSS includes shared pages more than once and includes local Ray services.
+  It is an aggregate process-tree observation, not unique resident memory or
+  remote-cluster cost. Physical read bytes are OS disk reads, not storage
+  request counts or object-store transfer bytes.
+
+Use the opt-in vector_search API when its workload-specific benchmark shows
+benefit. For batches, compare query-level parallelism (`num_tasks=1` with a
+query matrix) first. Small nprobes, compression, selective filters, uneven
+segments, real embeddings, remote storage and multi-node execution need
+separate measurements before choosing deployment settings. The benchmark
+supports these controls, but the above table does not measure them.