diff --git a/java/lance-jni/src/blocking_scanner.rs b/java/lance-jni/src/blocking_scanner.rs index 85fda0a97cf..6ad09fed920 100644 --- a/java/lance-jni/src/blocking_scanner.rs +++ b/java/lance-jni/src/blocking_scanner.rs @@ -532,6 +532,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 782e3941919..a7ee30cd25b 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()?; let search_effort = env .call_method(&java_obj, "getSearchEffort", "()D", &[])? @@ -335,6 +338,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")?, }; query.validate_search_effort()?; Ok(query) @@ -533,7 +538,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..665674e4ad8 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 (3 through 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 00edc18508d..26c47d1dc93 100644 --- a/java/src/main/java/org/lance/ipc/Query.java +++ b/java/src/main/java/org/lance/ipc/Query.java @@ -35,6 +35,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"); @@ -65,6 +67,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() { @@ -136,6 +140,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; } @@ -178,6 +194,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. @@ -401,6 +419,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 599399dd65a..be1469d96d9 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; @@ -56,6 +57,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)); + } assertEquals(0.5, defaultQuery.getSearchEffort()); assertEquals(Optional.empty(), defaultQuery.getSearchEffortOption()); @@ -146,6 +159,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 f03a98a43d7..cd854f07019 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; @@ -35,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; @@ -470,14 +472,16 @@ public void testCreateIvfSqIndexDistributively(@TempDir Path tempDir) throws Exc } } - @Test - public void testCreateIvfRqIndex(@TempDir Path tempDir) throws Exception { + @ParameterizedTest + @ValueSource(ints = {1, 3, 4, 5, 6, 7, 8, 9}) + 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); @@ -510,6 +514,38 @@ 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; + 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); + } + } + } } } } diff --git a/python/python/lance/dataset.py b/python/python/lance/dataset.py index d08fa1e5117..8f3173843ca 100644 --- a/python/python/lance/dataset.py +++ b/python/python/lance/dataset.py @@ -4539,6 +4539,11 @@ 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. Supports total widths 3 through 9: 3 (1+1+1), 4 (1+2+1), + 5 (1+2+2), 6 (1+4+1), 7 (1+4+2), 8 (1+4+3), and 9 (1+4+4). + This layout requires a layered-capable reader. Optional parameters for `IVF_HNSW_*`: max_level @@ -7206,6 +7211,8 @@ def nearest( approx_mode: Literal["fast", "normal", "accurate"] = "normal", distance_range: Optional[tuple[Optional[float], Optional[float]]] = None, search_effort: Optional[float] = None, + rq_precision: Literal["sign", "high", "full"] = "full", + rq_cascade_factor: Optional[int] = None, ) -> ScannerBuilder: """Configure nearest neighbor search. @@ -7243,7 +7250,20 @@ 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 2 total + bits for RQ3, 3 for RQ4/RQ5, and 5 for RQ6 through 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, @@ -7258,6 +7278,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, search_effort=search_effort, distance_range=distance_range, @@ -8499,6 +8521,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, search_effort: Optional[float] = None, ) -> dict: """Configure nearest neighbor search. @@ -8646,6 +8670,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, "search_effort": search_effort, "distance_range": distance_range, } @@ -8812,6 +8838,8 @@ def __init__( query_parallelism: Optional[int] = None, approx_mode: Literal["fast", "normal", "accurate"] = "normal", search_effort: Optional[float] = None, + rq_precision: Literal["sign", "high", "full"] = "full", + rq_cascade_factor: Optional[int] = None, ): self._inner = _build_vector_search_query( column, @@ -8825,6 +8853,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, search_effort=search_effort, ) diff --git a/python/python/tests/test_vector_index.py b/python/python/tests/test_vector_index.py index d8f4bff79f0..e9e79a4ba43 100644 --- a/python/python/tests/test_vector_index.py +++ b/python/python/tests/test_vector_index.py @@ -4250,3 +4250,61 @@ def test_fts_filter_vector_search(tmp_path): with pytest.raises(ValueError): scanner.to_table() + + +@pytest.mark.parametrize("num_bits", range(3, 10)) +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]) +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 in 3..=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 2a289857410..e2498cecb70 100644 --- a/python/src/dataset.rs +++ b/python/src/dataset.rs @@ -1552,6 +1552,8 @@ impl Dataset { ef, query_parallelism, approx_mode, + rq_precision, + rq_cascade_factor, search_effort, ) = vector_query_params_from_dict(nearest, default_k)?; @@ -1624,6 +1626,10 @@ 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 = s.search_effort(search_effort)?; s.use_index(use_index); if let Some((lower, upper)) = distance_range { @@ -5342,6 +5348,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; @@ -5735,6 +5744,8 @@ type VectorQueryParams = ( Option, i32, ApproxMode, + lance_index::vector::bq::layered::RQPrecision, + Option, f64, ); @@ -5888,7 +5899,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, @@ -5901,6 +5928,8 @@ fn vector_query_params_from_dict( ef, query_parallelism, approx_mode, + rq_precision, + rq_cascade_factor, search_effort, )) } @@ -5939,6 +5968,8 @@ impl PySearchFilter { ef, query_parallelism, approx_mode, + rq_precision, + rq_cascade_factor, search_effort, ) = vector_query_params_from_dict(query, default_k)?; @@ -5959,6 +5990,8 @@ impl PySearchFilter { query_parallelism, dist_q_c: 0.0, approx_mode, + rq_precision, + rq_cascade_factor, search_effort, }; vector_query