searchVectors accepts a source array for multi-file query

Widen the source / metadata / binary options to accept either a single
value or a same-length array. Each source is searched in parallel and
the global top-K is heap-merged. When more than one source is supplied,
each result carries a sourceIndex pointing back at its origin.

  // single source — unchanged
  searchVectors({ source: 'logs/today.parquet', query, topK })

  // multi-source — new
  searchVectors({
    source: dailyFiles,
    metadata: parsedMetas,  // optional, same length
    binary: prefetchedBins, // optional, same length
    query,
    topK,
  })

Single-source callers see no behavioral change. Length-mismatched
metadata or binary arrays throw. All sources must share the same metric
direction (mixed cosine/dot vs euclidean is rejected at merge time).

This is the foundation for sharded log queries — one parquet per day /
session / tenant, queried opportunistically as a group.

Author: platypii

src/searchVectors.js

@@ -6,12 +6,19 @@ import { l2Normalize, parseKvMetadata } from './utils.js'
 
 /**
  * @import { SearchResult, SearchVectorsOptions } from './types.js'
+ * @import { AsyncBuffer, FileMetaData } from 'hyparquet'
  */
 
 /**
  * Find the top-k nearest neighbors to a query vector.
  *
- * Three paths, in order of preference:
+ * `source` may be a single URL/path/AsyncBuffer or an array of them. With
+ * an array, each file is searched in parallel and results are heap-merged
+ * to the global top-K. Per-source `metadata` and `binary` (from
+ * `prefetchBinary`) may be passed as same-length arrays; pass `undefined`
+ * in any slot to fall back to the default behavior for that source.
+ *
+ * Three paths are picked per source, in order of preference:
  *   - Clustered + binary + rerank (file has centroids): phase 1 scans only
  *     the top-N nearest clusters' row ranges (skipping whole row groups),
  *     phase 2 fetches the candidate float32 vectors and reranks.
@@ -26,7 +33,7 @@ import { l2Normalize, parseKvMetadata } from './utils.js'
 export async function searchVectors({
   query,
   source,
-  metadata: providedMetadata,
+  metadata,
   topK = 10,
   metric,
   rerankFactor = 10,
@@ -39,6 +46,74 @@ export async function searchVectors({
   if (source === undefined || source === null) {
     throw new Error('searchVectors: `source` is required (URL, file path, or AsyncBuffer)')
   }
+  const sources = Array.isArray(source) ? source : [source]
+  if (sources.length === 0) {
+    throw new Error('searchVectors: `source` array is empty')
+  }
+  if (Array.isArray(metadata) && metadata.length !== sources.length) {
+    throw new Error(`searchVectors: \`metadata\` array length ${metadata.length} does not match \`source\` array length ${sources.length}`)
+  }
+  if (Array.isArray(binary) && binary.length !== sources.length) {
+    throw new Error(`searchVectors: \`binary\` array length ${binary.length} does not match \`source\` array length ${sources.length}`)
+  }
+  const metadatas = Array.isArray(metadata) ? metadata : sources.map(() => /** @type {FileMetaData | undefined} */ (metadata))
+  const binaries = Array.isArray(binary) ? binary : sources.map(() => /** @type {Uint8Array | undefined} */ (binary))
+
+  const multi = sources.length > 1
+  const perSource = await Promise.all(sources.map((src, i) => searchOne({
+    query,
+    source: src,
+    metadata: metadatas[i],
+    binary: binaries[i],
+    topK,
+    metric,
+    rerankFactor,
+    probe,
+    signal,
+    asyncBufferFactory,
+    compressors,
+    sourceIndex: multi ? i : undefined,
+  })))
+
+  if (!multi) return perSource[0].results
+
+  // Merge: each per-source array is already sorted best-first under the same
+  // direction (we assert below). Flatten, sort, slice — N is small.
+  const direction = perSource[0].direction
+  for (let i = 1; i < perSource.length; i += 1) {
+    if (perSource[i].direction !== direction) {
+      throw new Error('searchVectors: sources have inconsistent metric directions')
+    }
+  }
+  const merged = perSource.flatMap(p => p.results)
+  merged.sort((a, b) => direction * (a.score - b.score))
+  return merged.slice(0, topK)
+}
+
+/**
+ * Search a single source. Returns the sorted results plus the score
+ * direction (1 = ascending / lower-is-better for euclidean, -1 = descending
+ * for cosine/dot). The direction lets the caller merge cross-source results
+ * without re-parsing metadata.
+ *
+ * @param {object} opts
+ * @param {Float32Array | number[]} opts.query
+ * @param {string | AsyncBuffer} opts.source
+ * @param {FileMetaData} [opts.metadata]
+ * @param {Uint8Array} [opts.binary]
+ * @param {number} opts.topK
+ * @param {import('./types.js').DistanceMetric} [opts.metric]
+ * @param {number} opts.rerankFactor
+ * @param {number} [opts.probe]
+ * @param {AbortSignal} [opts.signal]
+ * @param {(options: { source: string, signal?: AbortSignal }) => Promise<AsyncBuffer>} [opts.asyncBufferFactory]
+ * @param {import('hyparquet').Compressors} [opts.compressors]
+ * @param {number} [opts.sourceIndex]
+ * @returns {Promise<{ results: SearchResult[], direction: 1 | -1 }>}
+ */
+async function searchOne({
+  query, source, metadata: providedMetadata, binary, topK, metric, rerankFactor, probe, signal, asyncBufferFactory, compressors, sourceIndex,
+}) {
   const file = typeof source === 'string'
     ? await (asyncBufferFactory ?? defaultAsyncBufferFactory)({ source, signal })
     : source
@@ -58,12 +133,18 @@ export async function searchVectors({
     scoringMetric = 'dot'
   }
 
-  if (meta.hasBinary && rerankFactor > 0) {
-    return searchRerank({
+  const results = meta.hasBinary && rerankFactor > 0
+    ? await searchRerank({
       file, metadata, meta, queryF32, scoringMetric, reportedMetric: requestedMetric, topK, rerankFactor, probe, binary, compressors,
     })
+    : await searchExact({
+      file, metadata, meta, queryF32, scoringMetric, reportedMetric: requestedMetric, topK, compressors,
+    })
+
+  if (sourceIndex !== undefined) {
+    for (const r of results) r.sourceIndex = sourceIndex
   }
-  return searchExact({
-    file, metadata, meta, queryF32, scoringMetric, reportedMetric: requestedMetric, topK, compressors,
-  })
+
+  const direction = requestedMetric === 'euclidean' ? 1 : -1
+  return { results, direction }
 }

src/types.d.ts

@@ -34,18 +34,34 @@ export interface ReadVectorsOptions {
 
 export interface SearchVectorsOptions {
   query: Float32Array | number[] // the query vector
-  source: string | AsyncBuffer // URL, file path, or an already-opened AsyncBuffer
-  metadata?: FileMetaData // optional pre-parsed parquet metadata; skips the footer fetch on every call. Reuse across queries for best throughput.
+
+  /**
+   * One source or an array of sources. Each source is a URL, file path, or
+   * an already-opened AsyncBuffer. With an array, each file is searched in
+   * parallel and the global top-K is returned with `result.sourceIndex` set.
+   * All sources must share the same dimension and metric.
+   */
+  source: string | AsyncBuffer | Array<string | AsyncBuffer>
+
+  /**
+   * Pre-parsed parquet metadata; skips the footer fetch on every call.
+   * Reuse across queries for best throughput. When `source` is an array,
+   * pass a same-length array of (FileMetaData | undefined) — undefined
+   * slots fall back to fetching the footer for that source.
+   */
+  metadata?: FileMetaData | Array<FileMetaData | undefined>
+
   topK?: number // number of nearest neighbors to return (default: 10)
   metric?: DistanceMetric // override the stored metric
 
   /**
    * Pre-fetched binary column from `prefetchBinary`. When provided, phase 1
    * Hamming scan runs entirely from memory and the binary parquet fetches
    * are skipped. The buffer must be `count × dim/8` bytes in the same row
-   * order as the file. Reuse across queries.
+   * order as the file. Reuse across queries. When `source` is an array,
+   * pass a same-length array of (Uint8Array | undefined).
    */
-  binary?: Uint8Array
+  binary?: Uint8Array | Array<Uint8Array | undefined>
 
   /**
    * When the file has a binary column, controls two-phase search:
@@ -87,6 +103,7 @@ export interface SearchResult {
   id: string | number
   score: number // similarity score (higher = better for cosine/dot, lower = better for euclidean)
   rowIndex: number // original row index in the source parquet
+  sourceIndex?: number // index into the `source` array; set only when multi-source search was used
   metadata?: Record<string, any>
 }
 

test/searchVectors.test.js

@@ -112,6 +112,61 @@ describe('searchVectors', () => {
     }
   })
 
+  it('accepts a single-element source array as equivalent to a bare source', async () => {
+    const dimension = 32
+    const vectors = makeVectors(80, dimension, 5)
+    const writer = fileWriter(TEST_FILE)
+    await writeVectors({ writer, vectors, dimension })
+
+    const query = vectors[12].vector
+    const bare = await searchVectors({ source: TEST_FILE, query, topK: 3 })
+    const wrapped = await searchVectors({ source: [TEST_FILE], query, topK: 3 })
+    expect(wrapped.map(r => r.id)).toEqual(bare.map(r => r.id))
+    // sourceIndex is only attached when truly multi-source (length > 1) to keep
+    // results identical between the bare and array-of-one forms.
+    expect(bare[0].sourceIndex).toBeUndefined()
+    expect(wrapped[0].sourceIndex).toBeUndefined()
+  })
+
+  it('merges results across multiple sources matching a single-file search', async () => {
+    const dimension = 32
+    const fileA = 'test/files/search-a.parquet'
+    const fileB = 'test/files/search-b.parquet'
+    const all = makeVectors(100, dimension, 9)
+    const halfA = all.slice(0, 50)
+    const halfB = all.slice(50)
+
+    await writeVectors({ writer: fileWriter(TEST_FILE), vectors: all, dimension, normalize: true })
+    await writeVectors({ writer: fileWriter(fileA), vectors: halfA, dimension, normalize: true })
+    await writeVectors({ writer: fileWriter(fileB), vectors: halfB, dimension, normalize: true })
+
+    try {
+      const query = all[37].vector
+      const combined = await searchVectors({ source: TEST_FILE, query, topK: 5 })
+      const split = await searchVectors({ source: [fileA, fileB], query, topK: 5 })
+
+      // Same IDs in the same order — vectors[37] lands in fileA so the top hit must be from source 0.
+      expect(split.map(r => r.id)).toEqual(combined.map(r => r.id))
+      expect(split[0].sourceIndex).toBe(0)
+      for (let i = 0; i < combined.length; i += 1) {
+        expect(split[i].score).toBeCloseTo(combined[i].score, 5)
+      }
+    } finally {
+      if (existsSync(fileA)) unlinkSync(fileA)
+      if (existsSync(fileB)) unlinkSync(fileB)
+    }
+  })
+
+  it('rejects metadata/binary arrays whose length does not match the source array', async () => {
+    const dimension = 16
+    const vectors = makeVectors(20, dimension, 1)
+    await writeVectors({ writer: fileWriter(TEST_FILE), vectors, dimension })
+
+    await expect(searchVectors({
+      source: [TEST_FILE, TEST_FILE], query: vectors[0].vector, topK: 1, metadata: [undefined],
+    })).rejects.toThrow(/metadata.*length/)
+  })
+
   it('produces the same top-1 with and without rerank', async () => {
     const dimension = 64
     const vectors = makeVectors(200, dimension, 23)