diff --git a/README.md b/README.md index ee6d1b5..216b625 100644 --- a/README.md +++ b/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 diff --git a/benchmarks/README.md b/benchmarks/README.md new file mode 100644 index 0000000..04e8110 --- /dev/null +++ b/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`. diff --git a/benchmarks/results/REPORT.md b/benchmarks/results/REPORT.md new file mode 100644 index 0000000..e761382 --- /dev/null +++ b/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. diff --git a/benchmarks/results/local-100k.jsonl b/benchmarks/results/local-100k.jsonl new file mode 100644 index 0000000..d9f4cdb --- /dev/null +++ b/benchmarks/results/local-100k.jsonl @@ -0,0 +1,7 @@ +{"backend": "pylance", "runner": "native", "workers": 4, "cpu_budget": 4, "cpu_affinity": [0, 1, 2, 3], "threads_per_process": "4", "dataset_version": 9, "rows": 100000, "segments": {"bench_idx": 8}, "queries_per_run": 5, "nearest": {"column": "vector", "k": 10, "nprobes": 32, "metric": "l2"}, "filter": null, "tasks": null, "first_run_seconds": 0.0851415004581213, "first_run_recall": 1.0, "warm_batch_p50_seconds": 0.03573167696595192, "warm_batch_p95_seconds": 0.03956762924790382, "queries_per_second": 137.99324810699196, "recall_min": 1.0, "local_cpu_seconds": 0.38000000000000034, "local_physical_read_bytes": 0, "local_peak_aggregate_rss_bytes": 418955264, "first_run_peak_aggregate_rss_bytes": 310300672, "durations_seconds": [0.03573167696595192, 0.033796364441514015, 0.03378709591925144, 0.03999499976634979, 0.03785814717411995], "execution_plan": null, "versions": {"daft": "0.7.14", "pylance": "12.0.0", "python": "3.11.2"}, "git_commit": "21d78f379d0d47ddebfc2044b32f4001b8501b28", "git_dirty": true} +{"backend": "daft-single", "runner": "native", "workers": 4, "cpu_budget": 4, "cpu_affinity": [0, 1, 2, 3], "threads_per_process": "4", "dataset_version": 9, "rows": 100000, "segments": {"bench_idx": 8}, "queries_per_run": 5, "nearest": {"column": "vector", "k": 10, "nprobes": 32, "metric": "l2"}, "filter": null, "tasks": null, "first_run_seconds": 0.5766022745519876, "first_run_recall": 1.0, "warm_batch_p50_seconds": 0.24319044686853886, "warm_batch_p95_seconds": 0.24695064388215543, "queries_per_second": 20.78861896931947, "recall_min": 1.0, "local_cpu_seconds": 2.49, "local_physical_read_bytes": 0, "local_peak_aggregate_rss_bytes": 496132096, "first_run_peak_aggregate_rss_bytes": 367427584, "durations_seconds": [0.24324788711965084, 0.24787633307278156, 0.22980329394340515, 0.24319044686853886, 0.2384631261229515], "execution_plan": null, "versions": {"daft": "0.7.14", "pylance": "12.0.0", "python": "3.11.2"}, "git_commit": "21d78f379d0d47ddebfc2044b32f4001b8501b28", "git_dirty": true} +{"backend": "segments", "runner": "ray", "workers": 1, "cpu_budget": 4, "cpu_affinity": [0, 1, 2, 3], "threads_per_process": "4", "dataset_version": 9, "rows": 100000, "segments": {"bench_idx": 8}, "queries_per_run": 5, "nearest": {"column": "vector", "k": 10, "nprobes": 32, "metric": "l2"}, "filter": null, "tasks": null, "first_run_seconds": 3.2385472264140844, "first_run_recall": 1.0, "warm_batch_p50_seconds": 0.5964183993637562, "warm_batch_p95_seconds": 0.6281368482857943, "queries_per_second": 8.33047423849619, "recall_min": 1.0, "local_cpu_seconds": 4.910000000000001, "local_physical_read_bytes": 0, "local_peak_aggregate_rss_bytes": 2125287424, "first_run_peak_aggregate_rss_bytes": 2002755584, "durations_seconds": [0.5964183993637562, 0.5695181880146265, 0.607533136382699, 0.6332877762615681, 0.5942721273750067], "execution_plan": "== Unoptimized Logical Plan ==\n\n* Sort: Sort by = (col(_query_id), ascending, nulls last), (col(_distance),\n| ascending, nulls last), (col(_rowid), ascending, nulls last)\n|\n* Limit: 10\n|\n* Sort: Sort by = (col(_distance), ascending, nulls last), (col(_rowid),\n| ascending, nulls last)\n|\n* PythonScanOperator: LanceVectorSearch(version=9, query=0)\n| File schema = id#Int64, _distance#Float32, _rowid#UInt64, _query_id#Int64\n| Partitioning keys = []\n| Output schema = id#Int64, _distance#Float32, _rowid#UInt64, _query_id#Int64\n\n\n== Optimized Logical Plan ==\n\n* Sort: Sort by = (col(_query_id), ascending, nulls last), (col(_distance),\n| ascending, nulls last), (col(_rowid), ascending, nulls last)\n| Stats = { Approx num rows = 0, Approx size bytes = 0 B, Accumulated selectivity\n| = 1.00 }\n|\n* TopN: Sort by = (col(_distance), ascending, nulls last), (col(_rowid),\n| ascending, nulls last), Num Rows = 10\n| Stats = { Approx num rows = 0, Approx size bytes = 0 B, Accumulated selectivity\n| = 1.00 }\n|\n* Num Scan Tasks = 8\n| Schema: {id#Int64, _distance#Float32, _rowid#UInt64, _query_id#Int64}\n| Scan Tasks: [\n| {daft_lance.lance_vector_search:_search_segments}\n| {daft_lance.lance_vector_search:_search_segments}\n| {daft_lance.lance_vector_search:_search_segments}\n| ...\n| {daft_lance.lance_vector_search:_search_segments}\n| {daft_lance.lance_vector_search:_search_segments}\n| {daft_lance.lance_vector_search:_search_segments}\n| ]\n| File schema = id#Int64, _distance#Float32, _rowid#UInt64, _query_id#Int64\n| Partitioning keys = []\n| Output schema = id#Int64, _distance#Float32, _rowid#UInt64, _query_id#Int64\n| Stats = { Approx num rows = 0, Approx size bytes = 0 B, Accumulated selectivity\n| = 1.00 }\n\n\n== Physical Plan ==\n\n* Sort\n| Sort by: col(3: _query_id), col(1: _distance), col(2: _rowid)\n| Descending = false, false, false\n| Nulls first = false, false, false\n|\n* TopN\n| Sort by: col(1: _distance), col(2: _rowid)\n| Descending = false, false\n| Nulls first = false, false\n| Num Rows = 10\n|\n* RayGather\n|\n* TopN\n| Sort by: col(1: _distance), col(2: _rowid)\n| Descending = false, false\n| Nulls first = false, false\n| Num Rows = 10\n| Offset = 0\n|\n* ScanTaskSource:\n| Num Scan Tasks = 8\n| Estimated Scan Bytes = 0\n| Source = LanceVectorSearch(version=9, query=0)\n| Schema: {id#Int64, _distance#Float32, _rowid#UInt64, _query_id#Int64}\n| Scan Tasks: [\n| {daft_lance.lance_vector_search:_search_segments}\n| {daft_lance.lance_vector_search:_search_segments}\n| {daft_lance.lance_vector_search:_search_segments}\n| ...\n| {daft_lance.lance_vector_search:_search_segments}\n| {daft_lance.lance_vector_search:_search_segments}\n| {daft_lance.lance_vector_search:_search_segments}\n| ]\n\n", "versions": {"daft": "0.7.14", "pylance": "12.0.0", "python": "3.11.2"}, "git_commit": "21d78f379d0d47ddebfc2044b32f4001b8501b28", "git_dirty": true} +{"backend": "segments", "runner": "ray", "workers": 2, "cpu_budget": 4, "cpu_affinity": [0, 1, 2, 3], "threads_per_process": "2", "dataset_version": 9, "rows": 100000, "segments": {"bench_idx": 8}, "queries_per_run": 5, "nearest": {"column": "vector", "k": 10, "nprobes": 32, "metric": "l2"}, "filter": null, "tasks": null, "first_run_seconds": 2.603702634572983, "first_run_recall": 1.0, "warm_batch_p50_seconds": 0.41350078396499157, "warm_batch_p95_seconds": 0.45352722853422167, "queries_per_second": 12.078187091058043, "recall_min": 1.0, "local_cpu_seconds": 4.4, "local_physical_read_bytes": 0, "local_peak_aggregate_rss_bytes": 2222669824, "first_run_peak_aggregate_rss_bytes": 2083549184, "durations_seconds": [0.41350078396499157, 0.443164125084877, 0.3826581407338381, 0.4561180043965578, 0.37440600246191025], "execution_plan": "== Unoptimized Logical Plan ==\n\n* Sort: Sort by = (col(_query_id), ascending, nulls last), (col(_distance),\n| ascending, nulls last), (col(_rowid), ascending, nulls last)\n|\n* Limit: 10\n|\n* Sort: Sort by = (col(_distance), ascending, nulls last), (col(_rowid),\n| ascending, nulls last)\n|\n* PythonScanOperator: LanceVectorSearch(version=9, query=0)\n| File schema = id#Int64, _distance#Float32, _rowid#UInt64, _query_id#Int64\n| Partitioning keys = []\n| Output schema = id#Int64, _distance#Float32, _rowid#UInt64, _query_id#Int64\n\n\n== Optimized Logical Plan ==\n\n* Sort: Sort by = (col(_query_id), ascending, nulls last), (col(_distance),\n| ascending, nulls last), (col(_rowid), ascending, nulls last)\n| Stats = { Approx num rows = 0, Approx size bytes = 0 B, Accumulated selectivity\n| = 1.00 }\n|\n* TopN: Sort by = (col(_distance), ascending, nulls last), (col(_rowid),\n| ascending, nulls last), Num Rows = 10\n| Stats = { Approx num rows = 0, Approx size bytes = 0 B, Accumulated selectivity\n| = 1.00 }\n|\n* Num Scan Tasks = 8\n| Schema: {id#Int64, _distance#Float32, _rowid#UInt64, _query_id#Int64}\n| Scan Tasks: [\n| {daft_lance.lance_vector_search:_search_segments}\n| {daft_lance.lance_vector_search:_search_segments}\n| {daft_lance.lance_vector_search:_search_segments}\n| ...\n| {daft_lance.lance_vector_search:_search_segments}\n| {daft_lance.lance_vector_search:_search_segments}\n| {daft_lance.lance_vector_search:_search_segments}\n| ]\n| File schema = id#Int64, _distance#Float32, _rowid#UInt64, _query_id#Int64\n| Partitioning keys = []\n| Output schema = id#Int64, _distance#Float32, _rowid#UInt64, _query_id#Int64\n| Stats = { Approx num rows = 0, Approx size bytes = 0 B, Accumulated selectivity\n| = 1.00 }\n\n\n== Physical Plan ==\n\n* Sort\n| Sort by: col(3: _query_id), col(1: _distance), col(2: _rowid)\n| Descending = false, false, false\n| Nulls first = false, false, false\n|\n* TopN\n| Sort by: col(1: _distance), col(2: _rowid)\n| Descending = false, false\n| Nulls first = false, false\n| Num Rows = 10\n|\n* RayGather\n|\n* TopN\n| Sort by: col(1: _distance), col(2: _rowid)\n| Descending = false, false\n| Nulls first = false, false\n| Num Rows = 10\n| Offset = 0\n|\n* ScanTaskSource:\n| Num Scan Tasks = 8\n| Estimated Scan Bytes = 0\n| Source = LanceVectorSearch(version=9, query=0)\n| Schema: {id#Int64, _distance#Float32, _rowid#UInt64, _query_id#Int64}\n| Scan Tasks: [\n| {daft_lance.lance_vector_search:_search_segments}\n| {daft_lance.lance_vector_search:_search_segments}\n| {daft_lance.lance_vector_search:_search_segments}\n| ...\n| {daft_lance.lance_vector_search:_search_segments}\n| {daft_lance.lance_vector_search:_search_segments}\n| {daft_lance.lance_vector_search:_search_segments}\n| ]\n\n", "versions": {"daft": "0.7.14", "pylance": "12.0.0", "python": "3.11.2"}, "git_commit": "21d78f379d0d47ddebfc2044b32f4001b8501b28", "git_dirty": true} +{"backend": "segments", "runner": "ray", "workers": 4, "cpu_budget": 4, "cpu_affinity": [0, 1, 2, 3], "threads_per_process": "1", "dataset_version": 9, "rows": 100000, "segments": {"bench_idx": 8}, "queries_per_run": 5, "nearest": {"column": "vector", "k": 10, "nprobes": 32, "metric": "l2"}, "filter": null, "tasks": null, "first_run_seconds": 2.243063896894455, "first_run_recall": 1.0, "warm_batch_p50_seconds": 0.33885527960956097, "warm_batch_p95_seconds": 0.3609825041145086, "queries_per_second": 14.785178808283973, "recall_min": 1.0, "local_cpu_seconds": 4.049999999999999, "local_physical_read_bytes": 0, "local_peak_aggregate_rss_bytes": 2254753792, "first_run_peak_aggregate_rss_bytes": 2090344448, "durations_seconds": [0.3194585535675287, 0.33885527960956097, 0.36513320729136467, 0.3230557609349489, 0.34437969140708447], "execution_plan": "== Unoptimized Logical Plan ==\n\n* Sort: Sort by = (col(_query_id), ascending, nulls last), (col(_distance),\n| ascending, nulls last), (col(_rowid), ascending, nulls last)\n|\n* Limit: 10\n|\n* Sort: Sort by = (col(_distance), ascending, nulls last), (col(_rowid),\n| ascending, nulls last)\n|\n* PythonScanOperator: LanceVectorSearch(version=9, query=0)\n| File schema = id#Int64, _distance#Float32, _rowid#UInt64, _query_id#Int64\n| Partitioning keys = []\n| Output schema = id#Int64, _distance#Float32, _rowid#UInt64, _query_id#Int64\n\n\n== Optimized Logical Plan ==\n\n* Sort: Sort by = (col(_query_id), ascending, nulls last), (col(_distance),\n| ascending, nulls last), (col(_rowid), ascending, nulls last)\n| Stats = { Approx num rows = 0, Approx size bytes = 0 B, Accumulated selectivity\n| = 1.00 }\n|\n* TopN: Sort by = (col(_distance), ascending, nulls last), (col(_rowid),\n| ascending, nulls last), Num Rows = 10\n| Stats = { Approx num rows = 0, Approx size bytes = 0 B, Accumulated selectivity\n| = 1.00 }\n|\n* Num Scan Tasks = 8\n| Schema: {id#Int64, _distance#Float32, _rowid#UInt64, _query_id#Int64}\n| Scan Tasks: [\n| {daft_lance.lance_vector_search:_search_segments}\n| {daft_lance.lance_vector_search:_search_segments}\n| {daft_lance.lance_vector_search:_search_segments}\n| ...\n| {daft_lance.lance_vector_search:_search_segments}\n| {daft_lance.lance_vector_search:_search_segments}\n| {daft_lance.lance_vector_search:_search_segments}\n| ]\n| File schema = id#Int64, _distance#Float32, _rowid#UInt64, _query_id#Int64\n| Partitioning keys = []\n| Output schema = id#Int64, _distance#Float32, _rowid#UInt64, _query_id#Int64\n| Stats = { Approx num rows = 0, Approx size bytes = 0 B, Accumulated selectivity\n| = 1.00 }\n\n\n== Physical Plan ==\n\n* Sort\n| Sort by: col(3: _query_id), col(1: _distance), col(2: _rowid)\n| Descending = false, false, false\n| Nulls first = false, false, false\n|\n* TopN\n| Sort by: col(1: _distance), col(2: _rowid)\n| Descending = false, false\n| Nulls first = false, false\n| Num Rows = 10\n|\n* RayGather\n|\n* TopN\n| Sort by: col(1: _distance), col(2: _rowid)\n| Descending = false, false\n| Nulls first = false, false\n| Num Rows = 10\n| Offset = 0\n|\n* ScanTaskSource:\n| Num Scan Tasks = 8\n| Estimated Scan Bytes = 0\n| Source = LanceVectorSearch(version=9, query=0)\n| Schema: {id#Int64, _distance#Float32, _rowid#UInt64, _query_id#Int64}\n| Scan Tasks: [\n| {daft_lance.lance_vector_search:_search_segments}\n| {daft_lance.lance_vector_search:_search_segments}\n| {daft_lance.lance_vector_search:_search_segments}\n| ...\n| {daft_lance.lance_vector_search:_search_segments}\n| {daft_lance.lance_vector_search:_search_segments}\n| {daft_lance.lance_vector_search:_search_segments}\n| ]\n\n", "versions": {"daft": "0.7.14", "pylance": "12.0.0", "python": "3.11.2"}, "git_commit": "21d78f379d0d47ddebfc2044b32f4001b8501b28", "git_dirty": true} +{"backend": "queries", "runner": "ray", "workers": 4, "cpu_budget": 4, "cpu_affinity": [0, 1, 2, 3], "threads_per_process": "1", "dataset_version": 9, "rows": 100000, "segments": {"bench_idx": 8}, "queries_per_run": 5, "nearest": {"column": "vector", "k": 10, "nprobes": 32, "metric": "l2"}, "filter": null, "tasks": null, "first_run_seconds": 2.1005610898137093, "first_run_recall": 1.0, "warm_batch_p50_seconds": 0.15925156511366367, "warm_batch_p95_seconds": 0.1639113999903202, "queries_per_second": 31.97557240280997, "recall_min": 1.0, "local_cpu_seconds": 1.7199999999999993, "local_physical_read_bytes": 0, "local_peak_aggregate_rss_bytes": 2431496192, "first_run_peak_aggregate_rss_bytes": 2261454848, "durations_seconds": [0.16030645743012428, 0.1505686528980732, 0.14690752141177654, 0.1648126356303692, 0.15925156511366367], "execution_plan": "== Unoptimized Logical Plan ==\n\n* Sort: Sort by = (col(_query_id), ascending, nulls last), (col(_distance),\n| ascending, nulls last), (col(_rowid), ascending, nulls last)\n|\n* Limit: 10\n|\n* Sort: Sort by = (col(_distance), ascending, nulls last), (col(_rowid),\n| ascending, nulls last)\n|\n* PythonScanOperator: LanceVectorSearch(version=9, query=0)\n| File schema = id#Int64, _distance#Float32, _rowid#UInt64, _query_id#Int64\n| Partitioning keys = []\n| Output schema = id#Int64, _distance#Float32, _rowid#UInt64, _query_id#Int64\n\n\n== Optimized Logical Plan ==\n\n* Sort: Sort by = (col(_query_id), ascending, nulls last), (col(_distance),\n| ascending, nulls last), (col(_rowid), ascending, nulls last)\n| Stats = { Approx num rows = 0, Approx size bytes = 0 B, Accumulated selectivity\n| = 1.00 }\n|\n* TopN: Sort by = (col(_distance), ascending, nulls last), (col(_rowid),\n| ascending, nulls last), Num Rows = 10\n| Stats = { Approx num rows = 0, Approx size bytes = 0 B, Accumulated selectivity\n| = 1.00 }\n|\n* Num Scan Tasks = 1\n| Schema: {id#Int64, _distance#Float32, _rowid#UInt64, _query_id#Int64}\n| Scan Tasks: [\n| {daft_lance.lance_vector_search:_search_segments}\n| ]\n| File schema = id#Int64, _distance#Float32, _rowid#UInt64, _query_id#Int64\n| Partitioning keys = []\n| Output schema = id#Int64, _distance#Float32, _rowid#UInt64, _query_id#Int64\n| Stats = { Approx num rows = 0, Approx size bytes = 0 B, Accumulated selectivity\n| = 1.00 }\n\n\n== Physical Plan ==\n\n* Sort\n| Sort by: col(3: _query_id), col(1: _distance), col(2: _rowid)\n| Descending = false, false, false\n| Nulls first = false, false, false\n|\n* TopN\n| Sort by: col(1: _distance), col(2: _rowid)\n| Descending = false, false\n| Nulls first = false, false\n| Num Rows = 10\n|\n* TopN\n| Sort by: col(1: _distance), col(2: _rowid)\n| Descending = false, false\n| Nulls first = false, false\n| Num Rows = 10\n| Offset = 0\n|\n* ScanTaskSource:\n| Num Scan Tasks = 1\n| Estimated Scan Bytes = 0\n| Source = LanceVectorSearch(version=9, query=0)\n| Schema: {id#Int64, _distance#Float32, _rowid#UInt64, _query_id#Int64}\n| Scan Tasks: [\n| {daft_lance.lance_vector_search:_search_segments}\n| ]\n\n", "versions": {"daft": "0.7.14", "pylance": "12.0.0", "python": "3.11.2"}, "git_commit": "21d78f379d0d47ddebfc2044b32f4001b8501b28", "git_dirty": true} +{"backend": "daft-single", "runner": "ray", "workers": 4, "cpu_budget": 4, "cpu_affinity": [0, 1, 2, 3], "threads_per_process": "4", "dataset_version": 9, "rows": 100000, "segments": {"bench_idx": 8}, "queries_per_run": 5, "nearest": {"column": "vector", "k": 10, "nprobes": 32, "metric": "l2"}, "filter": null, "tasks": null, "first_run_seconds": 2.020776331424713, "first_run_recall": 1.0, "warm_batch_p50_seconds": 0.3225444760173559, "warm_batch_p95_seconds": 0.3290407590568066, "queries_per_second": 15.708996188673416, "recall_min": 1.0, "local_cpu_seconds": 3.4499999999999997, "local_physical_read_bytes": 0, "local_peak_aggregate_rss_bytes": 2154016768, "first_run_peak_aggregate_rss_bytes": 1997156352, "durations_seconds": [0.31318376027047634, 0.3260476626455784, 0.3225444760173559, 0.29987984895706177, 0.3297890331596136], "execution_plan": null, "versions": {"daft": "0.7.14", "pylance": "12.0.0", "python": "3.11.2"}, "ray_version": "2.54.0", "source_sha256": {"benchmarks/vector_search.py": "8b8b4b84a436353062f3c542eaf9d2d92c77c4acfbcd0d3ba29f36d9daaf54f2", "daft_lance/lance_vector_search.py": "acc7d4cc3e70875b35b9eed9ba656899bf7efb94aa60d7a66742e763eb80fbb5", "daft_lance/_lance.py": "4f854d14838a299f0959b223d372c4c4a0fe577e9d3330da155c05ccf22e6820"}, "git_commit": "21d78f379d0d47ddebfc2044b32f4001b8501b28", "git_dirty": true} diff --git a/benchmarks/results/local-1m-single-query.jsonl b/benchmarks/results/local-1m-single-query.jsonl new file mode 100644 index 0000000..2611ba3 --- /dev/null +++ b/benchmarks/results/local-1m-single-query.jsonl @@ -0,0 +1,6 @@ +{"backend": "pylance", "runner": "native", "workers": 4, "cpu_budget": 4, "cpu_affinity": [0, 1, 2, 3], "threads_per_process": "4", "dataset_version": 9, "rows": 1000000, "segments": {"bench_idx": 8}, "queries_per_run": 1, "nearest": {"column": "vector", "k": 10, "nprobes": 32, "metric": "l2"}, "filter": null, "tasks": null, "first_run_seconds": 0.25192880257964134, "first_run_recall": 1.0, "warm_batch_p50_seconds": 0.03475245460867882, "warm_batch_p95_seconds": 0.03633422814309597, "queries_per_second": 28.65633335570346, "recall_min": 1.0, "local_cpu_seconds": 0.7999999999999998, "local_physical_read_bytes": 0, "local_peak_aggregate_rss_bytes": 1072095232, "first_run_peak_aggregate_rss_bytes": 885923840, "durations_seconds": [0.03476601466536522, 0.03418327867984772, 0.03451469913125038, 0.03556239977478981, 0.03458936884999275, 0.03517370671033859, 0.0333770327270031, 0.034738894551992416, 0.03509187884628773, 0.03696572408080101], "execution_plan": null, "versions": {"daft": "0.7.14", "pylance": "12.0.0", "python": "3.11.2"}, "ray_version": null, "source_sha256": {"benchmarks/vector_search.py": "79bcd4580732ef6b3c0bbfa78cd1aee500f7277b8521291b2952985d22cb59c6", "daft_lance/lance_vector_search.py": "5dfe9f473538e80db6afdeb72755b51223b40375753173f5d80dca7d8eeedb42", "daft_lance/_lance.py": "4f854d14838a299f0959b223d372c4c4a0fe577e9d3330da155c05ccf22e6820"}, "git_commit": "21d78f379d0d47ddebfc2044b32f4001b8501b28", "git_dirty": true} +{"backend": "segments", "runner": "ray", "workers": 1, "cpu_budget": 4, "cpu_affinity": [0, 1, 2, 3], "threads_per_process": "4", "dataset_version": 9, "rows": 1000000, "segments": {"bench_idx": 8}, "queries_per_run": 1, "nearest": {"column": "vector", "k": 10, "nprobes": 32, "metric": "l2"}, "filter": null, "tasks": null, "first_run_seconds": 2.5409005265682936, "first_run_recall": 1.0, "warm_batch_p50_seconds": 0.18046724423766136, "warm_batch_p95_seconds": 0.18243334535509348, "queries_per_second": 5.59467955308829, "recall_min": 1.0, "local_cpu_seconds": 2.82, "local_physical_read_bytes": 0, "local_peak_aggregate_rss_bytes": 2581250048, "first_run_peak_aggregate_rss_bytes": 2321862656, "durations_seconds": [0.17551222071051598, 0.18082896620035172, 0.180105522274971, 0.18165523745119572, 0.18100272119045258, 0.16940638422966003, 0.1741715930402279, 0.18194342032074928, 0.1799522116780281, 0.182834193110466], "execution_plan": "== Unoptimized Logical Plan ==\n\n* Sort: Sort by = (col(_query_id), ascending, nulls last), (col(_distance),\n| ascending, nulls last), (col(_rowid), ascending, nulls last)\n|\n* Limit: 10\n|\n* Sort: Sort by = (col(_distance), ascending, nulls last), (col(_rowid),\n| ascending, nulls last)\n|\n* PythonScanOperator: LanceVectorSearch(version=9, query=0)\n| File schema = id#Int64, _distance#Float32, _rowid#UInt64, _query_id#Int64\n| Partitioning keys = []\n| Output schema = id#Int64, _distance#Float32, _rowid#UInt64, _query_id#Int64\n\n\n== Optimized Logical Plan ==\n\n* Sort: Sort by = (col(_query_id), ascending, nulls last), (col(_distance),\n| ascending, nulls last), (col(_rowid), ascending, nulls last)\n| Stats = { Approx num rows = 0, Approx size bytes = 0 B, Accumulated selectivity\n| = 1.00 }\n|\n* TopN: Sort by = (col(_distance), ascending, nulls last), (col(_rowid),\n| ascending, nulls last), Num Rows = 10\n| Stats = { Approx num rows = 0, Approx size bytes = 0 B, Accumulated selectivity\n| = 1.00 }\n|\n* Num Scan Tasks = 8\n| Schema: {id#Int64, _distance#Float32, _rowid#UInt64, _query_id#Int64}\n| Scan Tasks: [\n| {daft_lance.lance_vector_search:_search_segments}\n| {daft_lance.lance_vector_search:_search_segments}\n| {daft_lance.lance_vector_search:_search_segments}\n| ...\n| {daft_lance.lance_vector_search:_search_segments}\n| {daft_lance.lance_vector_search:_search_segments}\n| {daft_lance.lance_vector_search:_search_segments}\n| ]\n| File schema = id#Int64, _distance#Float32, _rowid#UInt64, _query_id#Int64\n| Partitioning keys = []\n| Output schema = id#Int64, _distance#Float32, _rowid#UInt64, _query_id#Int64\n| Stats = { Approx num rows = 0, Approx size bytes = 0 B, Accumulated selectivity\n| = 1.00 }\n\n\n== Physical Plan ==\n\n* Sort\n| Sort by: col(3: _query_id), col(1: _distance), col(2: _rowid)\n| Descending = false, false, false\n| Nulls first = false, false, false\n|\n* TopN\n| Sort by: col(1: _distance), col(2: _rowid)\n| Descending = false, false\n| Nulls first = false, false\n| Num Rows = 10\n|\n* RayGather\n|\n* TopN\n| Sort by: col(1: _distance), col(2: _rowid)\n| Descending = false, false\n| Nulls first = false, false\n| Num Rows = 10\n| Offset = 0\n|\n* ScanTaskSource:\n| Num Scan Tasks = 8\n| Estimated Scan Bytes = 0\n| Source = LanceVectorSearch(version=9, query=0)\n| Schema: {id#Int64, _distance#Float32, _rowid#UInt64, _query_id#Int64}\n| Scan Tasks: [\n| {daft_lance.lance_vector_search:_search_segments}\n| {daft_lance.lance_vector_search:_search_segments}\n| {daft_lance.lance_vector_search:_search_segments}\n| ...\n| {daft_lance.lance_vector_search:_search_segments}\n| {daft_lance.lance_vector_search:_search_segments}\n| {daft_lance.lance_vector_search:_search_segments}\n| ]\n\n", "versions": {"daft": "0.7.14", "pylance": "12.0.0", "python": "3.11.2"}, "ray_version": "2.54.0", "source_sha256": {"benchmarks/vector_search.py": "79bcd4580732ef6b3c0bbfa78cd1aee500f7277b8521291b2952985d22cb59c6", "daft_lance/lance_vector_search.py": "5dfe9f473538e80db6afdeb72755b51223b40375753173f5d80dca7d8eeedb42", "daft_lance/_lance.py": "4f854d14838a299f0959b223d372c4c4a0fe577e9d3330da155c05ccf22e6820"}, "git_commit": "21d78f379d0d47ddebfc2044b32f4001b8501b28", "git_dirty": true} +{"backend": "segments", "runner": "ray", "workers": 2, "cpu_budget": 4, "cpu_affinity": [0, 1, 2, 3], "threads_per_process": "2", "dataset_version": 9, "rows": 1000000, "segments": {"bench_idx": 8}, "queries_per_run": 1, "nearest": {"column": "vector", "k": 10, "nprobes": 32, "metric": "l2"}, "filter": null, "tasks": null, "first_run_seconds": 2.1914080902934074, "first_run_recall": 1.0, "warm_batch_p50_seconds": 0.11732977163046598, "warm_batch_p95_seconds": 0.1380083255469799, "queries_per_second": 8.290980128518607, "recall_min": 1.0, "local_cpu_seconds": 2.510000000000001, "local_physical_read_bytes": 0, "local_peak_aggregate_rss_bytes": 2573103104, "first_run_peak_aggregate_rss_bytes": 2292535296, "durations_seconds": [0.1209481805562973, 0.1254721824079752, 0.1302979588508606, 0.1155436746776104, 0.11911586858332157, 0.11336985975503922, 0.11486527509987354, 0.11068514548242092, 0.14431680738925934, 0.11151506379246712], "execution_plan": "== Unoptimized Logical Plan ==\n\n* Sort: Sort by = (col(_query_id), ascending, nulls last), (col(_distance),\n| ascending, nulls last), (col(_rowid), ascending, nulls last)\n|\n* Limit: 10\n|\n* Sort: Sort by = (col(_distance), ascending, nulls last), (col(_rowid),\n| ascending, nulls last)\n|\n* PythonScanOperator: LanceVectorSearch(version=9, query=0)\n| File schema = id#Int64, _distance#Float32, _rowid#UInt64, _query_id#Int64\n| Partitioning keys = []\n| Output schema = id#Int64, _distance#Float32, _rowid#UInt64, _query_id#Int64\n\n\n== Optimized Logical Plan ==\n\n* Sort: Sort by = (col(_query_id), ascending, nulls last), (col(_distance),\n| ascending, nulls last), (col(_rowid), ascending, nulls last)\n| Stats = { Approx num rows = 0, Approx size bytes = 0 B, Accumulated selectivity\n| = 1.00 }\n|\n* TopN: Sort by = (col(_distance), ascending, nulls last), (col(_rowid),\n| ascending, nulls last), Num Rows = 10\n| Stats = { Approx num rows = 0, Approx size bytes = 0 B, Accumulated selectivity\n| = 1.00 }\n|\n* Num Scan Tasks = 8\n| Schema: {id#Int64, _distance#Float32, _rowid#UInt64, _query_id#Int64}\n| Scan Tasks: [\n| {daft_lance.lance_vector_search:_search_segments}\n| {daft_lance.lance_vector_search:_search_segments}\n| {daft_lance.lance_vector_search:_search_segments}\n| ...\n| {daft_lance.lance_vector_search:_search_segments}\n| {daft_lance.lance_vector_search:_search_segments}\n| {daft_lance.lance_vector_search:_search_segments}\n| ]\n| File schema = id#Int64, _distance#Float32, _rowid#UInt64, _query_id#Int64\n| Partitioning keys = []\n| Output schema = id#Int64, _distance#Float32, _rowid#UInt64, _query_id#Int64\n| Stats = { Approx num rows = 0, Approx size bytes = 0 B, Accumulated selectivity\n| = 1.00 }\n\n\n== Physical Plan ==\n\n* Sort\n| Sort by: col(3: _query_id), col(1: _distance), col(2: _rowid)\n| Descending = false, false, false\n| Nulls first = false, false, false\n|\n* TopN\n| Sort by: col(1: _distance), col(2: _rowid)\n| Descending = false, false\n| Nulls first = false, false\n| Num Rows = 10\n|\n* RayGather\n|\n* TopN\n| Sort by: col(1: _distance), col(2: _rowid)\n| Descending = false, false\n| Nulls first = false, false\n| Num Rows = 10\n| Offset = 0\n|\n* ScanTaskSource:\n| Num Scan Tasks = 8\n| Estimated Scan Bytes = 0\n| Source = LanceVectorSearch(version=9, query=0)\n| Schema: {id#Int64, _distance#Float32, _rowid#UInt64, _query_id#Int64}\n| Scan Tasks: [\n| {daft_lance.lance_vector_search:_search_segments}\n| {daft_lance.lance_vector_search:_search_segments}\n| {daft_lance.lance_vector_search:_search_segments}\n| ...\n| {daft_lance.lance_vector_search:_search_segments}\n| {daft_lance.lance_vector_search:_search_segments}\n| {daft_lance.lance_vector_search:_search_segments}\n| ]\n\n", "versions": {"daft": "0.7.14", "pylance": "12.0.0", "python": "3.11.2"}, "ray_version": "2.54.0", "source_sha256": {"benchmarks/vector_search.py": "79bcd4580732ef6b3c0bbfa78cd1aee500f7277b8521291b2952985d22cb59c6", "daft_lance/lance_vector_search.py": "5dfe9f473538e80db6afdeb72755b51223b40375753173f5d80dca7d8eeedb42", "daft_lance/_lance.py": "4f854d14838a299f0959b223d372c4c4a0fe577e9d3330da155c05ccf22e6820"}, "git_commit": "21d78f379d0d47ddebfc2044b32f4001b8501b28", "git_dirty": true} +{"backend": "segments", "runner": "ray", "workers": 4, "cpu_budget": 4, "cpu_affinity": [0, 1, 2, 3], "threads_per_process": "1", "dataset_version": 9, "rows": 1000000, "segments": {"bench_idx": 8}, "queries_per_run": 1, "nearest": {"column": "vector", "k": 10, "nprobes": 32, "metric": "l2"}, "filter": null, "tasks": null, "first_run_seconds": 1.865760125219822, "first_run_recall": 1.0, "warm_batch_p50_seconds": 0.09158205147832632, "warm_batch_p95_seconds": 0.10046986695379019, "queries_per_second": 10.830935482974454, "recall_min": 1.0, "local_cpu_seconds": 2.1799999999999997, "local_physical_read_bytes": 0, "local_peak_aggregate_rss_bytes": 2906271744, "first_run_peak_aggregate_rss_bytes": 2386391040, "durations_seconds": [0.08616053126752377, 0.0924088079482317, 0.09477577172219753, 0.09075529500842094, 0.09749267622828484, 0.08938949555158615, 0.10026966594159603, 0.10063366778194904, 0.08665031008422375, 0.08474506065249443], "execution_plan": "== Unoptimized Logical Plan ==\n\n* Sort: Sort by = (col(_query_id), ascending, nulls last), (col(_distance),\n| ascending, nulls last), (col(_rowid), ascending, nulls last)\n|\n* Limit: 10\n|\n* Sort: Sort by = (col(_distance), ascending, nulls last), (col(_rowid),\n| ascending, nulls last)\n|\n* PythonScanOperator: LanceVectorSearch(version=9, query=0)\n| File schema = id#Int64, _distance#Float32, _rowid#UInt64, _query_id#Int64\n| Partitioning keys = []\n| Output schema = id#Int64, _distance#Float32, _rowid#UInt64, _query_id#Int64\n\n\n== Optimized Logical Plan ==\n\n* Sort: Sort by = (col(_query_id), ascending, nulls last), (col(_distance),\n| ascending, nulls last), (col(_rowid), ascending, nulls last)\n| Stats = { Approx num rows = 0, Approx size bytes = 0 B, Accumulated selectivity\n| = 1.00 }\n|\n* TopN: Sort by = (col(_distance), ascending, nulls last), (col(_rowid),\n| ascending, nulls last), Num Rows = 10\n| Stats = { Approx num rows = 0, Approx size bytes = 0 B, Accumulated selectivity\n| = 1.00 }\n|\n* Num Scan Tasks = 8\n| Schema: {id#Int64, _distance#Float32, _rowid#UInt64, _query_id#Int64}\n| Scan Tasks: [\n| {daft_lance.lance_vector_search:_search_segments}\n| {daft_lance.lance_vector_search:_search_segments}\n| {daft_lance.lance_vector_search:_search_segments}\n| ...\n| {daft_lance.lance_vector_search:_search_segments}\n| {daft_lance.lance_vector_search:_search_segments}\n| {daft_lance.lance_vector_search:_search_segments}\n| ]\n| File schema = id#Int64, _distance#Float32, _rowid#UInt64, _query_id#Int64\n| Partitioning keys = []\n| Output schema = id#Int64, _distance#Float32, _rowid#UInt64, _query_id#Int64\n| Stats = { Approx num rows = 0, Approx size bytes = 0 B, Accumulated selectivity\n| = 1.00 }\n\n\n== Physical Plan ==\n\n* Sort\n| Sort by: col(3: _query_id), col(1: _distance), col(2: _rowid)\n| Descending = false, false, false\n| Nulls first = false, false, false\n|\n* TopN\n| Sort by: col(1: _distance), col(2: _rowid)\n| Descending = false, false\n| Nulls first = false, false\n| Num Rows = 10\n|\n* RayGather\n|\n* TopN\n| Sort by: col(1: _distance), col(2: _rowid)\n| Descending = false, false\n| Nulls first = false, false\n| Num Rows = 10\n| Offset = 0\n|\n* ScanTaskSource:\n| Num Scan Tasks = 8\n| Estimated Scan Bytes = 0\n| Source = LanceVectorSearch(version=9, query=0)\n| Schema: {id#Int64, _distance#Float32, _rowid#UInt64, _query_id#Int64}\n| Scan Tasks: [\n| {daft_lance.lance_vector_search:_search_segments}\n| {daft_lance.lance_vector_search:_search_segments}\n| {daft_lance.lance_vector_search:_search_segments}\n| ...\n| {daft_lance.lance_vector_search:_search_segments}\n| {daft_lance.lance_vector_search:_search_segments}\n| {daft_lance.lance_vector_search:_search_segments}\n| ]\n\n", "versions": {"daft": "0.7.14", "pylance": "12.0.0", "python": "3.11.2"}, "ray_version": "2.54.0", "source_sha256": {"benchmarks/vector_search.py": "79bcd4580732ef6b3c0bbfa78cd1aee500f7277b8521291b2952985d22cb59c6", "daft_lance/lance_vector_search.py": "5dfe9f473538e80db6afdeb72755b51223b40375753173f5d80dca7d8eeedb42", "daft_lance/_lance.py": "4f854d14838a299f0959b223d372c4c4a0fe577e9d3330da155c05ccf22e6820"}, "git_commit": "21d78f379d0d47ddebfc2044b32f4001b8501b28", "git_dirty": true} +{"backend": "segments", "runner": "ray", "workers": 8, "cpu_budget": 8, "cpu_affinity": [0, 1, 2, 3, 4, 5, 6, 7], "threads_per_process": "1", "dataset_version": 9, "rows": 1000000, "segments": {"bench_idx": 8}, "queries_per_run": 1, "nearest": {"column": "vector", "k": 10, "nprobes": 32, "metric": "l2"}, "filter": null, "tasks": null, "first_run_seconds": 1.7860618010163307, "first_run_recall": 1.0, "warm_batch_p50_seconds": 0.06881084106862545, "warm_batch_p95_seconds": 0.07283904692158102, "queries_per_second": 14.597775546630052, "recall_min": 1.0, "local_cpu_seconds": 2.1600000000000006, "local_physical_read_bytes": 0, "local_peak_aggregate_rss_bytes": 3248947200, "first_run_peak_aggregate_rss_bytes": 2896703488, "durations_seconds": [0.06942787952721119, 0.07402101345360279, 0.06642328761518002, 0.07086150720715523, 0.06753488443791866, 0.0713944211602211, 0.06819380261003971, 0.0632675550878048, 0.06440788879990578, 0.06950363889336586], "execution_plan": "== Unoptimized Logical Plan ==\n\n* Sort: Sort by = (col(_query_id), ascending, nulls last), (col(_distance),\n| ascending, nulls last), (col(_rowid), ascending, nulls last)\n|\n* Limit: 10\n|\n* Sort: Sort by = (col(_distance), ascending, nulls last), (col(_rowid),\n| ascending, nulls last)\n|\n* PythonScanOperator: LanceVectorSearch(version=9, query=0)\n| File schema = id#Int64, _distance#Float32, _rowid#UInt64, _query_id#Int64\n| Partitioning keys = []\n| Output schema = id#Int64, _distance#Float32, _rowid#UInt64, _query_id#Int64\n\n\n== Optimized Logical Plan ==\n\n* Sort: Sort by = (col(_query_id), ascending, nulls last), (col(_distance),\n| ascending, nulls last), (col(_rowid), ascending, nulls last)\n| Stats = { Approx num rows = 0, Approx size bytes = 0 B, Accumulated selectivity\n| = 1.00 }\n|\n* TopN: Sort by = (col(_distance), ascending, nulls last), (col(_rowid),\n| ascending, nulls last), Num Rows = 10\n| Stats = { Approx num rows = 0, Approx size bytes = 0 B, Accumulated selectivity\n| = 1.00 }\n|\n* Num Scan Tasks = 8\n| Schema: {id#Int64, _distance#Float32, _rowid#UInt64, _query_id#Int64}\n| Scan Tasks: [\n| {daft_lance.lance_vector_search:_search_segments}\n| {daft_lance.lance_vector_search:_search_segments}\n| {daft_lance.lance_vector_search:_search_segments}\n| ...\n| {daft_lance.lance_vector_search:_search_segments}\n| {daft_lance.lance_vector_search:_search_segments}\n| {daft_lance.lance_vector_search:_search_segments}\n| ]\n| File schema = id#Int64, _distance#Float32, _rowid#UInt64, _query_id#Int64\n| Partitioning keys = []\n| Output schema = id#Int64, _distance#Float32, _rowid#UInt64, _query_id#Int64\n| Stats = { Approx num rows = 0, Approx size bytes = 0 B, Accumulated selectivity\n| = 1.00 }\n\n\n== Physical Plan ==\n\n* Sort\n| Sort by: col(3: _query_id), col(1: _distance), col(2: _rowid)\n| Descending = false, false, false\n| Nulls first = false, false, false\n|\n* TopN\n| Sort by: col(1: _distance), col(2: _rowid)\n| Descending = false, false\n| Nulls first = false, false\n| Num Rows = 10\n|\n* RayGather\n|\n* TopN\n| Sort by: col(1: _distance), col(2: _rowid)\n| Descending = false, false\n| Nulls first = false, false\n| Num Rows = 10\n| Offset = 0\n|\n* ScanTaskSource:\n| Num Scan Tasks = 8\n| Estimated Scan Bytes = 0\n| Source = LanceVectorSearch(version=9, query=0)\n| Schema: {id#Int64, _distance#Float32, _rowid#UInt64, _query_id#Int64}\n| Scan Tasks: [\n| {daft_lance.lance_vector_search:_search_segments}\n| {daft_lance.lance_vector_search:_search_segments}\n| {daft_lance.lance_vector_search:_search_segments}\n| ...\n| {daft_lance.lance_vector_search:_search_segments}\n| {daft_lance.lance_vector_search:_search_segments}\n| {daft_lance.lance_vector_search:_search_segments}\n| ]\n\n", "versions": {"daft": "0.7.14", "pylance": "12.0.0", "python": "3.11.2"}, "ray_version": "2.54.0", "source_sha256": {"benchmarks/vector_search.py": "79bcd4580732ef6b3c0bbfa78cd1aee500f7277b8521291b2952985d22cb59c6", "daft_lance/lance_vector_search.py": "acc7d4cc3e70875b35b9eed9ba656899bf7efb94aa60d7a66742e763eb80fbb5", "daft_lance/_lance.py": "4f854d14838a299f0959b223d372c4c4a0fe577e9d3330da155c05ccf22e6820"}, "git_commit": "21d78f379d0d47ddebfc2044b32f4001b8501b28", "git_dirty": true} +{"backend": "daft-single", "runner": "ray", "workers": 4, "cpu_budget": 4, "cpu_affinity": [0, 1, 2, 3], "threads_per_process": "4", "dataset_version": 9, "rows": 1000000, "segments": {"bench_idx": 8}, "queries_per_run": 1, "nearest": {"column": "vector", "k": 10, "nprobes": 32, "metric": "l2"}, "filter": null, "tasks": null, "first_run_seconds": 2.154959099367261, "first_run_recall": 1.0, "warm_batch_p50_seconds": 0.2054679747670889, "warm_batch_p95_seconds": 0.22365041794255375, "queries_per_second": 4.811192755872849, "recall_min": 1.0, "local_cpu_seconds": 5.62, "local_physical_read_bytes": 0, "local_peak_aggregate_rss_bytes": 2955853824, "first_run_peak_aggregate_rss_bytes": 2504040448, "durations_seconds": [0.2200795505195856, 0.19801519997417927, 0.20659608207643032, 0.20924281887710094, 0.22657203674316406, 0.20270158164203167, 0.20433986745774746, 0.21296095103025436, 0.19480932131409645, 0.20316925831139088], "execution_plan": null, "versions": {"daft": "0.7.14", "pylance": "12.0.0", "python": "3.11.2"}, "ray_version": "2.54.0", "source_sha256": {"benchmarks/vector_search.py": "8b8b4b84a436353062f3c542eaf9d2d92c77c4acfbcd0d3ba29f36d9daaf54f2", "daft_lance/lance_vector_search.py": "acc7d4cc3e70875b35b9eed9ba656899bf7efb94aa60d7a66742e763eb80fbb5", "daft_lance/_lance.py": "4f854d14838a299f0959b223d372c4c4a0fe577e9d3330da155c05ccf22e6820"}, "git_commit": "21d78f379d0d47ddebfc2044b32f4001b8501b28", "git_dirty": true} diff --git a/benchmarks/results/local-1m.jsonl b/benchmarks/results/local-1m.jsonl new file mode 100644 index 0000000..4e9264b --- /dev/null +++ b/benchmarks/results/local-1m.jsonl @@ -0,0 +1,6 @@ +{"backend": "pylance", "runner": "native", "workers": 4, "cpu_budget": 4, "cpu_affinity": [0, 1, 2, 3], "threads_per_process": "4", "dataset_version": 9, "rows": 1000000, "segments": {"bench_idx": 8}, "queries_per_run": 5, "nearest": {"column": "vector", "k": 10, "nprobes": 32, "metric": "l2"}, "filter": null, "tasks": null, "first_run_seconds": 0.3945568297058344, "first_run_recall": 1.0, "warm_batch_p50_seconds": 0.17314535938203335, "warm_batch_p95_seconds": 0.17856472097337245, "queries_per_second": 28.684039067954952, "recall_min": 1.0, "local_cpu_seconds": 2.0199999999999996, "local_physical_read_bytes": 0, "local_peak_aggregate_rss_bytes": 1075986432, "first_run_peak_aggregate_rss_bytes": 880943104, "durations_seconds": [0.1719979103654623, 0.17314535938203335, 0.17966712825000286, 0.17259935289621353, 0.17415509186685085], "execution_plan": null, "versions": {"daft": "0.7.14", "pylance": "12.0.0", "python": "3.11.2"}, "git_commit": "21d78f379d0d47ddebfc2044b32f4001b8501b28", "git_dirty": true} +{"backend": "segments", "runner": "ray", "workers": 1, "cpu_budget": 4, "cpu_affinity": [0, 1, 2, 3], "threads_per_process": "4", "dataset_version": 9, "rows": 1000000, "segments": {"bench_idx": 8}, "queries_per_run": 5, "nearest": {"column": "vector", "k": 10, "nprobes": 32, "metric": "l2"}, "filter": null, "tasks": null, "first_run_seconds": 3.865571115165949, "first_run_recall": 1.0, "warm_batch_p50_seconds": 0.9383172523230314, "warm_batch_p95_seconds": 0.9571197457611561, "queries_per_second": 5.3440353752048715, "recall_min": 1.0, "local_cpu_seconds": 7.420000000000001, "local_physical_read_bytes": 0, "local_peak_aggregate_rss_bytes": 2908196864, "first_run_peak_aggregate_rss_bytes": 2497482752, "durations_seconds": [0.9089293479919434, 0.9383172523230314, 0.9397258721292019, 0.9614682141691446, 0.9296720586717129], "execution_plan": "== Unoptimized Logical Plan ==\n\n* Sort: Sort by = (col(_query_id), ascending, nulls last), (col(_distance),\n| ascending, nulls last), (col(_rowid), ascending, nulls last)\n|\n* Limit: 10\n|\n* Sort: Sort by = (col(_distance), ascending, nulls last), (col(_rowid),\n| ascending, nulls last)\n|\n* PythonScanOperator: LanceVectorSearch(version=9, query=0)\n| File schema = id#Int64, _distance#Float32, _rowid#UInt64, _query_id#Int64\n| Partitioning keys = []\n| Output schema = id#Int64, _distance#Float32, _rowid#UInt64, _query_id#Int64\n\n\n== Optimized Logical Plan ==\n\n* Sort: Sort by = (col(_query_id), ascending, nulls last), (col(_distance),\n| ascending, nulls last), (col(_rowid), ascending, nulls last)\n| Stats = { Approx num rows = 0, Approx size bytes = 0 B, Accumulated selectivity\n| = 1.00 }\n|\n* TopN: Sort by = (col(_distance), ascending, nulls last), (col(_rowid),\n| ascending, nulls last), Num Rows = 10\n| Stats = { Approx num rows = 0, Approx size bytes = 0 B, Accumulated selectivity\n| = 1.00 }\n|\n* Num Scan Tasks = 8\n| Schema: {id#Int64, _distance#Float32, _rowid#UInt64, _query_id#Int64}\n| Scan Tasks: [\n| {daft_lance.lance_vector_search:_search_segments}\n| {daft_lance.lance_vector_search:_search_segments}\n| {daft_lance.lance_vector_search:_search_segments}\n| ...\n| {daft_lance.lance_vector_search:_search_segments}\n| {daft_lance.lance_vector_search:_search_segments}\n| {daft_lance.lance_vector_search:_search_segments}\n| ]\n| File schema = id#Int64, _distance#Float32, _rowid#UInt64, _query_id#Int64\n| Partitioning keys = []\n| Output schema = id#Int64, _distance#Float32, _rowid#UInt64, _query_id#Int64\n| Stats = { Approx num rows = 0, Approx size bytes = 0 B, Accumulated selectivity\n| = 1.00 }\n\n\n== Physical Plan ==\n\n* Sort\n| Sort by: col(3: _query_id), col(1: _distance), col(2: _rowid)\n| Descending = false, false, false\n| Nulls first = false, false, false\n|\n* TopN\n| Sort by: col(1: _distance), col(2: _rowid)\n| Descending = false, false\n| Nulls first = false, false\n| Num Rows = 10\n|\n* RayGather\n|\n* TopN\n| Sort by: col(1: _distance), col(2: _rowid)\n| Descending = false, false\n| Nulls first = false, false\n| Num Rows = 10\n| Offset = 0\n|\n* ScanTaskSource:\n| Num Scan Tasks = 8\n| Estimated Scan Bytes = 0\n| Source = LanceVectorSearch(version=9, query=0)\n| Schema: {id#Int64, _distance#Float32, _rowid#UInt64, _query_id#Int64}\n| Scan Tasks: [\n| {daft_lance.lance_vector_search:_search_segments}\n| {daft_lance.lance_vector_search:_search_segments}\n| {daft_lance.lance_vector_search:_search_segments}\n| ...\n| {daft_lance.lance_vector_search:_search_segments}\n| {daft_lance.lance_vector_search:_search_segments}\n| {daft_lance.lance_vector_search:_search_segments}\n| ]\n\n", "versions": {"daft": "0.7.14", "pylance": "12.0.0", "python": "3.11.2"}, "git_commit": "21d78f379d0d47ddebfc2044b32f4001b8501b28", "git_dirty": true} +{"backend": "segments", "runner": "ray", "workers": 2, "cpu_budget": 4, "cpu_affinity": [0, 1, 2, 3], "threads_per_process": "2", "dataset_version": 9, "rows": 1000000, "segments": {"bench_idx": 8}, "queries_per_run": 5, "nearest": {"column": "vector", "k": 10, "nprobes": 32, "metric": "l2"}, "filter": null, "tasks": null, "first_run_seconds": 2.9262431748211384, "first_run_recall": 1.0, "warm_batch_p50_seconds": 0.5409161671996117, "warm_batch_p95_seconds": 0.5592106062918901, "queries_per_second": 9.201664908682663, "recall_min": 1.0, "local_cpu_seconds": 6.040000000000001, "local_physical_read_bytes": 0, "local_peak_aggregate_rss_bytes": 3037667328, "first_run_peak_aggregate_rss_bytes": 2621431808, "durations_seconds": [0.5318338833749294, 0.5446091238409281, 0.5366794802248478, 0.5628609769046307, 0.5409161671996117], "execution_plan": "== Unoptimized Logical Plan ==\n\n* Sort: Sort by = (col(_query_id), ascending, nulls last), (col(_distance),\n| ascending, nulls last), (col(_rowid), ascending, nulls last)\n|\n* Limit: 10\n|\n* Sort: Sort by = (col(_distance), ascending, nulls last), (col(_rowid),\n| ascending, nulls last)\n|\n* PythonScanOperator: LanceVectorSearch(version=9, query=0)\n| File schema = id#Int64, _distance#Float32, _rowid#UInt64, _query_id#Int64\n| Partitioning keys = []\n| Output schema = id#Int64, _distance#Float32, _rowid#UInt64, _query_id#Int64\n\n\n== Optimized Logical Plan ==\n\n* Sort: Sort by = (col(_query_id), ascending, nulls last), (col(_distance),\n| ascending, nulls last), (col(_rowid), ascending, nulls last)\n| Stats = { Approx num rows = 0, Approx size bytes = 0 B, Accumulated selectivity\n| = 1.00 }\n|\n* TopN: Sort by = (col(_distance), ascending, nulls last), (col(_rowid),\n| ascending, nulls last), Num Rows = 10\n| Stats = { Approx num rows = 0, Approx size bytes = 0 B, Accumulated selectivity\n| = 1.00 }\n|\n* Num Scan Tasks = 8\n| Schema: {id#Int64, _distance#Float32, _rowid#UInt64, _query_id#Int64}\n| Scan Tasks: [\n| {daft_lance.lance_vector_search:_search_segments}\n| {daft_lance.lance_vector_search:_search_segments}\n| {daft_lance.lance_vector_search:_search_segments}\n| ...\n| {daft_lance.lance_vector_search:_search_segments}\n| {daft_lance.lance_vector_search:_search_segments}\n| {daft_lance.lance_vector_search:_search_segments}\n| ]\n| File schema = id#Int64, _distance#Float32, _rowid#UInt64, _query_id#Int64\n| Partitioning keys = []\n| Output schema = id#Int64, _distance#Float32, _rowid#UInt64, _query_id#Int64\n| Stats = { Approx num rows = 0, Approx size bytes = 0 B, Accumulated selectivity\n| = 1.00 }\n\n\n== Physical Plan ==\n\n* Sort\n| Sort by: col(3: _query_id), col(1: _distance), col(2: _rowid)\n| Descending = false, false, false\n| Nulls first = false, false, false\n|\n* TopN\n| Sort by: col(1: _distance), col(2: _rowid)\n| Descending = false, false\n| Nulls first = false, false\n| Num Rows = 10\n|\n* RayGather\n|\n* TopN\n| Sort by: col(1: _distance), col(2: _rowid)\n| Descending = false, false\n| Nulls first = false, false\n| Num Rows = 10\n| Offset = 0\n|\n* ScanTaskSource:\n| Num Scan Tasks = 8\n| Estimated Scan Bytes = 0\n| Source = LanceVectorSearch(version=9, query=0)\n| Schema: {id#Int64, _distance#Float32, _rowid#UInt64, _query_id#Int64}\n| Scan Tasks: [\n| {daft_lance.lance_vector_search:_search_segments}\n| {daft_lance.lance_vector_search:_search_segments}\n| {daft_lance.lance_vector_search:_search_segments}\n| ...\n| {daft_lance.lance_vector_search:_search_segments}\n| {daft_lance.lance_vector_search:_search_segments}\n| {daft_lance.lance_vector_search:_search_segments}\n| ]\n\n", "versions": {"daft": "0.7.14", "pylance": "12.0.0", "python": "3.11.2"}, "git_commit": "21d78f379d0d47ddebfc2044b32f4001b8501b28", "git_dirty": true} +{"backend": "segments", "runner": "ray", "workers": 4, "cpu_budget": 4, "cpu_affinity": [0, 1, 2, 3], "threads_per_process": "1", "dataset_version": 9, "rows": 1000000, "segments": {"bench_idx": 8}, "queries_per_run": 5, "nearest": {"column": "vector", "k": 10, "nprobes": 32, "metric": "l2"}, "filter": null, "tasks": null, "first_run_seconds": 2.4935148023068905, "first_run_recall": 1.0, "warm_batch_p50_seconds": 0.43652975745499134, "warm_batch_p95_seconds": 0.458802667632699, "queries_per_second": 11.568039653435868, "recall_min": 1.0, "local_cpu_seconds": 5.529999999999999, "local_physical_read_bytes": 0, "local_peak_aggregate_rss_bytes": 3125792768, "first_run_peak_aggregate_rss_bytes": 2658480128, "durations_seconds": [0.4080206826329231, 0.43652975745499134, 0.40501311607658863, 0.4507464859634638, 0.4608167130500078], "execution_plan": "== Unoptimized Logical Plan ==\n\n* Sort: Sort by = (col(_query_id), ascending, nulls last), (col(_distance),\n| ascending, nulls last), (col(_rowid), ascending, nulls last)\n|\n* Limit: 10\n|\n* Sort: Sort by = (col(_distance), ascending, nulls last), (col(_rowid),\n| ascending, nulls last)\n|\n* PythonScanOperator: LanceVectorSearch(version=9, query=0)\n| File schema = id#Int64, _distance#Float32, _rowid#UInt64, _query_id#Int64\n| Partitioning keys = []\n| Output schema = id#Int64, _distance#Float32, _rowid#UInt64, _query_id#Int64\n\n\n== Optimized Logical Plan ==\n\n* Sort: Sort by = (col(_query_id), ascending, nulls last), (col(_distance),\n| ascending, nulls last), (col(_rowid), ascending, nulls last)\n| Stats = { Approx num rows = 0, Approx size bytes = 0 B, Accumulated selectivity\n| = 1.00 }\n|\n* TopN: Sort by = (col(_distance), ascending, nulls last), (col(_rowid),\n| ascending, nulls last), Num Rows = 10\n| Stats = { Approx num rows = 0, Approx size bytes = 0 B, Accumulated selectivity\n| = 1.00 }\n|\n* Num Scan Tasks = 8\n| Schema: {id#Int64, _distance#Float32, _rowid#UInt64, _query_id#Int64}\n| Scan Tasks: [\n| {daft_lance.lance_vector_search:_search_segments}\n| {daft_lance.lance_vector_search:_search_segments}\n| {daft_lance.lance_vector_search:_search_segments}\n| ...\n| {daft_lance.lance_vector_search:_search_segments}\n| {daft_lance.lance_vector_search:_search_segments}\n| {daft_lance.lance_vector_search:_search_segments}\n| ]\n| File schema = id#Int64, _distance#Float32, _rowid#UInt64, _query_id#Int64\n| Partitioning keys = []\n| Output schema = id#Int64, _distance#Float32, _rowid#UInt64, _query_id#Int64\n| Stats = { Approx num rows = 0, Approx size bytes = 0 B, Accumulated selectivity\n| = 1.00 }\n\n\n== Physical Plan ==\n\n* Sort\n| Sort by: col(3: _query_id), col(1: _distance), col(2: _rowid)\n| Descending = false, false, false\n| Nulls first = false, false, false\n|\n* TopN\n| Sort by: col(1: _distance), col(2: _rowid)\n| Descending = false, false\n| Nulls first = false, false\n| Num Rows = 10\n|\n* RayGather\n|\n* TopN\n| Sort by: col(1: _distance), col(2: _rowid)\n| Descending = false, false\n| Nulls first = false, false\n| Num Rows = 10\n| Offset = 0\n|\n* ScanTaskSource:\n| Num Scan Tasks = 8\n| Estimated Scan Bytes = 0\n| Source = LanceVectorSearch(version=9, query=0)\n| Schema: {id#Int64, _distance#Float32, _rowid#UInt64, _query_id#Int64}\n| Scan Tasks: [\n| {daft_lance.lance_vector_search:_search_segments}\n| {daft_lance.lance_vector_search:_search_segments}\n| {daft_lance.lance_vector_search:_search_segments}\n| ...\n| {daft_lance.lance_vector_search:_search_segments}\n| {daft_lance.lance_vector_search:_search_segments}\n| {daft_lance.lance_vector_search:_search_segments}\n| ]\n\n", "versions": {"daft": "0.7.14", "pylance": "12.0.0", "python": "3.11.2"}, "ray_version": "2.54.0", "source_sha256": {"benchmarks/vector_search.py": "79bcd4580732ef6b3c0bbfa78cd1aee500f7277b8521291b2952985d22cb59c6", "daft_lance/lance_vector_search.py": "5dfe9f473538e80db6afdeb72755b51223b40375753173f5d80dca7d8eeedb42", "daft_lance/_lance.py": "4f854d14838a299f0959b223d372c4c4a0fe577e9d3330da155c05ccf22e6820"}, "git_commit": "21d78f379d0d47ddebfc2044b32f4001b8501b28", "git_dirty": true} +{"backend": "queries", "runner": "ray", "workers": 4, "cpu_budget": 4, "cpu_affinity": [0, 1, 2, 3], "threads_per_process": "1", "dataset_version": 9, "rows": 1000000, "segments": {"bench_idx": 8}, "queries_per_run": 5, "nearest": {"column": "vector", "k": 10, "nprobes": 32, "metric": "l2"}, "filter": null, "tasks": null, "first_run_seconds": 3.007525434717536, "first_run_recall": 1.0, "warm_batch_p50_seconds": 0.26666103675961494, "warm_batch_p95_seconds": 0.2711159996688366, "queries_per_second": 18.761841245132235, "recall_min": 1.0, "local_cpu_seconds": 3.35, "local_physical_read_bytes": 0, "local_peak_aggregate_rss_bytes": 4583419904, "first_run_peak_aggregate_rss_bytes": 4231446528, "durations_seconds": [0.27193792164325714, 0.2678283117711544, 0.2653754875063896, 0.2606890629976988, 0.26666103675961494], "execution_plan": "== Unoptimized Logical Plan ==\n\n* Sort: Sort by = (col(_query_id), ascending, nulls last), (col(_distance),\n| ascending, nulls last), (col(_rowid), ascending, nulls last)\n|\n* Limit: 10\n|\n* Sort: Sort by = (col(_distance), ascending, nulls last), (col(_rowid),\n| ascending, nulls last)\n|\n* PythonScanOperator: LanceVectorSearch(version=9, query=0)\n| File schema = id#Int64, _distance#Float32, _rowid#UInt64, _query_id#Int64\n| Partitioning keys = []\n| Output schema = id#Int64, _distance#Float32, _rowid#UInt64, _query_id#Int64\n\n\n== Optimized Logical Plan ==\n\n* Sort: Sort by = (col(_query_id), ascending, nulls last), (col(_distance),\n| ascending, nulls last), (col(_rowid), ascending, nulls last)\n| Stats = { Approx num rows = 0, Approx size bytes = 0 B, Accumulated selectivity\n| = 1.00 }\n|\n* TopN: Sort by = (col(_distance), ascending, nulls last), (col(_rowid),\n| ascending, nulls last), Num Rows = 10\n| Stats = { Approx num rows = 0, Approx size bytes = 0 B, Accumulated selectivity\n| = 1.00 }\n|\n* Num Scan Tasks = 1\n| Schema: {id#Int64, _distance#Float32, _rowid#UInt64, _query_id#Int64}\n| Scan Tasks: [\n| {daft_lance.lance_vector_search:_search_segments}\n| ]\n| File schema = id#Int64, _distance#Float32, _rowid#UInt64, _query_id#Int64\n| Partitioning keys = []\n| Output schema = id#Int64, _distance#Float32, _rowid#UInt64, _query_id#Int64\n| Stats = { Approx num rows = 0, Approx size bytes = 0 B, Accumulated selectivity\n| = 1.00 }\n\n\n== Physical Plan ==\n\n* Sort\n| Sort by: col(3: _query_id), col(1: _distance), col(2: _rowid)\n| Descending = false, false, false\n| Nulls first = false, false, false\n|\n* TopN\n| Sort by: col(1: _distance), col(2: _rowid)\n| Descending = false, false\n| Nulls first = false, false\n| Num Rows = 10\n|\n* TopN\n| Sort by: col(1: _distance), col(2: _rowid)\n| Descending = false, false\n| Nulls first = false, false\n| Num Rows = 10\n| Offset = 0\n|\n* ScanTaskSource:\n| Num Scan Tasks = 1\n| Estimated Scan Bytes = 0\n| Source = LanceVectorSearch(version=9, query=0)\n| Schema: {id#Int64, _distance#Float32, _rowid#UInt64, _query_id#Int64}\n| Scan Tasks: [\n| {daft_lance.lance_vector_search:_search_segments}\n| ]\n\n", "versions": {"daft": "0.7.14", "pylance": "12.0.0", "python": "3.11.2"}, "ray_version": "2.54.0", "source_sha256": {"benchmarks/vector_search.py": "79bcd4580732ef6b3c0bbfa78cd1aee500f7277b8521291b2952985d22cb59c6", "daft_lance/lance_vector_search.py": "5dfe9f473538e80db6afdeb72755b51223b40375753173f5d80dca7d8eeedb42", "daft_lance/_lance.py": "4f854d14838a299f0959b223d372c4c4a0fe577e9d3330da155c05ccf22e6820"}, "git_commit": "21d78f379d0d47ddebfc2044b32f4001b8501b28", "git_dirty": true} +{"backend": "daft-single", "runner": "ray", "workers": 4, "cpu_budget": 4, "cpu_affinity": [0, 1, 2, 3], "threads_per_process": "4", "dataset_version": 9, "rows": 1000000, "segments": {"bench_idx": 8}, "queries_per_run": 5, "nearest": {"column": "vector", "k": 10, "nprobes": 32, "metric": "l2"}, "filter": null, "tasks": null, "first_run_seconds": 2.7965883295983076, "first_run_recall": 1.0, "warm_batch_p50_seconds": 0.973969854414463, "warm_batch_p95_seconds": 1.0046620029956101, "queries_per_second": 5.122165515131486, "recall_min": 1.0, "local_cpu_seconds": 12.729999999999999, "local_physical_read_bytes": 0, "local_peak_aggregate_rss_bytes": 3142033408, "first_run_peak_aggregate_rss_bytes": 2596118528, "durations_seconds": [0.9950273428112268, 0.973969854414463, 0.9562962613999844, 0.9483840502798557, 1.007070668041706], "execution_plan": null, "versions": {"daft": "0.7.14", "pylance": "12.0.0", "python": "3.11.2"}, "ray_version": "2.54.0", "source_sha256": {"benchmarks/vector_search.py": "8b8b4b84a436353062f3c542eaf9d2d92c77c4acfbcd0d3ba29f36d9daaf54f2", "daft_lance/lance_vector_search.py": "acc7d4cc3e70875b35b9eed9ba656899bf7efb94aa60d7a66742e763eb80fbb5", "daft_lance/_lance.py": "4f854d14838a299f0959b223d372c4c4a0fe577e9d3330da155c05ccf22e6820"}, "git_commit": "21d78f379d0d47ddebfc2044b32f4001b8501b28", "git_dirty": true} diff --git a/benchmarks/vector_search.py b/benchmarks/vector_search.py new file mode 100644 index 0000000..7b31a51 --- /dev/null +++ b/benchmarks/vector_search.py @@ -0,0 +1,374 @@ +"""Reproducible vector-search benchmark; see benchmarks/README.md.""" + +from __future__ import annotations + +import argparse +import json +import os +import platform +import subprocess +import threading +import time +from pathlib import Path +from typing import Any + + +def _prepare(args: argparse.Namespace) -> None: + import lance + import numpy as np + import pyarrow as pa + + from daft_lance import create_vector_index + + root = Path(args.root) + root.mkdir(parents=True, exist_ok=True) + uri = str(root / "data.lance") + if Path(uri).exists(): + raise ValueError("Dataset already exists; use a new root to avoid replacing benchmark data") + rng = np.random.default_rng(args.seed) + centers = rng.normal(size=(16, args.dim)).astype("float32") + for segment in range(args.segments): + start = args.rows * segment // args.segments + stop = args.rows * (segment + 1) // args.segments + labels = rng.integers(0, len(centers), stop - start) + vectors = centers[labels] + rng.normal(scale=0.5, size=(stop - start, args.dim)).astype("float32") + lance.write_dataset( + pa.table( + { + "id": np.arange(start, stop), + "bucket": rng.integers(0, 1000, stop - start), + "vector": pa.FixedSizeListArray.from_arrays(pa.array(vectors.ravel()), args.dim), + } + ), + uri, + mode="create" if segment == 0 else "append", + max_rows_per_file=stop - start, + ) + # Independent queries, drawn from the same mixture rather than stored rows. + query_rng = np.random.default_rng(args.seed + 1) + queries = centers[query_rng.integers(0, len(centers), args.queries)] + query_rng.normal( + scale=0.5, size=(args.queries, args.dim) + ).astype("float32") + np.save(root / "queries.npy", queries) + create_vector_index( + uri, + column="vector", + index_type=args.index_type, + name="bench_idx", + num_partitions=args.partitions, + num_sub_vectors=args.subvectors if args.index_type == "IVF_PQ" else None, + metric=args.metric, + sample_rate=2, + fragment_group_size=1, + ) + ds = lance.dataset(uri) + metadata = { + key: getattr(args, key) + for key in ("rows", "dim", "segments", "queries", "seed", "index_type", "partitions", "metric") + } + metadata.update(version=ds.version, actual_segments=len(ds.describe_indices()[0].segments)) + (root / "dataset.json").write_text(json.dumps(metadata, indent=2) + "\n") + print(json.dumps(metadata)) + + +class _Resources: + """Sample local driver descendants; remote workers are outside this metric.""" + + def __init__(self) -> None: + import psutil + + self.process = psutil.Process() + self.stop = threading.Event() + self.peak_rss = 0 + self.initial: dict[int, tuple[float, int]] = {} + self.cpu_seconds = 0.0 + self.read_bytes = 0 + self.thread = threading.Thread(target=self._sample, daemon=True) + + def _sample_once(self) -> None: + import psutil + + rss = 0 + cpu = 0.0 + read_bytes = 0 + for process in [self.process, *self.process.children(recursive=True)]: + try: + times = process.cpu_times() + current = times.user + times.system + reads = process.io_counters().read_bytes + initial_cpu, initial_reads = self.initial.setdefault(process.pid, (current, reads)) + cpu += current - initial_cpu + read_bytes += reads - initial_reads + rss += process.memory_info().rss + except (psutil.NoSuchProcess, psutil.AccessDenied): + continue + self.peak_rss = max(self.peak_rss, rss) + self.cpu_seconds = max(self.cpu_seconds, cpu) + self.read_bytes = max(self.read_bytes, read_bytes) + + def _sample(self) -> None: + while not self.stop.wait(0.02): + self._sample_once() + + def __enter__(self) -> _Resources: + self._sample_once() + self.thread.start() + return self + + def __exit__(self, *args: Any) -> None: + self.stop.set() + self.thread.join() + self._sample_once() + + +def _run(args: argparse.Namespace) -> None: + import hashlib + import io + + import lance + import numpy as np + import pyarrow as pa + + import daft + from daft_lance import read_lance, vector_search + + root = Path(args.root) + uri = args.uri or str(root / "data.lance") + queries = np.load(args.query_file or root / "queries.npy")[: args.queries] + if len(queries) == 0: + raise ValueError("Query file is empty") + ds = lance.dataset(uri) + version = ds.version + # The single-scan baselines select their index through Lance. Requiring one + # relevant index makes every backend use the same files without assuming + # Lance's default selection agrees with an explicit distributed selection. + field = ds.lance_schema.field(args.column) + indices = [ + idx + for idx in ds.describe_indices() + if str(idx.index_type).startswith("IVF_") and field is not None and field.id() in idx.fields + ] + if len(indices) != 1: + raise ValueError("Benchmark requires exactly one vector index for the selected column") + index_name = indices[0].name + if args.index_name is not None and args.index_name != index_name: + raise ValueError(f"Requested index {args.index_name!r} does not match the column's index {index_name!r}") + predicate = args.filter + if predicate is None and args.selectivity < 1: + predicate = f"bucket < {round(1000 * args.selectivity)}" + base_nearest = {"column": args.column, "k": args.k, "nprobes": args.nprobes, "metric": args.metric} + if args.refine_factor is not None: + base_nearest["refine_factor"] = args.refine_factor + # Explicit segments make Lance reject a metric mismatch instead of allowing + # a baseline to silently fall back to brute-force search. + ds.scanner( + columns=[args.id_column], + nearest={**base_nearest, "q": queries[0]}, + index_segments=[str(segment.uuid) for segment in indices[0].segments], + filter=predicate, + prefilter=True, + fast_search=True, + ).projected_schema + ray_module = None + if args.runner == "ray" and args.backend != "pylance": + import ray + + ray_module = ray + ray.init(num_cpus=args.workers, include_dashboard=False, object_store_memory=128 * 1024 * 1024) + daft.set_runner_ray() + + def search() -> list[list[int]]: + if args.backend in ("segments", "queries"): + df = vector_search( + uri, + version=version, + nearest={**base_nearest, "q": queries}, + index_name=index_name, + columns=[args.id_column], + filter=predicate, + num_tasks=1 if args.backend == "queries" else args.tasks, + ) + table = df.to_arrow() + return [ + table.filter(pa.compute.equal(table["_query_id"], query_id))[args.id_column].to_pylist() + for query_id in range(len(queries)) + ] + results = [] + for q in queries: + nearest = {**base_nearest, "q": q} + if args.backend == "pylance": + table = ds.scanner( + columns=[args.id_column, "_distance"], + nearest=nearest, + filter=predicate, + prefilter=True, + with_row_id=True, + ).to_table() + else: + df = read_lance(uri, version=version, default_scan_options={"nearest": nearest, "prefilter": True}) + if predicate is not None: + df = df.where(daft.sql_expr(predicate)) + table = df.select(args.id_column).to_arrow() + results.append(table[args.id_column].to_pylist()) + return results + + # Exact ground truth is outside all timed measurements and computed after + # the first query so it cannot prewarm the first-query cache observation. + with _Resources() as first_resources: + start = time.perf_counter() + first_results = search() + first_seconds = time.perf_counter() - start + truth = [] + for q in queries: + table = ds.scanner( + columns=[args.id_column], + nearest={"column": args.column, "q": q, "k": args.k, "metric": args.metric, "use_index": False}, + filter=predicate, + prefilter=True, + ).to_table() + truth.append(set(table[args.id_column].to_pylist())) + + def recall(result: list[list[int]]) -> float: + return float( + np.mean( + [len(set(got) & expected) / len(expected) if expected else 1.0 for got, expected in zip(result, truth)] + ) + ) + + if recall(first_results) < args.min_recall: + raise RuntimeError(f"Recall {recall(first_results):.4f} is below required {args.min_recall}") + for _ in range(args.warmup): + search() + durations = [] + recalls = [] + with _Resources() as resources: + for _ in range(args.repeats): + start = time.perf_counter() + result = search() + durations.append(time.perf_counter() - start) + recalls.append(recall(result)) + if min(recalls) < args.min_recall: + raise RuntimeError("A timed run failed the recall gate") + explanation = None + if args.backend in ("segments", "queries"): + output = io.StringIO() + vector_search( + uri, + version=version, + nearest={**base_nearest, "q": queries[0]}, + index_name=index_name, + columns=[args.id_column], + filter=predicate, + num_tasks=1 if args.backend == "queries" else args.tasks, + ).explain(show_all=True, file=output) + explanation = output.getvalue() + record = { + "backend": args.backend, + "runner": args.runner, + "workers": args.workers, + "cpu_budget": args.cpu_budget, + "cpu_affinity": sorted(os.sched_getaffinity(0)) if hasattr(os, "sched_getaffinity") else None, + "threads_per_process": os.environ["RAYON_NUM_THREADS"], + "dataset_version": version, + "index_name": index_name, + "rows": ds.count_rows(), + "segments": {idx.name: len(idx.segments) for idx in ds.describe_indices()}, + "queries_per_run": len(queries), + "nearest": base_nearest, + "filter": predicate, + "tasks": args.tasks, + "first_run_seconds": first_seconds, + "first_run_recall": recall(first_results), + "warm_batch_p50_seconds": float(np.percentile(durations, 50)), + "warm_batch_p95_seconds": float(np.percentile(durations, 95)), + "queries_per_second": len(queries) * len(durations) / sum(durations), + "recall_min": min(recalls), + "local_cpu_seconds": resources.cpu_seconds, + "local_physical_read_bytes": resources.read_bytes, + "local_peak_aggregate_rss_bytes": resources.peak_rss, + "first_run_peak_aggregate_rss_bytes": first_resources.peak_rss, + "durations_seconds": durations, + "execution_plan": explanation, + "versions": {"daft": daft.__version__, "pylance": lance.__version__, "python": platform.python_version()}, + "ray_version": ray_module.__version__ if ray_module is not None else None, + "source_sha256": { + str(path.relative_to(Path(__file__).resolve().parents[1])): hashlib.sha256(path.read_bytes()).hexdigest() + for path in [ + Path(__file__).resolve(), + Path(__file__).resolve().parents[1] / "daft_lance/lance_vector_search.py", + Path(__file__).resolve().parents[1] / "daft_lance/_lance.py", + ] + }, + "git_commit": subprocess.check_output(["git", "rev-parse", "HEAD"], text=True).strip(), + "git_dirty": bool(subprocess.check_output(["git", "status", "--porcelain"], text=True).strip()), + } + destination = Path(args.output) + destination.parent.mkdir(parents=True, exist_ok=True) + with destination.open("a") as output_file: + output_file.write(json.dumps(record) + "\n") + print(json.dumps({key: value for key, value in record.items() if key != "execution_plan"}, indent=2)) + if ray_module is not None: + ray_module.shutdown() + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("action", choices=["prepare", "run"]) + parser.add_argument("--root", required=True) + parser.add_argument("--rows", type=int, default=100_000) + parser.add_argument("--dim", type=int, default=128) + parser.add_argument("--segments", type=int, default=8) + parser.add_argument("--queries", type=int, default=20) + parser.add_argument("--seed", type=int, default=42) + parser.add_argument("--partitions", type=int, default=32) + parser.add_argument("--index-type", choices=["IVF_FLAT", "IVF_PQ"], default="IVF_FLAT") + parser.add_argument("--subvectors", type=int, default=16) + parser.add_argument("--backend", choices=["pylance", "daft-single", "segments", "queries"], default="pylance") + parser.add_argument("--runner", choices=["native", "ray"], default="native") + parser.add_argument("--workers", type=int, default=4) + parser.add_argument("--cpu-budget", type=int, default=4) + parser.add_argument("--tasks", type=int) + parser.add_argument("--k", type=int, default=10) + parser.add_argument("--nprobes", type=int, default=32) + parser.add_argument("--metric", choices=["l2", "cosine", "dot"], default="l2") + parser.add_argument("--refine-factor", type=int) + parser.add_argument("--selectivity", type=float, default=1) + parser.add_argument("--filter") + parser.add_argument("--warmup", type=int, default=2) + parser.add_argument("--repeats", type=int, default=10) + parser.add_argument("--min-recall", type=float, default=0.99) + parser.add_argument("--output", default="benchmark-results.jsonl") + parser.add_argument("--uri", help="Existing local or object-store table; prepare is not required") + parser.add_argument("--query-file", help="Independent query vectors stored as a NumPy matrix") + parser.add_argument("--column", default="vector") + parser.add_argument("--id-column", default="id") + parser.add_argument("--index-name") + args = parser.parse_args() + os.environ["DAFT_PROGRESS_BAR"] = "0" + os.environ["DAFT_ANALYTICS_ENABLED"] = "0" + if args.cpu_budget <= 0 or args.workers <= 0 or args.repeats <= 0 or args.queries <= 0: + parser.error("CPU budget, workers, repeats and queries must be positive") + if not 0 < args.selectivity <= 1 or not 0 <= args.min_recall <= 1: + parser.error("selectivity must be in (0, 1] and min-recall in [0, 1]") + if args.rows < args.segments or args.segments <= 0 or args.dim <= 0: + parser.error("Require rows >= segments > 0 and dim > 0") + if hasattr(os, "sched_setaffinity"): + available = sorted(os.sched_getaffinity(0)) + if args.cpu_budget > len(available): + parser.error("CPU budget exceeds this process's available CPU affinity") + os.sched_setaffinity(0, available[: args.cpu_budget]) + # A sequential single-task baseline can use the whole CPU budget. Giving + # it only budget/workers threads would manufacture a sharding speedup. + distributed = args.runner == "ray" and args.backend in ("segments", "queries") + threads = max(1, args.cpu_budget // args.workers) if distributed else args.cpu_budget + os.environ["RAYON_NUM_THREADS"] = str(threads) + os.environ["OMP_NUM_THREADS"] = str(threads) + os.environ["OPENBLAS_NUM_THREADS"] = str(threads) + if args.action == "prepare": + _prepare(args) + else: + _run(args) + + +if __name__ == "__main__": + main() diff --git a/daft_lance/__init__.py b/daft_lance/__init__.py index 7397441..137dcc3 100644 --- a/daft_lance/__init__.py +++ b/daft_lance/__init__.py @@ -13,6 +13,7 @@ optimize_indices, read_lance, update_columns_df, + vector_search, write_lance, ) from .lance_update_column import UpdateColumnsResult @@ -29,5 +30,6 @@ "read_lance", "take_blobs", "update_columns_df", + "vector_search", "write_lance", ] diff --git a/daft_lance/_lance.py b/daft_lance/_lance.py index c43e496..d435b27 100644 --- a/daft_lance/_lance.py +++ b/daft_lance/_lance.py @@ -20,6 +20,7 @@ from .lance_scan import LanceScanOperator from .lance_update_column import UpdateColumnsResult, update_columns_from_df, validate_update_arguments from .lance_vector_index import create_vector_index_internal +from .lance_vector_search import vector_search_internal from .namespace import validate_uri_or_namespace from .utils import construct_lance_dataset_handle @@ -30,6 +31,76 @@ from daft.checkpoint import CheckpointConfig +@PublicAPI +def vector_search( + uri: str | pathlib.Path | None = None, + io_config: IOConfig | None = None, + *, + nearest: dict[str, Any], + index_name: str | None = None, + columns: list[str] | None = None, + filter: str | None = None, + num_tasks: int | None = None, + version: int | str | None = None, + storage_options: dict[str, Any] | None = None, + index_cache_size: int | None = None, + table_id: list[str] | None = None, + namespace_impl: str | None = None, + namespace_properties: dict[str, str] | None = None, +) -> DataFrame: + """Search vector index segments in parallel and merge the global top-k. + + Returns a lazy DataFrame with ``_distance``, ``_rowid`` and ``_query_id``. + ``nearest`` contains Lance's ``column``, ``q`` and positive ``k`` options; + metric, nprobes and refine_factor are passed to Lance. ``q`` can be one + vector or a matrix: each query gets its own global k results, ordered by + distance then row ID. ``filter`` is applied before nearest search. + Equal-distance membership at the k boundary is chosen by Lance and may + vary with segment grouping; row-ID ordering applies to returned candidates. + + Workers read the pinned dataset version and only their assigned segments. + ``num_tasks`` caps segment groups per query (default: one per segment), + balanced by fragment row counts. + The runner determines concurrency; use the Ray runner for multiple workers. + Index coverage must be complete; use optimize_indices after appending data. + A single segment produces one task per query. Approximate index results can + differ from a monolithic search because candidate refinement runs per task. + + URI and namespace table entry points follow read_lance's conventions. + Output uses blob descriptors, like read_lance. Projection must retain the + result columns needed for the global merge (added by Lance automatically). + + Example: + >>> result = vector_search( # doctest: +SKIP + ... "table.lance", nearest={"column": "vector", "q": [0.1, 0.2], "k": 10} + ... ) + >>> result.explain() # doctest: +SKIP + >>> result.collect() # doctest: +SKIP + """ + if nearest.get("use_index") is False: + raise ValueError("vector_search requires use_index=True") + io_config = context.get_context().daft_planning_config.default_io_config if io_config is None else io_config + handle = construct_lance_dataset_handle( + uri, + io_config=io_config, + storage_options=storage_options, + index_cache_size=index_cache_size, + version=version, + table_id=table_id, + namespace_impl=namespace_impl, + namespace_properties=namespace_properties, + ) + return vector_search_internal( + handle.dataset, + handle.worker_open_context(), + nearest=nearest, + index_name=index_name, + columns=columns, + predicate=filter, + num_tasks=num_tasks, + ) + + @PublicAPI def read_lance( uri: str | pathlib.Path | None = None, diff --git a/daft_lance/lance_vector_search.py b/daft_lance/lance_vector_search.py new file mode 100644 index 0000000..46dcef0 --- /dev/null +++ b/daft_lance/lance_vector_search.py @@ -0,0 +1,209 @@ +from __future__ import annotations + +import pickle +from collections.abc import Iterator +from functools import lru_cache +from typing import Any + +import lance + +from daft.daft import PyPartitionField, PyPushdowns, PyRecordBatch, ScanOperatorHandle, ScanTask +from daft.dataframe import DataFrame +from daft.dependencies import pa +from daft.io.scan import ScanOperator +from daft.logical.builder import LogicalPlanBuilder +from daft.recordbatch import RecordBatch +from daft.schema import Schema + +from ._metadata import convert_lance_schema +from .namespace import DatasetOpenContext + + +@lru_cache(maxsize=1) +def _worker_dataset(serialized_context: bytes) -> lance.LanceDataset: + context: DatasetOpenContext = pickle.loads(serialized_context) + return context.open_pinned() + + +def _search_segments( + serialized_context: bytes, + segments: list[str], + nearest: dict[str, Any], + columns: list[str] | None, + predicate: str | None, + query_id: int, +) -> Iterator[PyRecordBatch]: + ds = _worker_dataset(serialized_context) + k = nearest["k"] + table = ds.scanner( + columns=columns, + nearest=nearest, + index_segments=segments, + filter=predicate, + prefilter=True, + fast_search=True, + with_row_id=True, + blob_handling="blobs_descriptions", + ).to_table() + table = table.sort_by([("_distance", "ascending"), ("_rowid", "ascending")]) + table = table.slice(0, k) + table = table.append_column("_query_id", pa.array([query_id] * table.num_rows, type=pa.int64())) + for batch in table.to_batches(): + yield RecordBatch.from_arrow_record_batches([batch], batch.schema)._recordbatch + + +class VectorSearchScanOperator(ScanOperator): + def __init__( + self, + open_context: DatasetOpenContext, + groups: list[list[str]], + nearest: dict[str, Any], + columns: list[str] | None, + predicate: str | None, + query_id: int, + schema: pa.Schema, + ) -> None: + self.open_context = open_context + self.groups = groups + self.nearest = nearest + self.columns = columns + self.predicate = predicate + self.query_id = query_id + self._schema = convert_lance_schema(schema) + + def schema(self) -> Schema: + return self._schema + + def name(self) -> str: + return "LanceVectorSearch" + + def as_pushdown_filter(self) -> None: + return None + + def display_name(self) -> str: + return f"LanceVectorSearch(version={self.open_context.version}, query={self.query_id})" + + def partitioning_keys(self) -> list[PyPartitionField]: + return [] + + def can_absorb_filter(self) -> bool: + return False + + def can_absorb_limit(self) -> bool: + return False + + def can_absorb_select(self) -> bool: + return False + + def multiline_display(self) -> list[str]: + return [self.display_name(), f"Index segment groups = {len(self.groups)}"] + + def to_scan_tasks(self, pushdowns: PyPushdowns) -> Iterator[ScanTask]: + for segments in self.groups: + yield ScanTask.python_factory_func_scan_task( + module=__name__, + func_name="_search_segments", + func_args=( + pickle.dumps(self.open_context), + segments, + self.nearest, + self.columns, + self.predicate, + self.query_id, + ), + schema=self._schema._schema, + num_rows=None, + size_bytes=None, + pushdowns=pushdowns, + stats=None, + source_name=self.display_name(), + ) + + +def vector_search_internal( + ds: lance.LanceDataset, + open_context: DatasetOpenContext, + *, + nearest: dict[str, Any], + index_name: str | None, + columns: list[str] | None, + predicate: str | None, + num_tasks: int | None, +) -> DataFrame: + import numpy as np + + if num_tasks is not None and (not isinstance(num_tasks, int) or isinstance(num_tasks, bool) or num_tasks <= 0): + raise ValueError("num_tasks must be a positive integer") + k = nearest.get("k") + if not isinstance(k, int) or isinstance(k, bool) or k <= 0: + raise ValueError("nearest['k'] must be a positive integer") + column = nearest.get("column") + if columns is not None: + columns = list(dict.fromkeys([*columns, "_distance"])) + vectors = np.array(nearest.get("q"), dtype=np.float32, copy=True) + if vectors.ndim == 1: + vectors = vectors.reshape(1, -1) + if vectors.ndim != 2 or not all(vectors.shape) or not np.isfinite(vectors).all(): + raise ValueError("nearest['q'] must be a finite vector or a non-empty matrix of query vectors") + if "_query_id" in ds.schema.names: + raise ValueError("_query_id is reserved for vector search results") + field = ds.lance_schema.field(column) if isinstance(column, str) else None + indices = [ + idx + for idx in ds.describe_indices() + if str(idx.index_type).startswith("IVF_") + and field is not None + and field.id() in idx.fields + and (index_name is None or idx.name == index_name) + ] + if len(indices) != 1: + raise ValueError("Select exactly one vector index for nearest['column']; specify index_name if ambiguous") + index = indices[0] + fragments = ds.get_fragments() + live = {fragment.fragment_id for fragment in fragments} + covered = {fid for segment in index.segments for fid in segment.fragment_ids} + if live - covered: + raise ValueError("vector_search requires complete index coverage; run optimize_indices or rebuild the index") + segments = [str(segment.uuid) for segment in index.segments if live.intersection(segment.fragment_ids)] + if not segments: + raise ValueError("vector_search requires a non-empty vector index") + count = min(num_tasks or len(segments), len(segments)) + if count == 1: + groups = [segments] + elif count == len(segments): + groups = [[segment] for segment in segments] + else: + groups = [[] for _ in range(count)] + rows = {fragment.fragment_id: fragment.metadata.num_rows for fragment in fragments} + weights = { + str(segment.uuid): sum(rows.get(fid, 0) for fid in segment.fragment_ids) for segment in index.segments + } + group_weights = [0] * count + for segment in sorted(segments, key=lambda uuid: weights[uuid], reverse=True): + group = min(range(count), key=lambda i: (group_weights[i], len(groups[i]))) + groups[group].append(segment) + group_weights[group] += weights[segment] + results: list[DataFrame] = [] + for query_id, vector in enumerate(vectors): + query = {**nearest, "q": vector} + # Building the scanner validates Lance-owned types, dimensions and + # parameters without executing a query on the coordinator. + schema = ds.scanner( + columns=columns, + nearest=query, + index_segments=segments, + filter=predicate, + prefilter=True, + fast_search=True, + with_row_id=True, + blob_handling="blobs_descriptions", + ).projected_schema + schema = schema.append(pa.field("_query_id", pa.int64())) + operator = VectorSearchScanOperator(open_context, groups, query, columns, predicate, query_id, schema) + handle = ScanOperatorHandle.from_python_scan_operator(operator) + candidates = DataFrame(LogicalPlanBuilder.from_tabular_scan(scan_operator=handle)) + results.append(candidates.sort(["_distance", "_rowid"]).limit(k)) + result = results[0] + for other in results[1:]: + result = result.concat(other) + return result.sort(["_query_id", "_distance", "_rowid"]) diff --git a/tests/io/lance/test_lance_distributed_vector_search.py b/tests/io/lance/test_lance_distributed_vector_search.py new file mode 100644 index 0000000..d353aa4 --- /dev/null +++ b/tests/io/lance/test_lance_distributed_vector_search.py @@ -0,0 +1,236 @@ +from __future__ import annotations + +import io +from pathlib import Path + +import lance +import numpy as np +import pyarrow as pa +import pytest + +from daft_lance import create_vector_index, vector_search + + +def _dataset(path: Path, metric: str = "l2", tied: bool = False) -> tuple[str, np.ndarray]: + vectors = np.random.default_rng(12).normal(size=(512, 8)).astype("float32") + if tied: + vectors[:] = 1 + uri = str(path) + lance.write_dataset( + pa.table( + { + "id": np.arange(len(vectors)), + "bucket": np.arange(len(vectors)) % 4, + "vector": pa.FixedSizeListArray.from_arrays(pa.array(vectors.ravel()), 8), + } + ), + uri, + max_rows_per_file=128, + ) + create_vector_index( + uri, + column="vector", + index_type="IVF_FLAT", + name="v_idx", + metric=metric, + num_partitions=1 if tied else 4, + sample_rate=2, + fragment_group_size=1, + ) + return uri, vectors + + +@pytest.mark.parametrize("metric", ["l2", "cosine", "dot"]) +@pytest.mark.parametrize("predicate", [None, "bucket = 1", "id < 3", "id < 0"]) +def test_global_topk_matches_reference(tmp_path: Path, metric: str, predicate: str | None) -> None: + uri, _ = _dataset(tmp_path / "search.lance", metric) + q = np.random.default_rng(17).normal(size=8).astype("float32") + nearest = {"column": "vector", "q": q, "k": 10, "nprobes": 4, "metric": metric, "refine_factor": 2} + result = vector_search(uri, nearest=nearest, columns=["id"], filter=predicate).to_arrow() + reference = ( + lance.dataset(uri) + .scanner(columns=["id", "_distance"], nearest=nearest, with_row_id=True, filter=predicate, prefilter=True) + .to_table() + .sort_by([("_distance", "ascending"), ("_rowid", "ascending")]) + ) + assert result["id"].to_pylist() == reference["id"].to_pylist() + assert result["_rowid"].to_pylist() == reference["_rowid"].to_pylist() + np.testing.assert_allclose(result["_distance"].to_numpy(), reference["_distance"].to_numpy(), atol=1e-6) + + +def test_plan_has_multiple_tasks_and_global_limit(tmp_path: Path) -> None: + uri, vectors = _dataset(tmp_path / "plan.lance") + df = vector_search(uri, nearest={"column": "vector", "q": vectors[3], "k": 7}) + output = io.StringIO() + df.explain(show_all=True, file=output) + assert "Num Scan Tasks = 4" in output.getvalue() + assert df.count_rows() == 7 + assert df.limit(2).count_rows() == 2 + # A downstream filter applies to the global top-k, not to individual shards. + import daft + + assert df.where(daft.col("id") < 0).to_arrow().num_rows == 0 + grouped = vector_search(uri, nearest={"column": "vector", "q": vectors[3], "k": 7}, num_tasks=2) + assert grouped.select("id").to_pydict() == df.select("id").to_pydict() + + +def test_batch_queries_have_independent_topk(tmp_path: Path) -> None: + uri, vectors = _dataset(tmp_path / "batch.lance") + query_vectors = vectors[[3, 301]] + result = vector_search(uri, nearest={"column": "vector", "q": query_vectors, "k": 5}).to_arrow() + assert result["_query_id"].to_pylist() == [0] * 5 + [1] * 5 + for query_id, q in enumerate(query_vectors): + expected = vector_search(uri, nearest={"column": "vector", "q": q, "k": 5}).to_arrow() + assert result.slice(query_id * 5, 5)["id"].to_pylist() == expected["id"].to_pylist() + + +def test_tie_boundary_keeps_bounded_candidates(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + from daft_lance import lance_vector_search + + uri, vectors = _dataset(tmp_path / "ties.lance", tied=True) + nearest = {"column": "vector", "q": vectors[0], "k": 5, "nprobes": 1} + ds = lance.dataset(uri) + segments = [str(segment.uuid) for segment in ds.describe_indices()[0].segments] + scanner = ds.scanner + calls: list[int] = [] + + def counted_scanner(**kwargs): + calls.append(kwargs["nearest"]["k"]) + return scanner(**kwargs) + + monkeypatch.setattr(lance_vector_search, "_worker_dataset", lambda _: ds) + monkeypatch.setattr(ds, "scanner", counted_scanner) + batches = list(lance_vector_search._search_segments(b"", segments, nearest, ["id", "_distance"], None, 0)) + assert calls == [5] + assert sum(len(batch) for batch in batches) == 5 + for tasks in (1, 2, 4): + result = vector_search(uri, nearest=nearest, num_tasks=tasks, columns=["id"]).to_arrow() + assert result.num_rows == 5 + assert len(set(result["id"].to_pylist())) == 5 + assert result["_distance"].to_pylist() == [0.0] * 5 + assert result["_rowid"].to_pylist() == sorted(result["_rowid"].to_pylist()) + + +@pytest.mark.parametrize("batch", [False, True]) +def test_query_buffer_is_copied_when_planning(tmp_path: Path, batch: bool) -> None: + uri, vectors = _dataset(tmp_path / "query-copy.lance") + queries = vectors[[3, 301]].copy() if batch else vectors[3].copy() + original = queries.copy() + pending = vector_search(uri, nearest={"column": "vector", "q": queries, "k": 5}, columns=["id"]) + queries[:] = vectors[400] + expected = vector_search(uri, nearest={"column": "vector", "q": original, "k": 5}, columns=["id"]) + assert pending.to_pydict() == expected.to_pydict() + + +def test_quoted_vector_column_matches_lance(tmp_path: Path) -> None: + uri, vectors = _dataset(tmp_path / "quoted.lance") + nearest = {"column": "`vector`", "q": vectors[3], "k": 5, "nprobes": 4} + result = vector_search(uri, nearest=nearest, columns=["id"]).to_arrow() + reference = lance.dataset(uri).scanner(nearest=nearest, columns=["id"]).to_table() + assert result["id"].to_pylist() == reference["id"].to_pylist() + + +def test_uneven_segments_are_balanced_by_rows(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + from types import SimpleNamespace + + from daft_lance import lance_vector_search + from daft_lance.namespace import DatasetOpenContext + + sizes = [256, 32, 256, 32] + vectors = np.random.default_rng(19).normal(size=(sum(sizes), 8)).astype("float32") + uri = str(tmp_path / "skew.lance") + start = 0 + for size in sizes: + lance.write_dataset( + pa.table( + { + "id": np.arange(start, start + size), + "vector": pa.FixedSizeListArray.from_arrays(pa.array(vectors[start : start + size].ravel()), 8), + } + ), + uri, + mode="create" if start == 0 else "append", + max_rows_per_file=size, + ) + start += size + create_vector_index( + uri, column="vector", index_type="IVF_FLAT", num_partitions=4, sample_rate=2, fragment_group_size=1 + ) + ds = lance.dataset(uri) + rows = {fragment.fragment_id: fragment.metadata.num_rows for fragment in ds.get_fragments()} + index = ds.describe_indices()[0] + # Manifest order is not a workload ordering guarantee. Arrange valid + # segments by fragment ID so round-robin groups both large segments together. + segments = sorted(index.segments, key=lambda segment: min(segment.fragment_ids)) + description = SimpleNamespace(index_type=index.index_type, fields=index.fields, name=index.name, segments=segments) + monkeypatch.setattr(ds, "describe_indices", lambda: [description]) + weights = { + str(segment.uuid): sum(rows[fid] for fid in segment.fragment_ids) + for segment in ds.describe_indices()[0].segments + } + operator = lance_vector_search.VectorSearchScanOperator + groups: list[list[str]] = [] + + def capture_groups(open_context, assigned, *args): + groups.extend(assigned) + return operator(open_context, assigned, *args) + + monkeypatch.setattr(lance_vector_search, "VectorSearchScanOperator", capture_groups) + nearest = {"column": "vector", "q": vectors[3], "k": 5, "nprobes": 4} + result = lance_vector_search.vector_search_internal( + ds, + DatasetOpenContext.from_dataset(ds, uri), + nearest=nearest, + index_name=None, + columns=["id"], + predicate=None, + num_tasks=2, + ).to_arrow() + assert [sum(weights[segment] for segment in group) for group in groups] == [288, 288] + reference = ds.scanner(nearest=nearest, columns=["id"]).to_table() + assert result["id"].to_pylist() == reference["id"].to_pylist() + + +def test_snapshot_is_pinned_and_partial_coverage_rejected(tmp_path: Path) -> None: + uri, vectors = _dataset(tmp_path / "pinned.lance") + nearest = {"column": "vector", "q": vectors[0], "k": 5} + df = vector_search(uri, nearest=nearest, columns=["id"]) + expected = df.to_pydict() + pending = vector_search(uri, nearest=nearest, columns=["id"]) + lance.write_dataset(lance.dataset(uri).to_table().slice(0, 1), uri, mode="append") + # A fresh plan sees uncovered data; the already planned DataFrame retains its snapshot. + with pytest.raises(ValueError, match="complete index coverage"): + vector_search(uri, nearest=nearest) + assert ( + vector_search(uri, nearest=nearest, version=lance.dataset(uri).version - 1, columns=["id"]).to_pydict() + == expected + ) + assert df.to_pydict() == expected + assert pending.to_pydict() == expected + + +def test_namespace_search(tmp_path: Path) -> None: + uri, vectors = _dataset(tmp_path / "ns.lance") + result = vector_search( + table_id=["ns"], + namespace_impl="dir", + namespace_properties={"root": str(tmp_path)}, + nearest={"column": "vector", "q": vectors[0], "k": 3}, + ) + assert result.select("id").to_pydict()["id"][0] == 0 + + +@pytest.mark.parametrize("nearest", [{"column": "vector", "q": [1] * 8, "k": 0}, {"column": "vector", "q": [], "k": 3}]) +def test_invalid_search_arguments(tmp_path: Path, nearest: dict) -> None: + uri, _ = _dataset(tmp_path / "invalid.lance") + with pytest.raises(ValueError): + vector_search(uri, nearest=nearest) + + +def test_missing_index_and_disabled_index_rejected(tmp_path: Path) -> None: + uri, vectors = _dataset(tmp_path / "missing.lance") + nearest = {"column": "vector", "q": vectors[0], "k": 3} + with pytest.raises(ValueError, match="Select exactly one"): + vector_search(uri, nearest=nearest, index_name="missing") + with pytest.raises(ValueError, match="use_index"): + vector_search(uri, nearest={**nearest, "use_index": False})