diff --git a/README.md b/README.md index f62b922..d8ff8bb 100644 --- a/README.md +++ b/README.md @@ -116,11 +116,12 @@ inline dataset definitions (same fields as | `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 | +| `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 | 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*: +single artifact. `npy`, `multivector`, and `parquet` are *components*, so a +config pairs them — one source per slot, row *i* of each landing on point *i*: ```yaml collection: @@ -138,6 +139,32 @@ Parquet sources accept three extra keys: `columns` (keep only these), `exclude` floats, which have no JSON form — by default such fields are simply absent). See [`examples/upload-laion-part.yaml`](examples/upload-laion-part.yaml). +#### Multivector (ColBERT-style) datasets + +A `multivector` source loads real per-point sub-vectors (e.g. one embedding +per token) from a directory of two files: `vectors.npy`, a flat 2-D float +array with every sub-vector from every point concatenated together, and +`offsets.npy`, a 1-D int array (`int32`/`int64`) of `num_points + 1` row +boundaries into it. Point `i`'s sub-vectors are +`vectors[offsets[i]:offsets[i+1]]`. The dense vector's `multivector:` block is +still required (for the comparator), but its `count` is ignored (arity comes +from `offsets.npy`), so points may have differing numbers of sub-vectors: + +```yaml +collection: + vectors: + - name: colbert + size: 128 + multivector: + comparator: max_sim + count: 1 # ignored for this source + source: + type: dataset + name: my-colbert-corpus + format: multivector + path: my-colbert-corpus # directory containing vectors.npy + offsets.npy +``` + #### Sharded datasets Corpora published as numbered parts are read as one row space with a `parts:` diff --git a/src/config/schema.rs b/src/config/schema.rs index 7041b31..4c2b715 100644 --- a/src/config/schema.rs +++ b/src/config/schema.rs @@ -103,6 +103,10 @@ collection: # # tar .tgz of vectors.npy + payloads.jsonl + tests.jsonl # # sparse CSR matrices # # npy one 2-D float .npy — dense vectors only + # # multivector directory of vectors.npy (flat sub-vectors) + + # # 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 # path: glove-25-angular/glove-25-angular.hdf5 # link: http://ann-benchmarks.com/glove-25-angular.hdf5 @@ -126,6 +130,13 @@ collection: # # 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. + # A ColBERT-style multivector dataset (`multivector:` above must be set): + # source: + # type: dataset + # name: colbert-corpus + # format: multivector + # path: colbert-corpus # directory containing vectors.npy + offsets.npy + # link: https://example.com/colbert-corpus.tgz # Sparse vectors (optional). Names must be unique across all vectors. sparse_vectors: diff --git a/src/dataset/config.rs b/src/dataset/config.rs index 5b91c45..0be8598 100644 --- a/src/dataset/config.rs +++ b/src/dataset/config.rs @@ -169,7 +169,7 @@ impl ResolvedDatasetConfig { } /// Formats accepted by `format:`, for error messages. -const KINDS: &str = "h5, tar, sparse, npy, parquet"; +const KINDS: &str = "h5, tar, sparse, npy, parquet, multivector"; #[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] #[serde(rename_all = "lowercase")] @@ -181,6 +181,9 @@ pub enum DatasetKind { Npy, /// A parquet file of payload rows: no vectors. Parquet, + /// A directory of `vectors.npy` (flat sub-vectors) + `offsets.npy` (row + /// boundaries per point): ColBERT-style multivectors, no payloads. + Multivector, } impl DatasetKind { diff --git a/src/dataset/reader.rs b/src/dataset/reader.rs index 94a868d..da64c94 100644 --- a/src/dataset/reader.rs +++ b/src/dataset/reader.rs @@ -7,7 +7,8 @@ use super::config::{DatasetConfig, DatasetKind}; use super::download::ensure_downloaded; use super::parts::PartitionedReader; use super::readers::{ - H5Reader, NpyReader, ParquetReader, QueryEntry, SparseReader, SparseVector, TarReader, + H5Reader, MultivectorReader, NpyReader, ParquetReader, QueryEntry, SparseReader, SparseVector, + TarReader, }; use super::registry::load_registry; @@ -17,6 +18,7 @@ enum DatasetReaderInner { Sparse(SparseReader), Npy(NpyReader), Parquet(ParquetReader), + Multivector(MultivectorReader), /// 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), @@ -74,6 +76,11 @@ impl DatasetReader { let n = reader.num_points(); (DatasetReaderInner::Parquet(reader), n) } + DatasetKind::Multivector => { + let reader = MultivectorReader::open(&local_path)?; + let n = reader.num_points(); + (DatasetReaderInner::Multivector(reader), n) + } }; Ok(DatasetReader { inner, num_points }) } @@ -84,7 +91,9 @@ impl DatasetReader { DatasetReaderInner::Tar(r) => r.vector_at(idx), DatasetReaderInner::Npy(r) => r.vector_at(idx), DatasetReaderInner::Partitioned(r) => r.vector_at(idx), - DatasetReaderInner::Sparse(_) | DatasetReaderInner::Parquet(_) => { + DatasetReaderInner::Sparse(_) + | DatasetReaderInner::Parquet(_) + | DatasetReaderInner::Multivector(_) => { bail!("dataset does not contain dense vectors") } } @@ -97,6 +106,14 @@ impl DatasetReader { } } + /// A point's sub-vectors from a `multivector` dataset (ColBERT-style). + pub fn multi_dense_vector(&self, idx: usize) -> Result>> { + match &self.inner { + DatasetReaderInner::Multivector(r) => r.vector_at(idx), + _ => bail!("dataset does not contain multivectors"), + } + } + pub fn payload_field(&self, idx: usize, field: &str) -> Result> { match &self.inner { DatasetReaderInner::Tar(r) => r.payload_field(idx, field), @@ -125,7 +142,8 @@ impl DatasetReader { // separate file, declared as its own source. DatasetReaderInner::Npy(_) | DatasetReaderInner::Parquet(_) - | DatasetReaderInner::Partitioned(_) => 0, + | DatasetReaderInner::Partitioned(_) + | DatasetReaderInner::Multivector(_) => 0, } } @@ -137,7 +155,8 @@ impl DatasetReader { DatasetReaderInner::Sparse(_) => bail!("sparse dataset has no dense queries"), DatasetReaderInner::Npy(_) | DatasetReaderInner::Parquet(_) - | DatasetReaderInner::Partitioned(_) => bail!("dataset has no query set"), + | DatasetReaderInner::Partitioned(_) + | DatasetReaderInner::Multivector(_) => bail!("dataset has no query set"), } } @@ -191,7 +210,8 @@ impl DatasetReader { DatasetReaderInner::Sparse(r) => r.query_ground_truth(idx), DatasetReaderInner::Npy(_) | DatasetReaderInner::Parquet(_) - | DatasetReaderInner::Partitioned(_) => bail!("dataset has no ground truth"), + | DatasetReaderInner::Partitioned(_) + | DatasetReaderInner::Multivector(_) => bail!("dataset has no ground truth"), } } } diff --git a/src/dataset/readers/mod.rs b/src/dataset/readers/mod.rs index 383c7d6..fb320be 100644 --- a/src/dataset/readers/mod.rs +++ b/src/dataset/readers/mod.rs @@ -1,6 +1,7 @@ mod binary; mod h5; mod jsonl; +mod multivector; mod npy; mod parquet; mod query; @@ -8,6 +9,7 @@ mod sparse; mod tar; pub use h5::H5Reader; +pub use multivector::MultivectorReader; pub use npy::{NpyReader, parse_npy_header}; pub use parquet::{ ParquetReader, parquet_footer_len, parquet_row_count, parquet_row_count_from_tail, diff --git a/src/dataset/readers/multivector.rs b/src/dataset/readers/multivector.rs new file mode 100644 index 0000000..b1682c7 --- /dev/null +++ b/src/dataset/readers/multivector.rs @@ -0,0 +1,275 @@ +//! Reads a ColBERT-style multivector dataset: a directory holding +//! `vectors.npy` (a flat `[total_subvectors, dim]` float array, in the same +//! format as the plain `npy` dataset) and `offsets.npy` (a 1-D int array of +//! length `num_points + 1` giving row boundaries into `vectors.npy`). +//! +//! Point `i`'s sub-vectors are `vectors[offsets[i]:offsets[i+1]]`, mirroring +//! how [`SparseReader`](super::SparseReader)'s CSR `index_pointer` addresses +//! ragged rows — except each unit here is a whole `dim`-wide row rather than +//! a single scalar. + +use std::fs::File; +use std::io::Read as _; +use std::path::Path; + +use anyhow::{Context, Result, bail}; + +use super::npy::{NpyMatrix, extract_quoted, parse_npy_header_str}; + +pub struct MultivectorReader { + vectors: NpyMatrix, + /// Row boundaries into `vectors`, length `num_points + 1`. + offsets: Vec, +} + +impl MultivectorReader { + pub fn open(path: &Path) -> Result { + let vectors = NpyMatrix::open(&path.join("vectors.npy"))?; + let offsets = read_offsets_npy(&path.join("offsets.npy"))?; + + if offsets.len() < 2 { + bail!( + "offsets.npy must have at least 2 entries (num_points + 1), got {}", + offsets.len() + ); + } + let last = *offsets.last().unwrap(); + if last < 0 || last as usize != vectors.rows() { + bail!( + "offsets.npy's last entry ({last}) does not match vectors.npy's row count ({})", + vectors.rows() + ); + } + + Ok(MultivectorReader { vectors, offsets }) + } + + pub fn num_points(&self) -> usize { + self.offsets.len() - 1 + } + + pub fn vector_at(&self, idx: usize) -> Result>> { + if idx + 1 >= self.offsets.len() { + bail!( + "index {idx} out of range (dataset has {} points)", + self.num_points() + ); + } + let start = self.offsets[idx]; + let end = self.offsets[idx + 1]; + if start < 0 || end < start { + bail!("offsets.npy is not non-decreasing at index {idx}"); + } + (start as usize..end as usize) + .map(|row| self.vectors.row(row)) + .collect() + } +} + +/// Minimal parser for a 1-D numeric `.npy` array (`int32`/`int64`, signed or +/// unsigned, or `float32`/`float64` downcast to `i64`), used for the +/// `offsets.npy` row-boundary array. Read in full rather than mmapped: it is +/// tiny (`num_points + 1` scalars) next to `vectors.npy`. +fn read_offsets_npy(path: &Path) -> Result> { + let mut file = + File::open(path).with_context(|| format!("failed to open {}", path.display()))?; + let mut buf = Vec::new(); + file.read_to_end(&mut buf) + .with_context(|| format!("failed to read {}", path.display()))?; + + let (header, header_end) = parse_npy_header_str(&buf) + .with_context(|| format!("failed to parse {}", path.display()))?; + + let descr = extract_quoted(header, "descr").context(".npy header missing 'descr'")?; + // Offsets are conceptually integers, but some exporters (e.g. numpy's + // default float dtype) write them as floats; downcast those to i64 rather + // than rejecting the file, since the values are still whole numbers. + let elem = match descr.as_str() { + " OffsetElem::I32, + " OffsetElem::U32, + " OffsetElem::I64, + " OffsetElem::U64, + " OffsetElem::F32, + " OffsetElem::F64, + other => bail!( + "unsupported offsets dtype {other:?} (expected int32/int64/uint32/uint64/float32/float64)" + ), + }; + let elem_size = elem.size(); + + if header.contains("'fortran_order': True") || header.contains("\"fortran_order\": true") { + bail!("offsets.npy is Fortran-ordered; expected C order"); + } + + let n = extract_1d_shape(header)?; + let data = &buf[header_end..]; + if data.len() < n * elem_size { + bail!( + "offsets.npy is truncated: need {} bytes of data, got {}", + n * elem_size, + data.len() + ); + } + + Ok((0..n) + .map(|i| elem.read(&data[i * elem_size..(i + 1) * elem_size])) + .collect()) +} + +#[derive(Debug, Clone, Copy)] +enum OffsetElem { + I32, + U32, + I64, + U64, + F32, + F64, +} + +impl OffsetElem { + fn size(self) -> usize { + match self { + OffsetElem::I32 | OffsetElem::U32 | OffsetElem::F32 => 4, + OffsetElem::I64 | OffsetElem::U64 | OffsetElem::F64 => 8, + } + } + + fn read(self, b: &[u8]) -> i64 { + match self { + OffsetElem::I32 => i32::from_le_bytes(b.try_into().unwrap()) as i64, + OffsetElem::U32 => u32::from_le_bytes(b.try_into().unwrap()) as i64, + OffsetElem::I64 => i64::from_le_bytes(b.try_into().unwrap()), + OffsetElem::U64 => u64::from_le_bytes(b.try_into().unwrap()) as i64, + OffsetElem::F32 => f32::from_le_bytes(b.try_into().unwrap()) as i64, + OffsetElem::F64 => f64::from_le_bytes(b.try_into().unwrap()) as i64, + } + } +} + +/// Extract the 1-D `(len,)` shape from a `.npy` header dict. +fn extract_1d_shape(header: &str) -> Result { + 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() != 1 { + bail!("expected a 1-D .npy array for offsets, got shape {dims:?}"); + } + Ok(dims[0]) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::dataset::fixtures::make_ramp_npy; + + fn make_offsets_npy(offsets: &[i64]) -> Vec { + make_offsets_npy_dtype(offsets, " Vec { + make_offsets_npy_dtype(offsets, " Vec, + ) -> Vec { + let mut header = format!( + "{{'descr': '{descr}', 'fortran_order': False, 'shape': ({},), }}", + offsets.len() + ); + while (10 + header.len() + 1) % 64 != 0 { + header.push(' '); + } + header.push('\n'); + let mut buf = Vec::new(); + buf.extend_from_slice(b"\x93NUMPY"); + buf.push(1); + buf.push(0); + buf.extend_from_slice(&(header.len() as u16).to_le_bytes()); + buf.extend_from_slice(header.as_bytes()); + for o in offsets { + buf.extend_from_slice(&encode(o)); + } + buf + } + + #[test] + fn reads_ragged_multivectors() { + let dir = tempfile::tempdir().unwrap(); + // 5 sub-vectors total, dim 3: rows 0..5. + std::fs::write(dir.path().join("vectors.npy"), make_ramp_npy(0, 5, 3)).unwrap(); + // Point 0 -> rows [0,2), point 1 -> [2,2) (empty), point 2 -> [2,5). + std::fs::write( + dir.path().join("offsets.npy"), + make_offsets_npy(&[0, 2, 2, 5]), + ) + .unwrap(); + + let reader = MultivectorReader::open(dir.path()).unwrap(); + assert_eq!(reader.num_points(), 3); + assert_eq!( + reader.vector_at(0).unwrap(), + vec![vec![0.0, 1.0, 2.0], vec![3.0, 4.0, 5.0]] + ); + assert!(reader.vector_at(1).unwrap().is_empty()); + assert_eq!(reader.vector_at(2).unwrap().len(), 3); + } + + /// Some exporters (e.g. numpy's default float dtype) write offsets as + /// floats; they must be downcast to i64, not rejected. + #[test] + fn reads_float_offsets() { + let dir = tempfile::tempdir().unwrap(); + std::fs::write(dir.path().join("vectors.npy"), make_ramp_npy(0, 5, 3)).unwrap(); + std::fs::write( + dir.path().join("offsets.npy"), + make_f32_offsets_npy(&[0, 2, 2, 5]), + ) + .unwrap(); + + let reader = MultivectorReader::open(dir.path()).unwrap(); + assert_eq!(reader.num_points(), 3); + assert_eq!( + reader.vector_at(0).unwrap(), + vec![vec![0.0, 1.0, 2.0], vec![3.0, 4.0, 5.0]] + ); + } + + #[test] + fn rejects_offsets_mismatched_with_vector_count() { + let dir = tempfile::tempdir().unwrap(); + std::fs::write(dir.path().join("vectors.npy"), make_ramp_npy(0, 5, 3)).unwrap(); + std::fs::write(dir.path().join("offsets.npy"), make_offsets_npy(&[0, 2, 4])).unwrap(); + + let err = match MultivectorReader::open(dir.path()) { + Ok(_) => panic!("expected an error"), + Err(e) => e.to_string(), + }; + assert!(err.contains("does not match"), "{err}"); + } + + #[test] + fn rejects_out_of_range_index() { + let dir = tempfile::tempdir().unwrap(); + std::fs::write(dir.path().join("vectors.npy"), make_ramp_npy(0, 5, 3)).unwrap(); + std::fs::write( + dir.path().join("offsets.npy"), + make_offsets_npy(&[0, 2, 2, 5]), + ) + .unwrap(); + + let reader = MultivectorReader::open(dir.path()).unwrap(); + assert!(reader.vector_at(3).is_err()); + } +} diff --git a/src/dataset/readers/npy.rs b/src/dataset/readers/npy.rs index a61a4bd..d614fcc 100644 --- a/src/dataset/readers/npy.rs +++ b/src/dataset/readers/npy.rs @@ -132,6 +132,35 @@ impl NpyReader { /// — which is what makes remote row counts a single ranged request rather than /// a download. pub fn parse_npy_header(buf: &[u8]) -> Result { + let (header, header_end) = parse_npy_header_str(buf)?; + + 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, + }) +} + +/// Parse the leading `.npy` magic/header framing, returning the header dict +/// string and the byte offset where the array data starts. Shared by +/// [`parse_npy_header`] (2-D float arrays) and the `offsets.npy` reader for +/// multivector datasets (1-D int arrays), which parse the returned dict +/// differently. +pub(crate) fn parse_npy_header_str(buf: &[u8]) -> Result<(&str, usize)> { if buf.len() < 10 || &buf[0..6] != b"\x93NUMPY" { bail!("not a .npy file (bad magic)"); } @@ -160,30 +189,11 @@ pub fn parse_npy_header(buf: &[u8]) -> Result { } 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, - }) + Ok((header, header_end)) } /// Extract a single-quoted string value for `key` from a `.npy` header dict. -fn extract_quoted(header: &str, key: &str) -> Option { +pub(crate) 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(); diff --git a/src/dataset/sources.rs b/src/dataset/sources.rs index f48e6d1..ccb8229 100644 --- a/src/dataset/sources.rs +++ b/src/dataset/sources.rs @@ -113,6 +113,16 @@ impl UploadDatasetSources { ) } + /// A point's sub-vectors from a `multivector` dataset (ColBERT-style). + pub fn multi_dense_vector(&self, slot: usize, idx: u64) -> Option>> { + let reader = self.vector.get(slot)?.as_ref()?; + Some( + reader + .multi_dense_vector(idx as usize) + .unwrap_or_else(|e| panic!("failed to read multivector at {idx}: {e}")), + ) + } + pub fn payload_value(&self, slot: usize, idx: u64) -> Option { let (reader, field) = self.payload.get(slot)?.as_ref()?; let value = reader diff --git a/src/generators/config.rs b/src/generators/config.rs index c425cec..a8948ee 100644 --- a/src/generators/config.rs +++ b/src/generators/config.rs @@ -166,6 +166,11 @@ impl ConfigGenerator { fn gen_dense(&self, i: usize, vc: &VectorConfig, idx: u64, rng: &mut impl Rng) -> Vector { let reader = &self.readers[i]; if let Some(mv) = &vc.multivector { + if let VectorSource::Dataset { .. } = &vc.source + && let Some(multi) = self.datasets.multi_dense_vector(i, idx) + { + return Vector::new_multi(multi); + } let multi: Vec<_> = (0..mv.count) .map(|_| self.gen_one_vector(vc, reader, i, idx, rng)) .collect(); @@ -581,6 +586,80 @@ collection: } } + /// A `format: multivector` dataset source must read each point's real, + /// ragged sub-vectors from `vectors.npy`/`offsets.npy`, not repeat one row + /// `multivector.count` times. + #[test] + fn reads_multivectors_from_a_dataset() { + use crate::dataset::fixtures::make_ramp_npy; + + let dir = tempfile::tempdir().unwrap(); + let mv_dir = dir.path().join("colbert"); + std::fs::create_dir(&mv_dir).unwrap(); + // 5 sub-vectors total, dim 4: point 0 -> rows [0,2), point 1 -> [2,5). + std::fs::write(mv_dir.join("vectors.npy"), make_ramp_npy(0, 5, 4)).unwrap(); + std::fs::write(mv_dir.join("offsets.npy"), make_offsets_npy(&[0i64, 2, 5])).unwrap(); + + let config: UploadConfig = serde_yaml::from_str( + " +collection: + name: t + vectors: + - name: m + size: 4 + multivector: { count: 1 } + source: + type: dataset + name: colbert + format: multivector + path: colbert +", + ) + .unwrap(); + config.validate().unwrap(); + + let generator = ConfigGenerator::new_with_datasets_dir(&config, dir.path()).unwrap(); + + let point0 = generator.make_point(0); + match named(&point0)["m"].vector.as_ref().unwrap() { + qdrant_client::qdrant::vector::Vector::MultiDense(m) => { + assert_eq!(m.vectors.len(), 2, "point 0 has 2 sub-vectors, not `count`"); + assert_eq!(m.vectors[0].data, vec![0.0, 1.0, 2.0, 3.0]); + assert_eq!(m.vectors[1].data, vec![4.0, 5.0, 6.0, 7.0]); + } + _ => panic!("expected multidense vector"), + } + + let point1 = generator.make_point(1); + match named(&point1)["m"].vector.as_ref().unwrap() { + qdrant_client::qdrant::vector::Vector::MultiDense(m) => { + assert_eq!(m.vectors.len(), 3, "point 1 has 3 sub-vectors"); + } + _ => panic!("expected multidense vector"), + } + } + + fn make_offsets_npy(offsets: &[i64]) -> Vec { + let mut header = format!( + "{{'descr': '