feat: support multivector type (#2005)

Signed-off-by: BubbleCal <bubble-cal@outlook.com>

docs/src/search.md

@@ -138,6 +138,36 @@ LanceDB supports binary vectors as a data type, and has the ability to search bi
         --8<-- "python/python/tests/docs/test_binary_vector.py:async_binary_vector"
         ```
 
+## Multivector type
+
+LanceDB supports multivector type, this is useful when you have multiple vectors for a single item (e.g. with ColBert and ColPali).
+
+You can index on a column with multivector type and search on it, the query can be single vector or multiple vectors. If the query is multiple vectors `mq`, the similarity (distance) from it to any multivector `mv` in the dataset, is defined as:
+
+![maxsim](assets/maxsim.png)
+
+where `sim` is the similarity function (e.g. cosine).
+
+For now, only `cosine` metric is supported for multivector search.
+
+=== "Python"
+
+    === "sync API"
+
+        ```python
+        --8<-- "python/python/tests/docs/test_multivector.py:imports"
+
+        --8<-- "python/python/tests/docs/test_multivector.py:sync_multivector"
+        ```
+
+    === "async API"
+
+        ```python
+        --8<-- "python/python/tests/docs/test_multivector.py:imports"
+
+        --8<-- "python/python/tests/docs/test_multivector.py:async_multivector"
+        ```
+
 ## Search with distance range
 
 You can also search for vectors within a specific distance range from the query vector. This is useful when you want to find vectors that are not just the nearest neighbors, but also those that are within a certain distance. This can be done by using the `distance_range` method.

docs/src/sql.md

@@ -18,7 +18,7 @@ import numpy as np
 
 uri = "data/sample-lancedb"
 data = [{"vector": row, "item": f"item {i}", "id": i}
-    for i, row in enumerate(np.random.random((10_000, 2)).astype('int'))]
+    for i, row in enumerate(np.random.random((10_000, 2)))]
 
 # Synchronous client
 db = lancedb.connect(uri)

python/python/lancedb/query.py

@@ -1741,12 +1741,14 @@ class AsyncQuery(AsyncQueryBase):
         a default `limit` of 10 will be used.
 
         Typically, a single vector is passed in as the query. However, you can also
-        pass in multiple vectors.  This can be useful if you want to find the nearest
-        vectors to multiple query vectors. This is not expected to be faster than
-        making multiple queries concurrently; it is just a convenience method.
-        If multiple vectors are passed in then an additional column `query_index`
-        will be added to the results.  This column will contain the index of the
-        query vector that the result is nearest to.
+        pass in multiple vectors. When multiple vectors are passed in, if the vector
+        column is with multivector type, then the vectors will be treated as a single
+        query. Or the vectors will be treated as multiple queries, this can be useful
+        if you want to find the nearest vectors to multiple query vectors.
+        This is not expected to be faster than making multiple queries concurrently;
+        it is just a convenience method. If multiple vectors are passed in then
+        an additional column `query_index` will be added to the results. This column
+        will contain the index of the query vector that the result is nearest to.
         """
         if query_vector is None:
             raise ValueError("query_vector can not be None")

python/python/lancedb/table.py

@@ -2856,6 +2856,8 @@ class AsyncTable:
             async_query = async_query.with_row_id()
 
         if query.vector:
+            # we need the schema to get the vector column type
+            # to determine whether the vectors is batch queries or not
             async_query = (
                 async_query.nearest_to(query.vector)
                 .distance_type(query.metric)

python/python/lancedb/util.py

@@ -223,7 +223,7 @@ def inf_vector_column_query(schema: pa.Schema) -> str:
     vector_col_count = 0
     for field_name in schema.names:
         field = schema.field(field_name)
-        if pa.types.is_fixed_size_list(field.type):
+        if is_vector_column(field.type):
             vector_col_count += 1
             if vector_col_count > 1:
                 raise ValueError(
@@ -231,7 +231,6 @@ def inf_vector_column_query(schema: pa.Schema) -> str:
                     "Please specify the vector column name "
                     "for vector search"
                 )
-                break
             elif vector_col_count == 1:
                 vector_col_name = field_name
     if vector_col_count == 0:
@@ -242,6 +241,29 @@ def inf_vector_column_query(schema: pa.Schema) -> str:
     return vector_col_name
 
 
+def is_vector_column(data_type: pa.DataType) -> bool:
+    """
+    Check if the column is a vector column.
+
+    Parameters
+    ----------
+    data_type : pa.DataType
+        The data type of the column.
+
+    Returns
+    -------
+    bool: True if the column is a vector column.
+    """
+    if pa.types.is_fixed_size_list(data_type) and (
+        pa.types.is_floating(data_type.value_type)
+        or pa.types.is_uint8(data_type.value_type)
+    ):
+        return True
+    elif pa.types.is_list(data_type):
+        return is_vector_column(data_type.value_type)
+    return False
+
+
 def infer_vector_column_name(
     schema: pa.Schema,
     query_type: str,

python/python/tests/test_query.py

@@ -68,6 +68,60 @@ async def table_struct_async(tmp_path) -> AsyncTable:
     return await conn.create_table("test_struct", table)
 
 
+@pytest.fixture
+def multivec_table() -> lancedb.table.Table:
+    db = lancedb.connect("memory://")
+    # Generate 256 rows of data
+    num_rows = 256
+
+    # Generate data for each column
+    vector_data = [
+        [[i, i + 1], [i + 2, i + 3]] for i in range(num_rows)
+    ]  # Adjust to match nested structure
+    id_data = list(range(1, num_rows + 1))
+    float_field_data = [float(i) for i in range(1, num_rows + 1)]
+
+    # Create the Arrow table
+    df = pa.table(
+        {
+            "vector": pa.array(
+                vector_data, type=pa.list_(pa.list_(pa.float32(), list_size=2))
+            ),
+            "id": pa.array(id_data),
+            "float_field": pa.array(float_field_data),
+        }
+    )
+    return db.create_table("test", df)
+
+
+@pytest_asyncio.fixture
+async def multivec_table_async(tmp_path) -> AsyncTable:
+    conn = await lancedb.connect_async(
+        "memory://", read_consistency_interval=timedelta(seconds=0)
+    )
+    # Generate 256 rows of data
+    num_rows = 256
+
+    # Generate data for each column
+    vector_data = [
+        [[i, i + 1], [i + 2, i + 3]] for i in range(num_rows)
+    ]  # Adjust to match nested structure
+    id_data = list(range(1, num_rows + 1))
+    float_field_data = [float(i) for i in range(1, num_rows + 1)]
+
+    # Create the Arrow table
+    df = pa.table(
+        {
+            "vector": pa.array(
+                vector_data, type=pa.list_(pa.list_(pa.float32(), list_size=2))
+            ),
+            "id": pa.array(id_data),
+            "float_field": pa.array(float_field_data),
+        }
+    )
+    return await conn.create_table("test_async", df)
+
+
 def test_cast(table):
     class TestModel(LanceModel):
         vector: Vector(2)
@@ -177,6 +231,62 @@ async def test_distance_range_async(table_async: AsyncTable):
     assert res["_distance"].to_pylist() == [min_dist, max_dist]
 
 
+def test_multivector(multivec_table: lancedb.table.Table):
+    # create index on multivector
+    multivec_table.create_index(
+        metric="cosine",
+        vector_column_name="vector",
+        index_type="IVF_PQ",
+        num_partitions=1,
+        num_sub_vectors=2,
+    )
+
+    # query with single vector
+    q = [1, 2]
+    rs = multivec_table.search(q).to_arrow()
+
+    # query with multiple vectors
+    q = [[1, 2], [1, 2]]
+    rs2 = multivec_table.search(q).to_arrow()
+    assert len(rs2) == len(rs)
+    for i in range(2):
+        assert rs2["_distance"][i].as_py() == rs["_distance"][i].as_py() * 2
+
+    # can't query with vector that dim not matched
+    with pytest.raises(Exception):
+        multivec_table.search([1, 2, 3]).to_arrow()
+    # can't query with vector list that some dim not matched
+    with pytest.raises(Exception):
+        multivec_table.search([[1, 2], [1, 2, 3]]).to_arrow()
+
+
+@pytest.mark.asyncio
+async def test_multivector_async(multivec_table_async: AsyncTable):
+    # create index on multivector
+    await multivec_table_async.create_index(
+        "vector",
+        config=IvfPq(distance_type="cosine", num_partitions=1, num_sub_vectors=2),
+    )
+
+    # query with single vector
+    q = [1, 2]
+    rs = await multivec_table_async.query().nearest_to(q).to_arrow()
+
+    # query with multiple vectors
+    q = [[1, 2], [1, 2]]
+    rs2 = await multivec_table_async.query().nearest_to(q).to_arrow()
+    assert len(rs2) == len(rs)
+    for i in range(2):
+        assert rs2["_distance"][i].as_py() == rs["_distance"][i].as_py() * 2
+
+    # can't query with vector that dim not matched
+    with pytest.raises(Exception):
+        await multivec_table_async.query().nearest_to([1, 2, 3]).to_arrow()
+    # can't query with vector list that some dim not matched
+    with pytest.raises(Exception):
+        await multivec_table_async.query().nearest_to([[1, 2], [1, 2, 3]]).to_arrow()
+
+
 def test_vector_query_with_no_limit(table):
     with pytest.raises(ValueError):
         LanceVectorQueryBuilder(table, [0, 0], "vector").limit(0).select(
@@ -448,11 +558,13 @@ async def test_query_to_pandas_flatten_async(table_struct_async: AsyncTable):
 
 @pytest.mark.asyncio
 async def test_query_to_polars_async(table_async: AsyncTable):
+    schema = await table_async.schema()
+    num_columns = len(schema.names)
     df = await table_async.query().to_polars()
-    assert df.shape == (2, 5)
+    assert df.shape == (2, num_columns)
 
     df = await table_async.query().where("id < 0").to_polars()
-    assert df.shape == (0, 5)
+    assert df.shape == (0, num_columns)
 
 
 @pytest.mark.asyncio

rust/lancedb/src/table.rs

@@ -18,7 +18,7 @@ use std::collections::HashMap;
 use std::path::Path;
 use std::sync::Arc;
 
-use arrow::array::AsArray;
+use arrow::array::{AsArray, FixedSizeListBuilder, Float32Builder};
 use arrow::datatypes::{Float32Type, UInt8Type};
 use arrow_array::{RecordBatchIterator, RecordBatchReader};
 use arrow_schema::{DataType, Field, Schema, SchemaRef};
@@ -1902,68 +1902,74 @@ impl TableInternal for NativeTable {
         options: QueryExecutionOptions,
     ) -> Result<Arc<dyn ExecutionPlan>> {
         let ds_ref = self.dataset.get().await?;
+        let mut column = query.column.clone();
+        let schema = ds_ref.schema();
 
+        let mut query_vector = query.query_vector.first().cloned();
         if query.query_vector.len() > 1 {
-            // If there are multiple query vectors, create a plan for each of them and union them.
-            let query_vecs = query.query_vector.clone();
-            let plan_futures = query_vecs
-                .into_iter()
-                .map(|query_vector| {
-                    let mut sub_query = query.clone();
-                    sub_query.query_vector = vec![query_vector];
-                    let options_ref = options.clone();
-                    async move { self.create_plan(&sub_query, options_ref).await }
-                })
-                .collect::<Vec<_>>();
-            let plans = futures::future::try_join_all(plan_futures).await?;
-            return Table::multi_vector_plan(plans);
+            if column.is_none() {
+                // Infer a vector column with the same dimension of the query vector.
+                let arrow_schema = Schema::from(ds_ref.schema());
+                column = Some(default_vector_column(
+                    &arrow_schema,
+                    Some(query.query_vector[0].len() as i32),
+                )?);
+            }
+            let vector_field = schema.field(column.as_ref().unwrap()).unwrap();
+            if let DataType::List(_) = vector_field.data_type() {
+                // it's multivector, then the vectors should be treated as single query
+                // concatenate the vectors into a FixedSizeList<FixedSizeList<_>>
+                // it's also possible to concatenate the vectors into a List<FixedSizeList<_>>,
+                // but FixedSizeList is more efficient and easier to construct
+                let vectors = query
+                    .query_vector
+                    .iter()
+                    .map(|arr| arr.as_ref())
+                    .collect::<Vec<_>>();
+                let dim = vectors[0].len();
+                let mut fsl_builder = FixedSizeListBuilder::with_capacity(
+                    Float32Builder::with_capacity(dim),
+                    dim as i32,
+                    vectors.len(),
+                );
+                for vec in vectors {
+                    fsl_builder
+                        .values()
+                        .append_slice(vec.as_primitive::<Float32Type>().values());
+                    fsl_builder.append(true);
+                }
+                query_vector = Some(Arc::new(fsl_builder.finish()));
+            } else {
+                // If there are multiple query vectors, create a plan for each of them and union them.
+                let query_vecs = query.query_vector.clone();
+                let plan_futures = query_vecs
+                    .into_iter()
+                    .map(|query_vector| {
+                        let mut sub_query = query.clone();
+                        sub_query.query_vector = vec![query_vector];
+                        let options_ref = options.clone();
+                        async move { self.create_plan(&sub_query, options_ref).await }
+                    })
+                    .collect::<Vec<_>>();
+                let plans = futures::future::try_join_all(plan_futures).await?;
+                return Table::multi_vector_plan(plans);
+            }
         }
 
         let mut scanner: Scanner = ds_ref.scan();
 
-        if let Some(query_vector) = query.query_vector.first() {
+        if let Some(query_vector) = query_vector {
             // If there is a vector query, default to limit=10 if unspecified
-            let column = if let Some(col) = query.column.as_ref() {
-                col.clone()
+            let column = if let Some(col) = column {
+                col
             } else {
                 // Infer a vector column with the same dimension of the query vector.
                 let arrow_schema = Schema::from(ds_ref.schema());
                 default_vector_column(&arrow_schema, Some(query_vector.len() as i32))?
             };
 
-            let field = ds_ref.schema().field(&column).ok_or(Error::Schema {
-                message: format!("Column {} not found in dataset schema", column),
-            })?;
-
-            let mut is_binary = false;
-            if let arrow_schema::DataType::FixedSizeList(element, dim) = field.data_type() {
-                match element.data_type() {
-                    e_type if e_type.is_floating() => {}
-                    e_type if *e_type == DataType::UInt8 => {
-                        is_binary = true;
-                    }
-                    _ => {
-                        return Err(Error::InvalidInput {
-                            message: format!(
-                                "The data type of the vector column '{}' is not a floating point type",
-                                column
-                            ),
-                        });
-                    }
-                }
-                if dim != query_vector.len() as i32 {
-                    return Err(Error::InvalidInput {
-                    message: format!(
-                        "The dimension of the query vector does not match with the dimension of the vector column '{}': \
-                            query dim={}, expected vector dim={}",
-                        column,
-                        query_vector.len(),
-                        dim,
-                    ),
-                });
-                }
-            }
-
+            let (_, element_type) = lance::index::vector::utils::get_vector_type(schema, &column)?;
+            let is_binary = matches!(element_type, DataType::UInt8);
             if is_binary {
                 let query_vector = arrow::compute::cast(&query_vector, &DataType::UInt8)?;
                 let query_vector = query_vector.as_primitive::<UInt8Type>();
@@ -1973,10 +1979,9 @@ impl TableInternal for NativeTable {
                     query.base.limit.unwrap_or(DEFAULT_TOP_K),
                 )?;
             } else {
-                let query_vector = query_vector.as_primitive::<Float32Type>();
                 scanner.nearest(
                     &column,
-                    query_vector,
+                    query_vector.as_ref(),
                     query.base.limit.unwrap_or(DEFAULT_TOP_K),
                 )?;
             }

rust/lancedb/src/utils.rs

@@ -108,13 +108,8 @@ pub(crate) fn default_vector_column(schema: &Schema, dim: Option<i32>) -> Result
     let candidates = schema
         .fields()
         .iter()
-        .filter_map(|field| match field.data_type() {
-            arrow_schema::DataType::FixedSizeList(f, d)
-                if (f.data_type().is_floating() || f.data_type() == &DataType::UInt8)
-                    && dim.map(|expect| *d == expect).unwrap_or(true) =>
-            {
-                Some(field.name())
-            }
+        .filter_map(|field| match inf_vector_dim(field) {
+            Some(d) if dim.is_none() || dim == Some(d) => Some(field.name()),
             _ => None,
         })
         .collect::<Vec<_>>();
@@ -138,6 +133,20 @@ pub(crate) fn default_vector_column(schema: &Schema, dim: Option<i32>) -> Result
     }
 }
 
+fn inf_vector_dim(field: &arrow_schema::Field) -> Option<i32> {
+    match field.data_type() {
+        arrow_schema::DataType::FixedSizeList(f, d) => {
+            if f.data_type().is_floating() || f.data_type() == &DataType::UInt8 {
+                Some(*d)
+            } else {
+                None
+            }
+        }
+        arrow_schema::DataType::List(f) => inf_vector_dim(f),
+        _ => None,
+    }
+}
+
 pub fn supported_btree_data_type(dtype: &DataType) -> bool {
     dtype.is_integer()
         || dtype.is_floating()
@@ -171,9 +180,10 @@ pub fn supported_fts_data_type(dtype: &DataType) -> bool {
 
 pub fn supported_vector_data_type(dtype: &DataType) -> bool {
     match dtype {
-        DataType::FixedSizeList(inner, _) => {
-            DataType::is_floating(inner.data_type()) || *inner.data_type() == DataType::UInt8
+        DataType::FixedSizeList(field, _) => {
+            field.data_type().is_floating() || field.data_type() == &DataType::UInt8
         }
+        DataType::List(field) => supported_vector_data_type(field.data_type()),
         _ => false,
     }
 }