From 436eafc0a3ce7b0d884935d6f80f1fbb6694622d Mon Sep 17 00:00:00 2001 From: Yang Cen <159225399+BubbleCal@users.noreply.github.com> Date: Fri, 9 Oct 2026 23:32:21 +0800 Subject: [PATCH 1/2] feat(bindings): expose layered RaBitQ build and query options --- java/lance-jni/src/blocking_scanner.rs | 7 +++ java/lance-jni/src/utils.rs | 8 ++- .../org/lance/index/vector/RQBuildParams.java | 13 +++++ java/src/main/java/org/lance/ipc/Query.java | 30 ++++++++++ .../main/java/org/lance/ipc/RQPrecision.java | 30 ++++++++++ java/src/test/java/org/lance/JNITest.java | 15 +++++ .../java/org/lance/index/VectorIndexTest.java | 32 +++++++++- python/python/lance/dataset.py | 30 +++++++++- python/python/tests/test_vector_index.py | 58 +++++++++++++++++++ python/src/dataset.rs | 37 +++++++++++- 10 files changed, 253 insertions(+), 7 deletions(-) create mode 100644 java/src/main/java/org/lance/ipc/RQPrecision.java diff --git a/java/lance-jni/src/blocking_scanner.rs b/java/lance-jni/src/blocking_scanner.rs index 9d65220aa53..33b06efc1e3 100644 --- a/java/lance-jni/src/blocking_scanner.rs +++ b/java/lance-jni/src/blocking_scanner.rs @@ -420,6 +420,13 @@ pub(crate) fn build_scanner_with_options<'a>( let approx_mode_str = env.get_string_from_method(&java_obj, "getApproxModeString")?; scanner.approx_mode(parse_approx_mode(&approx_mode_str)?); + scanner.rq_precision( + env.get_string_from_method(&java_obj, "getRqPrecisionString")? + .parse()?, + ); + if let Some(factor) = env.get_optional_u32_from_method(&java_obj, "getRqCascadeFactor")? { + scanner.rq_cascade_factor(factor)?; + } Ok(()) })?; diff --git a/java/lance-jni/src/utils.rs b/java/lance-jni/src/utils.rs index 5a3f795880c..fa8551ba662 100644 --- a/java/lance-jni/src/utils.rs +++ b/java/lance-jni/src/utils.rs @@ -315,6 +315,9 @@ pub fn get_query(env: &mut JNIEnv, query_obj: JObject) -> Result> .i()?; let approx_mode_str = env.get_string_from_method(&java_obj, "getApproxModeString")?; let approx_mode = parse_approx_mode(&approx_mode_str)?; + let rq_precision = env + .get_string_from_method(&java_obj, "getRqPrecisionString")? + .parse()?; Ok(Query { column, @@ -331,6 +334,8 @@ pub fn get_query(env: &mut JNIEnv, query_obj: JObject) -> Result> dist_q_c: 0.0, query_parallelism, approx_mode, + rq_precision, + rq_cascade_factor: env.get_optional_u32_from_method(&java_obj, "getRqCascadeFactor")?, }) })?; @@ -527,7 +532,8 @@ pub fn get_vector_index_params( "getRqParams", |env, rq_obj| { let num_bits = env.call_method(&rq_obj, "getNumBits", "()B", &[])?.b()? as u8; - Ok(RQBuildParams::new(num_bits)) + let layered = env.call_method(&rq_obj, "getLayered", "()Z", &[])?.z()?; + Ok(RQBuildParams::new(num_bits).with_layered(layered)) }, )?; diff --git a/java/src/main/java/org/lance/index/vector/RQBuildParams.java b/java/src/main/java/org/lance/index/vector/RQBuildParams.java index 49db5d9567e..e25a25a4b4c 100755 --- a/java/src/main/java/org/lance/index/vector/RQBuildParams.java +++ b/java/src/main/java/org/lance/index/vector/RQBuildParams.java @@ -18,13 +18,16 @@ /** Parameters for building a Rabit Quantizer (RQ) index stage. Defaults to 5 bits per dimension. */ public class RQBuildParams { private final byte numBits; + private final boolean layered; private RQBuildParams(Builder builder) { this.numBits = builder.numBits; + this.layered = builder.layered; } public static class Builder { private byte numBits = 5; + private boolean layered = false; public Builder() {} @@ -37,11 +40,21 @@ public Builder setNumBits(byte numBits) { return this; } + /** Opt into prefix planes (5, 7 or 9 bits); requires a layered-capable reader. */ + public Builder setLayered(boolean layered) { + this.layered = layered; + return this; + } + public RQBuildParams build() { return new RQBuildParams(this); } } + public boolean getLayered() { + return layered; + } + public byte getNumBits() { return numBits; } diff --git a/java/src/main/java/org/lance/ipc/Query.java b/java/src/main/java/org/lance/ipc/Query.java index 215865310df..ae503b81b35 100644 --- a/java/src/main/java/org/lance/ipc/Query.java +++ b/java/src/main/java/org/lance/ipc/Query.java @@ -33,6 +33,8 @@ public class Query { private final boolean useIndex; private final int queryParallelism; private final ApproxMode approxMode; + private final RQPrecision rqPrecision; + private final Optional rqCascadeFactor; private Query(Builder builder) { this.column = Preconditions.checkNotNull(builder.column, "Columns must be set"); @@ -54,6 +56,8 @@ private Query(Builder builder) { this.useIndex = builder.useIndex; this.queryParallelism = builder.queryParallelism; this.approxMode = builder.approxMode; + this.rqPrecision = builder.rqPrecision; + this.rqCascadeFactor = builder.rqCascadeFactor; } public String getColumn() { @@ -100,6 +104,18 @@ public int getQueryParallelism() { return queryParallelism; } + public Optional getRqCascadeFactor() { + return rqCascadeFactor; + } + + public RQPrecision getRqPrecision() { + return rqPrecision; + } + + public String getRqPrecisionString() { + return rqPrecision.toRustString(); + } + public ApproxMode getApproxMode() { return approxMode; } @@ -137,6 +153,8 @@ public static class Builder { private boolean useIndex = true; private int queryParallelism = 0; private ApproxMode approxMode = ApproxMode.NORMAL; + private RQPrecision rqPrecision = RQPrecision.FULL; + private Optional rqCascadeFactor = Optional.empty(); /** * Sets the column to be searched. @@ -296,6 +314,18 @@ public Builder setQueryParallelism(int queryParallelism) { * @param approxMode The approximate search mode to use for the query. * @return The Builder instance for method chaining. */ + /** Select the precision of a layered IVF_RQ query; default FULL. */ + /** Candidate oversampling; absent disables cascading. Rust validates positive values. */ + public Builder setRqCascadeFactor(int factor) { + this.rqCascadeFactor = Optional.of(factor); + return this; + } + + public Builder setRqPrecision(RQPrecision precision) { + this.rqPrecision = Preconditions.checkNotNull(precision, "RQPrecision must not be null"); + return this; + } + public Builder setApproxMode(ApproxMode approxMode) { this.approxMode = Preconditions.checkNotNull(approxMode, "ApproxMode must not be null"); return this; diff --git a/java/src/main/java/org/lance/ipc/RQPrecision.java b/java/src/main/java/org/lance/ipc/RQPrecision.java new file mode 100644 index 00000000000..b693c0f263a --- /dev/null +++ b/java/src/main/java/org/lance/ipc/RQPrecision.java @@ -0,0 +1,30 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.lance.ipc; + +/** Search precision for a layered IVF_RQ index. */ +public enum RQPrecision { + SIGN("sign"), + HIGH("high"), + FULL("full"); + private final String value; + + RQPrecision(String value) { + this.value = value; + } + + public String toRustString() { + return value; + } +} diff --git a/java/src/test/java/org/lance/JNITest.java b/java/src/test/java/org/lance/JNITest.java index 94db13d6dea..e1e5949f74a 100644 --- a/java/src/test/java/org/lance/JNITest.java +++ b/java/src/test/java/org/lance/JNITest.java @@ -23,6 +23,7 @@ import org.lance.index.vector.VectorIndexParams; import org.lance.ipc.ApproxMode; import org.lance.ipc.Query; +import org.lance.ipc.RQPrecision; import org.lance.test.JniTestHelper; import org.junit.jupiter.api.Test; @@ -54,6 +55,18 @@ public void testQuery() { Query defaultQuery = new Query.Builder().setColumn("column").setKey(new float[] {1.0f, 2.0f, 3.0f}).build(); assertEquals(ApproxMode.NORMAL, defaultQuery.getApproxMode()); + assertEquals(RQPrecision.FULL, defaultQuery.getRqPrecision()); + assertEquals(Optional.empty(), defaultQuery.getRqCascadeFactor()); + for (RQPrecision precision : RQPrecision.values()) { + Query query = + new Query.Builder() + .setColumn("vector") + .setKey(new float[] {1.0f, 2.0f}) + .setRqPrecision(precision) + .build(); + assertEquals(precision.toRustString(), query.getRqPrecisionString()); + JniTestHelper.parseQuery(Optional.of(query)); + } JniTestHelper.parseQuery( Optional.of( @@ -82,6 +95,8 @@ public void testIvfFlatIndexParams() { @Test public void testRqBuildParamsDefaultNumBits() { assertEquals((byte) 5, new RQBuildParams.Builder().build().getNumBits()); + assertEquals(false, new RQBuildParams.Builder().build().getLayered()); + assertEquals(true, new RQBuildParams.Builder().setLayered(true).build().getLayered()); } @Test diff --git a/java/src/test/java/org/lance/index/VectorIndexTest.java b/java/src/test/java/org/lance/index/VectorIndexTest.java index 6e5a7e8562d..c1b904ccb46 100755 --- a/java/src/test/java/org/lance/index/VectorIndexTest.java +++ b/java/src/test/java/org/lance/index/VectorIndexTest.java @@ -25,6 +25,7 @@ import org.lance.ipc.AsyncScanner; import org.lance.ipc.LanceScanner; import org.lance.ipc.Query; +import org.lance.ipc.RQPrecision; import org.lance.ipc.ScanOptions; import org.apache.arrow.memory.RootAllocator; @@ -446,14 +447,16 @@ public void testCreateIvfSqIndexDistributively(@TempDir Path tempDir) throws Exc } } - @Test - public void testCreateIvfRqIndex(@TempDir Path tempDir) throws Exception { + @ParameterizedTest + @ValueSource(ints = {1, 5, 7}) + public void testCreateIvfRqIndex(int bits, @TempDir Path tempDir) throws Exception { Path datasetPath = tempDir.resolve("ivf_rq_index"); try (TestVectorDataset testVectorDataset = new TestVectorDataset(datasetPath)) { try (Dataset dataset = testVectorDataset.create()) { IvfBuildParams ivf = new IvfBuildParams.Builder().setNumPartitions(2).build(); - RQBuildParams rq = new RQBuildParams.Builder().setNumBits((byte) 1).build(); + RQBuildParams rq = + new RQBuildParams.Builder().setNumBits((byte) bits).setLayered(bits > 1).build(); VectorIndexParams vectorIndexParams = VectorIndexParams.withIvfRqParams(DistanceType.L2, ivf, rq); @@ -486,6 +489,29 @@ public void testCreateIvfRqIndex(@TempDir Path tempDir) throws Exception { assertTrue( indexType == IndexType.VECTOR || indexType == IndexType.IVF_RQ, "IndexType for IVF_RQ index should be VECTOR or IVF_RQ but was " + indexType); + if (bits > 1) { + for (RQPrecision precision : RQPrecision.values()) { + Query.Builder query = + new Query.Builder() + .setColumn(TestVectorDataset.vectorColumnName) + .setKey(new float[32]) + .setK(5) + .setNprobes(2) + .setRqPrecision(precision); + if (precision == RQPrecision.FULL) { + query.setRqCascadeFactor(8); + } + try (LanceScanner scanner = + dataset.newScan(new ScanOptions.Builder().nearest(query.build()).build()); + ArrowReader reader = scanner.scanBatches()) { + int rows = 0; + while (reader.loadNextBatch()) { + rows += reader.getVectorSchemaRoot().getRowCount(); + } + assertEquals(5, rows); + } + } + } } } } diff --git a/python/python/lance/dataset.py b/python/python/lance/dataset.py index bb96abcf3ed..5a1b41750ee 100644 --- a/python/python/lance/dataset.py +++ b/python/python/lance/dataset.py @@ -4482,6 +4482,10 @@ def create_index( - num_bits The number of bits for RQ (Rabit Quantization). Default is 5. + - layered + Opt in to independently readable sign/high/low planes. Default is + False. Supported total widths are 5 (1+2+2), 7 (1+4+2), and 9 + (1+4+4); this layout requires a layered-capable reader. Optional parameters for `IVF_HNSW_*`: max_level @@ -7144,6 +7148,8 @@ def nearest( query_parallelism: Optional[int] = None, approx_mode: Literal["fast", "normal", "accurate"] = "normal", distance_range: Optional[tuple[Optional[float], Optional[float]]] = None, + rq_precision: Literal["sign", "high", "full"] = "full", + rq_cascade_factor: Optional[int] = None, ) -> ScannerBuilder: """Configure nearest neighbor search. @@ -7172,7 +7178,19 @@ def nearest( RQ-quantized indexes, such as IVF_RQ. Other index types ignore this setting. ``fast`` favors lower latency and may reduce recall, ``normal`` uses the default balance, and ``accurate`` favors higher - recall and may increase latency. + recall and may increase latency. Layered RQ uses the native angular + confidence policy in normal mode and a conservative estimator-difference + pruning bound in accurate mode. Neither requires raw-vector refinement. + rq_precision: {"sign", "high", "full"}, default "full" + Select a stored level of a layered IVF_RQ index. High uses 3 total + bits for RQ5 and 5 total bits for RQ7/RQ9. Non-layered indexes only + support full. The legacy approx_mode="fast" still scores binary codes. + rq_cascade_factor: int, optional + Positive oversampling factor for quantized candidate reranking. + Requires layered full precision and normal/accurate approx_mode. + Lower planes are fetched only for candidates; no raw-vector Take is + required. refine_factor may additionally refine with data-file vectors. + Range queries and row-id remapping use complete partition scoring. """ self._nearest = _build_vector_search_query( column, @@ -7187,6 +7205,8 @@ def nearest( use_index=use_index, ef=ef, query_parallelism=query_parallelism, + rq_precision=rq_precision, + rq_cascade_factor=rq_cascade_factor, approx_mode=approx_mode, distance_range=distance_range, ) @@ -8397,6 +8417,8 @@ def _build_vector_search_query( query_parallelism: Optional[int] = None, approx_mode: Literal["fast", "normal", "accurate"] = "normal", distance_range: Optional[tuple[Optional[float], Optional[float]]] = None, + rq_precision: Literal["sign", "high", "full"] = "full", + rq_cascade_factor: Optional[int] = None, ) -> dict: """Configure nearest neighbor search. @@ -8541,6 +8563,8 @@ def _build_vector_search_query( "ef": ef, "query_parallelism": query_parallelism, "approx_mode": approx_mode, + "rq_precision": rq_precision, + "rq_cascade_factor": rq_cascade_factor, "distance_range": distance_range, } @@ -8697,6 +8721,8 @@ def __init__( use_index: bool = True, ef: Optional[int] = None, query_parallelism: Optional[int] = None, + rq_precision: Literal["sign", "high", "full"] = "full", + rq_cascade_factor: Optional[int] = None, approx_mode: Literal["fast", "normal", "accurate"] = "normal", ): self._inner = _build_vector_search_query( @@ -8711,6 +8737,8 @@ def __init__( use_index=use_index, ef=ef, query_parallelism=query_parallelism, + rq_precision=rq_precision, + rq_cascade_factor=rq_cascade_factor, approx_mode=approx_mode, ) diff --git a/python/python/tests/test_vector_index.py b/python/python/tests/test_vector_index.py index ae51beca6c0..e0d595f192d 100644 --- a/python/python/tests/test_vector_index.py +++ b/python/python/tests/test_vector_index.py @@ -3745,3 +3745,61 @@ def test_fts_filter_vector_search(tmp_path): with pytest.raises(ValueError): scanner.to_table() + + +@pytest.mark.parametrize("num_bits", [5, 7, 9]) +def test_layered_rq_precision_and_cascade(tmp_path, num_bits): + rng = np.random.default_rng(472) + vectors = rng.normal(size=(512, 64)).astype(np.float32) + table = pa.table({"vector": pa.FixedSizeListArray.from_arrays(vectors.ravel(), 64)}) + dataset = lance.write_dataset(table, tmp_path, max_rows_per_file=128) + dataset.create_index( + "vector", + index_type="IVF_RQ", + num_bits=num_bits, + layered=True, + num_partitions=4, + ) + query = vectors[0] + nearest = {"column": "vector", "q": query, "k": 100, "nprobes": 4} + exact = dataset.to_table( + columns=[], with_row_id=True, nearest={**nearest, "use_index": False} + )["_rowid"].to_pylist() + full = None + for precision in ["sign", "high", "full"]: + result = dataset.to_table( + columns=[], with_row_id=True, nearest={**nearest, "rq_precision": precision} + ) + assert len(set(result["_rowid"].to_pylist()) & set(exact)) >= 50 + if precision == "full": + full = result + cold = lance.dataset(tmp_path, index_cache_size_bytes=0) + cascade = cold.to_table( + columns=[], with_row_id=True, nearest={**nearest, "rq_cascade_factor": 32} + ) + assert cascade.equals(full) + with pytest.raises(ValueError, match="rq_precision"): + dataset.to_table(nearest={**nearest, "rq_precision": "invalid"}) + with pytest.raises(ValueError, match="positive"): + dataset.to_table(nearest={**nearest, "rq_cascade_factor": 0}) + + +@pytest.mark.parametrize("num_bits", [1, 2, 3, 4, 6, 8]) +def test_layered_rq_rejects_unsupported_bits(tmp_path, num_bits): + vectors = np.zeros((64, 64), dtype=np.float32) + dataset = lance.write_dataset( + pa.table( + { + "vector": pa.FixedSizeListArray.from_arrays(vectors.ravel(), 64), + } + ), + tmp_path, + ) + with pytest.raises(ValueError, match="requires num_bits=5, 7 or 9"): + dataset.create_index( + "vector", + index_type="IVF_RQ", + num_bits=num_bits, + layered=True, + num_partitions=1, + ) diff --git a/python/src/dataset.rs b/python/src/dataset.rs index dbb6220d1ae..c2f5fb03a2b 100644 --- a/python/src/dataset.rs +++ b/python/src/dataset.rs @@ -1539,6 +1539,8 @@ impl Dataset { ef, query_parallelism, approx_mode, + rq_precision, + rq_cascade_factor, ) = vector_query_params_from_dict(nearest, default_k)?; let (_, element_type) = get_vector_type(self_.ds.schema(), &column) @@ -1590,7 +1592,7 @@ impl Dataset { }; scanner - .map(|s| { + .and_then(|s| { let mut s = s.minimum_nprobes(minimum_nprobes); if let Some(maximum_nprobes) = maximum_nprobes { s = s.maximum_nprobes(maximum_nprobes); @@ -1606,11 +1608,15 @@ impl Dataset { } s = s.query_parallelism(query_parallelism); s = s.approx_mode(approx_mode); + s.rq_precision(rq_precision); + if let Some(factor) = rq_cascade_factor { + s.rq_cascade_factor(factor)?; + } s.use_index(use_index); if let Some((lower, upper)) = distance_range { s.distance_range(lower, upper); } - s + Ok(s) }) .map_err(|err| PyValueError::new_err(err.to_string()))?; } @@ -5149,6 +5155,9 @@ fn prepare_vector_index_params( } // Parse PQ/RQ params + if let Some(value) = kwargs.get_item("layered")? { + rq_params.layered = value.extract()?; + } if let Some(n) = kwargs.get_item("num_bits")? { let num_bits: u8 = n.extract()?; pq_params.num_bits = num_bits as usize; @@ -5517,6 +5526,8 @@ type VectorQueryParams = ( Option, i32, ApproxMode, + lance_index::vector::bq::layered::RQPrecision, + Option, ); fn extract_query_parallelism(value: &Bound<'_, PyAny>) -> PyResult { @@ -5662,7 +5673,23 @@ fn vector_query_params_from_dict( let query_parallelism = vector_query_query_parallelism_from_dict(dict)?; let approx_mode = vector_query_approx_mode_from_dict(dict)?; + let rq_precision = dict + .get_item("rq_precision")? + .filter(|v| !v.is_none()) + .map(|v| v.extract::()) + .transpose()? + .map(|v| { + v.parse::() + .map_err(|e| PyValueError::new_err(e.to_string())) + }) + .transpose()? + .unwrap_or_default(); + let rq_cascade_factor = dict + .get_item("rq_cascade_factor")? + .filter(|v| !v.is_none()) + .map(|v| v.extract::()) + .transpose()?; Ok(( column, key, @@ -5675,6 +5702,8 @@ fn vector_query_params_from_dict( ef, query_parallelism, approx_mode, + rq_precision, + rq_cascade_factor, )) } @@ -5712,6 +5741,8 @@ impl PySearchFilter { ef, query_parallelism, approx_mode, + rq_precision, + rq_cascade_factor, ) = vector_query_params_from_dict(query, default_k)?; let metric_type = Some(metric_type_opt.unwrap_or(MetricType::L2)); @@ -5731,6 +5762,8 @@ impl PySearchFilter { query_parallelism, dist_q_c: 0.0, approx_mode, + rq_precision, + rq_cascade_factor, }; Ok(Self { From 99417ede8aa6d4b8c7eaa20802d3438506793dea Mon Sep 17 00:00:00 2001 From: Yang Cen <159225399+BubbleCal@users.noreply.github.com> Date: Sun, 11 Oct 2026 20:02:57 +0800 Subject: [PATCH 2/2] test(java): assert layered RaBitQ recall across all widths --- .../src/test/java/org/lance/index/VectorIndexTest.java | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/java/src/test/java/org/lance/index/VectorIndexTest.java b/java/src/test/java/org/lance/index/VectorIndexTest.java index a0bbff6b34a..cd854f07019 100755 --- a/java/src/test/java/org/lance/index/VectorIndexTest.java +++ b/java/src/test/java/org/lance/index/VectorIndexTest.java @@ -36,6 +36,7 @@ import org.junit.jupiter.params.ParameterizedTest; import org.junit.jupiter.params.provider.CsvSource; import org.junit.jupiter.params.provider.EnumSource; +import org.junit.jupiter.params.provider.ValueSource; import java.nio.file.Path; import java.util.Collections; @@ -529,10 +530,19 @@ public void testCreateIvfRqIndex(int bits, @TempDir Path tempDir) throws Excepti dataset.newScan(new ScanOptions.Builder().nearest(query.build()).build()); ArrowReader reader = scanner.scanBatches()) { int rows = 0; + int correct = 0; while (reader.loadNextBatch()) { rows += reader.getVectorSchemaRoot().getRowCount(); + IntVector ids = (IntVector) reader.getVectorSchemaRoot().getVector("i"); + for (int i = 0; i < ids.getValueCount(); i++) { + // Each of the five fragments has the same nearest vector at offset zero. + if (ids.get(i) % 80 == 0) { + correct++; + } + } } assertEquals(5, rows); + assertTrue(correct / 5.0 >= 0.5, "Expected recall >= 0.5 for " + precision); } } }