feat: component-wise and sharded dataset sources (npy + parquet, LAION-400M)

Cargo.toml

@@ -7,6 +7,7 @@ edition = "2024"
 
 [dependencies]
 anyhow = "1.0.102"
+bytes = "1.11"
 chrono = { version = "0.4.44", default-features = false, features = [
     "now",
     "std",
@@ -20,6 +21,15 @@ half = "2.7"
 hdf5-pure-rust = "0.3.10"
 indicatif = "0.18.4"
 memmap2 = "0.9.10"
+# `arrow` deliberately off: the record API is enough to read payload rows, and
+# skipping it keeps the dependency tree (and build time) an order smaller.
+parquet = { version = "59.1", default-features = false, features = [
+    "snap",
+    "brotli",
+    "flate2-zlib-rs",
+    "lz4",
+    "zstd",
+] }
 qdrant-client = { git = "https://github.com/qdrant/rust-client", branch = "dev" }
 rand = "0.10.1"
 rand_distr = "0.6.0"

README.md

@@ -24,9 +24,91 @@ each field is generated). The runtime flags (`-n`, `-b`, `-p`, `-t`, `--uri`,
 Upload configs can source dense vectors, sparse vectors, and payloads from
 inline dataset definitions (same fields as
 [vector-db-benchmark `datasets.json`](https://github.com/qdrant/vector-db-benchmark/blob/master/datasets/datasets.json)).
-Supported formats are `h5` (HDF5, pure-Rust reader — no system libraries), `tar`
-(`.tgz` with `vectors.npy` + optional `payloads.jsonl`), and `sparse` (CSR
-matrices).
+
+| `format` | Contents |
+|----------|----------|
+| `h5` | ann-benchmarks HDF5 bundle — `train`, optional `test`/`neighbors`. Pure-Rust reader, no system libraries |
+| `tar` | `.tgz` of `vectors.npy` + optional `payloads.jsonl` / `tests.jsonl` |
+| `sparse` | CSR matrices (`data.csr`, optional `queries.csr` / `results.gt`) |
+| `npy` | One 2-D float `.npy` — dense vectors only |
+| `parquet` | One parquet file — payload rows only |
+
+The first three are *bundles*: vectors, payloads, and queries all come out of a
+single artifact. `npy` and `parquet` are *components*, so a config pairs them —
+one source per slot, row *i* of each landing on point *i*:
+
+```yaml
+collection:
+  vectors:
+    - size: 512
+      source: { type: dataset, name: emb, format: npy, path: emb.npy }
+  payload:
+    source:
+      type: dataset
+      dataset: { name: meta, format: parquet, path: meta.parquet, exclude: [exif] }
+```
+
+Parquet sources accept three extra keys: `columns` (keep only these), `exclude`
+(drop these), and `fill_null` (a value substituted for nulls and for NaN/±inf
+floats, which have no JSON form — by default such fields are simply absent).
+See [`examples/upload-laion-part.yaml`](examples/upload-laion-part.yaml).
+
+#### Sharded datasets
+
+Corpora published as numbered parts are read as one row space with a `parts:`
+block, so point ids stay global across the whole set. `npy` and `parquet`
+sources support it; `{i}` is substituted with each part's number:
+
+```yaml
+source:
+  type: dataset
+  name: laion-400m-img-emb
+  format: npy
+  parts:
+    count: 410                 # parts 0..409; `start:` moves the first index
+    path: laion/img_emb_{i}.npy
+    link: https://deploy.laion.ai/.../img_emb_{i}.npy
+```
+
+Part row counts are **measured, never configured**. Both formats keep their
+shape at a known end of the file — the `.npy` header at the front, the parquet
+footer at the back — so bfb sizes every part with one ranged HTTP request each
+and downloads none of them (820 LAION parts in ~1.5 minutes). The result is
+cached in `datasets/.parts-index/<name>.json`, keyed on the parts spec, so later
+runs issue no requests at all.
+
+There is deliberately no "rows per part" setting. LAION-400M turns out to have
+seven distinct part sizes — 404 parts of 1,000,448 rows, one of 1,000,501, and
+five short ones (parts 8, 107, 220, 319 and 409, from 189,159 to 642,675 rows)
+— so a fixed guess would go wrong at part 8 and silently pair payloads with the
+wrong vectors across ~98% of the corpus. The host must support ranged requests;
+one that answers `200` to a `Range:` request is reported rather than silently
+downloaded.
+
+Because a point's id *is* its dataset row, `--offset` resumes an interrupted
+upload — it skips that many rows as well as ids, and `-n` is capped by what
+remains. See [`examples/upload-laion-400m.yaml`](examples/upload-laion-400m.yaml)
+for the full 410-part, ~409.7M-point corpus.
+
+##### Streaming a corpus larger than the disk
+
+Parts are downloaded as they are reached, and by default they accumulate.
+`cache: evict` streams instead — the next part is fetched in the background
+while the current one uploads, and parts already passed are deleted:
+
+```yaml
+source:
+  type: dataset
+  name: laion-400m-img-emb
+  format: npy
+  parts: { count: 410, path: laion/img_emb_{i}.npy, link: "https://…/img_emb_{i}.npy" }
+  cache: evict               # keep | evict (default: keep)
+```
+
+That holds peak disk to the few parts in flight (~4 GB for LAION) instead of
+the ~600 GB the whole corpus occupies. Eviction only ever removes parts bfb
+downloaded itself — a file staged in the datasets dir by hand is never deleted,
+and a part still being read is left until nothing references it.
 
 Use `format` for the dataset storage type in upload configs (`type` is reserved
 for the source kind). An optional local `datasets/datasets.json` registry is

examples/upload-laion-400m.yaml

@@ -0,0 +1,80 @@
+# The full LAION-400M benchmark corpus
+# (https://github.com/qdrant/laion-400m-benchmark) — 407,314,954 points across
+# 410 published parts, uploaded as one collection with global point ids.
+#
+#   bfb upload --file examples/upload-laion-400m.yaml -b 256 -p 16 -t 8 \
+#       --uri http://localhost:6334
+#
+# On first use bfb sizes every part with one ranged request each (~1.5 minutes
+# for all 820, no downloads) and caches the result in `datasets/.parts-index/`.
+# Part row counts are *not* uniform: 404 parts hold 1,000,448 rows, one holds
+# 1,000,501, and parts 8, 107, 220, 319 and 409 are short (189,159 to 642,675).
+# They are measured rather than assumed — a fixed size would go wrong at part 8
+# and misalign payloads against vectors for the rest of the corpus.
+#
+# Resuming: point ids are dataset rows, so an interrupted run continues with
+# `--offset <points already uploaded>`; `-n` is then capped by what is left.
+#
+# Disk: `cache: evict` streams the corpus — each part is fetched as it is
+# reached, the next one is fetched in the background while the current one
+# uploads, and parts already passed are deleted. Peak usage is a few parts
+# (~4 GB) rather than the ~600 GB the full corpus would occupy. Drop
+# `cache: evict` to keep every part for repeat runs.
+#
+# Eviction only ever deletes parts bfb downloaded; a file you staged in the
+# datasets dir yourself is left alone.
+
+collection:
+  name: laion
+  on_disk_payload: true
+
+  quantization:
+    type: binary
+    always_ram: true
+
+  hnsw:
+    m: 6
+    on_disk: false
+
+  optimizers:
+    default_segment_number: 2
+    # Bigger segments search faster, at the cost of slower indexing.
+    max_segment_size: 5000000
+
+  vectors:
+    - size: 512
+      distance: cosine
+      datatype: float16
+      on_disk: true
+      source:
+        type: dataset
+        name: laion-400m-img-emb
+        format: npy
+        parts:
+          count: 410           # parts 0..409 inclusive
+          path: laion/img_emb_{i}.npy
+          link: https://deploy.laion.ai/8f83b608504d46bb81708ec86e912220/embeddings/img_emb/img_emb_{i}.npy
+        cache: evict
+
+  payload:
+    source:
+      type: dataset
+      dataset:
+        name: laion-400m-metadata
+        format: parquet
+        parts:
+          count: 410
+          path: laion/metadata_{i}.parquet
+          link: https://deploy.laion.ai/8f83b608504d46bb81708ec86e912220/embeddings/metadata/metadata_{i}.parquet
+        cache: evict
+        # `exif` is the bulk of the metadata and useless as a filter; excluding
+        # it also skips decoding the column. The reference `upload.py` drops it.
+        exclude: [exif]
+        # `upload.py` does `df.fillna(0)`. Drop this line to leave fields with
+        # no value absent instead, which filters more predictably.
+        fill_null: 0
+
+  fields:
+    # Index-only declaration: the value comes from `payload.source` above.
+    - name: similarity
+      type: float

examples/upload-laion-part.yaml

@@ -0,0 +1,70 @@
+# One part of the LAION-400M benchmark corpus
+# (https://github.com/qdrant/laion-400m-benchmark), uploaded straight from the
+# files LAION publishes: dense vectors from a bare `.npy`, payloads from the
+# matching `.parquet`. No repacking step.
+#
+#   bfb upload --file examples/upload-laion-part.yaml -b 256 -p 16 -t 8 \
+#       --uri http://localhost:6334
+#
+# The two sources are row-aligned: row *i* of `img_emb_0.npy` is the same image
+# as row *i* of `metadata_0.parquet`, and becomes point *i*. Omit `-n` to upload
+# the whole part (~1,000,448 points — the parts are not round millions).
+#
+# Both files are downloaded on first use into `./datasets/` (override with
+# BFB_DATASETS_DIR) and total ~1.2 GB. For all 410 parts at once, see
+# `upload-laion-400m.yaml`, which streams them without keeping them on disk.
+#
+# The collection matches the reference benchmark's `upload.py`: fp16 vectors on
+# disk, binary quantization pinned in RAM, a small HNSW graph (m=6) and large
+# segments.
+
+collection:
+  name: laion
+  on_disk_payload: true
+
+  quantization:
+    type: binary
+    always_ram: true
+
+  hnsw:
+    m: 6
+    on_disk: false
+
+  optimizers:
+    default_segment_number: 2
+    # Bigger segments search faster, at the cost of slower indexing.
+    max_segment_size: 5000000
+
+  vectors:
+    - size: 512
+      distance: cosine
+      datatype: float16
+      on_disk: true
+      source:
+        type: dataset
+        name: laion-img-emb-0
+        format: npy
+        path: laion/img_emb_0.npy
+        link: https://deploy.laion.ai/8f83b608504d46bb81708ec86e912220/embeddings/img_emb/img_emb_0.npy
+
+  # Whole-payload source: every column becomes a payload field. Fields not
+  # listed under `fields` are still uploaded, just left unindexed.
+  payload:
+    source:
+      type: dataset
+      dataset:
+        name: laion-metadata-0
+        format: parquet
+        path: laion/metadata_0.parquet
+        link: https://deploy.laion.ai/8f83b608504d46bb81708ec86e912220/embeddings/metadata/metadata_0.parquet
+        # `exif` is the bulk of the metadata and is useless as a filter — the
+        # reference `upload.py` drops it too.
+        exclude: [exif]
+        # `upload.py` does `df.fillna(0)`; drop this line to leave fields with
+        # no value absent instead, which is usually what you want for filtering.
+        fill_null: 0
+
+  fields:
+    # Index-only declaration: the value comes from `payload.source` above.
+    - name: similarity
+      type: float

src/config/payload.rs

@@ -72,8 +72,9 @@ pub struct PayloadSource {
     #[serde(default, rename = "type")]
     pub kind: PayloadSourceKind,
     /// vector-db-benchmark dataset for payload values (`type: dataset`).
+    /// Boxed so a dataset definition does not bloat every source it can appear in.
     #[serde(default)]
-    pub dataset: Option<DatasetConfig>,
+    pub dataset: Option<Box<DatasetConfig>>,
     /// Payload field name inside the dataset schema / `payloads.jsonl`.
     #[serde(default)]
     pub field: Option<String>,

src/config/schema.rs

@@ -94,11 +94,34 @@ collection:
       # source:
       #   type: dataset          # inline dataset definition (vector-db-benchmark format)
       #   name: glove-25-angular
-      #   format: h5             # h5 | tar | sparse  (`type` alias accepted in nested `dataset:` maps)
+      #   format: h5             # dataset format (`type` alias accepted in nested `dataset:` maps):
+      #                          #   h5      ann-benchmarks bundle (train/test/neighbors)
+      #                          #   tar     .tgz of vectors.npy + payloads.jsonl + tests.jsonl
+      #                          #   sparse  CSR matrices
+      #                          #   npy     one 2-D float .npy — dense vectors only
+      #                          #   parquet one parquet file — payload rows only
       #   path: glove-25-angular/glove-25-angular.hdf5
       #   link: http://ann-benchmarks.com/glove-25-angular.hdf5
       #   vector_size: 25
       #   distance: cosine
+      # A sharded dataset (`npy` / `parquet` only) replaces `path`/`link` with a
+      # `parts` block; the files are read as one row space and `{i}` is
+      # substituted with each part's number. Row counts per part are measured,
+      # not configured — one ranged request per part, cached thereafter.
+      # source:
+      #   type: dataset
+      #   name: laion-400m-img-emb
+      #   format: npy
+      #   parts:
+      #     count: 410           # uint     required   number of parts
+      #     start: 0             # uint     default=0  index of the first part
+      #     path: laion/img_emb_{i}.npy     # string   required
+      #     link: https://host/img_emb_{i}.npy  # string  optional
+      #   cache: keep          # enum  default=keep  [keep | evict] (sharded only)
+      #                        #   evict deletes each downloaded part once the reader
+      #                        #   moves past it, and prefetches the next one, so a
+      #                        #   corpus larger than the disk can still be streamed.
+      #                        #   Only parts bfb downloaded are ever deleted.
 
   # Sparse vectors (optional). Names must be unique across all vectors.
   sparse_vectors:
@@ -136,6 +159,13 @@ collection:
     #     format: tar
     #     path: laion-small-clip/laion-small-clip
     #     link: https://example.com/laion-small-clip.tgz
+    # `format: parquet` reads payload rows from a parquet file, and accepts
+    # three extra keys (ignored by every other format):
+    #   columns: [url, similarity]   # list  optional  columns to keep (default: all)
+    #   exclude: [exif]              # list  default=[]  columns to drop (applied after `columns`)
+    #   fill_null: 0                 # any   optional  value substituted for nulls and for
+    #                                #   NaN/±inf floats, which have no JSON form. Omitted by
+    #                                #   default, leaving the payload field absent.
 
   # Payload field declarations (optional). Names must be unique. Each entry
   # generates a value and/or declares a field index.

src/config/vector.rs

@@ -80,9 +80,10 @@ pub enum VectorSource {
         strategy: FileStrategy,
     },
     /// vector-db-benchmark dataset (specified inline in the source definition).
+    /// Boxed to keep the enum from being sized by its largest variant.
     Dataset {
         #[serde(flatten)]
-        dataset: DatasetConfig,
+        dataset: Box<DatasetConfig>,
     },
 }
 
@@ -149,8 +150,9 @@ pub struct SparseSource {
     #[serde(default)]
     pub distribution: DistributionKind,
     /// vector-db-benchmark dataset (specified inline under `dataset`).
+    /// Boxed so a dataset definition does not bloat every source it can appear in.
     #[serde(default)]
-    pub dataset: Option<DatasetConfig>,
+    pub dataset: Option<Box<DatasetConfig>>,
 }
 
 impl Default for SparseSource {