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..02ef0f8 100644 --- a/README.md +++ b/README.md @@ -24,9 +24,91 @@ each field is generated). The runtime flags (`-n`, `-b`, `-p`, `-t`, `--uri`, Upload configs can source dense vectors, sparse vectors, and payloads from inline dataset definitions (same fields as [vector-db-benchmark `datasets.json`](https://github.com/qdrant/vector-db-benchmark/blob/master/datasets/datasets.json)). -Supported formats are `h5` (HDF5, pure-Rust reader — no system libraries), `tar` -(`.tgz` with `vectors.npy` + optional `payloads.jsonl`), and `sparse` (CSR -matrices). + +| `format` | Contents | +|----------|----------| +| `h5` | ann-benchmarks HDF5 bundle — `train`, optional `test`/`neighbors`. Pure-Rust reader, no system libraries | +| `tar` | `.tgz` of `vectors.npy` + optional `payloads.jsonl` / `tests.jsonl` | +| `sparse` | CSR matrices (`data.csr`, optional `queries.csr` / `results.gt`) | +| `npy` | One 2-D float `.npy` — dense vectors only | +| `parquet` | One parquet file — payload rows only | + +The first three are *bundles*: vectors, payloads, and queries all come out of a +single artifact. `npy` and `parquet` are *components*, so a config pairs them — +one source per slot, row *i* of each landing on point *i*: + +```yaml +collection: + vectors: + - size: 512 + source: { type: dataset, name: emb, format: npy, path: emb.npy } + payload: + source: + type: dataset + dataset: { name: meta, format: parquet, path: meta.parquet, exclude: [exif] } +``` + +Parquet sources accept three extra keys: `columns` (keep only these), `exclude` +(drop these), and `fill_null` (a value substituted for nulls and for NaN/±inf +floats, which have no JSON form — by default such fields are simply absent). +See [`examples/upload-laion-part.yaml`](examples/upload-laion-part.yaml). + +#### Sharded datasets + +Corpora published as numbered parts are read as one row space with a `parts:` +block, so point ids stay global across the whole set. `npy` and `parquet` +sources support it; `{i}` is substituted with each part's number: + +```yaml +source: + type: dataset + name: laion-400m-img-emb + format: npy + parts: + count: 410 # parts 0..409; `start:` moves the first index + path: laion/img_emb_{i}.npy + link: https://deploy.laion.ai/.../img_emb_{i}.npy +``` + +Part row counts are **measured, never configured**. Both formats keep their +shape at a known end of the file — the `.npy` header at the front, the parquet +footer at the back — so bfb sizes every part with one ranged HTTP request each +and downloads none of them (820 LAION parts in ~1.5 minutes). The result is +cached in `datasets/.parts-index/.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 diff --git a/examples/upload-laion-400m.yaml b/examples/upload-laion-400m.yaml new file mode 100644 index 0000000..bcfa5a5 --- /dev/null +++ b/examples/upload-laion-400m.yaml @@ -0,0 +1,80 @@ +# The full LAION-400M benchmark corpus +# (https://github.com/qdrant/laion-400m-benchmark) — 407,314,954 points across +# 410 published parts, uploaded as one collection with global point ids. +# +# bfb upload --file examples/upload-laion-400m.yaml -b 256 -p 16 -t 8 \ +# --uri http://localhost:6334 +# +# On first use bfb sizes every part with one ranged request each (~1.5 minutes +# for all 820, no downloads) and caches the result in `datasets/.parts-index/`. +# Part row counts are *not* uniform: 404 parts hold 1,000,448 rows, one holds +# 1,000,501, and parts 8, 107, 220, 319 and 409 are short (189,159 to 642,675). +# They are measured rather than assumed — a fixed size would go wrong at part 8 +# and misalign payloads against vectors for the rest of the corpus. +# +# Resuming: point ids are dataset rows, so an interrupted run continues with +# `--offset `; `-n` is then capped by what is left. +# +# Disk: `cache: evict` streams the corpus — each part is fetched as it is +# reached, the next one is fetched in the background while the current one +# uploads, and parts already passed are deleted. Peak usage is a few parts +# (~4 GB) rather than the ~600 GB the full corpus would occupy. Drop +# `cache: evict` to keep every part for repeat runs. +# +# Eviction only ever deletes parts bfb downloaded; a file you staged in the +# datasets dir yourself is left alone. + +collection: + name: laion + on_disk_payload: true + + quantization: + type: binary + always_ram: true + + hnsw: + m: 6 + on_disk: false + + optimizers: + default_segment_number: 2 + # Bigger segments search faster, at the cost of slower indexing. + max_segment_size: 5000000 + + vectors: + - size: 512 + distance: cosine + datatype: float16 + on_disk: true + source: + type: dataset + name: laion-400m-img-emb + format: npy + parts: + count: 410 # parts 0..409 inclusive + path: laion/img_emb_{i}.npy + link: https://deploy.laion.ai/8f83b608504d46bb81708ec86e912220/embeddings/img_emb/img_emb_{i}.npy + cache: evict + + payload: + source: + type: dataset + dataset: + name: laion-400m-metadata + format: parquet + parts: + count: 410 + path: laion/metadata_{i}.parquet + link: https://deploy.laion.ai/8f83b608504d46bb81708ec86e912220/embeddings/metadata/metadata_{i}.parquet + cache: evict + # `exif` is the bulk of the metadata and useless as a filter; excluding + # it also skips decoding the column. The reference `upload.py` drops it. + exclude: [exif] + # `upload.py` does `df.fillna(0)`. Drop this line to leave fields with + # no value absent instead, which filters more predictably. + fill_null: 0 + + fields: + # Index-only declaration: the value comes from `payload.source` above. + - name: similarity + type: float 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..b48a9e2 100644 --- a/src/config/schema.rs +++ b/src/config/schema.rs @@ -94,11 +94,34 @@ collection: # source: # type: dataset # inline dataset definition (vector-db-benchmark format) # name: glove-25-angular - # format: h5 # h5 | tar | sparse (`type` alias accepted in nested `dataset:` maps) + # format: h5 # dataset format (`type` alias accepted in nested `dataset:` maps): + # # h5 ann-benchmarks bundle (train/test/neighbors) + # # tar .tgz of vectors.npy + payloads.jsonl + tests.jsonl + # # sparse CSR matrices + # # npy one 2-D float .npy — dense vectors only + # # parquet one parquet file — payload rows only # path: glove-25-angular/glove-25-angular.hdf5 # link: http://ann-benchmarks.com/glove-25-angular.hdf5 # vector_size: 25 # distance: cosine + # A sharded dataset (`npy` / `parquet` only) replaces `path`/`link` with a + # `parts` block; the files are read as one row space and `{i}` is + # substituted with each part's number. Row counts per part are measured, + # not configured — one ranged request per part, cached thereafter. + # source: + # type: dataset + # name: laion-400m-img-emb + # format: npy + # parts: + # count: 410 # uint required number of parts + # start: 0 # uint default=0 index of the first part + # path: laion/img_emb_{i}.npy # string required + # link: https://host/img_emb_{i}.npy # string optional + # cache: keep # enum default=keep [keep | evict] (sharded only) + # # evict deletes each downloaded part once the reader + # # moves past it, and prefetches the next one, so a + # # corpus larger than the disk can still be streamed. + # # Only parts bfb downloaded are ever deleted. # Sparse vectors (optional). Names must be unique across all vectors. sparse_vectors: @@ -136,6 +159,13 @@ collection: # format: tar # path: laion-small-clip/laion-small-clip # link: https://example.com/laion-small-clip.tgz + # `format: parquet` reads payload rows from a parquet file, and accepts + # three extra keys (ignored by every other format): + # columns: [url, similarity] # list optional columns to keep (default: all) + # exclude: [exif] # list default=[] columns to drop (applied after `columns`) + # fill_null: 0 # any optional value substituted for nulls and for + # # NaN/±inf floats, which have no JSON form. Omitted by + # # default, leaving the payload field absent. # Payload field declarations (optional). Names must be unique. Each entry # generates a value and/or declares a field index. 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..5b91c45 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, @@ -22,12 +22,41 @@ 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)] 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, + /// 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 { @@ -40,19 +69,45 @@ 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)] pub distance: Option, #[allow(dead_code)] pub schema: Option>, + pub columns: Option>, + pub exclude: Vec, + pub fill_null: Option, + pub cache: CacheMode, } impl ResolvedDatasetConfig { @@ -64,16 +119,17 @@ 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 - ) - })?; - let path = inline - .path - .or_else(|| base.and_then(|b| b.path.clone())) - .with_context(|| format!("dataset {:?}: missing `path`", inline.name))?; + .with_context(|| format!("dataset {:?}: missing `format` ({KINDS})", 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 @@ -83,6 +139,7 @@ impl ResolvedDatasetConfig { kind, path, link, + parts, vector_size: inline .vector_size .or_else(|| base.and_then(|b| b.vector_size)), @@ -92,16 +149,49 @@ 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())), + cache: match inline.cache { + CacheMode::Keep => base.map(|b| b.cache).unwrap_or_default(), + explicit => explicit, + }, }) } } +/// 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 { @@ -118,14 +208,59 @@ 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.cache == CacheMode::Evict { bail!( - "dataset {:?} requires `format` (h5, tar, sparse)", + "dataset {:?}: `cache: evict` only applies to a sharded (`parts:`) dataset", self.name ); } if self.path.is_none() { - bail!("dataset {:?} requires `path`", self.name); + bail!("dataset {:?} requires `path` (or `parts`)", self.name); } Ok(()) } @@ -144,7 +279,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 +300,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(); @@ -184,4 +314,86 @@ 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}"); + } + + /// `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); + 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 dd751ab..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() @@ -140,19 +206,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..24f5b7e 100644 --- a/src/dataset/mod.rs +++ b/src/dataset/mod.rs @@ -1,5 +1,8 @@ mod config; mod download; +#[cfg(test)] +pub(crate) mod fixtures; +mod parts; mod payload; mod reader; mod readers; @@ -56,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..ceddb36 --- /dev/null +++ b/src/dataset/parts.rs @@ -0,0 +1,894 @@ +//! 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-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 — 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::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::{Arc, Mutex, RwLock, Weak}; + +use anyhow::{Context, Result, bail}; +use serde::{Deserialize, Serialize}; +use serde_json::Value; + +use super::config::{CacheMode, 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, + 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 { + 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(), + cache: config.cache, + guards: Mutex::new(HashMap::new()), + downloaded: Mutex::new(HashSet::new()), + prefetching: Mutex::new(HashSet::new()), + } + } + + 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. + /// + /// 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", + self.name, + target.display() + ) + })?; + 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)?), + 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: Arc, + 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<()>, + /// 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 { + 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: Arc::new(source), + entries, + starts, + open: RwLock::new(VecDeque::new()), + opening: Mutex::new(()), + pending_evict: Mutex::new(Vec::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 = 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) + } + + 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 + ); + } + + /// 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"); + 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 f8c6129..491c6e3 100644 --- a/src/dataset/reader.rs +++ b/src/dataset/reader.rs @@ -5,13 +5,19 @@ use serde_json::Value; use super::config::{DatasetConfig, DatasetKind}; use super::download::ensure_downloaded; -use super::readers::{H5Reader, SparseReader, TarReader}; +use super::parts::PartitionedReader; +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), + /// 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. @@ -24,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 => { @@ -41,6 +57,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 +80,11 @@ 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::Partitioned(r) => r.vector_at(idx), + DatasetReaderInner::Sparse(_) | DatasetReaderInner::Parquet(_) => { + bail!("dataset does not contain dense vectors") + } } } @@ -63,6 +98,8 @@ 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), + DatasetReaderInner::Partitioned(r) => r.payload_field(idx, field), _ => bail!("dataset does not contain payloads"), } } @@ -70,6 +107,8 @@ 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), + DatasetReaderInner::Partitioned(r) => r.payload_object(idx), _ => bail!("dataset does not contain payloads"), } } @@ -80,6 +119,11 @@ 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(_) + | DatasetReaderInner::Partitioned(_) => 0, } } @@ -89,6 +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(_) + | DatasetReaderInner::Partitioned(_) => bail!("dataset has no query set"), } } @@ -106,6 +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(_) + | DatasetReaderInner::Partitioned(_) => bail!("dataset has no ground truth"), } } } diff --git a/src/dataset/readers/mod.rs b/src/dataset/readers/mod.rs index c4a338a..0f29139 100644 --- a/src/dataset/readers/mod.rs +++ b/src/dataset/readers/mod.rs @@ -1,9 +1,15 @@ mod binary; mod h5; mod jsonl; +mod npy; +mod parquet; mod sparse; mod tar; pub use h5::H5Reader; +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/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) + } +} + +/// 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) +} + +/// 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 } @@ -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/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, 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(), )?);