diff --git a/python/python/lancedb/_lancedb.pyi b/python/python/lancedb/_lancedb.pyi index 2cf04a1b..2a326041 100644 --- a/python/python/lancedb/_lancedb.pyi +++ b/python/python/lancedb/_lancedb.pyi @@ -3,7 +3,17 @@ from typing import Dict, List, Optional, Tuple, Any, TypedDict, Union, Literal import pyarrow as pa -from .index import BTree, IvfFlat, IvfPq, Bitmap, LabelList, HnswPq, HnswSq, FTS +from .index import ( + BTree, + IvfFlat, + IvfPq, + IvfSq, + Bitmap, + LabelList, + HnswPq, + HnswSq, + FTS, +) from .io import StorageOptionsProvider from lance_namespace import ( ListNamespacesResponse, @@ -14,6 +24,9 @@ from lance_namespace import ( ) from .remote import ClientConfig +IvfHnswPq: type[HnswPq] = HnswPq +IvfHnswSq: type[HnswSq] = HnswSq + class Session: def __init__( self, @@ -131,7 +144,17 @@ class Table: async def create_index( self, column: str, - index: Union[IvfFlat, IvfPq, HnswPq, HnswSq, BTree, Bitmap, LabelList, FTS], + index: Union[ + IvfFlat, + IvfSq, + IvfPq, + HnswPq, + HnswSq, + BTree, + Bitmap, + LabelList, + FTS, + ], replace: Optional[bool], wait_timeout: Optional[object], *, diff --git a/python/python/lancedb/index.py b/python/python/lancedb/index.py index 27202c01..d5e4bfcd 100644 --- a/python/python/lancedb/index.py +++ b/python/python/lancedb/index.py @@ -376,6 +376,11 @@ class HnswSq: target_partition_size: Optional[int] = None +# Backwards-compatible aliases +IvfHnswPq = HnswPq +IvfHnswSq = HnswSq + + @dataclass class IvfFlat: """Describes an IVF Flat Index @@ -475,6 +480,36 @@ class IvfFlat: target_partition_size: Optional[int] = None +@dataclass +class IvfSq: + """Describes an IVF Scalar Quantization (SQ) index. + + This index applies scalar quantization to compress vectors and organizes the + quantized vectors into IVF partitions. It offers a balance between search + speed and storage efficiency while keeping good recall. + + Attributes + ---------- + distance_type: str, default "l2" + The distance metric used to train and search the index. Supported values + are "l2", "cosine", and "dot". + num_partitions: int, default sqrt(num_rows) + Number of IVF partitions to create. + max_iterations: int, default 50 + Maximum iterations for kmeans during partition training. + sample_rate: int, default 256 + Controls the number of training vectors: sample_rate * num_partitions. + target_partition_size: int, optional + Target size for each partition; adjusts the balance between speed and accuracy. + """ + + distance_type: Literal["l2", "cosine", "dot"] = "l2" + num_partitions: Optional[int] = None + max_iterations: int = 50 + sample_rate: int = 256 + target_partition_size: Optional[int] = None + + @dataclass class IvfPq: """Describes an IVF PQ Index @@ -661,6 +696,9 @@ class IvfRq: __all__ = [ "BTree", "IvfPq", + "IvfHnswPq", + "IvfHnswSq", + "IvfSq", "IvfRq", "IvfFlat", "HnswPq", diff --git a/python/python/lancedb/remote/table.py b/python/python/lancedb/remote/table.py index eee35caa..c97199e4 100644 --- a/python/python/lancedb/remote/table.py +++ b/python/python/lancedb/remote/table.py @@ -18,7 +18,7 @@ from lancedb._lancedb import ( UpdateResult, ) from lancedb.embeddings.base import EmbeddingFunctionConfig -from lancedb.index import FTS, BTree, Bitmap, HnswSq, IvfFlat, IvfPq, LabelList +from lancedb.index import FTS, BTree, Bitmap, HnswSq, IvfFlat, IvfPq, IvfSq, LabelList from lancedb.remote.db import LOOP import pyarrow as pa @@ -265,6 +265,8 @@ class RemoteTable(Table): num_sub_vectors=num_sub_vectors, num_bits=num_bits, ) + elif index_type == "IVF_SQ": + config = IvfSq(distance_type=metric, num_partitions=num_partitions) elif index_type == "IVF_HNSW_PQ": raise ValueError( "IVF_HNSW_PQ is not supported on LanceDB cloud." @@ -277,7 +279,7 @@ class RemoteTable(Table): else: raise ValueError( f"Unknown vector index type: {index_type}. Valid options are" - " 'IVF_FLAT', 'IVF_PQ', 'IVF_HNSW_PQ', 'IVF_HNSW_SQ'" + " 'IVF_FLAT', 'IVF_SQ', 'IVF_PQ', 'IVF_HNSW_PQ', 'IVF_HNSW_SQ'" ) LOOP.run( diff --git a/python/python/lancedb/table.py b/python/python/lancedb/table.py index 57af5e72..230daa89 100644 --- a/python/python/lancedb/table.py +++ b/python/python/lancedb/table.py @@ -44,7 +44,18 @@ import numpy as np from .common import DATA, VEC, VECTOR_COLUMN_NAME from .embeddings import EmbeddingFunctionConfig, EmbeddingFunctionRegistry -from .index import BTree, IvfFlat, IvfPq, Bitmap, IvfRq, LabelList, HnswPq, HnswSq, FTS +from .index import ( + BTree, + IvfFlat, + IvfPq, + IvfSq, + Bitmap, + IvfRq, + LabelList, + HnswPq, + HnswSq, + FTS, +) from .merge import LanceMergeInsertBuilder from .pydantic import LanceModel, model_to_dict from .query import ( @@ -2054,7 +2065,7 @@ class LanceTable(Table): index_cache_size: Optional[int] = None, num_bits: int = 8, index_type: Literal[ - "IVF_FLAT", "IVF_PQ", "IVF_RQ", "IVF_HNSW_SQ", "IVF_HNSW_PQ" + "IVF_FLAT", "IVF_SQ", "IVF_PQ", "IVF_RQ", "IVF_HNSW_SQ", "IVF_HNSW_PQ" ] = "IVF_PQ", max_iterations: int = 50, sample_rate: int = 256, @@ -2092,6 +2103,14 @@ class LanceTable(Table): sample_rate=sample_rate, target_partition_size=target_partition_size, ) + elif index_type == "IVF_SQ": + config = IvfSq( + distance_type=metric, + num_partitions=num_partitions, + max_iterations=max_iterations, + sample_rate=sample_rate, + target_partition_size=target_partition_size, + ) elif index_type == "IVF_PQ": config = IvfPq( distance_type=metric, @@ -3456,11 +3475,22 @@ class AsyncTable: if config is not None: if not isinstance( config, - (IvfFlat, IvfPq, IvfRq, HnswPq, HnswSq, BTree, Bitmap, LabelList, FTS), + ( + IvfFlat, + IvfSq, + IvfPq, + IvfRq, + HnswPq, + HnswSq, + BTree, + Bitmap, + LabelList, + FTS, + ), ): raise TypeError( - "config must be an instance of IvfPq, IvfRq, HnswPq, HnswSq, BTree," - " Bitmap, LabelList, or FTS, but got " + str(type(config)) + "config must be an instance of IvfSq, IvfPq, IvfRq, HnswPq, HnswSq," + " BTree, Bitmap, LabelList, or FTS, but got " + str(type(config)) ) try: await self._inner.create_index( diff --git a/python/python/lancedb/types.py b/python/python/lancedb/types.py index 6ca72de2..e7b185f2 100644 --- a/python/python/lancedb/types.py +++ b/python/python/lancedb/types.py @@ -18,12 +18,20 @@ AddMode = Literal["append", "overwrite"] CreateMode = Literal["create", "overwrite"] # Index type literals -VectorIndexType = Literal["IVF_FLAT", "IVF_PQ", "IVF_HNSW_SQ", "IVF_HNSW_PQ", "IVF_RQ"] +VectorIndexType = Literal[ + "IVF_FLAT", + "IVF_SQ", + "IVF_PQ", + "IVF_HNSW_SQ", + "IVF_HNSW_PQ", + "IVF_RQ", +] ScalarIndexType = Literal["BTREE", "BITMAP", "LABEL_LIST"] IndexType = Literal[ "IVF_PQ", "IVF_HNSW_PQ", "IVF_HNSW_SQ", + "IVF_SQ", "FTS", "BTREE", "BITMAP", diff --git a/python/python/tests/test_index.py b/python/python/tests/test_index.py index 4c594edf..b4097a8f 100644 --- a/python/python/tests/test_index.py +++ b/python/python/tests/test_index.py @@ -12,6 +12,9 @@ from lancedb.index import ( BTree, IvfFlat, IvfPq, + IvfSq, + IvfHnswPq, + IvfHnswSq, IvfRq, Bitmap, LabelList, @@ -229,6 +232,35 @@ async def test_create_hnswsq_index(some_table: AsyncTable): assert len(indices) == 1 +@pytest.mark.asyncio +async def test_create_hnswsq_alias_index(some_table: AsyncTable): + await some_table.create_index("vector", config=IvfHnswSq(num_partitions=5)) + indices = await some_table.list_indices() + assert len(indices) == 1 + assert indices[0].index_type in {"HnswSq", "IvfHnswSq"} + + +@pytest.mark.asyncio +async def test_create_hnswpq_alias_index(some_table: AsyncTable): + await some_table.create_index("vector", config=IvfHnswPq(num_partitions=5)) + indices = await some_table.list_indices() + assert len(indices) == 1 + assert indices[0].index_type in {"HnswPq", "IvfHnswPq"} + + +@pytest.mark.asyncio +async def test_create_ivfsq_index(some_table: AsyncTable): + await some_table.create_index("vector", config=IvfSq(num_partitions=10)) + indices = await some_table.list_indices() + assert len(indices) == 1 + assert indices[0].index_type == "IvfSq" + stats = await some_table.index_stats(indices[0].name) + assert stats.index_type == "IVF_SQ" + assert stats.distance_type == "l2" + assert stats.num_indexed_rows == await some_table.count_rows() + assert stats.num_unindexed_rows == 0 + + @pytest.mark.asyncio async def test_create_index_with_binary_vectors(binary_table: AsyncTable): await binary_table.create_index( diff --git a/python/src/index.rs b/python/src/index.rs index 82f61c2e..c93b23eb 100644 --- a/python/src/index.rs +++ b/python/src/index.rs @@ -1,7 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright The LanceDB Authors -use lancedb::index::vector::{IvfFlatIndexBuilder, IvfRqIndexBuilder}; +use lancedb::index::vector::{IvfFlatIndexBuilder, IvfRqIndexBuilder, IvfSqIndexBuilder}; use lancedb::index::{ scalar::{BTreeIndexBuilder, FtsIndexBuilder}, vector::{IvfHnswPqIndexBuilder, IvfHnswSqIndexBuilder, IvfPqIndexBuilder}, @@ -87,6 +87,21 @@ pub fn extract_index_params(source: &Option>) -> PyResult { + let params = source.extract::()?; + let distance_type = parse_distance_type(params.distance_type)?; + let mut ivf_sq_builder = IvfSqIndexBuilder::default() + .distance_type(distance_type) + .max_iterations(params.max_iterations) + .sample_rate(params.sample_rate); + if let Some(num_partitions) = params.num_partitions { + ivf_sq_builder = ivf_sq_builder.num_partitions(num_partitions); + } + if let Some(target_partition_size) = params.target_partition_size { + ivf_sq_builder = ivf_sq_builder.target_partition_size(target_partition_size); + } + Ok(LanceDbIndex::IvfSq(ivf_sq_builder)) + }, "IvfRq" => { let params = source.extract::()?; let distance_type = parse_distance_type(params.distance_type)?; @@ -142,7 +157,7 @@ pub fn extract_index_params(source: &Option>) -> PyResult Err(PyValueError::new_err(format!( - "Invalid index type '{}'. Must be one of BTree, Bitmap, LabelList, FTS, IvfPq, IvfHnswPq, or IvfHnswSq", + "Invalid index type '{}'. Must be one of BTree, Bitmap, LabelList, FTS, IvfPq, IvfSq, IvfHnswPq, or IvfHnswSq", not_supported ))), } @@ -186,6 +201,15 @@ struct IvfPqParams { target_partition_size: Option, } +#[derive(FromPyObject)] +struct IvfSqParams { + distance_type: String, + num_partitions: Option, + max_iterations: u32, + sample_rate: u32, + target_partition_size: Option, +} + #[derive(FromPyObject)] struct IvfRqParams { distance_type: String, diff --git a/rust/lancedb/src/index.rs b/rust/lancedb/src/index.rs index 83814ffa..2269ad85 100644 --- a/rust/lancedb/src/index.rs +++ b/rust/lancedb/src/index.rs @@ -13,7 +13,7 @@ use crate::{table::BaseTable, DistanceType, Error, Result}; use self::{ scalar::{BTreeIndexBuilder, BitmapIndexBuilder, LabelListIndexBuilder}, - vector::{IvfHnswPqIndexBuilder, IvfHnswSqIndexBuilder, IvfPqIndexBuilder}, + vector::{IvfHnswPqIndexBuilder, IvfHnswSqIndexBuilder, IvfPqIndexBuilder, IvfSqIndexBuilder}, }; pub mod scalar; @@ -54,6 +54,9 @@ pub enum Index { /// IVF index with Product Quantization IvfPq(IvfPqIndexBuilder), + /// IVF index with Scalar Quantization + IvfSq(IvfSqIndexBuilder), + /// IVF index with RabitQ Quantization IvfRq(IvfRqIndexBuilder), @@ -277,6 +280,8 @@ pub enum IndexType { // Vector #[serde(alias = "IVF_FLAT")] IvfFlat, + #[serde(alias = "IVF_SQ")] + IvfSq, #[serde(alias = "IVF_PQ")] IvfPq, #[serde(alias = "IVF_RQ")] @@ -301,6 +306,7 @@ impl std::fmt::Display for IndexType { fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result { match self { Self::IvfFlat => write!(f, "IVF_FLAT"), + Self::IvfSq => write!(f, "IVF_SQ"), Self::IvfPq => write!(f, "IVF_PQ"), Self::IvfRq => write!(f, "IVF_RQ"), Self::IvfHnswPq => write!(f, "IVF_HNSW_PQ"), @@ -323,6 +329,7 @@ impl std::str::FromStr for IndexType { "LABEL_LIST" | "LABELLIST" => Ok(Self::LabelList), "FTS" | "INVERTED" => Ok(Self::FTS), "IVF_FLAT" => Ok(Self::IvfFlat), + "IVF_SQ" => Ok(Self::IvfSq), "IVF_PQ" => Ok(Self::IvfPq), "IVF_RQ" => Ok(Self::IvfRq), "IVF_HNSW_PQ" => Ok(Self::IvfHnswPq), diff --git a/rust/lancedb/src/index/vector.rs b/rust/lancedb/src/index/vector.rs index 20f3a847..6268b36c 100644 --- a/rust/lancedb/src/index/vector.rs +++ b/rust/lancedb/src/index/vector.rs @@ -209,6 +209,38 @@ impl IvfFlatIndexBuilder { impl_ivf_params_setter!(); } +/// Builder for an IVF SQ index. +/// +/// This index compresses vectors using scalar quantization and groups them into IVF partitions. +/// It offers a balance between search performance and storage footprint. +#[derive(Debug, Clone)] +pub struct IvfSqIndexBuilder { + pub(crate) distance_type: DistanceType, + + // IVF + pub(crate) num_partitions: Option, + pub(crate) sample_rate: u32, + pub(crate) max_iterations: u32, + pub(crate) target_partition_size: Option, +} + +impl Default for IvfSqIndexBuilder { + fn default() -> Self { + Self { + distance_type: DistanceType::L2, + num_partitions: None, + sample_rate: 256, + max_iterations: 50, + target_partition_size: None, + } + } +} + +impl IvfSqIndexBuilder { + impl_distance_type_setter!(); + impl_ivf_params_setter!(); +} + /// Builder for an IVF PQ index. /// /// This index stores a compressed (quantized) copy of every vector. These vectors diff --git a/rust/lancedb/src/remote/table.rs b/rust/lancedb/src/remote/table.rs index f28770da..5f62c69b 100644 --- a/rust/lancedb/src/remote/table.rs +++ b/rust/lancedb/src/remote/table.rs @@ -1072,6 +1072,14 @@ impl BaseTable for RemoteTable { body["num_bits"] = serde_json::Value::Number(num_bits.into()); } } + Index::IvfSq(index) => { + body[INDEX_TYPE_KEY] = serde_json::Value::String("IVF_SQ".to_string()); + body[METRIC_TYPE_KEY] = + serde_json::Value::String(index.distance_type.to_string().to_lowercase()); + if let Some(num_partitions) = index.num_partitions { + body["num_partitions"] = serde_json::Value::Number(num_partitions.into()); + } + } Index::IvfHnswSq(index) => { body[INDEX_TYPE_KEY] = serde_json::Value::String("IVF_HNSW_SQ".to_string()); body[METRIC_TYPE_KEY] = diff --git a/rust/lancedb/src/table.rs b/rust/lancedb/src/table.rs index cb0b34be..14e04714 100644 --- a/rust/lancedb/src/table.rs +++ b/rust/lancedb/src/table.rs @@ -1946,6 +1946,25 @@ impl NativeTable { VectorIndexParams::with_ivf_flat_params(index.distance_type.into(), ivf_params); Ok(Box::new(lance_idx_params)) } + Index::IvfSq(index) => { + Self::validate_index_type(field, "IVF SQ", supported_vector_data_type)?; + let ivf_params = Self::build_ivf_params( + index.num_partitions, + index.target_partition_size, + index.sample_rate, + index.max_iterations, + ); + let sq_params = SQBuildParams { + sample_rate: index.sample_rate as usize, + ..Default::default() + }; + let lance_idx_params = VectorIndexParams::with_ivf_sq_params( + index.distance_type.into(), + ivf_params, + sq_params, + ); + Ok(Box::new(lance_idx_params)) + } Index::IvfPq(index) => { Self::validate_index_type(field, "IVF PQ", supported_vector_data_type)?; let dim = Self::get_vector_dimension(field)?; @@ -2053,6 +2072,7 @@ impl NativeTable { Index::LabelList(_) => IndexType::LabelList, Index::FTS(_) => IndexType::Inverted, Index::IvfFlat(_) + | Index::IvfSq(_) | Index::IvfPq(_) | Index::IvfRq(_) | Index::IvfHnswPq(_)