feat: support batch vector queries

python/python/benchmarks/test_search.py

@@ -13,6 +13,11 @@
 N_DIMS = 768
 NUM_ROWS = 100_000
 NEW_ROWS = 10_000
+BATCH_FLAT_KNN_DIM = 512
+BATCH_FLAT_KNN_K = 10
+BATCH_FLAT_KNN_BATCH_SIZE = 10_000
+BATCH_FLAT_KNN_QUERY_COUNT = 10
+BATCH_FLAT_KNN_ROWS = 1_000_000
 
 
 def find_or_clean(dataset_path: Path) -> Union[lance.LanceDataset, None]:
@@ -64,6 +69,42 @@ def create_table(num_rows, offset) -> pa.Table:
     )
 
 
+def create_flat_vector_table(num_rows: int, offset: int, dim: int) -> pa.Table:
+    rng = np.random.default_rng(seed=offset)
+    values = rng.random((num_rows, dim), dtype=np.float32)
+    vectors = pa.FixedSizeListArray.from_arrays(pa.array(values.ravel()), dim)
+    ids = pa.array(range(offset, offset + num_rows))
+    return pa.table({"vector": vectors, "id": ids})
+
+
+def create_batch_flat_knn_dataset(
+    data_dir: Path, num_rows: int, batch_size: int, dim: int
+) -> lance.LanceDataset:
+    tmp_path = data_dir / f"batch_flat_knn_{num_rows}_{batch_size}_{dim}"
+    dataset = find_or_clean(tmp_path)
+    if dataset:
+        return dataset
+
+    rows_remaining = num_rows
+    offset = 0
+    dataset = None
+    while rows_remaining > 0:
+        next_batch_length = min(rows_remaining, batch_size)
+        rows_remaining -= next_batch_length
+        table = create_flat_vector_table(next_batch_length, offset, dim)
+        if offset == 0:
+            dataset = lance.write_dataset(
+                table, tmp_path, data_storage_version="stable"
+            )
+        else:
+            dataset = lance.write_dataset(
+                table, tmp_path, mode="append", data_storage_version="stable"
+            )
+        offset += next_batch_length
+
+    return dataset
+
+
 def create_base_dataset(data_dir: Path) -> lance.LanceDataset:
     tmp_path = data_dir / "search_dataset"
     dataset = find_or_clean(tmp_path)
@@ -173,6 +214,70 @@ def test_knn_search(test_dataset, benchmark):
     assert result.num_rows > 0
 
 
+@pytest.mark.benchmark(group="batch_flat_knn")
+@pytest.mark.parametrize("mode", ["separate", "batch"])
+@pytest.mark.parametrize(
+    ("dim", "num_rows", "batch_size", "query_count", "rounds"),
+    [
+        (
+            BATCH_FLAT_KNN_DIM,
+            BATCH_FLAT_KNN_ROWS,
+            BATCH_FLAT_KNN_BATCH_SIZE,
+            BATCH_FLAT_KNN_QUERY_COUNT,
+            10,
+        )
+    ],
+    ids=["1m_rows_512d_m10"],
+)
+def test_batch_flat_knn(
+    data_dir: Path,
+    benchmark,
+    mode: str,
+    dim: int,
+    num_rows: int,
+    batch_size: int,
+    query_count: int,
+    rounds: int,
+):
+    dataset = create_batch_flat_knn_dataset(data_dir, num_rows, batch_size, dim)
+    query_table = dataset.to_table(columns=["vector"], limit=query_count)
+    query_values = np.asarray(
+        query_table["vector"].combine_chunks().values, dtype=np.float32
+    ).reshape(query_count, dim)
+
+    def separate_queries():
+        total_rows = 0
+        for query in query_values:
+            total_rows += dataset.to_table(
+                columns=[],
+                nearest={
+                    "column": "vector",
+                    "q": query,
+                    "k": BATCH_FLAT_KNN_K,
+                    "use_index": False,
+                },
+            ).num_rows
+        return total_rows
+
+    def batch_query():
+        return dataset.to_table(
+            columns=[],
+            nearest={
+                "column": "vector",
+                "q": query_values,
+                "k": BATCH_FLAT_KNN_K,
+                "use_index": False,
+            },
+        ).num_rows
+
+    if mode == "separate":
+        result = benchmark.pedantic(separate_queries, rounds=rounds, iterations=1)
+    else:
+        result = benchmark.pedantic(batch_query, rounds=rounds, iterations=1)
+
+    assert result == query_count * BATCH_FLAT_KNN_K
+
+
 @pytest.mark.benchmark(group="query_ann")
 def test_ann_no_refine(test_dataset, benchmark):
     q = pc.random(N_DIMS).cast(pa.float32())

python/python/lance/dataset.py

@@ -92,7 +92,7 @@
         pa.Array,
         pa.Scalar,
         np.ndarray,
-        Iterable[float],
+        Iterable[Union[float, Iterable[float]]],
     ]
 LANCE_COMMIT_MESSAGE_KEY = "__lance_commit_message"
 _BLOB_PANDAS_MODE_LAZY = "lazy"
@@ -1099,6 +1099,13 @@ def scanner(
                     "distance_range": (0.0, 1.0),
                 }
 
+            ``q`` may also be a 2-D array-like value for fixed-size vector columns.
+            In that case Lance runs a batch nearest-neighbor query, returns up to
+            ``k`` rows for each query vector, and adds ``query_index`` to identify the
+            source query for each result row. When ``use_index`` is true and a vector
+            index is available, each query vector is searched through the index
+            path; otherwise the flat batch path is used.
+
         batch_size: int, default None
             The maximum number of rows per batch.  In some cases batches can be
             smaller than this size.  Note: this can be overridden by
@@ -5989,6 +5996,12 @@ def nearest(
 
         Parameters
         ----------
+        q: QueryVectorLike
+            A single query vector or, for fixed-size vector columns, a 2-D array-like
+            batch of query vectors. Batch queries return up to ``k`` rows per query
+            and include ``query_index`` in the output. When ``use_index`` is true and
+            a vector index is available, each query vector is searched through the
+            index path; otherwise the flat batch path is used.
         query_parallelism: int, optional
             Maximum partition-search concurrency for a single vector query.
             The default is 0. Value 0 uses the automatic policy, which
@@ -7137,7 +7150,11 @@ def _build_vector_search_query(
     column: str
         The name of the vector column to search.
     q: QueryVectorLike
-        The query vector.
+        The query vector. For fixed-size vector columns, this may be a 2-D
+        array-like batch of query vectors. Batch queries return up to ``k`` rows per
+        query vector and include ``query_index`` in the output. When ``use_index``
+        is true and a vector index is available, each query vector is searched
+        through the index path; otherwise the flat batch path is used.
     k: int, optional
         The number of nearest neighbors to return.
     metric: str, optional

python/python/tests/test_vector_index.py

@@ -180,6 +180,108 @@ def test_flat(dataset):
     run(dataset)
 
 
+@pytest.mark.parametrize(
+    "queries",
+    [
+        np.random.randn(2, 128).astype(np.float32),
+        np.random.randn(1, 128).astype(np.float32),
+    ],
+    ids=["two_queries", "single_query"],
+)
+def test_batch_flat_query_matches_repeated_single_queries(dataset, queries):
+    k = 5
+    query_count = queries.shape[0]
+
+    batch = dataset.to_table(
+        columns=["id"],
+        nearest={
+            "column": "vector",
+            "q": queries,
+            "k": k,
+            "use_index": False,
+        },
+    )
+
+    assert batch.num_rows == query_count * k
+    assert batch.column_names == ["id", "_distance", "query_index"]
+    expected_query_index = sum([[i] * k for i in range(query_count)], [])
+    assert batch["query_index"].to_pylist() == expected_query_index
+
+    _assert_batch_matches_single_queries(
+        dataset,
+        queries,
+        k=k,
+        nearest_kwargs={"use_index": False},
+    )
+
+
+def _assert_batch_matches_single_queries(ds, queries, k, nearest_kwargs):
+    batch = ds.to_table(
+        columns=["id"],
+        nearest={
+            "column": "vector",
+            "q": queries,
+            "k": k,
+            **nearest_kwargs,
+        },
+    )
+    if "distance_range" in nearest_kwargs:
+        lo, hi = nearest_kwargs["distance_range"]
+        assert all(lo <= d < hi for d in batch["_distance"].to_pylist())
+
+    for query_index, query in enumerate(queries):
+        single = ds.to_table(
+            columns=["id"],
+            nearest={
+                "column": "vector",
+                "q": query,
+                "k": k,
+                **nearest_kwargs,
+            },
+        )
+        batch_slice = batch.filter(pc.field("query_index") == query_index)
+        assert batch_slice["id"].to_pylist() == single["id"].to_pylist()
+        np.testing.assert_allclose(
+            batch_slice["_distance"].to_numpy(),
+            single["_distance"].to_numpy(),
+        )
+
+
+def test_batch_flat_respects_distance_range(dataset):
+    queries = np.random.randn(2, 128).astype(np.float32)
+    _assert_batch_matches_single_queries(
+        dataset,
+        queries,
+        k=5,
+        nearest_kwargs={"use_index": False, "distance_range": (0.0, 50.0)},
+    )
+
+
+def test_batch_indexed_respects_distance_range(indexed_dataset):
+    queries = np.random.randn(2, 128).astype(np.float32)
+    _assert_batch_matches_single_queries(
+        indexed_dataset,
+        queries,
+        k=5,
+        nearest_kwargs={"distance_range": (0.0, 50.0)},
+    )
+
+
+def test_batch_fast_search_without_index_returns_empty_with_query_index(dataset):
+    queries = np.random.randn(2, 128).astype(np.float32)
+    batch = dataset.to_table(
+        columns=["id"],
+        nearest={
+            "column": "vector",
+            "q": queries,
+            "k": 5,
+        },
+        fast_search=True,
+    )
+    assert batch.num_rows == 0
+    assert "query_index" in batch.column_names
+
+
 def test_ann(indexed_dataset):
     run(indexed_dataset)
 

python/src/dataset.rs

@@ -1368,17 +1368,26 @@ impl Dataset {
                 query_parallelism,
             ) = vector_query_params_from_dict(nearest, default_k)?;
 
-            let (_, element_type) = get_vector_type(self_.ds.schema(), &column)
+            let (vector_type, element_type) = get_vector_type(self_.ds.schema(), &column)
                 .map_err(|e| PyValueError::new_err(e.to_string()))?;
-            let scanner = match element_type {
-                DataType::UInt8 => {
+            let is_batch_query = matches!(
+                q.data_type(),
+                DataType::List(_) | DataType::FixedSizeList(_, _)
+            ) && matches!(vector_type, DataType::FixedSizeList(_, _));
+            let scanner = match (is_batch_query, element_type) {
+                (true, DataType::UInt8) => {
+                    return Err(PyValueError::new_err(
+                        "Batch nearest is not supported for binary vector columns",
+                    ));
+                }
+                (false, DataType::UInt8) => {
                     let q = arrow::compute::cast(&q, &DataType::UInt8).map_err(|e| {
                         PyValueError::new_err(format!("Failed to cast q to binary vector: {}", e))
                     })?;
                     let q = q.as_primitive::<UInt8Type>();
                     scanner.nearest(&column, q, k)
                 }
-                _ => scanner.nearest(&column, &q, k),
+                (_, _) => scanner.nearest(&column, &q, k),
             };
             let distance_range: Option<(Option<f32>, Option<f32>)> =
                 if let Some(dr) = nearest.get_item("distance_range")? {