diff --git a/docs/src/js/interfaces/IvfRqOptions.md b/docs/src/js/interfaces/IvfRqOptions.md index 6b48cf653..298565a62 100644 --- a/docs/src/js/interfaces/IvfRqOptions.md +++ b/docs/src/js/interfaces/IvfRqOptions.md @@ -58,7 +58,8 @@ Number of bits per dimension for residual quantization. This value controls how much each residual component is compressed. The more bits, the more accurate the index will be but the slower search. Typical values -are small integers; the default is 1 bit per dimension. +are small integers. If unset, uses Lance's default (currently 5 bits per +dimension). Set to 1 for maximum compression. *** diff --git a/docs/web/indexing/vector-index.mdx b/docs/web/indexing/vector-index.mdx index b26e0c295..037d75f8d 100644 --- a/docs/web/indexing/vector-index.mdx +++ b/docs/web/indexing/vector-index.mdx @@ -58,7 +58,7 @@ Lance with `index_type`. | If your priority is... | Use this index | Why | Approx. compression ratio | Python config class | | :--- | :--- | :--- | :--- | :--- | -| Maximum compression | `IVF_RQ` | IVF indexing with RaBitQ quantization | Around `1/32` of raw size | `IvfRq` | +| Maximum compression | `IVF_RQ` | IVF indexing with RaBitQ quantization | Around `1/32` of raw size with `num_bits=1` | `IvfRq` | | Higher accuracy at small dimensions (`dimension <= 256`) | `IVF_PQ` | IVF indexing with product quantization | Usually `1/64` to `1/16` of raw size (depends on `num_sub_vectors`) | `IvfPq` | | Best recall/latency trade-off | `IVF_HNSW_SQ` | IVF-HNSW indexing with scalar | Typically a little larger than `1/4` of raw size | `IvfHnswSq` | | Highest recall / no quantization | `IVF_HNSW_FLAT` | IVF-HNSW indexing with no quantization; embedded library only | Around raw vector size plus HNSW graph overhead | `IvfHnswFlat` | @@ -377,7 +377,7 @@ which has been compressed to $4$ chunks $\times 8$-bit integers $= 32$ bits of q RaBitQ is an advanced quantization technique that outperforms PQ in several ways. It needs no codebook to train, estimates distances very quickly at query-time, -and crucially, quantizes each vector in (with some small overhead) just **one bit per dimension!** +and can quantize each vector in (with some small overhead) as little as **one bit per dimension!** In practice, RaBitQ compresses a $1024$-dimensional `float32` vector into just a few thousand bits, while maintaining good recall. The inner workings of RaBitQ quantization are rather mathematically dense. It generates a quantization codebook by @@ -385,14 +385,19 @@ applying a uniformly random, approximately distance-preserving orthogonal transf where $d$ is the dimensionality of the dataset. We defer the details, and an elegant theoretical error bound, to [the original paper.](https://arxiv.org/pdf/2405.12497) ##### Using RaBitQ {#using-rabitq} -Use RaBitQ quantization by selecting quantized index types ending in the suffix `RQ`. For example, call `create_index` with `index_type="IVF_RQ"`. +Use RaBitQ quantization by selecting quantized index types ending in the suffix `RQ`. For example, import `IvfRq` from `lancedb.index` and call `table.create_index("vector", config=IvfRq())`. Note that when using `IVF_RQ`, the dimension of the dataset must be a multiple of `8`. `num_bits` determines how many bits are used to quantize each dimension. -`1` is the standard RaBitQ setting. Increase to `2`, `4`, or `8` bits to achieve better recall for additional storage and query-time compute. +By default, Python, Rust, and Node leave this unset and inherit Lance's default, currently `5` bits. +Higher values improve recall at the cost of additional storage and query-time compute. +To retain the previous one-bit default and maximum compression, explicitly set `num_bits=1` (for example, `IvfRq(num_bits=1)` in Python). +Five-bit quantized vector codes use roughly five times the storage of one-bit codes; total index size also includes other overhead. -RaBitQ-quantized indexes computed with `num_bits >= 2` use a newer on-disk layout, and cannot be read by some older LanceDB versions. +RaBitQ-quantized indexes computed with `num_bits >= 2` use a newer on-disk layout and cannot be read by some older LanceDB versions. +The new five-bit default therefore affects reader compatibility as well as index size. +Upgrade readers to versions that support multi-bit IVF_RQ before creating indexes with the default settings, or explicitly set `num_bits=1` when sharing indexes with readers that support only one-bit IVF_RQ. See this [blog post](https://lancedb.com/blog/feature-rabitq-quantization/) for further discussion and benchmarking of LanceDB's RaBitQ implementation. @@ -410,7 +415,7 @@ among all integral multiples $\{j \times 10^5: -128 \leq j \leq 127 \}$. | Parameter | Description | | :--- | :--- | -| `num_bits` | Bits per dimension for quantization. Only applies to `IVF_PQ`/`IVF_HNSW_PQ` (default `8`) and `IVF_RQ` (default `1`, RaBitQ); not used by `IVF_FLAT`/`IVF_SQ`/`IVF_HNSW_FLAT`/`IVF_HNSW_SQ`. Higher values improve accuracy at the cost of additional storage. | +| `num_bits` | Bits per dimension for quantization. Only applies to `IVF_PQ`/`IVF_HNSW_PQ` (default `8`) and `IVF_RQ` (`IvfRq()` inherits Lance's default, currently `5`, RaBitQ); not used by `IVF_FLAT`/`IVF_SQ`/`IVF_HNSW_FLAT`/`IVF_HNSW_SQ`. Higher values improve accuracy at the cost of additional storage. | `max_iterations` and `sample_rate` also affect quantizer training, but since they apply to every IVF/HNSW index type (not just quantized ones), they're documented as general [Build-time Parameters](/indexing/vector-index#build-time-parameters) instead. All three are passed as keyword arguments to `create_index`, alongside `index_type`: diff --git a/nodejs/lancedb/indices.ts b/nodejs/lancedb/indices.ts index 43d1cf1b5..0e383e984 100644 --- a/nodejs/lancedb/indices.ts +++ b/nodejs/lancedb/indices.ts @@ -131,7 +131,8 @@ export interface IvfRqOptions { * * This value controls how much each residual component is compressed. The more * bits, the more accurate the index will be but the slower search. Typical values - * are small integers; the default is 1 bit per dimension. + * are small integers. If unset, uses Lance's default (currently 5 bits per + * dimension). Set to 1 for maximum compression. */ numBits?: number; diff --git a/python/python/lancedb/index.py b/python/python/lancedb/index.py index cf3bf8e4f..80179f4a7 100644 --- a/python/python/lancedb/index.py +++ b/python/python/lancedb/index.py @@ -882,8 +882,12 @@ class IvfRq: num_partitions: int, default sqrt(num_rows) Number of IVF partitions to create. - num_bits: int, default 1 + num_bits: int, optional Number of bits to encode each dimension in the RabitQ codebook. + If None, uses Lance's default (currently 5), matching Rust and Node. + Set to 1 to retain the previous Python default and maximum compression. + Five-bit vector codes use roughly five times the storage of one-bit codes, + in addition to other index overhead. max_iterations: int, default 50 Max iterations to train kmeans when computing IVF partitions. @@ -897,7 +901,7 @@ class IvfRq: distance_type: Literal["l2", "cosine", "dot"] = "l2" num_partitions: Optional[int] = None - num_bits: int = 1 + num_bits: Optional[int] = None max_iterations: int = 50 sample_rate: int = 256 target_partition_size: Optional[int] = None diff --git a/python/python/lancedb/table.py b/python/python/lancedb/table.py index 2ad7f821d..cc455f746 100644 --- a/python/python/lancedb/table.py +++ b/python/python/lancedb/table.py @@ -3754,6 +3754,11 @@ class LanceTable(Table): cfg_type = type(config).__name__ lance_index_type = index_type_map.get(cfg_type, "IVF_PQ") + quantization_kwargs = {} + num_bits = getattr(config, "num_bits", 8) + if num_bits is not None: + quantization_kwargs["num_bits"] = num_bits + dataset.create_index( column=column, index_type=lance_index_type, @@ -3762,12 +3767,12 @@ class LanceTable(Table): num_sub_vectors=getattr(config, "num_sub_vectors", None), replace=replace, accelerator=acc, - num_bits=getattr(config, "num_bits", 8), m=getattr(config, "m", 20), ef_construction=getattr(config, "ef_construction", 300), target_partition_size=getattr( config, "target_partition_size", None ), + **quantization_kwargs, ) self.checkout_latest() return diff --git a/python/python/tests/test_index.py b/python/python/tests/test_index.py index 0252a09ad..2ec14f6f8 100644 --- a/python/python/tests/test_index.py +++ b/python/python/tests/test_index.py @@ -421,13 +421,23 @@ async def test_create_4bit_ivfpq_index(some_table: AsyncTable): @pytest.mark.asyncio -async def test_create_ivfrq_index(some_table: AsyncTable): - await some_table.create_index("vector", config=IvfRq(num_bits=1)) +@pytest.mark.parametrize( + "config, expected_bits", + [ + pytest.param(IvfRq(), 5, id="default"), + pytest.param(IvfRq(num_bits=None), 5, id="none"), + pytest.param(IvfRq(num_bits=1), 1, id="one-bit"), + pytest.param(IvfRq(num_bits=4), 4, id="multi-bit"), + ], +) +async def test_create_ivfrq_index(some_table: AsyncTable, config, expected_bits): + await some_table.create_index("vector", config=config) indices = await some_table.list_indices() assert len(indices) == 1 assert indices[0].index_type == "IvfRq" assert indices[0].columns == ["vector"] assert indices[0].name == "vector_idx" + assert indices[0].index_details["compression"]["num_bits"] == expected_bits @pytest.mark.asyncio diff --git a/python/python/tests/test_remote_db.py b/python/python/tests/test_remote_db.py index 95da3eba2..049a5fc58 100644 --- a/python/python/tests/test_remote_db.py +++ b/python/python/tests/test_remote_db.py @@ -1593,7 +1593,10 @@ def test_remote_json_fts_config(server_version): assert request["block_size"] == 256 -def test_remote_create_index_new_api(caplog): +@pytest.mark.parametrize( + "rq_kwargs", [{}, {"num_bits": None}, {"num_bits": 1}, {"num_bits": 4}] +) +def test_remote_create_index_new_api(caplog, rq_kwargs): received_requests = [] def handler(request): @@ -1666,7 +1669,7 @@ def test_remote_create_index_new_api(caplog): table.create_index("category", config=BTree()) table.create_index("text", config=FTS(block_size=256)) # IvfRq via new API - table.create_index("vector", config=IvfRq(distance_type="l2")) + table.create_index("vector", config=IvfRq(distance_type="l2", **rq_kwargs)) table.create_index( "vector", config=IvfPq(distance_type="l2"), replace=False ) @@ -1691,6 +1694,10 @@ def test_remote_create_index_new_api(caplog): "vector", ] assert received_requests[2]["block_size"] == 256 + if rq_kwargs.get("num_bits") is None: + assert "num_bits" not in received_requests[3] + else: + assert received_requests[3]["num_bits"] == rq_kwargs["num_bits"] assert received_requests[4]["replace"] is False assert received_requests[5]["replace"] is False assert "replace is not supported" not in caplog.text diff --git a/python/python/tests/test_table.py b/python/python/tests/test_table.py index 7d92b0e6d..3b3ccd122 100644 --- a/python/python/tests/test_table.py +++ b/python/python/tests/test_table.py @@ -1751,6 +1751,28 @@ def test_create_index_async_returns_done_job(mem_db: DBConnection): job.cancel() +@pytest.mark.parametrize( + "rq_kwargs", [{}, {"num_bits": None}, {"num_bits": 1}, {"num_bits": 4}] +) +@patch("lancedb.table.LanceTable.to_lance") +def test_create_ivfrq_index_accelerator(mock_to_lance, mem_db: DBConnection, rq_kwargs): + from lancedb.index import IvfRq + + table = mem_db.create_table("test", [{"vector": [1.0] * 8}]) + torch = MagicMock() + torch.cuda.is_available.return_value = True + with patch.dict(sys.modules, {"torch": torch}): + table.create_index("vector", config=IvfRq(accelerator="cuda", **rq_kwargs)) + + kwargs = mock_to_lance.return_value.create_index.call_args.kwargs + assert kwargs["index_type"] == "IVF_RQ" + assert kwargs["accelerator"] == "cuda" + if rq_kwargs.get("num_bits") is None: + assert "num_bits" not in kwargs + else: + assert kwargs["num_bits"] == rq_kwargs["num_bits"] + + @pytest.mark.parametrize("method", ["create_index", "create_index_async"]) @pytest.mark.parametrize("wait_timeout", ["invalid", 1, []]) def test_create_index_invalid_wait_timeout(mem_db: DBConnection, method, wait_timeout): diff --git a/python/src/index.rs b/python/src/index.rs index d4e16a67e..c6c3fe17d 100644 --- a/python/src/index.rs +++ b/python/src/index.rs @@ -136,8 +136,10 @@ pub fn extract_index_params(source: &Option>) -> PyResult, - num_bits: u32, + num_bits: Option, max_iterations: u32, sample_rate: u32, target_partition_size: Option, diff --git a/rust/lancedb/src/index/vector.rs b/rust/lancedb/src/index/vector.rs index bce77a610..3e6794b54 100644 --- a/rust/lancedb/src/index/vector.rs +++ b/rust/lancedb/src/index/vector.rs @@ -362,6 +362,9 @@ impl IvfRqIndexBuilder { impl_distance_type_setter!(); impl_ivf_params_setter!(); + /// Set the number of bits used to encode each vector dimension. + /// + /// If unset, uses Lance's default (currently 5). Set to 1 for maximum compression. pub fn num_bits(mut self, num_bits: u32) -> Self { self.num_bits = Some(num_bits); self diff --git a/rust/lancedb/src/table/create_index.rs b/rust/lancedb/src/table/create_index.rs index 3ac759bdd..1795b778e 100644 --- a/rust/lancedb/src/table/create_index.rs +++ b/rust/lancedb/src/table/create_index.rs @@ -355,7 +355,10 @@ impl NativeTable { index.sample_rate, index.max_iterations, ); - let rq_params = RQBuildParams::new(index.num_bits.unwrap_or(1) as u8); + let rq_params = index + .num_bits + .map(|num_bits| RQBuildParams::new(num_bits as u8)) + .unwrap_or_default(); let lance_idx_params = VectorIndexParams::with_ivf_rq_params( index.distance_type.into(), ivf_params,