mirror of
https://github.com/lancedb/lancedb.git
synced 2026-10-26 21:46:15 +00:00
fix: inherit Lance defaults for IVF RQ quantization (#4585)
`IvfRq()` forced `num_bits=1` through Python's Rust binding, overriding Lance's five-bit default. The shared native index creation path also substituted one bit when the Rust or Node builder left the value unset. Make Python's bit count optional and forward only explicit values. Use `RQBuildParams::default()` for an unset native bit count, and omit unset values when dispatching to pylance for accelerator builds. Update the Python, Rust, Node, and vector-index documentation. New indexes with an unset bit count now inherit Lance's default, currently five bits, across Python, Rust, and Node. Five-bit vector codes use roughly five times the storage of one-bit codes, plus other index overhead. This also changes reader compatibility: indexes created with `num_bits >= 2` use a newer on-disk layout that some older LanceDB versions cannot read. Upgrade readers to versions that support multi-bit IVF_RQ, or explicitly set `num_bits=1` (for example, `IvfRq(num_bits=1)` in Python) when sharing indexes with readers limited to one-bit IVF_RQ. Validation: - Reproduced six regression failures before the fix: one-bit default metadata, rejected `None`, and incorrect remote/accelerator forwarding. - After merging `main` through `b762ec22`, rebuilt both native bindings against Lance 15.0.0-beta.1. Python validation passed with 106 tests and 15 optional Ollama skips. Coverage includes the full index suite, omitted and explicit bit counts, reading Lance-created IVF_RQ indexes, accelerator forwarding, CUDA-unavailability handling, timeout validation, remote requests, and table/namespace listing behavior. - Rust formatting and Python package Ruff lint/format checks pass. Node `pnpm lint`, `pnpm build` (including TypeScript checks), `pnpm run docs`, the full connection suite, and the existing IVF_RQ regression pass. Generated documentation adds no further changes. The merge preserves upstream's explicit RaBitQ heading anchor alongside the `IvfRq()` example and five-bit default/reader-compatibility guidance. Repository-wide `ruff check .` still reports 49 existing errors in unchanged documentation examples under `docs/web-tests`; root formatting also finds 11 unchanged documentation example files. The Python package checks pass. Fixes #4582 <!-- lance-gatekeeper-fix:v1 agent=6ef9406f0142b1ab66d53d9d56098223 generation=1 --> --------- Co-authored-by: Gatefixer <313497061+lancedb-gatefixer[bot]@users.noreply.github.com> Co-authored-by: Xuanwo <github@xuanwo.io>
This commit is contained in:
lancedb-gatefixer[bot]
co-authored by
Xuanwo
Xuanwo
parent
22fc2a2087
commit
66357275bb
@@ -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.
|
||||
|
||||
***
|
||||
|
||||
|
||||
@@ -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.
|
||||
|
||||
<Warning title="Reading multi-bit indexes across versions">
|
||||
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.
|
||||
</Warning>
|
||||
|
||||
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`:
|
||||
|
||||
|
||||
@@ -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;
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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):
|
||||
|
||||
+5
-3
@@ -136,8 +136,10 @@ pub fn extract_index_params(source: &Option<Bound<'_, PyAny>>) -> PyResult<Lance
|
||||
let mut ivf_rq_builder = IvfRqIndexBuilder::default()
|
||||
.distance_type(distance_type)
|
||||
.max_iterations(params.max_iterations)
|
||||
.sample_rate(params.sample_rate)
|
||||
.num_bits(params.num_bits);
|
||||
.sample_rate(params.sample_rate);
|
||||
if let Some(num_bits) = params.num_bits {
|
||||
ivf_rq_builder = ivf_rq_builder.num_bits(num_bits);
|
||||
}
|
||||
if let Some(num_partitions) = params.num_partitions {
|
||||
ivf_rq_builder = ivf_rq_builder.num_partitions(num_partitions);
|
||||
}
|
||||
@@ -265,7 +267,7 @@ struct IvfSqParams {
|
||||
struct IvfRqParams {
|
||||
distance_type: String,
|
||||
num_partitions: Option<u32>,
|
||||
num_bits: u32,
|
||||
num_bits: Option<u32>,
|
||||
max_iterations: u32,
|
||||
sample_rate: u32,
|
||||
target_partition_size: Option<u32>,
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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,
|
||||
|
||||
Reference in New Issue
Block a user