From 532186edbb5661eee5923a95fd9d8c7e0caf7b3d Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Sun, 6 Sep 2026 09:49:52 +0000 Subject: [PATCH] feat: FineWeb-10B parquet upload and search examples Add parquet vector_column/sparse_column support and zero-padded {i:04d} parts templates so bfb can ingest Qdrant/FineWeb-10B dense+sparse embeddings from the same shards, with parts.count to limit how much of the corpus is loaded. --- README.md | 31 ++- examples/search-fineweb.yaml | 58 +++++ examples/upload-fineweb-10b.yaml | 95 ++++++++ examples/upload-fineweb-part.yaml | 81 +++++++ src/config/examples.rs | 15 ++ src/config/schema.rs | 20 +- src/dataset/config.rs | 54 ++++- src/dataset/parts.rs | 91 +++++++- src/dataset/reader.rs | 9 +- src/dataset/readers/parquet.rs | 252 +++++++++++++++++++++- src/dataset/testdata/fineweb_tiny.parquet | Bin 0 -> 3080 bytes 11 files changed, 668 insertions(+), 38 deletions(-) create mode 100644 examples/search-fineweb.yaml create mode 100644 examples/upload-fineweb-10b.yaml create mode 100644 examples/upload-fineweb-part.yaml create mode 100644 src/dataset/testdata/fineweb_tiny.parquet diff --git a/README.md b/README.md index 8f3ab0a..1374bec 100644 --- a/README.md +++ b/README.md @@ -117,7 +117,7 @@ inline dataset definitions (same fields as | `sparse` | CSR matrices (`data.csr`, optional `queries.csr` / `results.gt`) | | `npy` | One 2-D float `.npy` — dense vectors only | | `multivector` | Directory of `vectors.npy` (flat sub-vectors) + `offsets.npy` (row boundaries per point). Late-interaction, ColBERT-style multivectors only | -| `parquet` | One parquet file — payload rows only | +| `parquet` | One parquet file — payload rows, and optionally dense/sparse vectors via `vector_column` / `sparse_column` | The first three are *bundles*: vectors, payloads, and queries all come out of a single artifact. `npy`, `multivector`, and `parquet` are *components*, so a @@ -134,10 +134,13 @@ collection: 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). +Parquet sources accept five extra keys: `columns` (keep only these), `exclude` +(drop these), `fill_null` (a value substituted for nulls and for NaN/±inf +floats, which have no JSON form — by default such fields are simply absent), +`vector_column` (list-of-floats column used as a dense vector source), and +`sparse_column` (`{indices, values}` struct column used as a sparse vector +source). See [`examples/upload-laion-part.yaml`](examples/upload-laion-part.yaml) +and [`examples/upload-fineweb-part.yaml`](examples/upload-fineweb-part.yaml). #### Multivector (ColBERT-style) datasets @@ -169,7 +172,8 @@ collection: 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: +sources support it; `{i}` (or zero-padded `{i:04d}`) is substituted with each +part's number: ```yaml source: @@ -182,6 +186,21 @@ source: link: https://deploy.laion.ai/.../img_emb_{i}.npy ``` +FineWeb-10B shards use zero-padded names — dial the subset with `parts.count`: + +```yaml +source: + type: dataset + name: fineweb-10b-dense + format: parquet + vector_column: dense_embedding + parts: + count: 1 # first N of train-part0 (raise toward 10000) + path: fineweb/train-part0/{i:04d}.parquet + link: https://huggingface.co/datasets/Qdrant/FineWeb-10B/resolve/refs%2Fconvert%2Fparquet/default/train-part0/{i:04d}.parquet +``` + +See [`examples/upload-fineweb-10b.yaml`](examples/upload-fineweb-10b.yaml). 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 diff --git a/examples/search-fineweb.yaml b/examples/search-fineweb.yaml new file mode 100644 index 0000000..8864f6f --- /dev/null +++ b/examples/search-fineweb.yaml @@ -0,0 +1,58 @@ +# Search workload for a FineWeb collection uploaded with +# `upload-fineweb-part` / `upload-fineweb-10b`. +# +# bfb search --example search-fineweb -n 10k -p 8 -t 4 \ +# --search-limit 10 --uri http://localhost:6334 +# +# Query vectors are random (MS MARCO query embeddings are not redistributed +# with the dataset — regenerate them from the HF scripts if you need recall +# against `queries/gt_*.parquet`). Filters exercise the FineWeb metadata +# indexes created on upload. + +collection: + name: fineweb + +requests: + - kind: dense + using: dense + size: 768 + datatype: float16 + source: random + + - kind: sparse + using: sparse + source: + type: random + vocab_size: 250048 + length: 100 + distribution: zipf + + - kind: dense + using: dense + size: 768 + datatype: float16 + source: random + filters: + - name: language + type: keyword + source: { type: random, cardinality: 1 } + + - kind: dense + using: dense + size: 768 + datatype: float16 + source: random + filters: + - name: language_score + type: float + source: { type: random, min: 0.65, max: 1.0 } + + - kind: dense + using: dense + size: 768 + datatype: float16 + source: random + filters: + - name: token_count + type: integer + source: { type: random, min: 50, max: 2000 } diff --git a/examples/upload-fineweb-10b.yaml b/examples/upload-fineweb-10b.yaml new file mode 100644 index 0000000..125c845 --- /dev/null +++ b/examples/upload-fineweb-10b.yaml @@ -0,0 +1,95 @@ +# Qdrant/FineWeb-10B corpus +# (https://huggingface.co/datasets/Qdrant/FineWeb-10B) — dense + sparse +# embeddings and FineWeb metadata in numbered parquet shards. +# +# bfb upload --example upload-fineweb-10b -b 64 -p 8 -t 4 \ +# --uri http://localhost:6334 +# +# Dial the subset with `parts.count` (and optional `parts.start`). The Hugging +# Face convert/parquet layout exposes train-part0 as 10_000 zero-padded shards +# `0000.parquet` … `9999.parquet` (~10B rows together with train-part1). The +# default below uploads a single shard — raise `count` toward 10000 for a +# larger slice (each shard is ~5 GB on disk). +# +# On first use bfb sizes every configured part with one ranged request each and +# caches the result in `datasets/.parts-index/`. Point ids are dataset rows, so +# `--offset` resumes an interrupted run. +# +# Disk: each shard is ~5 GB. Dense, sparse, and payload all read the same files; +# keep `cache: keep` (default) so one source does not delete a part another still +# needs. Drop sparse/payload and set `cache: evict` if you want a dense-only +# stream that holds only a few shards on disk. + +collection: + name: fineweb + on_disk_payload: true + + quantization: + type: turbo-2bit + always_ram: true + + hnsw: + on_disk: false + + optimizers: + default_segment_number: 2 + max_segment_size: 5000000 + + vectors: + - name: dense + size: 768 + distance: cosine + datatype: float16 + on_disk: true + source: + type: dataset + name: fineweb-10b-dense + format: parquet + vector_column: dense_embedding + parts: + # First N of train-part0 (0..9999). Default is one shard; raise for more. + count: 1 + path: fineweb/train-part0/{i:04d}.parquet + link: https://huggingface.co/datasets/Qdrant/FineWeb-10B/resolve/refs%2Fconvert%2Fparquet/default/train-part0/{i:04d}.parquet + + sparse_vectors: + - name: sparse + on_disk: true + source: + type: dataset + dataset: + name: fineweb-10b-sparse + format: parquet + sparse_column: sparse_embedding + parts: + count: 1 + path: fineweb/train-part0/{i:04d}.parquet + link: https://huggingface.co/datasets/Qdrant/FineWeb-10B/resolve/refs%2Fconvert%2Fparquet/default/train-part0/{i:04d}.parquet + + payload: + source: + type: dataset + dataset: + name: fineweb-10b-payload + format: parquet + exclude: [dense_embedding, sparse_embedding, text] + parts: + count: 1 + path: fineweb/train-part0/{i:04d}.parquet + link: https://huggingface.co/datasets/Qdrant/FineWeb-10B/resolve/refs%2Fconvert%2Fparquet/default/train-part0/{i:04d}.parquet + + fields: + - name: id + type: keyword + - name: dump + type: keyword + - name: language + type: keyword + - name: url + type: keyword + - name: date + type: datetime + - name: language_score + type: float + - name: token_count + type: integer diff --git a/examples/upload-fineweb-part.yaml b/examples/upload-fineweb-part.yaml new file mode 100644 index 0000000..31a4e74 --- /dev/null +++ b/examples/upload-fineweb-part.yaml @@ -0,0 +1,81 @@ +# One parquet shard of Qdrant/FineWeb-10B +# (https://huggingface.co/datasets/Qdrant/FineWeb-10B). +# +# bfb upload --example upload-fineweb-part -b 64 -p 8 -t 4 \ +# --uri http://localhost:6334 +# +# Dense + sparse embeddings and FineWeb metadata live in the same parquet. +# `vector_column` / `sparse_column` tell bfb which columns are vectors; the +# payload source drops those (and the full `text`) so only filterable metadata +# is uploaded. Omit `-n` to upload the whole shard (~1M points; size varies). +# +# The file is downloaded on first use into `./datasets/` (override with +# BFB_DATASETS_DIR), ~5 GB. For N shards at once, see `upload-fineweb-10b.yaml`. + +collection: + name: fineweb + on_disk_payload: true + + quantization: + type: turbo-2bit + always_ram: true + + hnsw: + on_disk: false + + optimizers: + default_segment_number: 2 + max_segment_size: 2000000 + + vectors: + - name: dense + size: 768 + distance: cosine + datatype: float16 + on_disk: true + source: + type: dataset + name: fineweb-part0-dense + format: parquet + path: fineweb/train-part0/0000.parquet + link: https://huggingface.co/datasets/Qdrant/FineWeb-10B/resolve/refs%2Fconvert%2Fparquet/default/train-part0/0000.parquet + vector_column: dense_embedding + + sparse_vectors: + - name: sparse + on_disk: true + source: + type: dataset + dataset: + name: fineweb-part0-sparse + format: parquet + path: fineweb/train-part0/0000.parquet + link: https://huggingface.co/datasets/Qdrant/FineWeb-10B/resolve/refs%2Fconvert%2Fparquet/default/train-part0/0000.parquet + sparse_column: sparse_embedding + + payload: + source: + type: dataset + dataset: + name: fineweb-part0-payload + format: parquet + path: fineweb/train-part0/0000.parquet + link: https://huggingface.co/datasets/Qdrant/FineWeb-10B/resolve/refs%2Fconvert%2Fparquet/default/train-part0/0000.parquet + # Embeddings are uploaded as vectors; `text` dominates the row size. + exclude: [dense_embedding, sparse_embedding, text] + + fields: + - name: id + type: keyword + - name: dump + type: keyword + - name: language + type: keyword + - name: url + type: keyword + - name: date + type: datetime + - name: language_score + type: float + - name: token_count + type: integer diff --git a/src/config/examples.rs b/src/config/examples.rs index 458952b..9e0c44c 100644 --- a/src/config/examples.rs +++ b/src/config/examples.rs @@ -90,6 +90,16 @@ pub static EXAMPLES: &[Example] = &[ Upload, "Full LAION-400M corpus (~410 parts, streamed with cache: evict)" ), + example!( + "upload-fineweb-part", + Upload, + "One FineWeb-10B parquet shard: dense + sparse + metadata" + ), + example!( + "upload-fineweb-10b", + Upload, + "FineWeb-10B subset via parts.count (parquet vectors + metadata)" + ), example!( "serverless-upload", Upload, @@ -105,6 +115,11 @@ pub static EXAMPLES: &[Example] = &[ Search, "Measure recall against a dataset query set + ground truth" ), + example!( + "search-fineweb", + Search, + "Dense/sparse/filtered search against an uploaded FineWeb collection" + ), example!( "scroll-config", Scroll, diff --git a/src/config/schema.rs b/src/config/schema.rs index 4c2b715..83cd3ad 100644 --- a/src/config/schema.rs +++ b/src/config/schema.rs @@ -107,23 +107,24 @@ collection: # # offsets.npy (row boundaries per point) — ColBERT-style # # multivectors; requires `multivector:` above (`count` is # # ignored — arity comes from `offsets.npy`) - # # parquet one parquet file — payload rows only + # # parquet payload rows + optional vector_column/sparse_column # 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. + # `parts` block; the files are read as one row space and `{i}` (or + # zero-padded `{i:04d}`) 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 + # start: 0 # uint default=0 first part index + # path: laion/img_emb_{i}.npy # string required `{i}` or `{i:04d}` # 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 @@ -175,13 +176,16 @@ collection: # 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): + # five 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. - + # vector_column: dense_embedding # optional list-of-floats dense vector column + # sparse_column: sparse_embedding # optional {indices, values} sparse vector column + # When `vector_column` / `sparse_column` are set and `columns` is omitted, + # only those vector columns are decoded (payload text is skipped). # Payload field declarations (optional). Names must be unique. Each entry # generates a value and/or declares a field index. fields: diff --git a/src/dataset/config.rs b/src/dataset/config.rs index 0be8598..d9dbd5a 100644 --- a/src/dataset/config.rs +++ b/src/dataset/config.rs @@ -42,6 +42,15 @@ pub struct DatasetConfig { /// Omitted by default, which leaves the payload field absent. #[serde(default)] pub fill_null: Option, + /// `parquet` only: column holding a dense vector (list of floats). + /// When set, this dataset answers [`DatasetReader::dense_vector`]. + #[serde(default)] + pub vector_column: Option, + /// `parquet` only: column holding a sparse vector + /// (`{indices: list, values: list}`). + /// When set, this dataset answers [`DatasetReader::sparse_vector`]. + #[serde(default)] + pub sparse_column: Option, /// What to do with downloaded parts once the upload has moved past them. #[serde(default)] pub cache: CacheMode, @@ -71,10 +80,10 @@ 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. +/// `path` and `link` are templates containing `{i}` or `{i:04d}`, substituted +/// with each part's number (`{i:04d}` zero-pads to width 4). 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 { @@ -107,6 +116,8 @@ pub struct ResolvedDatasetConfig { pub columns: Option>, pub exclude: Vec, pub fill_null: Option, + pub vector_column: Option, + pub sparse_column: Option, pub cache: CacheMode, } @@ -160,6 +171,12 @@ impl ResolvedDatasetConfig { fill_null: inline .fill_null .or_else(|| base.and_then(|b| b.fill_null.clone())), + vector_column: inline + .vector_column + .or_else(|| base.and_then(|b| b.vector_column.clone())), + sparse_column: inline + .sparse_column + .or_else(|| base.and_then(|b| b.sparse_column.clone())), cache: match inline.cache { CacheMode::Keep => base.map(|b| b.cache).unwrap_or_default(), explicit => explicit, @@ -179,7 +196,8 @@ pub enum DatasetKind { Sparse, /// A standalone 2-D float `.npy` array: dense vectors, no payloads. Npy, - /// A parquet file of payload rows: no vectors. + /// A parquet file of payload rows, and optionally dense/sparse vectors + /// via `vector_column` / `sparse_column`. Parquet, /// A directory of `vectors.npy` (flat sub-vectors) + `offsets.npy` (row /// boundaries per point): ColBERT-style multivectors, no payloads. @@ -215,6 +233,16 @@ impl DatasetConfig { bail!("dataset {:?} requires `format` ({KINDS})", self.name); }; + if (self.vector_column.is_some() || self.sparse_column.is_some()) + && !matches!(kind, DatasetKind::Parquet) + { + bail!( + "dataset {:?}: `vector_column` / `sparse_column` are only supported \ + for `format: parquet`", + self.name + ); + } + if let Some(parts) = &self.parts { if self.path.is_some() || self.link.is_some() { bail!( @@ -238,18 +266,21 @@ impl DatasetConfig { } // 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}") { + // `{i:04d}` (zero-padded) also counts — it contains the `{i` marker. + if parts.count > 1 && !parts.path.contains("{i") { bail!( - "dataset {:?}: `parts.path` must contain `{{i}}` to distinguish parts", + "dataset {:?}: `parts.path` must contain `{{i}}` or `{{i:04d}}` \ + to distinguish parts", self.name ); } if let Some(link) = &parts.link && parts.count > 1 - && !link.contains("{i}") + && !link.contains("{i") { bail!( - "dataset {:?}: `parts.link` must contain `{{i}}` to distinguish parts", + "dataset {:?}: `parts.link` must contain `{{i}}` or `{{i:04d}}` \ + to distinguish parts", self.name ); } @@ -357,6 +388,11 @@ mod tests { parts_config(template("laion/img_emb.npy", 1), DatasetKind::Npy) .validate_inline() .unwrap(); + + // Zero-padded `{i:04d}` is accepted. + parts_config(template("p/{i:04d}.parquet", 10), DatasetKind::Parquet) + .validate_inline() + .unwrap(); } #[test] diff --git a/src/dataset/parts.rs b/src/dataset/parts.rs index ceddb36..c4b79ce 100644 --- a/src/dataset/parts.rs +++ b/src/dataset/parts.rs @@ -81,6 +81,8 @@ pub struct PartSource { columns: Option>, exclude: Vec, fill_null: Option, + vector_column: Option, + sparse_column: Option, cache: CacheMode, /// One lock per part, so fetching part *n+1* in the background never blocks /// a reader that wants part *n*. @@ -102,6 +104,8 @@ impl PartSource { columns: config.columns.clone(), exclude: config.exclude.clone(), fill_null: config.fill_null.clone(), + vector_column: config.vector_column.clone(), + sparse_column: config.sparse_column.clone(), cache: config.cache, guards: Mutex::new(HashMap::new()), downloaded: Mutex::new(HashSet::new()), @@ -229,6 +233,8 @@ impl PartSource { self.columns.as_deref(), &self.exclude, self.fill_null.as_ref(), + self.vector_column.as_deref(), + self.sparse_column.as_deref(), )?)), other => bail!("`parts:` is not supported for format {other:?} (use npy or parquet)"), }) @@ -304,9 +310,44 @@ impl PartSource { } } -/// Substitute a part number into a `{i}` template. +/// Substitute a part number into a `{i}` / `{i:04d}` template. +/// +/// `{i}` expands to the decimal index (`7` → `"7"`). `{i:0Nd}` zero-pads to +/// width N (`7` → `"0007"` for `{i:04d}`), which HuggingFace convert/parquet +/// shards and similar layouts need. fn expand(template: &str, index: usize) -> String { - template.replace("{i}", &index.to_string()) + let mut out = String::with_capacity(template.len() + 8); + let mut rest = template; + while let Some(start) = rest.find("{i") { + out.push_str(&rest[..start]); + let after = &rest[start + 2..]; + if let Some(stripped) = after.strip_prefix('}') { + out.push_str(&index.to_string()); + rest = stripped; + } else if let Some(end) = after.find('}') { + out.push_str(&format_index(index, &after[..end])); + rest = &after[end + 1..]; + } else { + // Unclosed `{i` — leave it literally so a typo stays visible. + out.push_str("{i"); + rest = after; + } + } + out.push_str(rest); + out +} + +/// Format `index` according to a `{i…}` body such as `:04d`. Unknown specs +/// fall back to plain decimal so a typo does not invent a second naming scheme. +fn format_index(index: usize, spec: &str) -> String { + if let Some(width) = spec + .strip_prefix(':') + .and_then(|s| s.strip_suffix('d')) + .and_then(|w| w.parse::().ok()) + { + return format!("{index:0width$}"); + } + index.to_string() } /// Keep a dataset name usable as a file name. @@ -407,7 +448,14 @@ impl PartReader { 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"), + PartReader::Parquet(r) => r.dense_vector_at(idx), + } + } + + fn sparse_vector_at(&self, idx: usize) -> Result> { + match self { + PartReader::Parquet(r) => r.sparse_vector_at(idx), + PartReader::Npy(_) => bail!("npy parts do not contain sparse vectors"), } } @@ -575,6 +623,11 @@ impl PartitionedReader { self.reader_for(slot)?.vector_at(local) } + pub fn sparse_vector_at(&self, idx: usize) -> Result> { + let (slot, local) = self.locate(idx)?; + self.reader_for(slot)?.sparse_vector_at(local) + } + pub fn payload_object(&self, idx: usize) -> Result> { let (slot, local) = self.locate(idx)?; self.reader_for(slot)?.payload_object(local) @@ -890,5 +943,37 @@ mod tests { 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"); + assert_eq!(expand("p/{i:04d}.parquet", 7), "p/0007.parquet"); + assert_eq!(expand("p/{i:04d}.parquet", 0), "p/0000.parquet"); + assert_eq!(expand("p/{i:02d}/x_{i}.npy", 3), "p/03/x_3.npy"); + } + + #[test] + fn reads_dense_vectors_across_padded_parquet_parts() { + let dir = tempfile::tempdir().unwrap(); + let fixture = Path::new(env!("CARGO_MANIFEST_DIR")) + .join("src/dataset/testdata/fineweb_tiny.parquet"); + std::fs::copy(&fixture, dir.path().join("0000.parquet")).unwrap(); + std::fs::copy(&fixture, dir.path().join("0001.parquet")).unwrap(); + + let config = DatasetConfig { + name: "fineweb-parts".to_string(), + kind: Some(DatasetKind::Parquet), + parts: Some(PartsConfig { + count: 2, + start: 0, + path: "{i:04d}.parquet".to_string(), + link: None, + }), + vector_column: Some("dense_embedding".to_string()), + ..Default::default() + }; + let resolved = DatasetConfig::resolve(config, &Default::default()).unwrap(); + let reader = PartitionedReader::open(dir.path(), &resolved).unwrap(); + assert_eq!(reader.num_points(), 10); + assert_eq!(reader.vector_at(0).unwrap()[0], 0.0); + // First row of the second part. + assert_eq!(reader.vector_at(5).unwrap()[0], 0.0); + assert_eq!(reader.vector_at(7).unwrap()[0], 2.0); } } diff --git a/src/dataset/reader.rs b/src/dataset/reader.rs index da64c94..7ad6b20 100644 --- a/src/dataset/reader.rs +++ b/src/dataset/reader.rs @@ -72,6 +72,8 @@ impl DatasetReader { config.columns.as_deref(), &config.exclude, config.fill_null.as_ref(), + config.vector_column.as_deref(), + config.sparse_column.as_deref(), )?; let n = reader.num_points(); (DatasetReaderInner::Parquet(reader), n) @@ -90,10 +92,9 @@ impl DatasetReader { DatasetReaderInner::H5(r) => r.vector_at(idx), DatasetReaderInner::Tar(r) => r.vector_at(idx), DatasetReaderInner::Npy(r) => r.vector_at(idx), + DatasetReaderInner::Parquet(r) => r.dense_vector_at(idx), DatasetReaderInner::Partitioned(r) => r.vector_at(idx), - DatasetReaderInner::Sparse(_) - | DatasetReaderInner::Parquet(_) - | DatasetReaderInner::Multivector(_) => { + DatasetReaderInner::Sparse(_) | DatasetReaderInner::Multivector(_) => { bail!("dataset does not contain dense vectors") } } @@ -102,6 +103,8 @@ impl DatasetReader { pub fn sparse_vector(&self, idx: usize) -> Result> { match &self.inner { DatasetReaderInner::Sparse(r) => r.vector_at(idx), + DatasetReaderInner::Parquet(r) => r.sparse_vector_at(idx), + DatasetReaderInner::Partitioned(r) => r.sparse_vector_at(idx), _ => bail!("dataset does not contain sparse vectors"), } } diff --git a/src/dataset/readers/parquet.rs b/src/dataset/readers/parquet.rs index 78eed64..4f97472 100644 --- a/src/dataset/readers/parquet.rs +++ b/src/dataset/readers/parquet.rs @@ -1,7 +1,9 @@ -//! Parquet payload reading. +//! Parquet payload and (optional) vector reading. //! -//! Payloads only: a parquet part carries no vectors, so a config pairs it with -//! a separate vector source (typically a `.npy` of the same row count). +//! By default a parquet part is payloads only, paired with a separate vector +//! source (typically a `.npy` of the same row count) — that is the LAION shape. +//! FineWeb-style corpora put dense and/or sparse embeddings in the same file; +//! set `vector_column` / `sparse_column` on the dataset config to read those. //! //! Access is a *streaming cursor* rather than a decode-the-whole-file cache. //! Upload walks point ids in order, so the reader keeps one live row iterator @@ -43,6 +45,10 @@ pub struct ParquetReader { columns: Option>, /// Value substituted for nulls and non-finite floats. `None` ⇒ omit the field. fill_null: Option, + /// Column holding a dense vector (list of floats), if any. + vector_column: Option, + /// Column holding a sparse vector struct, if any. + sparse_column: Option, cursor: Mutex, } @@ -56,11 +62,17 @@ struct Cursor { impl ParquetReader { /// Open `path`, keeping `columns` (default: all) minus `exclude`. + /// + /// When `vector_column` / `sparse_column` are set and `columns` is omitted, + /// the projection collapses to just those vector columns so payload text is + /// not decoded on a vectors-only read. pub fn open( path: &Path, columns: Option<&[String]>, exclude: &[String], fill_null: Option<&Value>, + vector_column: Option<&str>, + sparse_column: Option<&str>, ) -> Result { let file = File::open(path).with_context(|| format!("failed to open {}", path.display()))?; @@ -85,7 +97,40 @@ impl ParquetReader { .iter() .map(|f| f.name().to_string()) .collect(); - let columns = resolve_projection(&available, columns, exclude, path)?; + + for name in vector_column.into_iter().chain(sparse_column) { + if !available.iter().any(|c| c == name) { + bail!( + "{} has no column {name:?}; available columns: {}", + path.display(), + available.join(", ") + ); + } + } + + // Vectors-only read: if the caller named vector columns and did not + // ask for a payload projection, decode just those columns. + let effective_columns: Option> = if columns.is_none() + && exclude.is_empty() + && (vector_column.is_some() || sparse_column.is_some()) + { + let mut cols = Vec::new(); + if let Some(name) = vector_column { + cols.push(name.to_string()); + } + if let Some(name) = sparse_column + && !cols.iter().any(|c| c == name) + { + cols.push(name.to_string()); + } + Some(cols) + } else { + None + }; + let columns = match effective_columns.as_deref() { + Some(cols) => resolve_projection(&available, Some(cols), exclude, path)?, + None => resolve_projection(&available, columns, exclude, path)?, + }; Ok(ParquetReader { path: path.to_path_buf(), @@ -93,6 +138,8 @@ impl ParquetReader { group_starts, columns, fill_null: fill_null.cloned(), + vector_column: vector_column.map(str::to_string), + sparse_column: sparse_column.map(str::to_string), cursor: Mutex::new(Cursor::default()), }) } @@ -101,6 +148,56 @@ impl ParquetReader { self.num_rows } + /// Dense vector from `vector_column` at row `idx`. + pub fn dense_vector_at(&self, idx: usize) -> Result> { + let column = self.vector_column.as_deref().with_context(|| { + format!( + "{}: dense vectors require `vector_column` on the parquet dataset", + self.path.display() + ) + })?; + let row = self + .payload_object(idx)? + .with_context(|| format!("{}: row {idx} is past the end", self.path.display()))?; + let value = row.get(column).with_context(|| { + format!( + "{}: row {idx} has no dense vector column {column:?}", + self.path.display() + ) + })?; + json_to_dense(value).with_context(|| { + format!( + "{}: row {idx} column {column:?} is not a list of floats", + self.path.display() + ) + }) + } + + /// Sparse vector from `sparse_column` at row `idx`. + pub fn sparse_vector_at(&self, idx: usize) -> Result> { + let column = self.sparse_column.as_deref().with_context(|| { + format!( + "{}: sparse vectors require `sparse_column` on the parquet dataset", + self.path.display() + ) + })?; + let row = self + .payload_object(idx)? + .with_context(|| format!("{}: row {idx} is past the end", self.path.display()))?; + let value = row.get(column).with_context(|| { + format!( + "{}: row {idx} has no sparse vector column {column:?}", + self.path.display() + ) + })?; + json_to_sparse(value).with_context(|| { + format!( + "{}: row {idx} column {column:?} is not `{{indices, values}}`", + self.path.display() + ) + }) + } + /// The whole payload object for a row. pub fn payload_object(&self, idx: usize) -> Result> { if idx >= self.num_rows { @@ -369,6 +466,66 @@ fn number(value: f64) -> Option { Number::from_f64(value).map(Value::Number) } +/// Decode a JSON array of numbers into an `f32` dense vector. +fn json_to_dense(value: &Value) -> Result> { + let Value::Array(items) = value else { + bail!("expected a JSON array"); + }; + items + .iter() + .map(|item| match item { + Value::Number(n) => n + .as_f64() + .map(|f| f as f32) + .context("non-finite float in dense vector"), + _ => bail!("dense vector element is not a number"), + }) + .collect() +} + +/// Decode `{indices: [...], values: [...]}` into sparse index/value pairs. +fn json_to_sparse(value: &Value) -> Result> { + let object = value + .as_object() + .context("expected a JSON object with indices and values")?; + let indices = object + .get("indices") + .and_then(Value::as_array) + .context("sparse vector missing `indices` array")?; + let values = object + .get("values") + .and_then(Value::as_array) + .context("sparse vector missing `values` array")?; + if indices.len() != values.len() { + bail!( + "sparse vector indices ({}) and values ({}) lengths differ", + indices.len(), + values.len() + ); + } + indices + .iter() + .zip(values) + .map(|(idx, val)| { + let index = match idx { + Value::Number(n) => n + .as_u64() + .context("sparse index is not an unsigned integer")? + as u32, + _ => bail!("sparse index is not a number"), + }; + let value = match val { + Value::Number(n) => n + .as_f64() + .map(|f| f as f32) + .context("non-finite float in sparse values")?, + _ => bail!("sparse value is not a number"), + }; + Ok((index, value)) + }) + .collect() +} + /// 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 { @@ -437,7 +594,7 @@ mod tests { use crate::dataset::fixtures::write_parquet; fn open(path: &Path) -> ParquetReader { - ParquetReader::open(path, None, &[], None).unwrap() + ParquetReader::open(path, None, &[], None, None, None).unwrap() } #[test] @@ -480,7 +637,7 @@ mod tests { 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 reader = ParquetReader::open(&path, None, &[], Some(&Value::from(0)), None, None).unwrap(); let hole = reader.payload_object(3).unwrap().unwrap(); assert_eq!(hole["similarity"], 0); assert_eq!(hole["caption"], 0); @@ -493,7 +650,7 @@ mod tests { write_parquet(&path, 0, 10, 64); let reader = - ParquetReader::open(&path, None, std::slice::from_ref(&"url".to_string()), None) + ParquetReader::open(&path, None, std::slice::from_ref(&"url".to_string()), None, None, None) .unwrap(); let row = reader.payload_object(1).unwrap().unwrap(); assert!(!row.as_object().unwrap().contains_key("url")); @@ -507,7 +664,7 @@ mod tests { 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 reader = ParquetReader::open(&path, Some(&wanted), &[], None, None, 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"]); @@ -519,7 +676,7 @@ mod tests { let path = dir.path().join("m.parquet"); write_parquet(&path, 0, 10, 64); - let err = ParquetReader::open(&path, None, &["exif".to_string()], None) + let err = ParquetReader::open(&path, None, &["exif".to_string()], None, None, None) .map(|_| ()) .unwrap_err(); let message = err.to_string(); @@ -563,4 +720,81 @@ mod tests { assert_eq!(open(&path).payload_object(10).unwrap(), None); } + + fn fineweb_tiny() -> std::path::PathBuf { + Path::new(env!("CARGO_MANIFEST_DIR")).join("src/dataset/testdata/fineweb_tiny.parquet") + } + + #[test] + fn reads_dense_vectors_from_vector_column() { + let path = fineweb_tiny(); + let reader = + ParquetReader::open(&path, None, &[], None, Some("dense_embedding"), None).unwrap(); + assert_eq!(reader.num_points(), 5); + let v = reader.dense_vector_at(0).unwrap(); + assert_eq!(v.len(), 8); + assert_eq!(v[0], 0.0); + assert_eq!(v[7], 7.0); + let v2 = reader.dense_vector_at(2).unwrap(); + assert_eq!(v2[0], 2.0); + assert_eq!(v2[1], 3.0); + } + + #[test] + fn reads_sparse_vectors_from_sparse_column() { + let path = fineweb_tiny(); + let reader = + ParquetReader::open(&path, None, &[], None, None, Some("sparse_embedding")).unwrap(); + let pairs = reader.sparse_vector_at(0).unwrap(); + assert_eq!(pairs, vec![(1, 0.1), (3, 0.2), (5, 0.3)]); + assert_eq!(reader.sparse_vector_at(2).unwrap(), vec![(7, 1.0)]); + } + + #[test] + fn dense_vector_requires_vector_column() { + let path = fineweb_tiny(); + let reader = ParquetReader::open(&path, None, &[], None, None, None).unwrap(); + let err = reader.dense_vector_at(0).unwrap_err().to_string(); + assert!(err.contains("vector_column"), "{err}"); + } + + #[test] + fn exclude_keeps_payload_without_embeddings() { + let path = fineweb_tiny(); + let exclude = vec![ + "dense_embedding".to_string(), + "sparse_embedding".to_string(), + "text".to_string(), + ]; + let reader = ParquetReader::open(&path, None, &exclude, None, None, None).unwrap(); + let row = reader.payload_object(1).unwrap().unwrap(); + let obj = row.as_object().unwrap(); + assert!(!obj.contains_key("dense_embedding")); + assert!(!obj.contains_key("text")); + assert_eq!(obj["token_count"], 101); + assert_eq!(obj["dump"], "CC-MAIN-2016-01"); + } + + /// Optional smoke against a real FineWeb shard (set `FINEWEB_SAMPLE` to the + /// path of a `*.parquet` from Qdrant/FineWeb-10B). + #[test] + fn reads_a_real_fineweb_sample_when_configured() { + let Ok(path) = std::env::var("FINEWEB_SAMPLE") else { + return; + }; + let path = Path::new(&path); + if !path.exists() { + return; + } + let dense = ParquetReader::open(path, None, &[], None, Some("dense_embedding"), None) + .unwrap(); + let v = dense.dense_vector_at(0).unwrap(); + assert_eq!(v.len(), 768, "mGTE dense vectors are 768-d"); + + let sparse = + ParquetReader::open(path, None, &[], None, None, Some("sparse_embedding")).unwrap(); + let pairs = sparse.sparse_vector_at(0).unwrap(); + assert!(!pairs.is_empty()); + assert!(pairs.windows(2).all(|w| w[0].0 < w[1].0) || pairs.len() == 1); + } } diff --git a/src/dataset/testdata/fineweb_tiny.parquet b/src/dataset/testdata/fineweb_tiny.parquet new file mode 100644 index 0000000000000000000000000000000000000000..a992cd9e64bb0edc220f37581d6d187b28d144a5 GIT binary patch literal 3080 zcmb7HO>7%Q6rSB|9mh#imv&a`%7TFC%0N;(j*}*l8+SH|od_p5R%^#rrN}=U$IZrZ zV>@Xh4n9EXp%REgMSG$`LPA0uIH4YT0E7@q5hsv(LqZ6Nwuh?RD#UxUw(HnVfl+qG zGyC56z4yK6NfKevVJgv&1-clfHcDO~B;xNQg!n64g~>`KpPlqfjoAqqo`lJZ(+EtG zZ;BEs(-VW>WTzHtagGsJ-!muQl+v{GOdUN6n7u9$W}DE}buCDMHA2tUNnKKw^aw=n z>6XaqVXHp}hB(&Dg|xCM*7_U8nyuN$#MR;S0Osg3=6idB@Xmw4OYi=<^worTy{86Aig!&uBdT$TMg8+ume%dhv5A}tq!WCC8Kn8wt zmS|*#%#t7pk-6#|nXk?RXP7Pyf>Nf)iddajL{bHw+9!^!9=8BO@85!8pu&+Zsn_m3 zxBc?8{vIqJ!SWd_cR=!S>JQU?&gwJ-ukg|TcJm8@=G7OKG+FT=zQ(tQ} z$B4_q?R|>5dOB9^u0R`|vod!CNT#QU7TEmakk>OjHstXHkHL2ZCofJTICbJWvfs_8 zv>bF~cRaO^62(%=tmzvm+D~ZTNl_8=*^?yfgh;3a+uri#x?WhxY*q?oLTzx;Z=?#j zN-C$Xlro!JIvEv}BU!yr(pU72w4Tl83ppxMOV|$(qmm*S`BIsR1iD0`qSCD!`i5R8 z8&C1IsMt%z)Yd*<36Hz;g={{fmpUSe5Jes5Ny56;@j}AFI-b?L7sx9aUf|ik*o}x)M*Va_^ zc4}pn+E8WxzMgDA;+MgM6y>La`om%Mejh@%+$aX6Je7LC3YpENHarP}SrF`AoZj4i=~<4wkPn3iKf#&T>r$Imp4 zu?74DhV!_>zP`dPW1fKRIGbfGjIXsNL|jJW8JzG1W2Yk{Zubs+SUS)7LU`Dq=D3W7 zqQhzK$e0i9FQ0qE?H&ia*&qjj;BbSJIS@+oB$wAb=KdnukGr?sZhSEl<@P0f2Z6!B z&@}{XNm~KAca+qwA^oXaQ~5Eb>qI8hDTz3{1L)$3r>K_eeFDTuaZz zj11(NiJW%3#~Sg^JSYBQHWVu(2Ghb&6rL~hpg-4ZhvUz21J2J=fg-`+`WMq1(PF|I z+hq%wzq^q?TWc8A;^liDN@_W3qyY{j1yPo8FqMr=J9GAhX((?1ab5FY;m4?2V(ZS^)h$+%g%7~148^)j1~lX7?F zgGDPPpg!4nu#k*R>?XVuSCX880Ywh;FsU$`g>!YhX$~^Y4V&<=H8fA!TQ<_h#KrVT ztdixev4i~%lA|bmm`Z>MpgZW5VA_uPy#O$i9e}hbw@u(|>c#f}#aJqiW}tJee~B^x xzmYFh^zu-qSR6`iZEfxhZ0lR4{AOXmJ95f%YS?id{ygaTL(W^_HirKJ{sC}RkbeLG literal 0 HcmV?d00001