From 81f48e6acb8d65bb984c4280a50c1b8bcad18956 Mon Sep 17 00:00:00 2001 From: generall Date: Sun, 26 Jul 2026 14:24:36 +0200 Subject: [PATCH 1/3] feat: read dense vectors from `.npy` and payloads from `.parquet` MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The dataset layer assumed a dataset is a *bundle* — one artifact holding vectors, payloads and queries together, which is what h5/tgz/csr are. Corpora published as separate files per component (LAION-400M ships `img_emb_*.npy` alongside `metadata_*.parquet`) had no way in: the only `.npy` reader was buried inside the tar format, keyed off the literal name `vectors.npy`, and there was no parquet reader at all. Add two component-scoped formats, `npy` (dense vectors only) and `parquet` (payload rows only). The upload config already keeps a separate source per slot, so pairing them needs no new plumbing — row *i* of each file lands on point *i*: 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] } The `.npy` reader moves out of `tar.rs` into `readers/npy.rs`, shared by both formats. Its header parser now accepts a short prefix of the file rather than the whole mapping, so a row count can later be had from a ranged request. Parquet is read through the record API with `arrow` deliberately off — payload rows need no columnar machinery, and skipping it keeps the dependency tree an order smaller. Access is a streaming cursor plus a ring of recent rows, not a decode-the-file cache: upload walks ids in order, and materializing a whole LAION metadata part up front would cost hundreds of MB before the first batch. Reads behind the ring rewind only to the containing row group. Values with no JSON form — nulls, NaN, ±inf, non-UTF-8 bytes — leave the field absent by default; `fill_null` substitutes a value instead, matching what the reference `upload.py` gets from `df.fillna(0)`. `columns`/`exclude` project at the parquet level, so dropping `exif` also skips decoding it. Boxing `DatasetConfig` at its three embedding sites keeps the config enums from being sized by a variant that only appears in dataset-backed runs. Verified against the published LAION files: `img_emb_0.npy` and `metadata_0.parquet` both report 1,000,448 rows, and payload rows decode to the expected string/int/float fields. Co-Authored-By: Claude Opus 5 (1M context) --- Cargo.lock | 275 ++++++++++++++++++- Cargo.toml | 10 + README.md | 31 ++- examples/upload-laion-part.yaml | 70 +++++ src/config/payload.rs | 3 +- src/config/schema.rs | 14 +- src/config/vector.rs | 6 +- src/dataset/config.rs | 67 +++-- src/dataset/download.rs | 17 +- src/dataset/fixtures.rs | 119 +++++++++ src/dataset/mod.rs | 2 + src/dataset/reader.rs | 35 ++- src/dataset/readers/mod.rs | 4 + src/dataset/readers/npy.rs | 295 ++++++++++++++++++++ src/dataset/readers/parquet.rs | 459 ++++++++++++++++++++++++++++++++ src/dataset/readers/tar.rs | 234 ++-------------- src/dataset/upload.rs | 8 +- src/generators/config.rs | 61 +++++ 18 files changed, 1446 insertions(+), 264 deletions(-) create mode 100644 examples/upload-laion-part.yaml create mode 100644 src/dataset/fixtures.rs create mode 100644 src/dataset/readers/npy.rs create mode 100644 src/dataset/readers/parquet.rs diff --git a/Cargo.lock b/Cargo.lock index e40ff8a..2a65cae 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -8,6 +8,20 @@ version = "2.0.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "320119579fcad9c21884f5c4861d16174d0e06250625266f50fe6898340abefa" +[[package]] +name = "ahash" +version = "0.8.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5a15f179cd60c4584b8a8c596927aadc462e27f2ca70c04e0071964a73ba7a75" +dependencies = [ + "cfg-if", + "const-random", + "getrandom 0.3.4", + "once_cell", + "version_check", + "zerocopy", +] + [[package]] name = "aho-corasick" version = "1.1.4" @@ -17,6 +31,30 @@ dependencies = [ "memchr", ] +[[package]] +name = "alloc-no-stdlib" +version = "2.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cc7bb162ec39d46ab1ca8c77bf72e890535becd1751bb45f64c597edb4c8c6b3" + +[[package]] +name = "alloc-stdlib" +version = "0.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0e76a019e91224d279006ff972f1e984179a6e9feb050adba6ce8274aef23195" +dependencies = [ + "alloc-no-stdlib", +] + +[[package]] +name = "android_system_properties" +version = "0.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "819e7219dbd41043ac279b19830f2efc897156490d7fd6ea916720117ee66311" +dependencies = [ + "libc", +] + [[package]] name = "anstream" version = "1.0.0" @@ -173,6 +211,7 @@ name = "bfb" version = "0.1.1" dependencies = [ "anyhow", + "bytes", "chrono", "clap", "clap_autocomplete", @@ -185,6 +224,7 @@ dependencies = [ "indicatif", "indicatif-log-bridge", "memmap2", + "parquet", "qdrant-client", "rand 0.10.1", "rand_distr", @@ -215,6 +255,27 @@ dependencies = [ "objc2", ] +[[package]] +name = "brotli" +version = "8.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5cc91aac060a7a1e25823bdccbfb6af1875b88f17c6daac97894eed8207166b3" +dependencies = [ + "alloc-no-stdlib", + "alloc-stdlib", + "brotli-decompressor", +] + +[[package]] +name = "brotli-decompressor" +version = "5.0.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3a32acac15fe1967bc3986b2a6347dffc965602354ea6f450ad07e8bfd253583" +dependencies = [ + "alloc-no-stdlib", + "alloc-stdlib", +] + [[package]] name = "bumpalo" version = "3.20.2" @@ -268,7 +329,9 @@ version = "0.4.44" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c673075a2e0e5f4a1dde27ce9dee1ea4558c7ffe648f576438a20ca1d2acc4b0" dependencies = [ + "iana-time-zone", "num-traits", + "windows-link", ] [[package]] @@ -369,6 +432,26 @@ dependencies = [ "windows-sys 0.61.2", ] +[[package]] +name = "const-random" +version = "0.1.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "87e00182fe74b066627d63b85fd550ac2998d4b0bd86bfed477a0ae4c7c71359" +dependencies = [ + "const-random-macro", +] + +[[package]] +name = "const-random-macro" +version = "0.1.16" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9d839f2a20b0aee515dc581a6172f2321f96cab76c1a38a4c584a194955390e" +dependencies = [ + "getrandom 0.2.17", + "once_cell", + "tiny-keccak", +] + [[package]] name = "core-foundation" version = "0.10.1" @@ -782,6 +865,7 @@ checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ "cfg-if", "crunchy", + "num-traits", "zerocopy", ] @@ -949,6 +1033,30 @@ dependencies = [ "tracing", ] +[[package]] +name = "iana-time-zone" +version = "0.1.65" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e31bc9ad994ba00e440a8aa5c9ef0ec67d5cb5e5cb0cc7f8b744a35b389cc470" +dependencies = [ + "android_system_properties", + "core-foundation-sys", + "iana-time-zone-haiku", + "js-sys", + "log", + "wasm-bindgen", + "windows-core", +] + +[[package]] +name = "iana-time-zone-haiku" +version = "0.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f31827a206f56af32e590ba56d5d2d085f558508192593743f16b2306495269f" +dependencies = [ + "cc", +] + [[package]] name = "icu_collections" version = "2.2.0" @@ -1282,6 +1390,15 @@ version = "0.1.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "112b39cec0b298b6c1999fee3e31427f74f676e4cb9879ed1a121b43661a4154" +[[package]] +name = "lz4_flex" +version = "0.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7ef0d4ed8669f8f8826eb00dc878084aa8f253506c4fd5e8f58f5bce72ddb97e" +dependencies = [ + "twox-hash", +] + [[package]] name = "matchit" version = "0.8.4" @@ -1374,6 +1491,16 @@ dependencies = [ "winapi", ] +[[package]] +name = "num-bigint" +version = "0.4.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c89e69e7e0f03bea5ef08013795c25018e101932225a656383bd384495ecc367" +dependencies = [ + "num-integer", + "num-traits", +] + [[package]] name = "num-complex" version = "0.4.6" @@ -1458,6 +1585,36 @@ dependencies = [ "windows-link", ] +[[package]] +name = "parquet" +version = "59.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5302d4da74d6596a1f11f9928767995b53bca657cbeea1e4e8c5074f8a1157dd" +dependencies = [ + "ahash", + "brotli", + "bytes", + "chrono", + "flate2", + "half", + "hashbrown 0.17.0", + "lz4_flex", + "num-bigint", + "num-integer", + "num-traits", + "paste", + "seq-macro", + "snap", + "twox-hash", + "zstd", +] + +[[package]] +name = "paste" +version = "1.0.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "57c0d7b74b563b49d38dae00a0c37d4d6de9b432382b2892f0574ddcae73fd0a" + [[package]] name = "percent-encoding" version = "2.3.2" @@ -1926,7 +2083,7 @@ dependencies = [ "security-framework", "security-framework-sys", "webpki-root-certs", - "windows-sys 0.52.0", + "windows-sys 0.61.2", ] [[package]] @@ -2012,6 +2169,12 @@ version = "1.0.28" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8a7852d02fc848982e0c167ef163aaff9cd91dc640ba85e263cb1ce46fae51cd" +[[package]] +name = "seq-macro" +version = "0.3.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1bc711410fbe7399f390ca1c3b60ad0f53f80e95c5eb935e52268a0e2cd49acc" + [[package]] name = "serde" version = "1.0.228" @@ -2118,6 +2281,12 @@ version = "1.15.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "67b1b7a3b5fe4f1376887184045fcf45c69e92af734b7aaddc05fb777b6fbd03" +[[package]] +name = "snap" +version = "1.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "199905e6153d6405f9728fe44daace35f8f837bbf830bb6e85fbd5828709a886" + [[package]] name = "socket2" version = "0.6.3" @@ -2255,6 +2424,15 @@ dependencies = [ "syn", ] +[[package]] +name = "tiny-keccak" +version = "2.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2c9d3793400a45f954c52e73d068316d76b6f4e36977e3fcebb13a2721e80237" +dependencies = [ + "crunchy", +] + [[package]] name = "tinystr" version = "0.8.3" @@ -2472,6 +2650,12 @@ version = "0.2.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e421abadd41a4225275504ea4d6566923418b7f05506fbc9c0fe86ba7396114b" +[[package]] +name = "twox-hash" +version = "2.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8464ec13c3691491391d9fce00f6416c9a48e46972f72d7865688be2080192c9" + [[package]] name = "unicode-ident" version = "1.0.24" @@ -2579,6 +2763,12 @@ dependencies = [ "wasm-bindgen", ] +[[package]] +name = "version_check" +version = "0.9.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0b928f33d975fc6ad9f86c8f283853ad26bdd5b10b7f1542aa2fa15e2289105a" + [[package]] name = "walkdir" version = "2.5.0" @@ -2784,7 +2974,7 @@ version = "0.1.11" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c2a7b1c03c876122aa43f3020e6c3c3ee5c05081c9a00739faf7503aeba10d22" dependencies = [ - "windows-sys 0.52.0", + "windows-sys 0.61.2", ] [[package]] @@ -2793,12 +2983,65 @@ version = "0.4.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "712e227841d057c1ee1cd2fb22fa7e5a5461ae8e48fa2ca79ec42cfc1931183f" +[[package]] +name = "windows-core" +version = "0.62.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b8e83a14d34d0623b51dce9581199302a221863196a1dde71a7663a4c2be9deb" +dependencies = [ + "windows-implement", + "windows-interface", + "windows-link", + "windows-result", + "windows-strings", +] + +[[package]] +name = "windows-implement" +version = "0.60.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "053e2e040ab57b9dc951b72c264860db7eb3b0200ba345b4e4c3b14f67855ddf" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "windows-interface" +version = "0.59.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3f316c4a2570ba26bbec722032c4099d8c8bc095efccdc15688708623367e358" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "windows-link" version = "0.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f0805222e57f7521d6a62e36fa9163bc891acd422f971defe97d64e70d0a4fe5" +[[package]] +name = "windows-result" +version = "0.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7781fa89eaf60850ac3d2da7af8e5242a5ea78d1a11c49bf2910bb5a73853eb5" +dependencies = [ + "windows-link", +] + +[[package]] +name = "windows-strings" +version = "0.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7837d08f69c77cf6b07689544538e017c1bfcf57e34b4c0ff58e6c2cd3b37091" +dependencies = [ + "windows-link", +] + [[package]] name = "windows-sys" version = "0.52.0" @@ -3179,3 +3422,31 @@ name = "zmij" version = "1.0.21" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b8848ee67ecc8aedbaf3e4122217aff892639231befc6a1b58d29fff4c2cabaa" + +[[package]] +name = "zstd" +version = "0.13.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e91ee311a569c327171651566e07972200e76fcfe2242a4fa446149a3881c08a" +dependencies = [ + "zstd-safe", +] + +[[package]] +name = "zstd-safe" +version = "7.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8f49c4d5f0abb602a93fb8736af2a4f4dd9512e36f7f570d66e65ff867ed3b9d" +dependencies = [ + "zstd-sys", +] + +[[package]] +name = "zstd-sys" +version = "2.0.16+zstd.1.5.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "91e19ebc2adc8f83e43039e79776e3fda8ca919132d68a1fed6a5faca2683748" +dependencies = [ + "cc", + "pkg-config", +] diff --git a/Cargo.toml b/Cargo.toml index 6c8e82b..06d45f2 100644 --- a/Cargo.toml +++ b/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" diff --git a/README.md b/README.md index 5374cb5..db047d9 100644 --- a/README.md +++ b/README.md @@ -24,9 +24,34 @@ 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). 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 diff --git a/examples/upload-laion-part.yaml b/examples/upload-laion-part.yaml new file mode 100644 index 0000000..f76df6f --- /dev/null +++ b/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 diff --git a/src/config/payload.rs b/src/config/payload.rs index 7532df8..6457839 100644 --- a/src/config/payload.rs +++ b/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, + pub dataset: Option>, /// Payload field name inside the dataset schema / `payloads.jsonl`. #[serde(default)] pub field: Option, diff --git a/src/config/schema.rs b/src/config/schema.rs index f5188a3..cd0aa12 100644 --- a/src/config/schema.rs +++ b/src/config/schema.rs @@ -94,7 +94,12 @@ 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 @@ -136,6 +141,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. diff --git a/src/config/vector.rs b/src/config/vector.rs index dc71e86..2f56fe5 100644 --- a/src/config/vector.rs +++ b/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, }, } @@ -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, + pub dataset: Option>, } impl Default for SparseSource { diff --git a/src/dataset/config.rs b/src/dataset/config.rs index d2cd384..e8dd235 100644 --- a/src/dataset/config.rs +++ b/src/dataset/config.rs @@ -9,7 +9,7 @@ use serde::{Deserialize, Serialize}; /// A single dataset entry (inline in upload config, or from an optional local /// `datasets.json` registry). -#[derive(Debug, Clone, Serialize, Deserialize)] +#[derive(Debug, Clone, Default, Serialize, Deserialize)] #[serde(deny_unknown_fields)] pub struct DatasetConfig { pub name: String, @@ -28,6 +28,16 @@ pub struct DatasetConfig { pub distance: Option, #[serde(default)] pub schema: Option>, + /// `parquet` only: columns to keep. `None` ⇒ every column. + #[serde(default)] + pub columns: Option>, + /// `parquet` only: columns to drop (applied after `columns`). + #[serde(default)] + pub exclude: Vec, + /// `parquet` only: value substituted for nulls and non-finite floats. + /// Omitted by default, which leaves the payload field absent. + #[serde(default)] + pub fill_null: Option, } impl DatasetConfig { @@ -53,6 +63,9 @@ pub struct ResolvedDatasetConfig { pub distance: Option, #[allow(dead_code)] pub schema: Option>, + pub columns: Option>, + pub exclude: Vec, + pub fill_null: Option, } impl ResolvedDatasetConfig { @@ -64,12 +77,7 @@ impl ResolvedDatasetConfig { let kind = inline .kind .or_else(|| base.and_then(|b| b.kind)) - .with_context(|| { - format!( - "dataset {:?}: missing `format` (h5, tar, sparse)", - inline.name - ) - })?; + .with_context(|| format!("dataset {:?}: missing `format` ({KINDS})", inline.name))?; let path = inline .path .or_else(|| base.and_then(|b| b.path.clone())) @@ -92,16 +100,45 @@ impl ResolvedDatasetConfig { schema: inline .schema .or_else(|| base.and_then(|b| b.schema.clone())), + columns: inline + .columns + .or_else(|| base.and_then(|b| b.columns.clone())), + exclude: if inline.exclude.is_empty() { + base.map(|b| b.exclude.clone()).unwrap_or_default() + } else { + inline.exclude + }, + fill_null: inline + .fill_null + .or_else(|| base.and_then(|b| b.fill_null.clone())), }) } } +/// Formats accepted by `format:`, for error messages. +const KINDS: &str = "h5, tar, sparse, npy, parquet"; + #[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] #[serde(rename_all = "lowercase")] pub enum DatasetKind { H5, Tar, Sparse, + /// A standalone 2-D float `.npy` array: dense vectors, no payloads. + Npy, + /// A parquet file of payload rows: no vectors. + Parquet, +} + +impl DatasetKind { + /// Is the dataset a single file, rather than a directory of extracted + /// files? Single-file datasets are installed as-is on download. + pub fn is_single_file(self) -> bool { + matches!( + self, + DatasetKind::H5 | DatasetKind::Npy | DatasetKind::Parquet + ) + } } impl DatasetConfig { @@ -119,10 +156,7 @@ impl DatasetConfig { bail!("dataset source requires `name`"); } if self.kind.is_none() { - bail!( - "dataset {:?} requires `format` (h5, tar, sparse)", - self.name - ); + bail!("dataset {:?} requires `format` ({KINDS})", self.name); } if self.path.is_none() { bail!("dataset {:?} requires `path`", self.name); @@ -144,7 +178,7 @@ mod tests { link: Some("http://ann-benchmarks.com/glove-25-angular.hdf5".to_string()), vector_size: Some(25), distance: Some("cosine".to_string()), - schema: None, + ..Default::default() }; let resolved = DatasetConfig::resolve(inline, &HashMap::new()).unwrap(); @@ -165,18 +199,13 @@ mod tests { link: Some("http://example.com/glove.hdf5".to_string()), vector_size: Some(100), distance: Some("cosine".to_string()), - schema: None, + ..Default::default() }, ); let inline = DatasetConfig { name: "glove-100-angular".to_string(), - kind: None, - path: None, - link: None, - vector_size: None, - distance: None, - schema: None, + ..Default::default() }; let resolved = DatasetConfig::resolve(inline, ®istry).unwrap(); diff --git a/src/dataset/download.rs b/src/dataset/download.rs index dd751ab..a559b74 100644 --- a/src/dataset/download.rs +++ b/src/dataset/download.rs @@ -140,19 +140,12 @@ fn install_download( return Ok(()); } - match kind { - DatasetKind::H5 => { - tmp.persist(target) - .with_context(|| format!("failed to install download at {}", target.display()))?; - } - DatasetKind::Tar | DatasetKind::Sparse => { - bail!( - "dataset archive at {link} must end with .tgz or .tar.gz for type {:?}", - kind - ); - } + if kind.is_single_file() { + tmp.persist(target) + .with_context(|| format!("failed to install download at {}", target.display()))?; + return Ok(()); } - Ok(()) + bail!("dataset archive at {link} must end with .tgz or .tar.gz for format {kind:?}") } #[cfg(test)] diff --git a/src/dataset/fixtures.rs b/src/dataset/fixtures.rs new file mode 100644 index 0000000..46de471 --- /dev/null +++ b/src/dataset/fixtures.rs @@ -0,0 +1,119 @@ +//! Dataset files built on the fly for tests: real `.npy` and `.parquet` bytes, +//! so the readers are exercised against the formats rather than against mocks. + +use std::fs::File; +use std::path::Path; +use std::sync::Arc; + +use parquet::data_type::{BoolType, ByteArray, ByteArrayType, DoubleType, Int64Type}; +use parquet::file::properties::WriterProperties; +use parquet::file::writer::SerializedFileWriter; +use parquet::schema::parser::parse_message_type; + +/// Build a minimal little-endian `.npy` v1.0 buffer. +pub fn make_npy(descr: &str, rows: usize, cols: usize, data: &[u8]) -> Vec { + let dict = + format!("{{'descr': '{descr}', 'fortran_order': False, 'shape': ({rows}, {cols}), }}"); + let mut header = dict.into_bytes(); + // Pad with spaces so the total (magic + version + len field + header) is a + // multiple of 64, and terminate with a newline. + let unpadded = 10 + header.len() + 1; + header.extend(std::iter::repeat_n(b' ', (64 - unpadded % 64) % 64)); + header.push(b'\n'); + + let mut out = Vec::new(); + out.extend_from_slice(b"\x93NUMPY"); + out.extend_from_slice(&[1, 0]); // version 1.0 + out.extend_from_slice(&(header.len() as u16).to_le_bytes()); + out.extend_from_slice(&header); + out.extend_from_slice(data); + out +} + +/// A `rows` x `cols` f32 `.npy` whose value at (r, c) is `base + r * cols + c`. +pub fn make_ramp_npy(base: usize, rows: usize, cols: usize) -> Vec { + let bytes: Vec = (0..rows * cols) + .flat_map(|x| ((base + x) as f32).to_le_bytes()) + .collect(); + make_npy(" = (start..end).collect(); + let mut group = writer.next_row_group().unwrap(); + + let ids: Vec = span.iter().map(|i| base + *i as i64).collect(); + let mut col = group.next_column().unwrap().unwrap(); + col.typed::() + .write_batch(&ids, None, None) + .unwrap(); + col.close().unwrap(); + + let sims: Vec = span + .iter() + .map(|i| if i % 3 == 0 { f64::NAN } else { *i as f64 }) + .collect(); + let mut col = group.next_column().unwrap().unwrap(); + col.typed::() + .write_batch(&sims, None, None) + .unwrap(); + col.close().unwrap(); + + let urls: Vec = span + .iter() + .map(|i| ByteArray::from(format!("http://example.com/{i}").as_str())) + .collect(); + let mut col = group.next_column().unwrap().unwrap(); + col.typed::() + .write_batch(&urls, None, None) + .unwrap(); + col.close().unwrap(); + + let flags: Vec = span.iter().map(|i| i % 2 == 0).collect(); + let mut col = group.next_column().unwrap().unwrap(); + col.typed::() + .write_batch(&flags, None, None) + .unwrap(); + col.close().unwrap(); + + // Optional column: definition level 0 = null, 1 = present. + let defs: Vec = span.iter().map(|i| i16::from(i % 3 != 0)).collect(); + let captions: Vec = span + .iter() + .filter(|i| *i % 3 != 0) + .map(|i| ByteArray::from(format!("caption {i}").as_str())) + .collect(); + let mut col = group.next_column().unwrap().unwrap(); + col.typed::() + .write_batch(&captions, Some(&defs), None) + .unwrap(); + col.close().unwrap(); + + group.close().unwrap(); + } + writer.close().unwrap(); +} diff --git a/src/dataset/mod.rs b/src/dataset/mod.rs index 6c9adc6..fb962ed 100644 --- a/src/dataset/mod.rs +++ b/src/dataset/mod.rs @@ -1,5 +1,7 @@ mod config; mod download; +#[cfg(test)] +pub(crate) mod fixtures; mod payload; mod reader; mod readers; diff --git a/src/dataset/reader.rs b/src/dataset/reader.rs index f8c6129..32fbee6 100644 --- a/src/dataset/reader.rs +++ b/src/dataset/reader.rs @@ -5,13 +5,15 @@ use serde_json::Value; use super::config::{DatasetConfig, DatasetKind}; use super::download::ensure_downloaded; -use super::readers::{H5Reader, SparseReader, TarReader}; +use super::readers::{H5Reader, NpyReader, ParquetReader, SparseReader, TarReader}; use super::registry::load_registry; enum DatasetReaderInner { H5(H5Reader), Tar(TarReader), Sparse(SparseReader), + Npy(NpyReader), + Parquet(ParquetReader), } /// Random access to points from a vector-db-benchmark dataset. @@ -41,6 +43,21 @@ impl DatasetReader { let n = reader.num_points(); (DatasetReaderInner::Sparse(reader), n) } + DatasetKind::Npy => { + let reader = NpyReader::open(&local_path)?; + let n = reader.num_points(); + (DatasetReaderInner::Npy(reader), n) + } + DatasetKind::Parquet => { + let reader = ParquetReader::open( + &local_path, + config.columns.as_deref(), + &config.exclude, + config.fill_null.as_ref(), + )?; + let n = reader.num_points(); + (DatasetReaderInner::Parquet(reader), n) + } }; Ok(DatasetReader { inner, num_points }) } @@ -49,7 +66,10 @@ impl DatasetReader { match &self.inner { DatasetReaderInner::H5(r) => r.vector_at(idx), DatasetReaderInner::Tar(r) => r.vector_at(idx), - DatasetReaderInner::Sparse(_) => bail!("dataset does not contain dense vectors"), + DatasetReaderInner::Npy(r) => r.vector_at(idx), + DatasetReaderInner::Sparse(_) | DatasetReaderInner::Parquet(_) => { + bail!("dataset does not contain dense vectors") + } } } @@ -63,6 +83,7 @@ impl DatasetReader { pub fn payload_field(&self, idx: usize, field: &str) -> Result> { match &self.inner { DatasetReaderInner::Tar(r) => r.payload_field(idx, field), + DatasetReaderInner::Parquet(r) => r.payload_field(idx, field), _ => bail!("dataset does not contain payloads"), } } @@ -70,6 +91,7 @@ impl DatasetReader { pub fn payload_object(&self, idx: usize) -> Result> { match &self.inner { DatasetReaderInner::Tar(r) => r.payload_object(idx), + DatasetReaderInner::Parquet(r) => r.payload_object(idx), _ => bail!("dataset does not contain payloads"), } } @@ -80,6 +102,9 @@ impl DatasetReader { DatasetReaderInner::H5(r) => r.num_queries(), DatasetReaderInner::Tar(r) => r.num_queries(), DatasetReaderInner::Sparse(r) => r.num_queries(), + // Component formats hold corpus rows only; a query set is a + // separate file, declared as its own source. + DatasetReaderInner::Npy(_) | DatasetReaderInner::Parquet(_) => 0, } } @@ -89,6 +114,9 @@ impl DatasetReader { DatasetReaderInner::H5(r) => r.query_at(idx), DatasetReaderInner::Tar(r) => r.query_at(idx), DatasetReaderInner::Sparse(_) => bail!("sparse dataset has no dense queries"), + DatasetReaderInner::Npy(_) | DatasetReaderInner::Parquet(_) => { + bail!("dataset has no query set") + } } } @@ -106,6 +134,9 @@ impl DatasetReader { DatasetReaderInner::H5(r) => r.neighbors_at(idx), DatasetReaderInner::Tar(r) => r.query_ground_truth(idx), DatasetReaderInner::Sparse(r) => r.query_ground_truth(idx), + DatasetReaderInner::Npy(_) | DatasetReaderInner::Parquet(_) => { + bail!("dataset has no ground truth") + } } } } diff --git a/src/dataset/readers/mod.rs b/src/dataset/readers/mod.rs index c4a338a..f2ce770 100644 --- a/src/dataset/readers/mod.rs +++ b/src/dataset/readers/mod.rs @@ -1,9 +1,13 @@ mod binary; mod h5; mod jsonl; +mod npy; +mod parquet; mod sparse; mod tar; pub use h5::H5Reader; +pub use npy::NpyReader; +pub use parquet::ParquetReader; pub use sparse::SparseReader; pub use tar::TarReader; diff --git a/src/dataset/readers/npy.rs b/src/dataset/readers/npy.rs new file mode 100644 index 0000000..a61a4bd --- /dev/null +++ b/src/dataset/readers/npy.rs @@ -0,0 +1,295 @@ +//! `.npy` (NumPy array) reading. +//! +//! Two consumers share this module: the `tar` bundle format, whose payload is a +//! `vectors.npy` inside an extracted archive, and the standalone `npy` dataset +//! format, which points straight at one such file. +//! +//! Rows are served from an mmap, so access is lock-free and resident memory is +//! reclaimable page cache rather than committed RAM. + +use std::fs::File; +use std::path::Path; + +use anyhow::{Context, Result, bail}; +use half::f16; +use memmap2::Mmap; + +/// Element type of a `.npy` array (little-endian float). +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum Dtype { + F16, + F32, + F64, +} + +impl Dtype { + pub fn size(self) -> usize { + match self { + Dtype::F16 => 2, + Dtype::F32 => 4, + Dtype::F64 => 8, + } + } +} + +/// Shape and element layout read from a `.npy` header. +#[derive(Debug, Clone, Copy)] +pub struct NpyLayout { + pub dtype: Dtype, + pub num_points: usize, + pub dim: usize, + /// Byte offset of the raw array data within the file. + pub data_offset: usize, +} + +/// A 2-D float `.npy` array served from an mmap. +pub struct NpyMatrix { + mmap: Mmap, + layout: NpyLayout, +} + +impl NpyMatrix { + pub fn open(path: &Path) -> Result { + let file = + File::open(path).with_context(|| format!("failed to open {}", path.display()))?; + let mmap = unsafe { Mmap::map(&file) } + .with_context(|| format!("failed to mmap {}", path.display()))?; + + let layout = parse_npy_header(&mmap) + .with_context(|| format!("failed to parse {}", path.display()))?; + let needed = layout.data_offset + layout.num_points * layout.dim * layout.dtype.size(); + if mmap.len() < needed { + bail!( + "{} is truncated: {} bytes, need {needed}", + path.display(), + mmap.len() + ); + } + + Ok(NpyMatrix { mmap, layout }) + } + + pub fn rows(&self) -> usize { + self.layout.num_points + } + + pub fn row(&self, idx: usize) -> Result> { + if idx >= self.layout.num_points { + bail!( + "index {idx} out of range (array has {} rows)", + self.layout.num_points + ); + } + let row_bytes = self.layout.dim * self.layout.dtype.size(); + let start = self.layout.data_offset + idx * row_bytes; + let bytes = &self.mmap[start..start + row_bytes]; + Ok(match self.layout.dtype { + Dtype::F16 => bytes + .chunks_exact(2) + .map(|b| f16::from_le_bytes([b[0], b[1]]).to_f32()) + .collect(), + Dtype::F32 => bytes + .chunks_exact(4) + .map(|b| f32::from_le_bytes(b.try_into().unwrap())) + .collect(), + Dtype::F64 => bytes + .chunks_exact(8) + .map(|b| f64::from_le_bytes(b.try_into().unwrap()) as f32) + .collect(), + }) + } +} + +/// A standalone `.npy` file used as a dense-vector dataset source. +/// +/// Vectors only: it carries no payloads, query set, or ground truth, so those +/// accessors are rejected by [`DatasetReader`](crate::dataset::DatasetReader). +pub struct NpyReader { + matrix: NpyMatrix, +} + +impl NpyReader { + pub fn open(path: &Path) -> Result { + Ok(NpyReader { + matrix: NpyMatrix::open(path)?, + }) + } + + pub fn num_points(&self) -> usize { + self.matrix.rows() + } + + pub fn vector_at(&self, idx: usize) -> Result> { + self.matrix.row(idx) + } +} + +/// Minimal parser for the `.npy` format (v1/v2 headers) sufficient for the 2-D +/// float arrays shipped by vector-db-benchmark and by embedding dumps such as +/// LAION's `img_emb_*.npy`. +/// +/// Only the leading header is read, so `buf` may be a short prefix of the file +/// — which is what makes remote row counts a single ranged request rather than +/// a download. +pub fn parse_npy_header(buf: &[u8]) -> Result { + if buf.len() < 10 || &buf[0..6] != b"\x93NUMPY" { + bail!("not a .npy file (bad magic)"); + } + let major = buf[6]; + // Header length field: 2 bytes (v1) or 4 bytes (v2+), little-endian. + let (header_len, header_start) = if major >= 2 { + if buf.len() < 12 { + bail!("truncated .npy header"); + } + ( + u32::from_le_bytes(buf[8..12].try_into().unwrap()) as usize, + 12, + ) + } else { + ( + u16::from_le_bytes(buf[8..10].try_into().unwrap()) as usize, + 10, + ) + }; + let header_end = header_start + header_len; + if buf.len() < header_end { + bail!( + "truncated .npy header: need {header_end} bytes, got {}", + buf.len() + ); + } + let header = std::str::from_utf8(&buf[header_start..header_end]) + .context(".npy header is not valid UTF-8")?; + + let descr = extract_quoted(header, "descr").context(".npy header missing 'descr'")?; + let dtype = match descr.as_str() { + " Dtype::F16, + " Dtype::F32, + " Dtype::F64, + other => bail!("unsupported .npy dtype {other:?} (expected float16/32/64)"), + }; + + if header.contains("'fortran_order': True") || header.contains("\"fortran_order\": true") { + bail!(".npy array is Fortran-ordered; expected C order"); + } + + let (num_points, dim) = extract_shape(header)?; + Ok(NpyLayout { + dtype, + num_points, + dim, + data_offset: header_end, + }) +} + +/// Extract a single-quoted string value for `key` from a `.npy` header dict. +fn extract_quoted(header: &str, key: &str) -> Option { + let after_key = &header[header.find(&format!("'{key}'"))?..]; + let after_colon = &after_key[after_key.find(':')? + 1..]; + let bytes = after_colon.as_bytes(); + let mut i = 0; + while i < bytes.len() && bytes[i].is_ascii_whitespace() { + i += 1; + } + if i >= bytes.len() || (bytes[i] != b'\'' && bytes[i] != b'"') { + return None; + } + let quote = bytes[i]; + i += 1; + let start = i; + while i < bytes.len() && bytes[i] != quote { + i += 1; + } + Some(after_colon[start..i].to_string()) +} + +/// Extract the 2-D `(rows, cols)` shape tuple from a `.npy` header dict. +fn extract_shape(header: &str) -> Result<(usize, usize)> { + let after_key = &header[header + .find("'shape'") + .context(".npy header missing 'shape'")?..]; + let open = after_key.find('(').context("malformed 'shape'")?; + let close = after_key[open..].find(')').context("malformed 'shape'")? + open; + let dims: Vec = after_key[open + 1..close] + .split(',') + .map(str::trim) + .filter(|s| !s.is_empty()) + .map(|s| s.parse::()) + .collect::>() + .context("malformed 'shape' dims")?; + if dims.len() != 2 { + bail!("expected a 2-D .npy array, got shape {dims:?}"); + } + Ok((dims[0], dims[1])) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::dataset::fixtures::{make_npy, make_ramp_npy}; + + #[test] + fn reads_f32_rows() { + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("v.npy"); + std::fs::write(&path, make_ramp_npy(0, 2, 3)).unwrap(); + + let reader = NpyReader::open(&path).unwrap(); + assert_eq!(reader.num_points(), 2); + assert_eq!(reader.vector_at(1).unwrap(), vec![3.0, 4.0, 5.0]); + } + + #[test] + fn reads_f16_rows() { + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("v.npy"); + let values = [0.5f32, -1.0, 2.0, 4.0]; + let bytes: Vec = values + .iter() + .flat_map(|&v| f16::from_f32(v).to_le_bytes()) + .collect(); + std::fs::write(&path, make_npy(" = (0..4u32).flat_map(|x| x.to_le_bytes()).collect(); + let err = parse_npy_header(&make_npy(", + /// Projected column names, in file order. `None` ⇒ every column. + columns: Option>, + /// Value substituted for nulls and non-finite floats. `None` ⇒ omit the field. + fill_null: Option, + cursor: Mutex, +} + +#[derive(Default)] +struct Cursor { + iter: Option>, + /// Global index of the row `iter` will yield next. + next: usize, + ring: VecDeque<(usize, Value)>, +} + +impl ParquetReader { + /// Open `path`, keeping `columns` (default: all) minus `exclude`. + pub fn open( + path: &Path, + columns: Option<&[String]>, + exclude: &[String], + fill_null: Option<&Value>, + ) -> Result { + let file = + File::open(path).with_context(|| format!("failed to open {}", path.display()))?; + let reader = SerializedFileReader::new(file) + .with_context(|| format!("failed to read parquet metadata from {}", path.display()))?; + let metadata = reader.metadata(); + + let num_rows = metadata.file_metadata().num_rows().max(0) as usize; + let mut group_starts = Vec::with_capacity(metadata.num_row_groups() + 1); + let mut running = 0usize; + for i in 0..metadata.num_row_groups() { + group_starts.push(running); + running += metadata.row_group(i).num_rows().max(0) as usize; + } + group_starts.push(running); + + let available: Vec = metadata + .file_metadata() + .schema_descr() + .root_schema() + .get_fields() + .iter() + .map(|f| f.name().to_string()) + .collect(); + let columns = resolve_projection(&available, columns, exclude, path)?; + + Ok(ParquetReader { + path: path.to_path_buf(), + num_rows, + group_starts, + columns, + fill_null: fill_null.cloned(), + cursor: Mutex::new(Cursor::default()), + }) + } + + pub fn num_points(&self) -> usize { + self.num_rows + } + + /// The whole payload object for a row. + pub fn payload_object(&self, idx: usize) -> Result> { + if idx >= self.num_rows { + return Ok(None); + } + let mut cursor = self.cursor.lock().unwrap(); + + if let Some((_, value)) = cursor.ring.iter().find(|(i, _)| *i == idx) { + return Ok(Some(value.clone())); + } + // Behind the ring (or nothing decoded yet): restart at the row group + // holding `idx` so a rewind costs one group, not the whole file. + if cursor.iter.is_none() || idx < cursor.next { + let group = self.group_of(idx); + cursor.iter = Some(self.open_iter(group)?); + cursor.next = self.group_starts[group]; + cursor.ring.clear(); + } + + while cursor.next <= idx { + let row = cursor + .iter + .as_mut() + .expect("iterator was just installed") + .next() + .with_context(|| { + format!( + "{} ended at row {} while reading row {idx}", + self.path.display(), + cursor.next + ) + })? + .with_context(|| format!("failed to decode {} row {idx}", self.path.display()))?; + + let value = self.row_to_value(&row); + let at = cursor.next; + cursor.ring.push_back((at, value)); + if cursor.ring.len() > RING_ROWS { + cursor.ring.pop_front(); + } + cursor.next += 1; + } + + Ok(cursor + .ring + .back() + .filter(|(i, _)| *i == idx) + .map(|(_, value)| value.clone())) + } + + pub fn payload_field(&self, idx: usize, field: &str) -> Result> { + Ok(self + .payload_object(idx)? + .and_then(|object| object.get(field).cloned())) + } + + /// Index of the row group containing global row `idx`. + fn group_of(&self, idx: usize) -> usize { + match self.group_starts.binary_search(&idx) { + Ok(exact) => exact.min(self.group_starts.len().saturating_sub(2)), + Err(next) => next - 1, + } + } + + /// A row iterator positioned at the first row of row group `group`. + fn open_iter(&self, group: usize) -> Result> { + let file = File::open(&self.path) + .with_context(|| format!("failed to open {}", self.path.display()))?; + let options = ReadOptionsBuilder::new() + .with_predicate(Box::new(move |_, i| i >= group)) + .build(); + let reader = SerializedFileReader::new_with_options(file, options) + .with_context(|| format!("failed to open {}", self.path.display()))?; + + let projection = self + .columns + .as_ref() + .map(|names| project_schema(reader.metadata().file_metadata().schema(), names)) + .transpose()?; + + RowIter::from_file_into(Box::new(reader)) + .project(projection) + .with_context(|| format!("failed to project columns of {}", self.path.display())) + } + + fn row_to_value(&self, row: &Row) -> Value { + let mut object = JsonMap::new(); + for (name, field) in row.get_column_iter() { + match self.field_to_value(field) { + Some(value) => { + object.insert(name.clone(), value); + } + None => { + if let Some(fill) = &self.fill_null { + object.insert(name.clone(), fill.clone()); + } + } + } + } + Value::Object(object) + } + + /// `None` for values with no JSON representation (null, NaN, ±inf, + /// non-UTF-8 bytes) — the caller then omits the field or substitutes + /// `fill_null`. + fn field_to_value(&self, field: &Field) -> Option { + Some(match field { + Field::Null => return None, + Field::Bool(v) => Value::Bool(*v), + Field::Byte(v) => Value::Number((*v).into()), + Field::Short(v) => Value::Number((*v).into()), + Field::Int(v) => Value::Number((*v).into()), + Field::Long(v) => Value::Number((*v).into()), + Field::UByte(v) => Value::Number((*v).into()), + Field::UShort(v) => Value::Number((*v).into()), + Field::UInt(v) => Value::Number((*v).into()), + Field::ULong(v) => Value::Number((*v).into()), + Field::Date(v) => Value::Number((*v).into()), + Field::TimeMillis(v) => Value::Number((*v).into()), + Field::TimeMicros(v) => Value::Number((*v).into()), + Field::TimestampMillis(v) => Value::Number((*v).into()), + Field::TimestampMicros(v) => Value::Number((*v).into()), + Field::Float16(v) => number(f32::from(*v) as f64)?, + Field::Float(v) => number(*v as f64)?, + Field::Double(v) => number(*v)?, + Field::Decimal(v) => number(decimal_to_f64(v))?, + Field::Str(v) => Value::String(v.clone()), + Field::Bytes(v) => Value::String(std::str::from_utf8(v.data()).ok()?.to_string()), + Field::Group(row) => self.row_to_value(row), + Field::ListInternal(list) => self.list_to_value(list), + Field::MapInternal(map) => self.map_to_value(map), + }) + } + + fn list_to_value(&self, list: &List) -> Value { + Value::Array( + list.elements() + .iter() + .map(|element| self.field_to_value(element).unwrap_or(Value::Null)) + .collect(), + ) + } + + fn map_to_value(&self, map: &Map) -> Value { + let mut object = JsonMap::new(); + for (key, value) in map.entries() { + // JSON object keys are strings; render the key's scalar form. + let key = match self.field_to_value(key) { + Some(Value::String(s)) => s, + Some(other) => other.to_string(), + None => continue, + }; + object.insert(key, self.field_to_value(value).unwrap_or(Value::Null)); + } + Value::Object(object) + } +} + +/// JSON has no NaN or infinity — those become "no value", same as null. +fn number(value: f64) -> Option { + Number::from_f64(value).map(Value::Number) +} + +/// Best-effort decimal → float. Payload filters are numeric, and Qdrant has no +/// fixed-point payload type, so precision beyond f64 has nowhere to go anyway. +fn decimal_to_f64(decimal: &parquet::data_type::Decimal) -> f64 { + let bytes = decimal.data(); + // Big-endian two's complement, sign-extended into i128. + let mut unscaled: i128 = if bytes.first().is_some_and(|b| b & 0x80 != 0) { + -1 + } else { + 0 + }; + for byte in bytes { + unscaled = (unscaled << 8) | i128::from(*byte); + } + unscaled as f64 / 10f64.powi(decimal.scale()) +} + +/// Resolve `columns` / `exclude` against the file's actual columns. +fn resolve_projection( + available: &[String], + columns: Option<&[String]>, + exclude: &[String], + path: &Path, +) -> Result>> { + for name in columns.unwrap_or_default().iter().chain(exclude) { + if !available.contains(name) { + bail!( + "{} has no column {name:?}; available columns: {}", + path.display(), + available.join(", ") + ); + } + } + + if columns.is_none() && exclude.is_empty() { + return Ok(None); + } + let kept: Vec = available + .iter() + .filter(|name| columns.is_none_or(|wanted| wanted.contains(name))) + .filter(|name| !exclude.contains(name)) + .cloned() + .collect(); + if kept.is_empty() { + bail!("{}: `columns`/`exclude` select no columns", path.display()); + } + Ok(Some(kept)) +} + +/// Rebuild the root group type with only `names`, preserving file order. +fn project_schema(schema: &SchemaType, names: &[String]) -> Result { + let fields = schema + .get_fields() + .iter() + .filter(|field| names.iter().any(|name| name == field.name())) + .cloned() + .collect(); + SchemaType::group_type_builder(schema.name()) + .with_fields(fields) + .build() + .context("failed to build parquet projection") +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::dataset::fixtures::write_parquet; + + fn open(path: &Path) -> ParquetReader { + ParquetReader::open(path, None, &[], None).unwrap() + } + + #[test] + fn reads_rows_in_order_across_row_groups() { + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("m.parquet"); + write_parquet(&path, 0, 250, 64); + + let reader = open(&path); + assert_eq!(reader.num_points(), 250); + for i in 0..250 { + let row = reader.payload_object(i).unwrap().unwrap(); + assert_eq!(row["id"], i as i64); + assert_eq!(row["url"], format!("http://example.com/{i}")); + assert_eq!(row["nsfw"], i % 2 == 0); + } + } + + /// NaN and null have no JSON form; by default the field is simply absent, + /// which Qdrant treats as "no value" rather than a bogus 0. + #[test] + fn omits_nan_and_null_by_default() { + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("m.parquet"); + write_parquet(&path, 0, 10, 64); + + let reader = open(&path); + let hole = reader.payload_object(3).unwrap().unwrap(); + assert!(!hole.as_object().unwrap().contains_key("similarity")); + assert!(!hole.as_object().unwrap().contains_key("caption")); + + let full = reader.payload_object(4).unwrap().unwrap(); + assert_eq!(full["similarity"], 4.0); + assert_eq!(full["caption"], "caption 4"); + } + + #[test] + fn fill_null_substitutes_a_value() { + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("m.parquet"); + write_parquet(&path, 0, 10, 64); + + let reader = ParquetReader::open(&path, None, &[], Some(&Value::from(0))).unwrap(); + let hole = reader.payload_object(3).unwrap().unwrap(); + assert_eq!(hole["similarity"], 0); + assert_eq!(hole["caption"], 0); + } + + #[test] + fn exclude_drops_a_column() { + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("m.parquet"); + write_parquet(&path, 0, 10, 64); + + let reader = + ParquetReader::open(&path, None, std::slice::from_ref(&"url".to_string()), None) + .unwrap(); + let row = reader.payload_object(1).unwrap().unwrap(); + assert!(!row.as_object().unwrap().contains_key("url")); + assert_eq!(row["id"], 1); + } + + #[test] + fn columns_keeps_only_the_listed_ones() { + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("m.parquet"); + write_parquet(&path, 0, 10, 64); + + let wanted = vec!["id".to_string(), "similarity".to_string()]; + let reader = ParquetReader::open(&path, Some(&wanted), &[], None).unwrap(); + let row = reader.payload_object(1).unwrap().unwrap(); + let keys: Vec<&String> = row.as_object().unwrap().keys().collect(); + assert_eq!(keys, vec!["id", "similarity"]); + } + + #[test] + fn unknown_column_names_are_rejected_with_the_available_ones() { + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("m.parquet"); + write_parquet(&path, 0, 10, 64); + + let err = ParquetReader::open(&path, None, &["exif".to_string()], None) + .map(|_| ()) + .unwrap_err(); + let message = err.to_string(); + assert!(message.contains("exif"), "{message}"); + assert!(message.contains("similarity"), "{message}"); + } + + /// Backward seeks must still return the right row (they rewind to the + /// containing row group), and forward jumps must not skip rows. + #[test] + fn serves_out_of_order_reads() { + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("m.parquet"); + write_parquet(&path, 0, 500, 64); + + let reader = open(&path); + for i in [499usize, 0, 250, 251, 12, 499, 13] { + assert_eq!(reader.payload_object(i).unwrap().unwrap()["id"], i as i64); + } + } + + #[test] + fn reads_a_single_field() { + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("m.parquet"); + write_parquet(&path, 0, 10, 64); + + let reader = open(&path); + assert_eq!( + reader.payload_field(4, "caption").unwrap().unwrap(), + "caption 4" + ); + assert_eq!(reader.payload_field(3, "caption").unwrap(), None); + } + + #[test] + fn out_of_range_row_is_none() { + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("m.parquet"); + write_parquet(&path, 0, 10, 64); + + assert_eq!(open(&path).payload_object(10).unwrap(), None); + } +} + diff --git a/src/dataset/readers/tar.rs b/src/dataset/readers/tar.rs index f1d56db..a7950e3 100644 --- a/src/dataset/readers/tar.rs +++ b/src/dataset/readers/tar.rs @@ -1,38 +1,15 @@ -use std::fs::File; use std::path::Path; -use anyhow::{Context, Result, bail}; -use half::f16; -use memmap2::Mmap; +use anyhow::{Context, Result}; use serde_json::Value; use super::jsonl::JsonlStore; +use super::npy::NpyMatrix; -/// Element type of a `vectors.npy` array (little-endian float). -#[derive(Debug, Clone, Copy)] -enum Dtype { - F16, - F32, - F64, -} - -impl Dtype { - fn size(self) -> usize { - match self { - Dtype::F16 => 2, - Dtype::F32 => 4, - Dtype::F64 => 8, - } - } -} - +/// An extracted `.tgz` bundle: `vectors.npy` plus optional `payloads.jsonl` and +/// `tests.jsonl` (ann-filtering-benchmark-datasets layout). pub struct TarReader { - mmap: Mmap, - dtype: Dtype, - num_points: usize, - dim: usize, - /// Byte offset of the raw array data within the mmap. - data_offset: usize, + vectors: NpyMatrix, payloads: Option, /// Query set + ground truth (`tests.jsonl`), if present. Each line has the /// shape `{ "query": [..], "conditions": {..}, "closest_ids": [..], @@ -42,21 +19,7 @@ pub struct TarReader { impl TarReader { pub fn open(path: &Path) -> Result { - let vectors_path = path.join("vectors.npy"); - let file = File::open(&vectors_path) - .with_context(|| format!("failed to open {}", vectors_path.display()))?; - let mmap = unsafe { Mmap::map(&file).context("failed to mmap vectors.npy")? }; - - let layout = parse_npy_header(&mmap) - .with_context(|| format!("failed to parse {}", vectors_path.display()))?; - let needed = layout.data_offset + layout.num_points * layout.dim * layout.dtype.size(); - if mmap.len() < needed { - bail!( - "{} is truncated: {} bytes, need {needed}", - vectors_path.display(), - mmap.len() - ); - } + let vectors = NpyMatrix::open(&path.join("vectors.npy"))?; let payloads_path = path.join("payloads.jsonl"); let payloads = if payloads_path.exists() { @@ -73,44 +36,18 @@ impl TarReader { }; Ok(TarReader { - mmap, - dtype: layout.dtype, - num_points: layout.num_points, - dim: layout.dim, - data_offset: layout.data_offset, + vectors, payloads, queries, }) } pub fn num_points(&self) -> usize { - self.num_points + self.vectors.rows() } pub fn vector_at(&self, idx: usize) -> Result> { - if idx >= self.num_points { - bail!( - "index {idx} out of range (dataset has {} points)", - self.num_points - ); - } - let row_bytes = self.dim * self.dtype.size(); - let start = self.data_offset + idx * row_bytes; - let bytes = &self.mmap[start..start + row_bytes]; - Ok(match self.dtype { - Dtype::F16 => bytes - .chunks_exact(2) - .map(|b| f16::from_le_bytes([b[0], b[1]]).to_f32()) - .collect(), - Dtype::F32 => bytes - .chunks_exact(4) - .map(|b| f32::from_le_bytes(b.try_into().unwrap())) - .collect(), - Dtype::F64 => bytes - .chunks_exact(8) - .map(|b| f64::from_le_bytes(b.try_into().unwrap()) as f32) - .collect(), - }) + self.vectors.row(idx) } pub fn payload_field(&self, idx: usize, field: &str) -> Result> { @@ -180,138 +117,19 @@ fn parse_f32_array(value: &Value) -> Option> { .collect::>>() } -struct NpyLayout { - dtype: Dtype, - num_points: usize, - dim: usize, - data_offset: usize, -} - -/// Minimal parser for the `.npy` format (v1/v2 headers) sufficient for the 2-D -/// float `vectors.npy` files shipped by vector-db-benchmark. -fn parse_npy_header(buf: &[u8]) -> Result { - if buf.len() < 10 || &buf[0..6] != b"\x93NUMPY" { - bail!("not a .npy file (bad magic)"); - } - let major = buf[6]; - // Header length field: 2 bytes (v1) or 4 bytes (v2+), little-endian. - let (header_len, header_start) = if major >= 2 { - if buf.len() < 12 { - bail!("truncated .npy header"); - } - ( - u32::from_le_bytes(buf[8..12].try_into().unwrap()) as usize, - 12, - ) - } else { - ( - u16::from_le_bytes(buf[8..10].try_into().unwrap()) as usize, - 10, - ) - }; - let header_end = header_start + header_len; - if buf.len() < header_end { - bail!("truncated .npy header"); - } - let header = std::str::from_utf8(&buf[header_start..header_end]) - .context(".npy header is not valid UTF-8")?; - - let descr = extract_quoted(header, "descr").context(".npy header missing 'descr'")?; - let dtype = match descr.as_str() { - " Dtype::F16, - " Dtype::F32, - " Dtype::F64, - other => bail!("unsupported vectors.npy dtype {other:?} (expected float16/32/64)"), - }; - - if header.contains("'fortran_order': True") || header.contains("\"fortran_order\": true") { - bail!("vectors.npy is Fortran-ordered; expected C order"); - } - - let (num_points, dim) = extract_shape(header)?; - Ok(NpyLayout { - dtype, - num_points, - dim, - data_offset: header_end, - }) -} - -/// Extract a single-quoted string value for `key` from a `.npy` header dict. -fn extract_quoted(header: &str, key: &str) -> Option { - let after_key = &header[header.find(&format!("'{key}'"))?..]; - let after_colon = &after_key[after_key.find(':')? + 1..]; - let bytes = after_colon.as_bytes(); - let mut i = 0; - while i < bytes.len() && bytes[i].is_ascii_whitespace() { - i += 1; - } - if i >= bytes.len() || (bytes[i] != b'\'' && bytes[i] != b'"') { - return None; - } - let quote = bytes[i]; - i += 1; - let start = i; - while i < bytes.len() && bytes[i] != quote { - i += 1; - } - Some(after_colon[start..i].to_string()) -} - -/// Extract the 2-D `(rows, cols)` shape tuple from a `.npy` header dict. -fn extract_shape(header: &str) -> Result<(usize, usize)> { - let after_key = &header[header - .find("'shape'") - .context(".npy header missing 'shape'")?..]; - let open = after_key.find('(').context("malformed 'shape'")?; - let close = after_key[open..].find(')').context("malformed 'shape'")? + open; - let dims: Vec = after_key[open + 1..close] - .split(',') - .map(str::trim) - .filter(|s| !s.is_empty()) - .map(|s| s.parse::()) - .collect::>() - .context("malformed 'shape' dims")?; - if dims.len() != 2 { - bail!("expected 2-D vectors.npy, got shape {dims:?}"); - } - Ok((dims[0], dims[1])) -} - #[cfg(test)] mod tests { use super::*; + use crate::dataset::fixtures::make_ramp_npy; - /// Build a minimal little-endian `.npy` v1.0 buffer. - fn make_npy(descr: &str, rows: usize, cols: usize, data: &[u8]) -> Vec { - let dict = - format!("{{'descr': '{descr}', 'fortran_order': False, 'shape': ({rows}, {cols}), }}"); - let mut header = dict.into_bytes(); - // Pad with spaces so the total (magic + version + len field + header) - // is a multiple of 64, and terminate with a newline. - let unpadded = 10 + header.len() + 1; - header.extend(std::iter::repeat_n(b' ', (64 - unpadded % 64) % 64)); - header.push(b'\n'); - - let mut out = Vec::new(); - out.extend_from_slice(b"\x93NUMPY"); - out.extend_from_slice(&[1, 0]); // version 1.0 - out.extend_from_slice(&(header.len() as u16).to_le_bytes()); - out.extend_from_slice(&header); - out.extend_from_slice(data); - out - } - - fn write_dataset(dir: &Path, npy: &[u8]) { - std::fs::write(dir.join("vectors.npy"), npy).unwrap(); + fn write_dataset(dir: &Path) { + std::fs::write(dir.join("vectors.npy"), make_ramp_npy(0, 2, 3)).unwrap(); } #[test] - fn reads_f32_vectors() { + fn reads_vectors() { let dir = tempfile::tempdir().unwrap(); - let values: Vec = (0..6).map(|x| x as f32).collect(); - let bytes: Vec = values.iter().flat_map(|v| v.to_le_bytes()).collect(); - write_dataset(dir.path(), &make_npy(" = (0..6).map(|x| x as f32).collect(); - let bytes: Vec = values.iter().flat_map(|v| v.to_le_bytes()).collect(); - write_dataset(dir.path(), &make_npy(" = (0..6).map(|x| x as f32).collect(); - let bytes: Vec = values.iter().flat_map(|v| v.to_le_bytes()).collect(); - write_dataset(dir.path(), &make_npy(" = values - .iter() - .flat_map(|&v| f16::from_f32(v).to_le_bytes()) - .collect(); - write_dataset(dir.path(), &make_npy(" Vec { let mut out = Vec::new(); for vector in &config.collection.vectors { if let crate::config::VectorSource::Dataset { dataset } = &vector.source { - out.push(dataset.clone()); + out.push((**dataset).clone()); } } for sparse in &config.collection.sparse_vectors { if let Some(dataset) = &sparse.source.dataset { - out.push(dataset.clone()); + out.push((**dataset).clone()); } } for field in &config.collection.fields { if let Some(dataset) = field.source.as_ref().and_then(|s| s.dataset.as_ref()) { - out.push(dataset.clone()); + out.push((**dataset).clone()); } } if let Some(dataset) = config @@ -32,7 +32,7 @@ pub fn collect_dataset_configs(config: &UploadConfig) -> Vec { .as_ref() .and_then(|s| s.dataset.as_ref()) { - out.push(dataset.clone()); + out.push((**dataset).clone()); } out } diff --git a/src/generators/config.rs b/src/generators/config.rs index 6e4ccda..c425cec 100644 --- a/src/generators/config.rs +++ b/src/generators/config.rs @@ -438,6 +438,67 @@ collection: assert_eq!(server.join().unwrap(), 1); } + /// The LAION shape: dense vectors from a bare `.npy`, the whole payload + /// object from a parquet of the same row count. Both are indexed by point + /// id, so row *i* of each file must land on point *i*. + #[test] + fn pairs_npy_vectors_with_parquet_payloads() { + use crate::dataset::fixtures::{make_ramp_npy, write_parquet}; + + let dir = tempfile::tempdir().unwrap(); + std::fs::write(dir.path().join("emb.npy"), make_ramp_npy(0, 8, 4)).unwrap(); + write_parquet(&dir.path().join("meta.parquet"), 0, 8, 4); + + let config: UploadConfig = serde_yaml::from_str( + " +collection: + name: t + vectors: + - size: 4 + source: + type: dataset + name: emb + format: npy + path: emb.npy + payload: + source: + type: dataset + dataset: + name: meta + format: parquet + path: meta.parquet + exclude: [url] + fields: + - name: similarity + type: float +", + ) + .unwrap(); + config.validate().unwrap(); + + let generator = ConfigGenerator::new_with_datasets_dir(&config, dir.path()).unwrap(); + + let point = generator.make_point(5); + let data = match point.vectors.unwrap().vectors_options.unwrap() { + VectorsOptions::Vector(v) => match v.vector.unwrap() { + qdrant_client::qdrant::vector::Vector::Dense(d) => d.data, + _ => panic!("expected a dense vector"), + }, + _ => panic!("expected the unnamed default vector"), + }; + assert_eq!(data, vec![20.0, 21.0, 22.0, 23.0], "row 5 of the .npy"); + + let payload: serde_json::Value = + serde_json::to_value(point.payload.into_iter().collect::>()).unwrap(); + assert_eq!(payload["id"], 5, "row 5 of the parquet"); + assert_eq!(payload["similarity"], 5.0); + assert_eq!(payload["caption"], "caption 5"); + assert!( + payload.get("url").is_none(), + "`exclude` must drop the column: {payload}" + ); + } + fn named(point: &PointStruct) -> HashMap { match point.vectors.clone().unwrap().vectors_options.unwrap() { VectorsOptions::Vectors(nv) => nv.vectors, From bcc75660d8306b15c31a0b2fd14ca96023f8d15a Mon Sep 17 00:00:00 2001 From: generall Date: Sun, 26 Jul 2026 14:36:38 +0200 Subject: [PATCH 2/3] feat: read numbered dataset parts as one row space MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Corpora too large to publish as one artifact ship as numbered parts, and there was no way to point a source at more than one file. Uploading LAION-400M meant 409 separate invocations, each of which would then have collided on point ids: a point's id doubles as its dataset row, so `--offset` shifted the row index too and reading part *k* at offset *k*×1M ran off the end of that part. Add a `parts:` block to dataset sources — `count`, optional `start`, and `{i}` templates for `path`/`link` — presenting the family as one contiguous row space. Point ids stay global across the whole corpus, and `--offset` becomes a resume switch rather than a hazard: it skips rows and ids together, and `-n` is now capped by the rows *remaining* instead of the corpus size (an offset past the end is an error rather than a silently empty run). Row counts per part are measured, never configured. A `rows_per_part` setting would have looked reasonable and been wrong: LAION's parts are 1,000,448 rows, except part 408 at 1,000,501 and part 409 at 518,720, so a fixed guess misaligns payloads against vectors near the end and silently drops ~700k points. 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 each part is sized with one ranged request and no download. `TailChunkReader` presents a fetched suffix to the parquet metadata reader as if the whole file were there. Sizing results are cached in `datasets/.parts-index/.json`, keyed on a hash of the parts spec so a changed spec re-measures. Rather than validating the cache with per-run ETag requests, each part's real row count is checked against the manifest when it is opened — free, since opening reads it anyway, and it catches the stale-sidecar case at the moment it would corrupt ids. A host that answers 200 to a `Range:` request is reported, not silently streamed: the failure mode this guards against is downloading 600 GB to compute 410 integers. Verified against the published LAION files: sizing 5 `img_emb_*.npy` and 5 `metadata_*.parquet` parts takes ~4s over the network, downloads nothing, and both families report an identical 5,002,240 rows. Co-Authored-By: Claude Opus 5 (1M context) --- README.md | 33 ++ examples/upload-laion-400m.yaml | 70 ++++ src/config/schema.rs | 13 + src/dataset/config.rs | 150 ++++++- src/dataset/download.rs | 78 +++- src/dataset/mod.rs | 103 +++++ src/dataset/parts.rs | 677 ++++++++++++++++++++++++++++++++ src/dataset/reader.rs | 33 +- src/dataset/readers/mod.rs | 6 +- src/dataset/readers/parquet.rs | 111 +++++- src/dataset/upload.rs | 88 ++++- src/main.rs | 1 + 12 files changed, 1337 insertions(+), 26 deletions(-) create mode 100644 examples/upload-laion-400m.yaml create mode 100644 src/dataset/parts.rs diff --git a/README.md b/README.md index db047d9..039447e 100644 --- a/README.md +++ b/README.md @@ -53,6 +53,39 @@ Parquet sources accept three extra keys: `columns` (keep only these), `exclude` 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. The result is cached in +`datasets/.parts-index/.json`, keyed on the parts spec, so later runs +issue no requests at all. (Assuming a uniform part size would be wrong in +practice: LAION's parts are 1,000,448 rows except part 408 at 1,000,501 and +part 409 at 518,720, so any fixed guess misaligns payloads against vectors near +the end 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. + 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 still supported for name-only shorthand. diff --git a/examples/upload-laion-400m.yaml b/examples/upload-laion-400m.yaml new file mode 100644 index 0000000..82a3701 --- /dev/null +++ b/examples/upload-laion-400m.yaml @@ -0,0 +1,70 @@ +# The full LAION-400M benchmark corpus +# (https://github.com/qdrant/laion-400m-benchmark) — ~409.7M 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 (a few seconds, +# no downloads) and caches the result in `datasets/.parts-index/`. Part row +# counts are *not* uniform — most are 1,000,448 rows, part 408 is 1,000,501 and +# part 409 is 518,720 — so they are measured rather than assumed. +# +# Resuming: point ids are dataset rows, so an interrupted run continues with +# `--offset `; `-n` is then capped by what is left. +# +# Disk: parts are downloaded as they are reached. Without `cache:` they all +# accumulate (~600 GB); see the `cache: evict` note at the bottom. + +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 + + 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 + # `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 diff --git a/src/config/schema.rs b/src/config/schema.rs index cd0aa12..82383b7 100644 --- a/src/config/schema.rs +++ b/src/config/schema.rs @@ -104,6 +104,19 @@ collection: # 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 # Sparse vectors (optional). Names must be unique across all vectors. sparse_vectors: diff --git a/src/dataset/config.rs b/src/dataset/config.rs index e8dd235..8b7bde2 100644 --- a/src/dataset/config.rs +++ b/src/dataset/config.rs @@ -22,6 +22,10 @@ pub struct DatasetConfig { pub path: Option, #[serde(default)] pub link: Option, + /// Sharded dataset: a numbered family of files read as one row space. + /// Mutually exclusive with `path` / `link`. + #[serde(default)] + pub parts: Option, #[serde(default)] pub vector_size: Option, #[serde(default)] @@ -50,13 +54,35 @@ impl DatasetConfig { } } +/// A numbered family of files making up one dataset. +/// +/// `path` and `link` are templates containing `{i}`, substituted with each +/// part's number. Part row counts are always measured rather than configured — +/// see [`crate::dataset::parts`] for why a "rows per part" setting would be +/// actively wrong. +#[derive(Debug, Clone, Default, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct PartsConfig { + /// Number of parts; indices run `start .. start + count`. + pub count: usize, + #[serde(default)] + pub start: usize, + /// Path template resolved against the datasets dir, e.g. `laion/img_emb_{i}.npy`. + pub path: String, + /// Download template, e.g. `https://host/img_emb_{i}.npy`. + #[serde(default)] + pub link: Option, +} + /// Fully-resolved dataset configuration. #[derive(Debug, Clone)] pub struct ResolvedDatasetConfig { pub name: String, pub kind: DatasetKind, + /// Path of the single dataset file/directory. Empty when `parts` is set. pub path: String, pub link: Option, + pub parts: Option, #[allow(dead_code)] pub vector_size: Option, #[allow(dead_code)] @@ -78,10 +104,16 @@ impl ResolvedDatasetConfig { .kind .or_else(|| base.and_then(|b| b.kind)) .with_context(|| format!("dataset {:?}: missing `format` ({KINDS})", inline.name))?; - let path = inline - .path - .or_else(|| base.and_then(|b| b.path.clone())) - .with_context(|| format!("dataset {:?}: missing `path`", inline.name))?; + let parts = inline.parts.or_else(|| base.and_then(|b| b.parts.clone())); + // A sharded dataset locates its files through the `parts` templates, so + // the single-file `path` is not required (and must not be set). + let path = match &parts { + Some(_) => String::new(), + None => inline + .path + .or_else(|| base.and_then(|b| b.path.clone())) + .with_context(|| format!("dataset {:?}: missing `path`", inline.name))?, + }; let link = inline.link.or_else(|| base.and_then(|b| b.link.clone())); // Whether the file actually exists is checked by `ensure_downloaded`, which // resolves `path` against the datasets dir; checking it here (relative to the @@ -91,6 +123,7 @@ impl ResolvedDatasetConfig { kind, path, link, + parts, vector_size: inline .vector_size .or_else(|| base.and_then(|b| b.vector_size)), @@ -155,11 +188,53 @@ impl DatasetConfig { if self.name.is_empty() { bail!("dataset source requires `name`"); } - if self.kind.is_none() { + let Some(kind) = self.kind else { bail!("dataset {:?} requires `format` ({KINDS})", self.name); + }; + + if let Some(parts) = &self.parts { + if self.path.is_some() || self.link.is_some() { + bail!( + "dataset {:?} sets both `parts` and `path`/`link`; \ + a sharded dataset locates its files through `parts.path` / `parts.link`", + self.name + ); + } + if !matches!(kind, DatasetKind::Npy | DatasetKind::Parquet) { + bail!( + "dataset {:?}: `parts` is only supported for `format: npy` or \ + `format: parquet`, not {kind:?}", + self.name + ); + } + if parts.count == 0 { + bail!( + "dataset {:?}: `parts.count` must be greater than 0", + self.name + ); + } + // Without the placeholder every part resolves to the same file, which + // would look like a working upload of `count` copies of part one. + if parts.count > 1 && !parts.path.contains("{i}") { + bail!( + "dataset {:?}: `parts.path` must contain `{{i}}` to distinguish parts", + self.name + ); + } + if let Some(link) = &parts.link + && parts.count > 1 + && !link.contains("{i}") + { + bail!( + "dataset {:?}: `parts.link` must contain `{{i}}` to distinguish parts", + self.name + ); + } + return Ok(()); } + if self.path.is_none() { - bail!("dataset {:?} requires `path`", self.name); + bail!("dataset {:?} requires `path` (or `parts`)", self.name); } Ok(()) } @@ -213,4 +288,67 @@ mod tests { assert_eq!(resolved.vector_size, Some(100)); assert!(resolved.link.is_some()); } + + fn parts_config(parts: PartsConfig, kind: DatasetKind) -> DatasetConfig { + DatasetConfig { + name: "sharded".to_string(), + kind: Some(kind), + parts: Some(parts), + ..Default::default() + } + } + + fn template(path: &str, count: usize) -> PartsConfig { + PartsConfig { + count, + start: 0, + path: path.to_string(), + link: None, + } + } + + #[test] + fn parts_resolve_without_a_single_file_path() { + let config = parts_config(template("laion/img_emb_{i}.npy", 410), DatasetKind::Npy); + config.validate_inline().unwrap(); + let resolved = DatasetConfig::resolve(config, &HashMap::new()).unwrap(); + assert_eq!(resolved.parts.unwrap().count, 410); + assert!(resolved.path.is_empty()); + } + + /// Without `{i}` every part resolves to the same file — which would look + /// like a successful upload of `count` copies of part one. + #[test] + fn parts_path_must_distinguish_parts() { + let config = parts_config(template("laion/img_emb.npy", 410), DatasetKind::Npy); + let err = config.validate_inline().unwrap_err().to_string(); + assert!(err.contains("{i}"), "{err}"); + + // A single part needs no placeholder. + parts_config(template("laion/img_emb.npy", 1), DatasetKind::Npy) + .validate_inline() + .unwrap(); + } + + #[test] + fn parts_and_path_are_mutually_exclusive() { + let mut config = parts_config(template("p_{i}.npy", 2), DatasetKind::Npy); + config.path = Some("p.npy".to_string()); + let err = config.validate_inline().unwrap_err().to_string(); + assert!(err.contains("both `parts` and `path`"), "{err}"); + } + + #[test] + fn parts_are_rejected_for_bundle_formats() { + let config = parts_config(template("d_{i}", 3), DatasetKind::Tar); + let err = config.validate_inline().unwrap_err().to_string(); + assert!(err.contains("only supported for"), "{err}"); + } + + #[test] + fn parts_count_must_be_positive() { + let config = parts_config(template("p_{i}.npy", 0), DatasetKind::Npy); + let err = config.validate_inline().unwrap_err().to_string(); + assert!(err.contains("greater than 0"), "{err}"); + } } diff --git a/src/dataset/download.rs b/src/dataset/download.rs index a559b74..d76b6b8 100644 --- a/src/dataset/download.rs +++ b/src/dataset/download.rs @@ -1,7 +1,7 @@ use std::collections::hash_map::DefaultHasher; use std::fs::{self, File}; use std::hash::{Hash, Hasher}; -use std::io; +use std::io::{self, Read}; use std::path::{Path, PathBuf}; use anyhow::{Context, Result, bail}; @@ -90,14 +90,80 @@ pub fn ensure_downloaded(datasets_dir: &Path, config: &ResolvedDatasetConfig) -> Ok(target) } -/// Download `url` into a temporary file inside `dir`. The file is deleted when the -/// returned handle is dropped, so an aborted download leaves nothing behind. -fn download_to_temp(url: &str, dir: &Path) -> Result { - let agent = ureq::Agent::new_with_config( +/// Download `url` to `target`, creating parent directories as needed. +/// +/// Staged next to the target and renamed into place, so an interrupted download +/// never leaves a truncated file that a later run would happily reuse. +pub fn download_file_to(url: &str, target: &Path) -> Result<()> { + let parent = target.parent().unwrap_or(Path::new(".")); + fs::create_dir_all(parent).with_context(|| format!("failed to create {}", parent.display()))?; + + println!("Downloading {url}..."); + let tmp = download_to_temp(url, parent)?; + tmp.persist(target) + .with_context(|| format!("failed to install download at {}", target.display()))?; + Ok(()) +} + +/// A byte range fetched from a remote file, plus the file's full length as +/// reported in `Content-Range`. +pub struct RangeResponse { + pub body: Vec, + pub total_len: u64, +} + +/// Fetch `range` (a `Range:` header value such as `bytes=0-511` or `bytes=-65536`) +/// from `url`. +/// +/// This is what lets a remote part be *sized* without being downloaded: both the +/// `.npy` header and the parquet footer live at a known end of the file. A server +/// that ignores the range is rejected rather than silently streaming gigabytes. +pub fn fetch_range(url: &str, range: &str) -> Result { + let response = agent() + .get(url) + .header("Range", range) + .call() + .with_context(|| format!("failed to fetch {range} of {url}"))?; + + if response.status() != 206 { + bail!( + "{url} answered {} to a `Range: {range}` request; \ + ranged requests are required to size dataset parts without downloading them", + response.status() + ); + } + + // `Content-Range: bytes -/` + let total_len = response + .headers() + .get("content-range") + .and_then(|value| value.to_str().ok()) + .and_then(|value| value.rsplit_once('/')) + .and_then(|(_, total)| total.trim().parse::().ok()) + .with_context(|| format!("{url} returned a Content-Range without a total length"))?; + + let mut body = Vec::new(); + response + .into_body() + .into_reader() + .read_to_end(&mut body) + .with_context(|| format!("failed to read {range} of {url}"))?; + + Ok(RangeResponse { body, total_len }) +} + +fn agent() -> ureq::Agent { + ureq::Agent::new_with_config( ureq::config::Config::builder() .user_agent("Mozilla/5.0") .build(), - ); + ) +} + +/// Download `url` into a temporary file inside `dir`. The file is deleted when the +/// returned handle is dropped, so an aborted download leaves nothing behind. +fn download_to_temp(url: &str, dir: &Path) -> Result { + let agent = agent(); let response = agent .get(url) .call() diff --git a/src/dataset/mod.rs b/src/dataset/mod.rs index fb962ed..24f5b7e 100644 --- a/src/dataset/mod.rs +++ b/src/dataset/mod.rs @@ -2,6 +2,7 @@ mod config; mod download; #[cfg(test)] pub(crate) mod fixtures; +mod parts; mod payload; mod reader; mod readers; @@ -58,4 +59,106 @@ pub(crate) mod test_http { (url, handle) } + + /// What a [`serve_ranges`] server actually did. + pub(crate) struct ServeStats { + pub requests: usize, + /// Total body bytes written back, so a test can assert that sizing a + /// dataset did *not* amount to downloading it. + pub bytes_served: usize, + } + + /// Serve `files` (by path, e.g. `"p_0.npy"`) over HTTP with `Range:` + /// support, for exactly `max_requests` requests. Returns the base URL and a + /// handle yielding the request/byte counts. + pub(crate) fn serve_ranges( + files: Vec<(String, Vec)>, + max_requests: usize, + ) -> (String, JoinHandle) { + let listener = TcpListener::bind("127.0.0.1:0").unwrap(); + let base = format!("http://{}", listener.local_addr().unwrap()); + + let handle = std::thread::spawn(move || { + let mut stats = ServeStats { + requests: 0, + bytes_served: 0, + }; + for _ in 0..max_requests { + let Ok((mut stream, _)) = listener.accept() else { + break; + }; + let mut buf = [0u8; 2048]; + let read = stream.read(&mut buf).unwrap_or(0); + let request = String::from_utf8_lossy(&buf[..read]).to_string(); + + let path = request + .lines() + .next() + .and_then(|line| line.split_whitespace().nth(1)) + .unwrap_or("/") + .trim_start_matches('/') + .to_string(); + let Some((_, body)) = files.iter().find(|(name, _)| *name == path) else { + let _ = stream.write_all( + b"HTTP/1.1 404 Not Found\r\nContent-Length: 0\r\nConnection: close\r\n\r\n", + ); + stats.requests += 1; + continue; + }; + + let range = request + .lines() + .find(|line| line.to_ascii_lowercase().starts_with("range:")) + .and_then(|line| { + line.split_once('=') + .map(|(_, spec)| spec.trim().to_string()) + }); + + let total = body.len(); + let (status, start, end) = match range.as_deref().map(parse_range) { + // `bytes=-N`: the last N bytes. + Some((None, Some(suffix))) => { + (206, total.saturating_sub(suffix), total.saturating_sub(1)) + } + // `bytes=A-B` (B optional, and clamped to the real end). + Some((Some(from), to)) => (206, from, to.unwrap_or(total - 1).min(total - 1)), + _ => (200, 0, total.saturating_sub(1)), + }; + let slice = &body[start..=end.min(total - 1)]; + + let mut header = format!( + "HTTP/1.1 {status} {}\r\nContent-Length: {}\r\nAccept-Ranges: bytes\r\n", + if status == 206 { + "Partial Content" + } else { + "OK" + }, + slice.len() + ); + if status == 206 { + header.push_str(&format!( + "Content-Range: bytes {start}-{}/{total}\r\n", + start + slice.len() - 1 + )); + } + header.push_str("Connection: close\r\n\r\n"); + + let _ = stream.write_all(header.as_bytes()); + let _ = stream.write_all(slice); + let _ = stream.flush(); + stats.requests += 1; + stats.bytes_served += slice.len(); + } + stats + }); + + (base, handle) + } + + /// Parse a `Range` header spec into `(start, end)`; a bare `-N` suffix + /// range comes back as `(None, Some(N))`. + fn parse_range(spec: &str) -> (Option, Option) { + let (from, to) = spec.split_once('-').unwrap_or((spec, "")); + (from.trim().parse().ok(), to.trim().parse().ok()) + } } diff --git a/src/dataset/parts.rs b/src/dataset/parts.rs new file mode 100644 index 0000000..d65cb94 --- /dev/null +++ b/src/dataset/parts.rs @@ -0,0 +1,677 @@ +//! Sharded datasets: a numbered family of files presented as one row space. +//! +//! Corpora too large to publish as a single artifact ship as parts — +//! LAION-400M is 410 `img_emb_{i}.npy` / `metadata_{i}.parquet` pairs. A +//! `parts:` block turns those into one logical dataset, so point ids stay +//! global and `--offset` resumes an interrupted run at the right row. +//! +//! # Sizing +//! +//! Mapping a global row to a part needs every part's row count up front, and +//! the counts are *not* uniform: LAION's parts are 1,000,448 rows except part +//! 408 (1,000,501) and the last (518,720). A configured "rows per part" would +//! therefore be wrong for the tail of the corpus, silently misaligning payloads +//! against vectors, so the counts are always measured instead. +//! +//! Measuring is cheap because both formats keep their shape at a known end of +//! the file: the `.npy` header is the first ~128 bytes, and the parquet footer +//! the last few KB. One ranged request per part sizes the whole corpus without +//! downloading any of it, and the result is cached in a sidecar so later runs +//! do no requests at all. + +use std::collections::VecDeque; +use std::collections::hash_map::DefaultHasher; +use std::hash::{Hash, Hasher}; +use std::path::{Path, PathBuf}; +use std::sync::atomic::{AtomicUsize, Ordering}; +use std::sync::{Mutex, RwLock}; + +use anyhow::{Context, Result, bail}; +use serde::{Deserialize, Serialize}; +use serde_json::Value; + +use super::config::{DatasetKind, PartsConfig, ResolvedDatasetConfig}; +use super::download::fetch_range; +use super::readers::{NpyReader, ParquetReader, parquet_row_count, parse_npy_header}; + +/// Parts kept open at once. Upload walks rows in order, so two is enough to +/// straddle a boundary; holding more would pin every part's decode buffers. +const OPEN_PARTS: usize = 2; + +/// Prefix of a `.npy` requested when sizing a remote part. Far more than the +/// ~128 bytes a 2-D header occupies, and still a single small request. +const NPY_HEADER_PROBE: usize = 4096; + +/// Tail of a parquet file requested when sizing a remote part. Comfortably +/// covers a typical footer (LAION's are ~9.5 KB) in one request; a larger +/// footer costs one more. +const PARQUET_TAIL_PROBE: usize = 64 * 1024; + +/// Concurrent sizing requests. Enough to hide per-request latency across a few +/// hundred parts without hammering the host. +const PROBE_CONCURRENCY: usize = 16; + +/// One part's measured size. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct PartEntry { + /// The part's own number, as substituted into the path/link templates. + pub index: usize, + pub rows: usize, + pub bytes: u64, +} + +/// Cached sizing result, invalidated by any change to the parts spec. +#[derive(Debug, Serialize, Deserialize)] +struct PartsManifest { + key: String, + parts: Vec, +} + +/// Everything needed to locate, fetch and open one part. +pub struct PartSource { + datasets_dir: PathBuf, + name: String, + kind: DatasetKind, + parts: PartsConfig, + columns: Option>, + exclude: Vec, + fill_null: Option, +} + +impl PartSource { + pub fn new(datasets_dir: &Path, config: &ResolvedDatasetConfig, parts: PartsConfig) -> Self { + PartSource { + datasets_dir: datasets_dir.to_path_buf(), + name: config.name.clone(), + kind: config.kind, + parts, + columns: config.columns.clone(), + exclude: config.exclude.clone(), + fill_null: config.fill_null.clone(), + } + } + + fn local_path(&self, index: usize) -> PathBuf { + self.datasets_dir.join(expand(&self.parts.path, index)) + } + + fn link(&self, index: usize) -> Option { + self.parts.link.as_deref().map(|tpl| expand(tpl, index)) + } + + /// Identity of the parts spec: a different spec must not reuse a manifest. + fn key(&self) -> String { + let mut hasher = DefaultHasher::new(); + format!("{:?}", self.kind).hash(&mut hasher); + self.parts.path.hash(&mut hasher); + self.parts.link.hash(&mut hasher); + self.parts.start.hash(&mut hasher); + self.parts.count.hash(&mut hasher); + format!("{:016x}", hasher.finish()) + } + + fn manifest_path(&self) -> PathBuf { + self.datasets_dir + .join(".parts-index") + .join(format!("{}.json", sanitize(&self.name))) + } + + fn indices(&self) -> impl Iterator + use<> { + let start = self.parts.start; + start..start + self.parts.count + } + + /// Ensure part `index` is present locally, downloading it if needed. + pub fn ensure_downloaded(&self, index: usize) -> Result { + let target = self.local_path(index); + if target.exists() { + return Ok(target); + } + let link = self.link(index).with_context(|| { + format!( + "dataset {:?} part {index} is missing at {} and no `parts.link` is configured", + self.name, + target.display() + ) + })?; + super::download::download_file_to(&link, &target)?; + Ok(target) + } + + fn open_reader(&self, path: &Path) -> Result { + Ok(match self.kind { + DatasetKind::Npy => PartReader::Npy(NpyReader::open(path)?), + DatasetKind::Parquet => PartReader::Parquet(Box::new(ParquetReader::open( + path, + self.columns.as_deref(), + &self.exclude, + self.fill_null.as_ref(), + )?)), + other => bail!("`parts:` is not supported for format {other:?} (use npy or parquet)"), + }) + } + + /// Measure one part's row count, without downloading it if it is remote. + fn probe(&self, index: usize) -> Result { + let local = self.local_path(index); + let (rows, bytes) = if local.exists() { + let bytes = std::fs::metadata(&local) + .with_context(|| format!("failed to stat {}", local.display()))? + .len(); + (measure_local(self.kind, &local)?, bytes) + } else { + let link = self.link(index).with_context(|| { + format!( + "dataset {:?} part {index} is missing at {} and no `parts.link` is configured", + self.name, + local.display() + ) + })?; + measure_remote(self.kind, &link)? + }; + Ok(PartEntry { index, rows, bytes }) + } + + /// Size every part, using the cached manifest when the spec is unchanged. + fn measure_all(&self) -> Result> { + let key = self.key(); + let manifest_path = self.manifest_path(); + if let Some(cached) = read_manifest(&manifest_path, &key) { + return Ok(cached); + } + + println!( + "Sizing {} parts of dataset {:?}...", + self.parts.count, self.name + ); + let indices: Vec = self.indices().collect(); + let results: Mutex>>> = + Mutex::new((0..indices.len()).map(|_| None).collect()); + + let next = AtomicUsize::new(0); + std::thread::scope(|scope| { + for _ in 0..PROBE_CONCURRENCY.min(indices.len()) { + scope.spawn(|| { + loop { + let slot = next.fetch_add(1, Ordering::Relaxed); + let Some(&index) = indices.get(slot) else { + break; + }; + let probed = self.probe(index); + results.lock().unwrap()[slot] = Some(probed); + } + }); + } + }); + + let entries = results + .into_inner() + .unwrap() + .into_iter() + .map(|slot| slot.expect("every slot is filled before the scope ends")) + .collect::>>()?; + + let total: usize = entries.iter().map(|e| e.rows).sum(); + println!( + "Dataset {:?}: {} parts, {total} rows total", + self.name, self.parts.count + ); + write_manifest(&manifest_path, &key, &entries); + Ok(entries) + } +} + +/// Substitute a part number into a `{i}` template. +fn expand(template: &str, index: usize) -> String { + template.replace("{i}", &index.to_string()) +} + +/// Keep a dataset name usable as a file name. +fn sanitize(name: &str) -> String { + name.chars() + .map(|c| { + if c.is_alphanumeric() || c == '-' || c == '_' { + c + } else { + '_' + } + }) + .collect() +} + +fn measure_local(kind: DatasetKind, path: &Path) -> Result { + match kind { + DatasetKind::Npy => { + use std::io::Read; + let mut file = std::fs::File::open(path) + .with_context(|| format!("failed to open {}", path.display()))?; + let mut head = vec![0u8; NPY_HEADER_PROBE]; + let read = file + .read(&mut head) + .with_context(|| format!("failed to read {}", path.display()))?; + head.truncate(read); + Ok(parse_npy_header(&head) + .with_context(|| format!("failed to parse {}", path.display()))? + .num_points) + } + DatasetKind::Parquet => parquet_row_count(path), + other => bail!("`parts:` is not supported for format {other:?} (use npy or parquet)"), + } +} + +fn measure_remote(kind: DatasetKind, link: &str) -> Result<(usize, u64)> { + match kind { + DatasetKind::Npy => { + let response = fetch_range(link, &format!("bytes=0-{}", NPY_HEADER_PROBE - 1))?; + let layout = parse_npy_header(&response.body) + .with_context(|| format!("failed to parse the .npy header of {link}"))?; + Ok((layout.num_points, response.total_len)) + } + DatasetKind::Parquet => { + let mut response = fetch_range(link, &format!("bytes=-{PARQUET_TAIL_PROBE}"))?; + // The footer length lives in the last 8 bytes; if the footer runs + // past what we fetched, ask for exactly as much as it needs. + if let Some(needed) = super::readers::parquet_footer_len(&response.body)? + && needed > response.body.len() + { + response = fetch_range(link, &format!("bytes=-{needed}"))?; + } + let rows = + super::readers::parquet_row_count_from_tail(&response.body, response.total_len) + .with_context(|| format!("failed to parse the parquet footer of {link}"))?; + Ok((rows, response.total_len)) + } + other => bail!("`parts:` is not supported for format {other:?} (use npy or parquet)"), + } +} + +fn read_manifest(path: &Path, key: &str) -> Option> { + let text = std::fs::read_to_string(path).ok()?; + let manifest: PartsManifest = serde_json::from_str(&text).ok()?; + (manifest.key == key).then_some(manifest.parts) +} + +/// Best-effort: an unwritable cache costs a re-probe next run, nothing more. +fn write_manifest(path: &Path, key: &str, parts: &[PartEntry]) { + let manifest = PartsManifest { + key: key.to_string(), + parts: parts.to_vec(), + }; + let Ok(text) = serde_json::to_string(&manifest) else { + return; + }; + if let Some(parent) = path.parent() { + let _ = std::fs::create_dir_all(parent); + } + let _ = std::fs::write(path, text); +} + +/// One open part. The parquet reader carries a decode ring, so it is boxed to +/// keep the enum from sizing every `npy` part by it too. +pub enum PartReader { + Npy(NpyReader), + Parquet(Box), +} + +impl PartReader { + fn num_points(&self) -> usize { + match self { + PartReader::Npy(r) => r.num_points(), + PartReader::Parquet(r) => r.num_points(), + } + } + + fn vector_at(&self, idx: usize) -> Result> { + match self { + PartReader::Npy(r) => r.vector_at(idx), + PartReader::Parquet(_) => bail!("parquet parts do not contain dense vectors"), + } + } + + fn payload_object(&self, idx: usize) -> Result> { + match self { + PartReader::Parquet(r) => r.payload_object(idx), + PartReader::Npy(_) => bail!("npy parts do not contain payloads"), + } + } + + fn payload_field(&self, idx: usize, field: &str) -> Result> { + match self { + PartReader::Parquet(r) => r.payload_field(idx, field), + PartReader::Npy(_) => bail!("npy parts do not contain payloads"), + } + } +} + +/// A family of parts addressed as one contiguous row space. +pub struct PartitionedReader { + source: PartSource, + entries: Vec, + /// Global index of each part's first row, plus a final total. + starts: Vec, + open: RwLock)>>, + /// Serializes the open-a-new-part path so concurrent readers crossing a + /// boundary fetch it once. Held *outside* `open`, so readers still working + /// on the previous part are never blocked behind a download. + opening: Mutex<()>, +} + +impl PartitionedReader { + pub fn open(datasets_dir: &Path, config: &ResolvedDatasetConfig) -> Result { + let parts = config + .parts + .clone() + .expect("PartitionedReader requires a `parts:` block"); + let source = PartSource::new(datasets_dir, config, parts); + let entries = source.measure_all()?; + + let mut starts = Vec::with_capacity(entries.len() + 1); + let mut running = 0usize; + for entry in &entries { + starts.push(running); + running += entry.rows; + } + starts.push(running); + + Ok(PartitionedReader { + source, + entries, + starts, + open: RwLock::new(VecDeque::new()), + opening: Mutex::new(()), + }) + } + + pub fn num_points(&self) -> usize { + *self.starts.last().unwrap_or(&0) + } + + /// Split a global row index into (slot in `entries`, row within the part). + fn locate(&self, idx: usize) -> Result<(usize, usize)> { + if idx >= self.num_points() { + bail!( + "row {idx} is past the end of dataset {:?} ({} rows)", + self.source.name, + self.num_points() + ); + } + let slot = match self.starts.binary_search(&idx) { + Ok(exact) => exact, + Err(next) => next - 1, + }; + Ok((slot, idx - self.starts[slot])) + } + + fn reader_for(&self, slot: usize) -> Result> { + let index = self.entries[slot].index; + if let Some((_, reader)) = self + .open + .read() + .unwrap() + .iter() + .find(|(open, _)| *open == index) + { + return Ok(reader.clone()); + } + + let _guard = self.opening.lock().unwrap(); + // Another thread may have opened it while we waited for the guard. + if let Some((_, reader)) = self + .open + .read() + .unwrap() + .iter() + .find(|(open, _)| *open == index) + { + return Ok(reader.clone()); + } + + let path = self.source.ensure_downloaded(index)?; + let reader = self.source.open_reader(&path)?; + + // A stale sidecar would misalign every row after this part; catching it + // here costs nothing, since opening already read the real count. + let expected = self.entries[slot].rows; + if reader.num_points() != expected { + bail!( + "{} holds {} rows but the cached parts index says {expected}; \ + delete {} and re-run to re-measure", + path.display(), + reader.num_points(), + self.source.manifest_path().display() + ); + } + + let reader = std::sync::Arc::new(reader); + let mut open = self.open.write().unwrap(); + open.push_back((index, reader.clone())); + while open.len() > OPEN_PARTS { + open.pop_front(); + } + Ok(reader) + } + + pub fn vector_at(&self, idx: usize) -> Result> { + let (slot, local) = self.locate(idx)?; + self.reader_for(slot)?.vector_at(local) + } + + pub fn payload_object(&self, idx: usize) -> Result> { + let (slot, local) = self.locate(idx)?; + self.reader_for(slot)?.payload_object(local) + } + + pub fn payload_field(&self, idx: usize, field: &str) -> Result> { + let (slot, local) = self.locate(idx)?; + self.reader_for(slot)?.payload_field(local, field) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::dataset::config::DatasetConfig; + use crate::dataset::fixtures::{make_ramp_npy, write_parquet}; + + fn resolved(dir: &Path, kind: DatasetKind, path: &str, count: usize) -> ResolvedDatasetConfig { + let _ = dir; + let config = DatasetConfig { + name: "test-parts".to_string(), + kind: Some(kind), + parts: Some(PartsConfig { + count, + start: 0, + path: path.to_string(), + link: None, + }), + ..Default::default() + }; + DatasetConfig::resolve(config, &Default::default()).unwrap() + } + + /// Parts of *different* sizes must still map global row -> (part, row) + /// correctly; this is exactly where a `rows_per_part` assumption breaks. + #[test] + fn maps_global_rows_across_uneven_parts() { + let dir = tempfile::tempdir().unwrap(); + // 4 + 7 + 2 rows: no uniform part size exists. + for (i, rows) in [4usize, 7, 2].iter().enumerate() { + let path = dir.path().join(format!("p_{i}.npy")); + std::fs::write(&path, make_ramp_npy(i * 100, *rows, 2)).unwrap(); + } + + let config = resolved(dir.path(), DatasetKind::Npy, "p_{i}.npy", 3); + let reader = PartitionedReader::open(dir.path(), &config).unwrap(); + assert_eq!(reader.num_points(), 13); + + // First row of each part carries that part's base value. + assert_eq!(reader.vector_at(0).unwrap(), vec![0.0, 1.0]); + assert_eq!(reader.vector_at(4).unwrap(), vec![100.0, 101.0]); + assert_eq!(reader.vector_at(11).unwrap(), vec![200.0, 201.0]); + // Last row overall. + assert_eq!(reader.vector_at(12).unwrap(), vec![202.0, 203.0]); + assert!(reader.vector_at(13).is_err(), "past the end must not wrap"); + } + + #[test] + fn reads_payload_rows_across_parquet_parts() { + let dir = tempfile::tempdir().unwrap(); + write_parquet(&dir.path().join("m_0.parquet"), 0, 5, 2); + write_parquet(&dir.path().join("m_1.parquet"), 1000, 3, 2); + + let config = resolved(dir.path(), DatasetKind::Parquet, "m_{i}.parquet", 2); + let reader = PartitionedReader::open(dir.path(), &config).unwrap(); + assert_eq!(reader.num_points(), 8); + + assert_eq!(reader.payload_object(1).unwrap().unwrap()["id"], 1); + assert_eq!(reader.payload_object(5).unwrap().unwrap()["id"], 1000); + assert_eq!(reader.payload_object(7).unwrap().unwrap()["id"], 1002); + } + + /// The sidecar must survive a re-open, and a spec change must invalidate it. + #[test] + fn caches_and_invalidates_the_parts_index() { + let dir = tempfile::tempdir().unwrap(); + for i in 0..2 { + std::fs::write( + dir.path().join(format!("p_{i}.npy")), + make_ramp_npy(0, 3, 2), + ) + .unwrap(); + } + + let two = resolved(dir.path(), DatasetKind::Npy, "p_{i}.npy", 2); + assert_eq!( + PartitionedReader::open(dir.path(), &two) + .unwrap() + .num_points(), + 6 + ); + let manifest = dir.path().join(".parts-index").join("test-parts.json"); + assert!(manifest.exists(), "sizing result must be cached"); + + // Re-opening with the same spec reuses it. + assert_eq!( + PartitionedReader::open(dir.path(), &two) + .unwrap() + .num_points(), + 6 + ); + + // A different part count is a different spec: the stale entry must not + // be reused, or the second part would go missing. + let one = resolved(dir.path(), DatasetKind::Npy, "p_{i}.npy", 1); + assert_eq!( + PartitionedReader::open(dir.path(), &one) + .unwrap() + .num_points(), + 3 + ); + } + + /// A part whose real size no longer matches the sidecar must be reported, + /// not silently used — every later row would land on the wrong point. + #[test] + fn rejects_a_part_that_no_longer_matches_the_cached_size() { + let dir = tempfile::tempdir().unwrap(); + std::fs::write(dir.path().join("p_0.npy"), make_ramp_npy(0, 6, 2)).unwrap(); + let config = resolved(dir.path(), DatasetKind::Npy, "p_{i}.npy", 1); + PartitionedReader::open(dir.path(), &config).unwrap(); + + // Same spec (so the sidecar is reused), fewer rows on disk. + std::fs::write(dir.path().join("p_0.npy"), make_ramp_npy(0, 3, 2)).unwrap(); + let reader = PartitionedReader::open(dir.path(), &config).unwrap(); + let err = reader.vector_at(0).unwrap_err().to_string(); + assert!(err.contains("parts index"), "{err}"); + } + + /// The point of measuring over `Range`: a remote corpus is sized without + /// being downloaded. Both formats must manage it in one small request each. + #[test] + fn sizes_remote_parts_without_downloading_them() { + let npy_a = make_ramp_npy(0, 900, 16); + let npy_b = make_ramp_npy(0, 350, 16); + let corpus_bytes = npy_a.len() + npy_b.len(); + let (base, server) = crate::dataset::test_http::serve_ranges( + vec![ + ("p_0.npy".to_string(), npy_a), + ("p_1.npy".to_string(), npy_b), + ], + 2, + ); + + let dir = tempfile::tempdir().unwrap(); + let config = DatasetConfig { + name: "remote-parts".to_string(), + kind: Some(DatasetKind::Npy), + parts: Some(PartsConfig { + count: 2, + start: 0, + path: "p_{i}.npy".to_string(), + link: Some(format!("{base}/p_{{i}}.npy")), + }), + ..Default::default() + }; + let config = DatasetConfig::resolve(config, &Default::default()).unwrap(); + + let reader = PartitionedReader::open(dir.path(), &config).unwrap(); + assert_eq!(reader.num_points(), 1250, "900 + 350 rows"); + + let stats = server.join().unwrap(); + assert_eq!(stats.requests, 2, "one ranged request per part"); + // The cost is a fixed header probe per part, not a function of how big + // the parts are — that is what makes sizing 410 LAION parts viable. + assert!( + stats.bytes_served <= 2 * NPY_HEADER_PROBE && stats.bytes_served < corpus_bytes, + "sizing served {} bytes of a {corpus_bytes}-byte corpus", + stats.bytes_served + ); + // Nothing was written to the datasets dir either. + assert!(!dir.path().join("p_0.npy").exists()); + } + + /// A remote parquet part is sized from its footer, which lives at the *end* + /// of the file — so this exercises the suffix-range path specifically. + #[test] + fn sizes_a_remote_parquet_part_from_its_footer() { + let dir = tempfile::tempdir().unwrap(); + let built = dir.path().join("built.parquet"); + write_parquet(&built, 0, 3000, 256); + let body = std::fs::read(&built).unwrap(); + std::fs::remove_file(&built).unwrap(); + let total = body.len(); + + let (base, server) = + crate::dataset::test_http::serve_ranges(vec![("m_0.parquet".to_string(), body)], 1); + + let config = DatasetConfig { + name: "remote-parquet".to_string(), + kind: Some(DatasetKind::Parquet), + parts: Some(PartsConfig { + count: 1, + start: 0, + path: "m_{i}.parquet".to_string(), + link: Some(format!("{base}/m_{{i}}.parquet")), + }), + ..Default::default() + }; + let config = DatasetConfig::resolve(config, &Default::default()).unwrap(); + + let reader = PartitionedReader::open(dir.path(), &config).unwrap(); + assert_eq!(reader.num_points(), 3000); + + let stats = server.join().unwrap(); + assert_eq!(stats.requests, 1); + assert!( + stats.bytes_served <= PARQUET_TAIL_PROBE && stats.bytes_served < total, + "the footer request served {} bytes of a {total}-byte file", + stats.bytes_served + ); + } + + #[test] + fn expands_templates() { + assert_eq!(expand("img_emb_{i}.npy", 0), "img_emb_0.npy"); + assert_eq!(expand("a/{i}/b_{i}.npy", 7), "a/7/b_7.npy"); + } +} diff --git a/src/dataset/reader.rs b/src/dataset/reader.rs index 32fbee6..491c6e3 100644 --- a/src/dataset/reader.rs +++ b/src/dataset/reader.rs @@ -5,6 +5,7 @@ use serde_json::Value; use super::config::{DatasetConfig, DatasetKind}; use super::download::ensure_downloaded; +use super::parts::PartitionedReader; use super::readers::{H5Reader, NpyReader, ParquetReader, SparseReader, TarReader}; use super::registry::load_registry; @@ -14,6 +15,9 @@ enum DatasetReaderInner { Sparse(SparseReader), Npy(NpyReader), Parquet(ParquetReader), + /// A `parts:` family read as one row space; the part format is `npy` or + /// `parquet`, so it answers the same accessors as those two. + Partitioned(PartitionedReader), } /// Random access to points from a vector-db-benchmark dataset. @@ -26,6 +30,16 @@ impl DatasetReader { pub fn open(datasets_dir: &Path, config: &DatasetConfig) -> Result { let registry = load_registry(datasets_dir)?; let config = DatasetConfig::resolve(config.clone(), ®istry)?; + + if config.parts.is_some() { + let reader = PartitionedReader::open(datasets_dir, &config)?; + let n = reader.num_points(); + return Ok(DatasetReader { + inner: DatasetReaderInner::Partitioned(reader), + num_points: n, + }); + } + let local_path = ensure_downloaded(datasets_dir, &config)?; let (inner, num_points) = match config.kind { DatasetKind::H5 => { @@ -67,6 +81,7 @@ impl DatasetReader { DatasetReaderInner::H5(r) => r.vector_at(idx), DatasetReaderInner::Tar(r) => r.vector_at(idx), DatasetReaderInner::Npy(r) => r.vector_at(idx), + DatasetReaderInner::Partitioned(r) => r.vector_at(idx), DatasetReaderInner::Sparse(_) | DatasetReaderInner::Parquet(_) => { bail!("dataset does not contain dense vectors") } @@ -84,6 +99,7 @@ impl DatasetReader { match &self.inner { DatasetReaderInner::Tar(r) => r.payload_field(idx, field), DatasetReaderInner::Parquet(r) => r.payload_field(idx, field), + DatasetReaderInner::Partitioned(r) => r.payload_field(idx, field), _ => bail!("dataset does not contain payloads"), } } @@ -92,6 +108,7 @@ impl DatasetReader { match &self.inner { DatasetReaderInner::Tar(r) => r.payload_object(idx), DatasetReaderInner::Parquet(r) => r.payload_object(idx), + DatasetReaderInner::Partitioned(r) => r.payload_object(idx), _ => bail!("dataset does not contain payloads"), } } @@ -104,7 +121,9 @@ impl DatasetReader { DatasetReaderInner::Sparse(r) => r.num_queries(), // Component formats hold corpus rows only; a query set is a // separate file, declared as its own source. - DatasetReaderInner::Npy(_) | DatasetReaderInner::Parquet(_) => 0, + DatasetReaderInner::Npy(_) + | DatasetReaderInner::Parquet(_) + | DatasetReaderInner::Partitioned(_) => 0, } } @@ -114,9 +133,9 @@ impl DatasetReader { DatasetReaderInner::H5(r) => r.query_at(idx), DatasetReaderInner::Tar(r) => r.query_at(idx), DatasetReaderInner::Sparse(_) => bail!("sparse dataset has no dense queries"), - DatasetReaderInner::Npy(_) | DatasetReaderInner::Parquet(_) => { - bail!("dataset has no query set") - } + DatasetReaderInner::Npy(_) + | DatasetReaderInner::Parquet(_) + | DatasetReaderInner::Partitioned(_) => bail!("dataset has no query set"), } } @@ -134,9 +153,9 @@ impl DatasetReader { DatasetReaderInner::H5(r) => r.neighbors_at(idx), DatasetReaderInner::Tar(r) => r.query_ground_truth(idx), DatasetReaderInner::Sparse(r) => r.query_ground_truth(idx), - DatasetReaderInner::Npy(_) | DatasetReaderInner::Parquet(_) => { - bail!("dataset has no ground truth") - } + DatasetReaderInner::Npy(_) + | DatasetReaderInner::Parquet(_) + | DatasetReaderInner::Partitioned(_) => bail!("dataset has no ground truth"), } } } diff --git a/src/dataset/readers/mod.rs b/src/dataset/readers/mod.rs index f2ce770..0f29139 100644 --- a/src/dataset/readers/mod.rs +++ b/src/dataset/readers/mod.rs @@ -7,7 +7,9 @@ mod sparse; mod tar; pub use h5::H5Reader; -pub use npy::NpyReader; -pub use parquet::ParquetReader; +pub use npy::{NpyReader, parse_npy_header}; +pub use parquet::{ + ParquetReader, parquet_footer_len, parquet_row_count, parquet_row_count_from_tail, +}; pub use sparse::SparseReader; pub use tar::TarReader; diff --git a/src/dataset/readers/parquet.rs b/src/dataset/readers/parquet.rs index bb0d1f8..78eed64 100644 --- a/src/dataset/readers/parquet.rs +++ b/src/dataset/readers/parquet.rs @@ -20,7 +20,10 @@ use std::path::{Path, PathBuf}; use std::sync::Mutex; use anyhow::{Context, Result, bail}; -use parquet::file::reader::{FileReader, SerializedFileReader}; +use bytes::Bytes; +use parquet::errors::ParquetError; +use parquet::file::metadata::ParquetMetaDataReader; +use parquet::file::reader::{ChunkReader, FileReader, Length, SerializedFileReader}; use parquet::file::serialized_reader::ReadOptionsBuilder; use parquet::record::reader::RowIter; use parquet::record::{Field, List, Map, Row}; @@ -256,6 +259,111 @@ impl ParquetReader { } } +/// Row count of a local parquet file, read from its footer. +pub fn parquet_row_count(path: &Path) -> Result { + let file = File::open(path).with_context(|| format!("failed to open {}", path.display()))?; + let reader = SerializedFileReader::new(file) + .with_context(|| format!("failed to read parquet metadata from {}", path.display()))?; + Ok(reader.metadata().file_metadata().num_rows().max(0) as usize) +} + +/// Length of the thrift footer described by the last 8 bytes of a parquet file, +/// *including* those 8 bytes — i.e. the smallest tail that holds the metadata. +/// +/// `None` when `tail` is too short to hold even the trailer. +pub fn parquet_footer_len(tail: &[u8]) -> Result> { + if tail.len() < FOOTER_TRAILER { + return Ok(None); + } + let trailer = &tail[tail.len() - FOOTER_TRAILER..]; + if &trailer[4..] != PARQUET_MAGIC { + bail!("not a parquet file (missing PAR1 trailer)"); + } + let len = u32::from_le_bytes(trailer[..4].try_into().unwrap()) as usize; + Ok(Some(len + FOOTER_TRAILER)) +} + +/// Row count read from the *tail* of a parquet file of `total_len` bytes. +/// +/// Parquet keeps its metadata at the end, so this sizes a remote part from one +/// ranged request instead of a download. `tail` must reach back at least +/// [`parquet_footer_len`] bytes from the end. +pub fn parquet_row_count_from_tail(tail: &[u8], total_len: u64) -> Result { + let needed = + parquet_footer_len(tail)?.context("tail is too short to hold a parquet trailer")?; + if needed > tail.len() { + bail!( + "parquet footer is {needed} bytes but only {} were fetched", + tail.len() + ); + } + if (tail.len() as u64) > total_len { + bail!( + "tail of {} bytes exceeds the file length {total_len}", + tail.len() + ); + } + + let reader = TailChunkReader { + total_len, + start: total_len - tail.len() as u64, + tail: Bytes::copy_from_slice(tail), + }; + let metadata = ParquetMetaDataReader::new() + .parse_and_finish(&reader) + .context("failed to decode the parquet footer")?; + Ok(metadata.file_metadata().num_rows().max(0) as usize) +} + +const FOOTER_TRAILER: usize = 8; +const PARQUET_MAGIC: &[u8] = b"PAR1"; + +/// Presents the last `tail.len()` bytes of a file as if the whole file were +/// available, so the parquet metadata reader — which only ever seeks from the +/// end — can work against a ranged response. +struct TailChunkReader { + total_len: u64, + start: u64, + tail: Bytes, +} + +impl TailChunkReader { + fn slice_from(&self, start: u64) -> parquet::errors::Result { + if start < self.start || start > self.total_len { + return Err(ParquetError::General(format!( + "parquet metadata read at {start} falls outside the fetched tail \ + [{}, {})", + self.start, self.total_len + ))); + } + Ok(self.tail.slice((start - self.start) as usize..)) + } +} + +impl Length for TailChunkReader { + fn len(&self) -> u64 { + self.total_len + } +} + +impl ChunkReader for TailChunkReader { + type T = bytes::buf::Reader; + + fn get_read(&self, start: u64) -> parquet::errors::Result { + self.slice_from(start).map(bytes::Buf::reader) + } + + fn get_bytes(&self, start: u64, length: usize) -> parquet::errors::Result { + let slice = self.slice_from(start)?; + if slice.len() < length { + return Err(ParquetError::General(format!( + "parquet metadata read of {length} bytes at {start} runs past the fetched tail" + ))); + } + Ok(slice.slice(..length)) + } +} + /// JSON has no NaN or infinity — those become "no value", same as null. fn number(value: f64) -> Option { Number::from_f64(value).map(Value::Number) @@ -456,4 +564,3 @@ mod tests { assert_eq!(open(&path).payload_object(10).unwrap(), None); } } - diff --git a/src/dataset/upload.rs b/src/dataset/upload.rs index c73c5c3..662b5a2 100644 --- a/src/dataset/upload.rs +++ b/src/dataset/upload.rs @@ -1,6 +1,6 @@ use std::path::Path; -use anyhow::Result; +use anyhow::{Result, bail}; use crate::config::UploadConfig; @@ -56,16 +56,98 @@ pub fn dataset_point_limit(config: &UploadConfig, datasets_dir: &Path) -> Result /// When dataset sources are present and `-n` is omitted, the full dataset (up /// to the smallest source) is uploaded. When `-n` is set, it is capped by that /// limit. Without dataset sources the legacy default of 100_000 applies. +/// +/// A point's id doubles as its dataset row, so `--offset` skips that many rows +/// too — which is what makes it a resume switch for a long upload. The rows +/// left to read are therefore counted from `offset`, not from zero. pub fn resolve_num_vectors( requested: Option, + offset: usize, config: &UploadConfig, datasets_dir: &Path, ) -> Result { let limit = dataset_point_limit(config, datasets_dir)?; + if let Some(limit) = limit + && offset >= limit + { + bail!( + "--offset {offset} starts past the end of the dataset ({limit} rows); \ + nothing would be uploaded" + ); + } Ok(match (requested, limit) { - (Some(n), Some(limit)) => n.min(limit), + (Some(n), Some(limit)) => n.min(limit - offset), (Some(n), None) => n, - (None, Some(limit)) => limit, + (None, Some(limit)) => limit - offset, (None, None) => 100_000, }) } + +#[cfg(test)] +mod tests { + use crate::config::UploadConfig; + use crate::dataset::fixtures::make_ramp_npy; + + fn dataset_config(dir: &std::path::Path) -> UploadConfig { + std::fs::write(dir.join("v.npy"), make_ramp_npy(0, 10, 4)).unwrap(); + serde_yaml::from_str( + " +collection: + vectors: + - size: 4 + source: { type: dataset, name: v, format: npy, path: v.npy } +", + ) + .unwrap() + } + + /// `--offset` skips dataset rows as well as ids, so the rows *remaining* + /// are what bound the run — otherwise resuming a part-way upload would read + /// past the end of the corpus. + #[test] + fn offset_reduces_the_rows_left_to_upload() { + let dir = tempfile::tempdir().unwrap(); + let config = dataset_config(dir.path()); + + let all = super::resolve_num_vectors(None, 0, &config, dir.path()).unwrap(); + assert_eq!(all, 10); + + let resumed = super::resolve_num_vectors(None, 6, &config, dir.path()).unwrap(); + assert_eq!(resumed, 4, "only rows 6..10 are left"); + + let capped = super::resolve_num_vectors(Some(100), 6, &config, dir.path()).unwrap(); + assert_eq!(capped, 4, "-n cannot exceed what remains"); + + let under = super::resolve_num_vectors(Some(2), 6, &config, dir.path()).unwrap(); + assert_eq!(under, 2); + } + + #[test] + fn offset_past_the_end_is_an_error_rather_than_an_empty_run() { + let dir = tempfile::tempdir().unwrap(); + let config = dataset_config(dir.path()); + + let err = super::resolve_num_vectors(None, 10, &config, dir.path()) + .unwrap_err() + .to_string(); + assert!(err.contains("past the end"), "{err}"); + } + + /// Without dataset sources there is nothing to bound, and `--offset` only + /// shifts ids — the legacy behaviour must be untouched. + #[test] + fn offset_does_not_bound_generated_data() { + let dir = tempfile::tempdir().unwrap(); + let config: UploadConfig = + serde_yaml::from_str("collection:\n vectors:\n - size: 4\n").unwrap(); + + assert_eq!( + super::resolve_num_vectors(Some(5), 1_000_000, &config, dir.path()).unwrap(), + 5 + ); + assert_eq!( + super::resolve_num_vectors(None, 1_000_000, &config, dir.path()).unwrap(), + 100_000 + ); + } +} diff --git a/src/main.rs b/src/main.rs index 9638084..c3889e6 100644 --- a/src/main.rs +++ b/src/main.rs @@ -80,6 +80,7 @@ async fn run_upload( let mut args = args; args.num_vectors = Some(dataset::resolve_num_vectors( args.num_vectors, + args.offset, &config, &dataset::default_datasets_dir(), )?); From 76ddb56e7c3efdab15b0f4d4c66e02e7f8a724bc Mon Sep 17 00:00:00 2001 From: generall Date: Sun, 26 Jul 2026 14:45:14 +0200 Subject: [PATCH 3/3] feat: stream sharded datasets larger than the disk MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit With `parts:`, a run downloads each part as it reaches it and keeps every one — so uploading LAION-400M needs ~600 GB of scratch space for a corpus that is only ever read once, front to back. The reference `upload.py` avoids this by fetching a part, uploading it and deleting it; bfb had no equivalent. Add `cache: keep | evict` to sharded sources. Under `evict`, a part is deleted once the reader has moved past it, holding peak disk to the few parts in flight (~4 GB for LAION rather than ~600 GB). Two rules keep that safe: * only parts bfb downloaded itself are ever deleted, tracked as they are fetched — a file staged in the datasets dir by hand is never touched, since bfb may have no way to get it back; * a part dropped from the LRU is deleted only once no reader still references it. The parquet reader reopens its file by path when a read rewinds, so unlinking under a live reader would break it; pending evictions hold a `Weak` and are swept when the last reference goes. Opening a part now also starts fetching the next one in the background, so crossing a part boundary does not stall the upload on a ~1 GB download. Each part gets its own download lock, so a prefetch of part n+1 neither duplicates nor blocks a reader that wants part n; whoever arrives second waits and finds the file already there. `cache: evict` is rejected on a non-sharded dataset rather than accepted as a no-op — it is a setting about deleting files, and silence would be the wrong default there. Sizing the real corpus turned up a sharper argument for measuring part rows than the one in the previous commit. LAION-400M has seven distinct part sizes, not two: 404 parts of 1,000,448 rows, one of 1,000,501, and five short parts — 8, 107, 220, 319 and 409 — holding 189,159 to 642,675 rows. The first short one is part *8*, so a plausible-looking fixed size would have misaligned payloads against vectors across ~98% of the corpus and overcounted it by 2.9M rows. The docs previously claimed only the tail was irregular; corrected here. All 820 parts size in ~1.5 minutes, and both families agree at 407,314,954 rows. Co-Authored-By: Claude Opus 5 (1M context) --- README.md | 38 ++++- examples/upload-laion-400m.yaml | 26 +++- src/config/schema.rs | 5 + src/dataset/config.rs | 45 ++++++ src/dataset/parts.rs | 253 +++++++++++++++++++++++++++++--- 5 files changed, 334 insertions(+), 33 deletions(-) diff --git a/README.md b/README.md index 039447e..02ef0f8 100644 --- a/README.md +++ b/README.md @@ -73,19 +73,43 @@ source: 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. The result is cached in -`datasets/.parts-index/.json`, keyed on the parts spec, so later runs -issue no requests at all. (Assuming a uniform part size would be wrong in -practice: LAION's parts are 1,000,448 rows except part 408 at 1,000,501 and -part 409 at 518,720, so any fixed guess misaligns payloads against vectors near -the end of the corpus.) The host must support ranged requests; one that answers -`200` to a `Range:` request is reported rather than silently downloaded. +and downloads none of them (820 LAION parts in ~1.5 minutes). The result is +cached in `datasets/.parts-index/.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 still supported for name-only shorthand. diff --git a/examples/upload-laion-400m.yaml b/examples/upload-laion-400m.yaml index 82a3701..bcfa5a5 100644 --- a/examples/upload-laion-400m.yaml +++ b/examples/upload-laion-400m.yaml @@ -1,20 +1,28 @@ # The full LAION-400M benchmark corpus -# (https://github.com/qdrant/laion-400m-benchmark) — ~409.7M points across 410 -# published parts, uploaded as one collection with global point ids. +# (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 (a few seconds, -# no downloads) and caches the result in `datasets/.parts-index/`. Part row -# counts are *not* uniform — most are 1,000,448 rows, part 408 is 1,000,501 and -# part 409 is 518,720 — so they are measured rather than assumed. +# 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 `; `-n` is then capped by what is left. # -# Disk: parts are downloaded as they are reached. Without `cache:` they all -# accumulate (~600 GB); see the `cache: evict` note at the bottom. +# 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 @@ -46,6 +54,7 @@ collection: 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: @@ -57,6 +66,7 @@ collection: 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] diff --git a/src/config/schema.rs b/src/config/schema.rs index 82383b7..b48a9e2 100644 --- a/src/config/schema.rs +++ b/src/config/schema.rs @@ -117,6 +117,11 @@ collection: # 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: diff --git a/src/dataset/config.rs b/src/dataset/config.rs index 8b7bde2..5b91c45 100644 --- a/src/dataset/config.rs +++ b/src/dataset/config.rs @@ -42,6 +42,21 @@ pub struct DatasetConfig { /// Omitted by default, which leaves the payload field absent. #[serde(default)] pub fill_null: Option, + /// What to do with downloaded parts once the upload has moved past them. + #[serde(default)] + pub cache: CacheMode, +} + +/// Disk policy for a sharded dataset's downloaded parts. +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "lowercase")] +pub enum CacheMode { + /// Keep every part that gets downloaded (default). + #[default] + Keep, + /// Delete each part once the reader has moved past it, bounding disk use to + /// the few parts in flight. Only ever deletes files bfb downloaded itself. + Evict, } impl DatasetConfig { @@ -92,6 +107,7 @@ pub struct ResolvedDatasetConfig { pub columns: Option>, pub exclude: Vec, pub fill_null: Option, + pub cache: CacheMode, } impl ResolvedDatasetConfig { @@ -144,6 +160,10 @@ impl ResolvedDatasetConfig { fill_null: inline .fill_null .or_else(|| base.and_then(|b| b.fill_null.clone())), + cache: match inline.cache { + CacheMode::Keep => base.map(|b| b.cache).unwrap_or_default(), + explicit => explicit, + }, }) } } @@ -233,6 +253,12 @@ impl DatasetConfig { return Ok(()); } + if self.cache == CacheMode::Evict { + bail!( + "dataset {:?}: `cache: evict` only applies to a sharded (`parts:`) dataset", + self.name + ); + } if self.path.is_none() { bail!("dataset {:?} requires `path` (or `parts`)", self.name); } @@ -345,6 +371,25 @@ mod tests { assert!(err.contains("only supported for"), "{err}"); } + /// `cache: evict` has nothing to act on without `parts`, so accepting it + /// there would be a silent no-op on a setting about deleting files. + #[test] + fn evict_requires_a_sharded_dataset() { + let config = DatasetConfig { + name: "single".to_string(), + kind: Some(DatasetKind::Npy), + path: Some("v.npy".to_string()), + cache: CacheMode::Evict, + ..Default::default() + }; + let err = config.validate_inline().unwrap_err().to_string(); + assert!(err.contains("only applies to a sharded"), "{err}"); + + let mut sharded = parts_config(template("p_{i}.npy", 3), DatasetKind::Npy); + sharded.cache = CacheMode::Evict; + sharded.validate_inline().unwrap(); + } + #[test] fn parts_count_must_be_positive() { let config = parts_config(template("p_{i}.npy", 0), DatasetKind::Npy); diff --git a/src/dataset/parts.rs b/src/dataset/parts.rs index d65cb94..ceddb36 100644 --- a/src/dataset/parts.rs +++ b/src/dataset/parts.rs @@ -8,29 +8,34 @@ //! # Sizing //! //! Mapping a global row to a part needs every part's row count up front, and -//! the counts are *not* uniform: LAION's parts are 1,000,448 rows except part -//! 408 (1,000,501) and the last (518,720). A configured "rows per part" would -//! therefore be wrong for the tail of the corpus, silently misaligning payloads -//! against vectors, so the counts are always measured instead. +//! the counts are *not* uniform. LAION-400M has 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, holding between 189,159 and 642,675 rows. +//! +//! That is why a configured "rows per part" is not offered. It would look +//! reasonable and be wrong from part 8 onward: every later part would be +//! addressed at the wrong offset, silently pairing payloads with the wrong +//! vectors across ~98% of the corpus, and overcounting the total by 2.9M rows. +//! The counts are therefore always measured. //! //! Measuring is cheap because both formats keep their shape at a known end of //! the file: the `.npy` header is the first ~128 bytes, and the parquet footer //! the last few KB. One ranged request per part sizes the whole corpus without -//! downloading any of it, and the result is cached in a sidecar so later runs -//! do no requests at all. +//! downloading any of it — 820 parts of LAION in ~1.5 minutes — and the result +//! is cached in a sidecar so later runs do no requests at all. -use std::collections::VecDeque; use std::collections::hash_map::DefaultHasher; +use std::collections::{HashMap, HashSet, VecDeque}; use std::hash::{Hash, Hasher}; use std::path::{Path, PathBuf}; use std::sync::atomic::{AtomicUsize, Ordering}; -use std::sync::{Mutex, RwLock}; +use std::sync::{Arc, Mutex, RwLock, Weak}; use anyhow::{Context, Result, bail}; use serde::{Deserialize, Serialize}; use serde_json::Value; -use super::config::{DatasetKind, PartsConfig, ResolvedDatasetConfig}; +use super::config::{CacheMode, DatasetKind, PartsConfig, ResolvedDatasetConfig}; use super::download::fetch_range; use super::readers::{NpyReader, ParquetReader, parquet_row_count, parse_npy_header}; @@ -76,6 +81,15 @@ pub struct PartSource { columns: Option>, exclude: Vec, fill_null: Option, + cache: CacheMode, + /// One lock per part, so fetching part *n+1* in the background never blocks + /// a reader that wants part *n*. + guards: Mutex>>>, + /// Parts bfb downloaded itself. Eviction consults this and nothing else, so + /// a file the user placed in the datasets dir is never deleted. + downloaded: Mutex>, + /// Parts with a prefetch in flight, so one is not started twice. + prefetching: Mutex>, } impl PartSource { @@ -88,6 +102,10 @@ impl PartSource { columns: config.columns.clone(), exclude: config.exclude.clone(), fill_null: config.fill_null.clone(), + cache: config.cache, + guards: Mutex::new(HashMap::new()), + downloaded: Mutex::new(HashSet::new()), + prefetching: Mutex::new(HashSet::new()), } } @@ -122,11 +140,23 @@ impl PartSource { } /// Ensure part `index` is present locally, downloading it if needed. + /// + /// Guarded per part, so a prefetch already fetching this part is waited on + /// rather than duplicated — and a prefetch of a *different* part does not + /// hold anyone up. pub fn ensure_downloaded(&self, index: usize) -> Result { let target = self.local_path(index); if target.exists() { return Ok(target); } + + let guard = self.guard_for(index); + let _held = guard.lock().unwrap(); + // Whoever held the guard may have just finished fetching it. + if target.exists() { + return Ok(target); + } + let link = self.link(index).with_context(|| { format!( "dataset {:?} part {index} is missing at {} and no `parts.link` is configured", @@ -135,9 +165,62 @@ impl PartSource { ) })?; super::download::download_file_to(&link, &target)?; + self.downloaded.lock().unwrap().insert(index); Ok(target) } + fn guard_for(&self, index: usize) -> Arc> { + self.guards + .lock() + .unwrap() + .entry(index) + .or_default() + .clone() + } + + fn is_valid(&self, index: usize) -> bool { + (self.parts.start..self.parts.start + self.parts.count).contains(&index) + } + + /// Start fetching part `index` in the background, if it is worth doing: + /// upload spends minutes on a part, which is ample time to have the next + /// one on disk before it is reached. + fn prefetch(self: &Arc, index: usize) { + if !self.is_valid(index) || self.link(index).is_none() || self.local_path(index).exists() { + return; + } + if !self.prefetching.lock().unwrap().insert(index) { + return; + } + + let source = Arc::clone(self); + std::thread::spawn(move || { + if let Err(e) = source.ensure_downloaded(index) { + // Not fatal: the reader will retry (and report) when it gets there. + tracing::warn!("prefetch of part {index} failed: {e:#}"); + } + source.prefetching.lock().unwrap().remove(&index); + }); + } + + /// Delete a part bfb downloaded, once nothing is reading it. + /// + /// Files that were already in the datasets dir are left alone — the whole + /// point of tracking `downloaded` is that eviction can never destroy data + /// bfb cannot fetch again. + fn evict(&self, index: usize) { + if self.cache != CacheMode::Evict { + return; + } + if !self.downloaded.lock().unwrap().remove(&index) { + return; + } + let path = self.local_path(index); + if let Err(e) = std::fs::remove_file(&path) { + tracing::warn!("failed to evict {}: {e}", path.display()); + } + } + fn open_reader(&self, path: &Path) -> Result { Ok(match self.kind { DatasetKind::Npy => PartReader::Npy(NpyReader::open(path)?), @@ -345,15 +428,19 @@ impl PartReader { /// A family of parts addressed as one contiguous row space. pub struct PartitionedReader { - source: PartSource, + source: Arc, entries: Vec, /// Global index of each part's first row, plus a final total. starts: Vec, - open: RwLock)>>, + open: RwLock)>>, /// Serializes the open-a-new-part path so concurrent readers crossing a /// boundary fetch it once. Held *outside* `open`, so readers still working /// on the previous part are never blocked behind a download. opening: Mutex<()>, + /// Parts dropped from the LRU and awaiting deletion. A reader may still + /// hold the `Arc` — the parquet reader reopens its file by path on a + /// rewind — so the file is only removed once every reference is gone. + pending_evict: Mutex)>>, } impl PartitionedReader { @@ -374,11 +461,12 @@ impl PartitionedReader { starts.push(running); Ok(PartitionedReader { - source, + source: Arc::new(source), entries, starts, open: RwLock::new(VecDeque::new()), opening: Mutex::new(()), + pending_evict: Mutex::new(Vec::new()), }) } @@ -402,7 +490,7 @@ impl PartitionedReader { Ok((slot, idx - self.starts[slot])) } - fn reader_for(&self, slot: usize) -> Result> { + fn reader_for(&self, slot: usize) -> Result> { let index = self.entries[slot].index; if let Some((_, reader)) = self .open @@ -442,15 +530,46 @@ impl PartitionedReader { ); } - let reader = std::sync::Arc::new(reader); - let mut open = self.open.write().unwrap(); - open.push_back((index, reader.clone())); - while open.len() > OPEN_PARTS { - open.pop_front(); + let reader = Arc::new(reader); + { + let mut open = self.open.write().unwrap(); + open.push_back((index, reader.clone())); + let mut pending = self.pending_evict.lock().unwrap(); + while open.len() > OPEN_PARTS { + if let Some((dropped, evicted)) = open.pop_front() { + pending.push((dropped, Arc::downgrade(&evicted))); + } + } } + + // Fetch the next part while this one is being uploaded, and reclaim the + // disk of any part nothing is reading any more. + self.source.prefetch(self.next_part_index(slot)); + self.sweep_evictions(); Ok(reader) } + /// Part number following the one at `slot`, or a value outside the range + /// when this is the last part (`prefetch` ignores it). + fn next_part_index(&self, slot: usize) -> usize { + self.entries + .get(slot + 1) + .map(|entry| entry.index) + .unwrap_or(usize::MAX) + } + + /// Delete the files of parts nothing holds a reader for any more. + fn sweep_evictions(&self) { + let mut pending = self.pending_evict.lock().unwrap(); + pending.retain(|(index, reader)| { + if reader.strong_count() > 0 { + return true; // still in use; try again next time round + } + self.source.evict(*index); + false + }); + } + pub fn vector_at(&self, idx: usize) -> Result> { let (slot, local) = self.locate(idx)?; self.reader_for(slot)?.vector_at(local) @@ -669,6 +788,104 @@ mod tests { ); } + /// Serve `count` npy parts and return the config that reads them. + /// The server thread is left running; tests here assert on the filesystem. + fn remote_parts(dir: &Path, count: usize, cache: CacheMode) -> ResolvedDatasetConfig { + let files: Vec<(String, Vec)> = (0..count) + .map(|i| (format!("p_{i}.npy"), make_ramp_npy(i * 100, 4, 2))) + .collect(); + let (base, _server) = crate::dataset::test_http::serve_ranges(files, 1000); + let _ = dir; + + let config = DatasetConfig { + name: format!("evictable-{count}-{cache:?}"), + kind: Some(DatasetKind::Npy), + parts: Some(PartsConfig { + count, + start: 0, + path: "p_{i}.npy".to_string(), + link: Some(format!("{base}/p_{{i}}.npy")), + }), + cache, + ..Default::default() + }; + DatasetConfig::resolve(config, &Default::default()).unwrap() + } + + /// With `cache: evict`, a part downloaded by bfb is deleted once the reader + /// has moved past it — this is what keeps a 600 GB corpus off a 1 TB disk. + #[test] + fn evicts_downloaded_parts_once_the_reader_moves_on() { + let dir = tempfile::tempdir().unwrap(); + let config = remote_parts(dir.path(), 4, CacheMode::Evict); + let reader = PartitionedReader::open(dir.path(), &config).unwrap(); + + // Parts are 4 rows each; walk into parts 0, 1, then 2. + reader.vector_at(0).unwrap(); + assert!(dir.path().join("p_0.npy").exists()); + reader.vector_at(4).unwrap(); + assert!(dir.path().join("p_0.npy").exists(), "still within the LRU"); + + reader.vector_at(8).unwrap(); + assert!( + !dir.path().join("p_0.npy").exists(), + "part 0 fell out of the LRU and should have been reclaimed" + ); + assert!(dir.path().join("p_2.npy").exists(), "the live part stays"); + } + + #[test] + fn keeps_parts_by_default() { + let dir = tempfile::tempdir().unwrap(); + let config = remote_parts(dir.path(), 4, CacheMode::Keep); + let reader = PartitionedReader::open(dir.path(), &config).unwrap(); + + for row in [0, 4, 8] { + reader.vector_at(row).unwrap(); + } + assert!(dir.path().join("p_0.npy").exists(), "default keeps parts"); + } + + /// Eviction must never delete a file the user put there: bfb cannot get it + /// back, and `parts.link` may not even point at it. + #[test] + fn never_evicts_a_file_it_did_not_download() { + let dir = tempfile::tempdir().unwrap(); + let config = remote_parts(dir.path(), 4, CacheMode::Evict); + // Pre-place part 0 as if the user had staged it themselves. + std::fs::write(dir.path().join("p_0.npy"), make_ramp_npy(0, 4, 2)).unwrap(); + + let reader = PartitionedReader::open(dir.path(), &config).unwrap(); + for row in [0, 4, 8] { + reader.vector_at(row).unwrap(); + } + assert!( + dir.path().join("p_0.npy").exists(), + "a user-supplied part must survive eviction" + ); + } + + /// Opening a part starts fetching the next one, so the upload does not + /// stall on a download every time it crosses a boundary. + #[test] + fn prefetches_the_next_part() { + let dir = tempfile::tempdir().unwrap(); + let config = remote_parts(dir.path(), 3, CacheMode::Keep); + let reader = PartitionedReader::open(dir.path(), &config).unwrap(); + + reader.vector_at(0).unwrap(); + + let next = dir.path().join("p_1.npy"); + let deadline = std::time::Instant::now() + std::time::Duration::from_secs(10); + while !next.exists() && std::time::Instant::now() < deadline { + std::thread::sleep(std::time::Duration::from_millis(10)); + } + assert!( + next.exists(), + "reading part 0 should have prefetched part 1" + ); + } + #[test] fn expands_templates() { assert_eq!(expand("img_emb_{i}.npy", 0), "img_emb_0.npy");