feat: add support for colbert-style multivector datasets

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:`

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:

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 {

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<Vec<Vec<f32>>> {
+        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<Option<Value>> {
         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"),
         }
     }
 }

src/dataset/readers/mod.rs

@@ -1,13 +1,15 @@
 mod binary;
 mod h5;
 mod jsonl;
+mod multivector;
 mod npy;
 mod parquet;
 mod query;
 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,