diff --git a/docs/src/format/index/index.md b/docs/src/format/index/index.md index 01970133208..6121011e750 100644 --- a/docs/src/format/index/index.md +++ b/docs/src/format/index/index.md @@ -100,10 +100,11 @@ Index segments are created and updated through a transactional process: - `name`: The index name (must match existing segments if adding to an existing index) - `fields`: The columns the index depends on: the keyed column(s) it is searched on, followed by any merely-carried columns named in `covering_fields`. `fields[0]` is always a keyed column. - - `covering_fields`: The trailing subset of `fields` whose values the index carries but is not - keyed on, letting a query that only projects those columns be answered without a fragment take. - Empty for an index that carries no extra columns. Declaring a column here does not by itself - make it servable -- see [Serving carried columns](#serving-carried-columns). + - `covering_fields`: The trailing subset of `fields` whose values the index carries, letting a + query that only projects those columns be answered without a fragment take. Usually these are + columns the index is not keyed on, but a keyed column may also be carried. Empty for an index + that carries no extra columns. Declaring a column here does not by itself make it servable -- + see [Serving carried columns](#serving-carried-columns). - `fragment_bitmap`: The set of fragment IDs covered by this segment - `index_details`: Index-specific configuration and parameters - `version`: The format version of this index type @@ -141,18 +142,36 @@ fragments that would have been covered by that segment. carries. It does not establish that the segment's storage holds their values. **The segment's storage schema is authoritative.** Before answering a query from a -carried column, an engine must confirm that column is present in the storage it opened, -and fall back to a take against the base table when it is not. A segment whose -declaration names a column its storage does not hold is a legal state, not corruption: -a maintenance operation that cannot carry the payload through a rebuild is permitted to -withdraw it and leave the declaration standing. - -!!! note "Current state" - - No index builder writes carried values yet, so today every declaration is ahead of - its storage. Engines that read `covering_fields` must therefore treat it purely as a - declaration and serve every column from the base table until they have verified the - storage themselves. This is transitional; the rule above is not. +carried column, an engine must confirm that column is present and bound to the declared +logical field in the storage it opened, and fall back to a take against the base table +when it cannot. A segment whose declaration names a column its storage does not hold is +a legal state, not corruption: a maintenance operation that cannot carry the payload +through a rebuild is permitted to withdraw it and leave the declaration standing. + +!!! note "Capability varies by segment" + + Whether a segment's storage holds a declared column depends on the index type, on the + writer that produced the segment, and on what later maintenance did to it, so one + logical index may hold values for some of its segments and not others. An engine + therefore verifies each selected segment rather than inferring capability from the + index type, the writer version, or the declaration alone, and serves from the base + table every column it cannot verify. + +!!! note "A keyed column may also be carried" + + `covering_fields` usually names columns the index is *not* keyed on, but an index is + permitted to carry a column it is also keyed on -- for instance a vector index that + keeps full-precision vectors so a refine pass can re-rank without a base-table take. + The id then appears twice in `fields`, once as `fields[0]` and again as the trailing + carried entry, and once in `covering_fields`. This is the only case in which an id + repeats in `fields`. + + A reader must therefore take the carried set from `covering_fields` directly, and + never derive it by subtracting the keyed prefix from `fields`: that set difference + silently drops a column that is both. The trailing-subset rule is stated over + `covering_fields` and is unaffected: `fields[0]` remains the column the index is + searched on, and an engine serves the repeated id from storage like any other + carried column. ## Loading an index @@ -175,9 +194,9 @@ The `IndexMetadata` message contains important information about the index segme - `fields`: the columns the index depends on: the keyed column(s) the index is searched on, followed by any columns it merely carries, as named in `covering_fields`. `fields[0]` is always a keyed column. - `covering_fields`: the trailing subset of `fields` whose values the index carries alongside its own - data but is not keyed on. Empty for an index that carries no extra columns. This declaration is - not authoritative for what the segment can serve -- see - [Serving carried columns](#serving-carried-columns). + data -- usually columns it is not keyed on, though a keyed column may also be carried. Empty for an + index that carries no extra columns. This declaration is not authoritative for what the segment can + serve -- see [Serving carried columns](#serving-carried-columns). - `fragment_bitmap`: the set of fragment IDs covered by this index segment. - `index_details`: a protobuf `Any` message that contains index-specific details, such as index type, parameters, and storage format. This allows different index types to store their own metadata. diff --git a/docs/src/format/index/vector/index.md b/docs/src/format/index/vector/index.md index 91d73a23a16..709b9a58cc7 100644 --- a/docs/src/format/index/vector/index.md +++ b/docs/src/format/index/vector/index.md @@ -161,6 +161,19 @@ the Arrow schema of the Lance file varies depending on the quantization method u !!! note All partitions are stored in the same file, and partitions must be written in order. +Every quantization format below lists only its internal columns. When a V3 IVF +writer materializes carried values, it appends one trailing column per carried +field after them, named and typed exactly as in the dataset schema. This physical +payload may be a subset of the manifest's `covering_fields` declaration (see +[Index Metadata](../index.md)). A reader returns only columns whose physical +schema and source field ids it verifies across every selected segment; all other +projected columns come from a base-table take. + +A reader discovers carried columns by exclusion, not by position: any column in the +auxiliary file's schema that is not one of the quantizer's internal columns is a +carried column. Writers append them in trailing order, but a reader must not depend +on that ordering to identify them. + ##### FLAT No quantization applied - stores original vectors in their full precision: @@ -229,6 +242,17 @@ Contains RabitQ-specific metadata in JSON format (only present for RQ quantizati This includes the rotation matrix position, number of bits, and packing information. See the RQ metadata specification in the "storage_metadata" section below. +##### "covering_field_ids" + +The *source dataset* field ids of the storage file's physical carried columns, +comma separated in physical schema order (only present when the storage carries +values). Arrow fields carry no Lance field id, so names and types alone cannot +prove which logical column a payload came from. Readers use these ids to bind +physical values to the segment's `covering_fields` declaration, and treat missing, +malformed, ambiguous, or mismatched metadata as no servable carried capability. +Distributed merges use the same identity to reject shards whose columns match by +name and type but come from different fields. + ##### "storage_metadata" Contains quantizer-specific metadata as a list of JSON strings. diff --git a/docs/src/guide/performance.md b/docs/src/guide/performance.md index 8d7f0c5401f..e749f68a1d0 100644 --- a/docs/src/guide/performance.md +++ b/docs/src/guide/performance.md @@ -529,3 +529,17 @@ Set `LANCE_DISABLE_AMX=1` to take the AMX paths out of service without rebuildin A/B measurement, or to get the previous behaviour back. Because it also moves partition assignment back to the approximate path, an index built with it set is not equivalent to one built without it; compare recall, not just build time. + +#### Covering Columns + +A vector index can carry the values of extra columns beside its vectors, so a query whose +projection they satisfy is answered from the index without a take against the base table. +That trade only pays off where the search settles into a single global top-k heap, because +the covering read is then bounded by the query's survivors. Every HNSW index, and any query +with `query_parallelism` above one, emits results per partition instead, so covering reads +scale with the number of partitions probed rather than with `k`, and are issued serially. + +On those shapes covering is *slower* than the base-table take it exists to avoid — roughly +2.8x a plain index's latency warm at nprobe 32, and 1.18x cold. Prefer IVF_PQ or IVF_FLAT +when declaring covering columns; index creation logs a warning when covering is combined +with HNSW. diff --git a/java/lance-jni/src/utils.rs b/java/lance-jni/src/utils.rs index c9d2d2005b1..12835f8b112 100644 --- a/java/lance-jni/src/utils.rs +++ b/java/lance-jni/src/utils.rs @@ -24,7 +24,7 @@ use lance_linalg::distance::DistanceType; use crate::error::{Error, Result}; use crate::ffi::JNIEnvExt; -use crate::traits::FromJObjectWithEnv; +use crate::traits::{FromJObjectWithEnv, import_vec_from_method}; use lance_index::vector::{ApproxMode, Query}; use std::collections::HashMap; use std::str::FromStr; @@ -261,6 +261,11 @@ pub fn build_compaction_options( } // Convert from Java Optional to Rust Option +// +// This builds a `Query` directly rather than through `Scanner::nearest`, and is used +// only by `JniTestHelper.parseQuery` to check the Java-side field marshalling. Java's +// real search path is `blocking_scanner.rs`, which configures a `Scanner`, so it picks +// up every plan-derived setting -- including the covering projection below. pub fn get_query(env: &mut JNIEnv, query_obj: JObject) -> Result> { let query = env.get_optional(&query_obj, |env, java_obj| { let column = env.get_string_from_method(&java_obj, "getColumn")?; @@ -305,6 +310,11 @@ pub fn get_query(env: &mut JNIEnv, query_obj: JObject) -> Result> dist_q_c: 0.0, query_parallelism, approx_mode, + // Not a user-settable search parameter: the covering projection is derived + // per plan from what the scan reads, so it stays `None` (materialize + // whatever the index declares) at the binding boundary. Java's real search + // path resolves it in `Scanner`; see the note on this function. + covering_projection: None, }) })?; @@ -509,12 +519,36 @@ pub fn get_vector_index_params( stages.push(StageParams::RQ(rq_params)); } + // Covering ("included") columns: names of extra dataset columns stored inline in + // the index. Empty when absent. The core validates them when the index is built. + let covering_columns: Vec = import_vec_from_method( + env, + &vector_index_params_obj, + "getCoveringColumns", + |env, elem| Ok(env.get_string(&JString::from(elem))?.into()), + )?; + + // Carry full-precision vectors in the index so `refine` re-ranks from index + // storage instead of taking them from the base table. Deliberately not part of + // the covering columns above: the core records it by naming the indexed column + // itself, which the covering setter rejects. The core validates it at build. + let store_vectors_for_refine = env + .call_method( + &vector_index_params_obj, + "isStoreVectorsForRefine", + "()Z", + &[], + )? + .z()?; + Ok(VectorIndexParams { metric_type: distance_type, stages, version: IndexFileVersion::V3, skip_transpose: false, runtime_hints: Default::default(), + covering_columns, + store_vectors_for_refine, }) }, )?; diff --git a/java/src/main/java/org/lance/index/Index.java b/java/src/main/java/org/lance/index/Index.java index 620474ff787..bf4b1cf4418 100644 --- a/java/src/main/java/org/lance/index/Index.java +++ b/java/src/main/java/org/lance/index/Index.java @@ -16,6 +16,7 @@ import com.google.common.base.MoreObjects; import java.time.Instant; +import java.util.ArrayList; import java.util.Arrays; import java.util.Collections; import java.util.List; @@ -55,8 +56,19 @@ private Index( Long sizeBytes, IndexType indexType) { this.uuid = uuid; - this.fields = fields; - this.coveringFields = coveringFields; + // Both lists are defensively copied and unmodifiable so a caller cannot mutate them + // after construction: equals/hashCode would silently drift from the index files that + // actually carry the payload, and `coveringFields` is documented as the trailing slice + // of `fields`, so mutating either alone breaks that relationship. Null-coalesced so an + // explicit null from the builder reads as "none" rather than blowing up at first use. + this.fields = + fields == null + ? Collections.emptyList() + : Collections.unmodifiableList(new ArrayList<>(fields)); + this.coveringFields = + coveringFields == null + ? Collections.emptyList() + : Collections.unmodifiableList(new ArrayList<>(coveringFields)); this.name = name; this.datasetVersion = datasetVersion; this.fragments = fragments; @@ -89,7 +101,7 @@ public List fields() { *

These ids also appear in {@link #fields()} — that is deliberate, so that every consumer * reading {@code fields()} as the index's dependency set also covers them with no change. * - * @return the covering field IDs + * @return the covering field IDs, as an unmodifiable list */ public List coveringFields() { return coveringFields; diff --git a/java/src/main/java/org/lance/index/IndexDescription.java b/java/src/main/java/org/lance/index/IndexDescription.java index b1e2689a00c..09a182c192b 100755 --- a/java/src/main/java/org/lance/index/IndexDescription.java +++ b/java/src/main/java/org/lance/index/IndexDescription.java @@ -69,7 +69,14 @@ public String getName() { return name; } - /** Field ids that this index is built on. */ + /** + * Field ids that this index is built on -- the columns it can answer queries for. + * + *

This is the index's keyed prefix only. An index may additionally carry values for + * columns it is not keyed on; those are deliberately absent here, because the index cannot be + * searched on them. They stay reachable per segment via {@link Index#coveringFields()} on the + * entries of {@link #getMetadata()}. + */ public List getFieldIds() { return fieldIds; } diff --git a/java/src/main/java/org/lance/index/vector/VectorIndexParams.java b/java/src/main/java/org/lance/index/vector/VectorIndexParams.java index e8928943b48..451c454dcf7 100644 --- a/java/src/main/java/org/lance/index/vector/VectorIndexParams.java +++ b/java/src/main/java/org/lance/index/vector/VectorIndexParams.java @@ -17,6 +17,9 @@ import com.google.common.base.MoreObjects; +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; import java.util.Optional; /** Parameters for creating a vector index. */ @@ -27,6 +30,8 @@ public class VectorIndexParams { private final Optional hnswParams; private final Optional sqParams; private final Optional rqParams; + private final List coveringColumns; + private final boolean storeVectorsForRefine; private VectorIndexParams(Builder builder) { this.distanceType = builder.distanceType; @@ -35,6 +40,11 @@ private VectorIndexParams(Builder builder) { this.hnswParams = builder.hnswParams; this.sqParams = builder.sqParams; this.rqParams = builder.rqParams; + this.coveringColumns = + builder.coveringColumns == null + ? Collections.emptyList() + : Collections.unmodifiableList(new ArrayList<>(builder.coveringColumns)); + this.storeVectorsForRefine = builder.storeVectorsForRefine; validate(); } @@ -179,6 +189,8 @@ public static class Builder { private Optional hnswParams = Optional.empty(); private Optional sqParams = Optional.empty(); private Optional rqParams = Optional.empty(); + private List coveringColumns = Collections.emptyList(); + private boolean storeVectorsForRefine = false; /** * Create a new builder to create a vector index. @@ -235,6 +247,47 @@ public Builder setRqParams(RQBuildParams rqParams) { return this; } + /** + * Set the columns to cover ("include") in the index. Their values are stored inline in the + * index so a query projecting only covered columns is answered from the index without a take + * from the base table. Each must name a top-level, non-key column of the dataset; the columns + * are validated by the core when the index is built. Empty by default (no covering). + * + * @param coveringColumns the covering column names + * @return Builder + */ + public Builder setCoveringColumns(List coveringColumns) { + this.coveringColumns = coveringColumns; + return this; + } + + /** + * Carry full-precision vectors in the index storage so a query using {@code refineFactor} + * re-ranks from the index instead of taking the vectors from the base table. Off by default: it + * stores a second copy of the vector column, which is the largest column in most datasets. + * + *

Not expressible through {@link #setCoveringColumns}: the build pipeline rewrites the + * indexed column in place, so a covering column under that name would serve values that are not + * the user's vectors. The copy taken here predates those transforms. + * + *

Requires the V3 index file format, and is rejected on flat quantizers (which already store + * the vectors), on multivector columns, on a nested indexed column, and when precomputed + * shuffle buffers are used. The core validates this when the index is built. + * + *

While such an index exists the indexed column can no longer be renamed, nor have its + * nullability changed: the carried copy lives in index storage under that column's name and + * nothing renames it, so altering the column is refused rather than leaving an index whose + * optimize and compaction paths are permanently broken. Drop the index, alter the column, then + * rebuild. Dropping the column itself stays allowed -- that removes the index along with it. + * + * @param storeVectorsForRefine whether to carry full-precision vectors for refine + * @return Builder + */ + public Builder setStoreVectorsForRefine(boolean storeVectorsForRefine) { + this.storeVectorsForRefine = storeVectorsForRefine; + return this; + } + public VectorIndexParams build() { return new VectorIndexParams(this); } @@ -268,6 +321,26 @@ public Optional getRqParams() { return rqParams; } + /** + * Get the covering ("included") columns whose values are stored inline in the index. Empty when + * the index has no covering columns. + * + * @return the covering column names, as an unmodifiable list + */ + public List getCoveringColumns() { + return coveringColumns; + } + + /** + * Whether the index carries full-precision vectors for {@code refine}. See {@link + * Builder#setStoreVectorsForRefine(boolean)}. + * + * @return true when carried refine vectors were requested + */ + public boolean isStoreVectorsForRefine() { + return storeVectorsForRefine; + } + @Override public String toString() { return MoreObjects.toStringHelper(this) @@ -277,6 +350,8 @@ public String toString() { .add("hnswParams", hnswParams.orElse(null)) .add("sqParams", sqParams.orElse(null)) .add("rqParams", rqParams.orElse(null)) + .add("coveringColumns", coveringColumns) + .add("storeVectorsForRefine", storeVectorsForRefine) .toString(); } } diff --git a/java/src/test/java/org/lance/index/VectorIndexTest.java b/java/src/test/java/org/lance/index/VectorIndexTest.java index 7a9e066a1c5..cd81ffe2ab4 100755 --- a/java/src/test/java/org/lance/index/VectorIndexTest.java +++ b/java/src/test/java/org/lance/index/VectorIndexTest.java @@ -22,20 +22,36 @@ import org.lance.index.vector.SQBuildParams; import org.lance.index.vector.VectorIndexParams; import org.lance.index.vector.VectorTrainer; - +import org.lance.ipc.Query; +import org.lance.ipc.ScanOptions; + +import org.apache.arrow.dataset.scanner.Scanner; +import org.apache.arrow.vector.IntVector; +import org.apache.arrow.vector.VarCharVector; +import org.apache.arrow.vector.VectorSchemaRoot; +import org.apache.arrow.vector.ipc.ArrowReader; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.io.TempDir; import org.junit.jupiter.params.ParameterizedTest; import org.junit.jupiter.params.provider.EnumSource; +import java.nio.charset.StandardCharsets; import java.nio.file.Path; +import java.util.ArrayList; +import java.util.Arrays; import java.util.Collections; +import java.util.HashMap; +import java.util.HashSet; import java.util.List; +import java.util.Map; import java.util.Optional; +import java.util.Set; +import java.util.UUID; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertFalse; import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertThrows; import static org.junit.jupiter.api.Assertions.assertTrue; public class VectorIndexTest { @@ -347,6 +363,408 @@ public void testCreateIvfRqIndex(@TempDir Path tempDir) throws Exception { } } + /** + * {@code coveringColumns} must be a defensive, unmodifiable copy: mutating the list the caller + * passed to the builder (or the list the getter returns) must not change the params, so the JNI + * read at build time sees exactly what was requested. + */ + @Test + public void testCoveringColumnsIsDefensivelyCopied() { + IvfBuildParams ivf = new IvfBuildParams.Builder().setNumPartitions(2).build(); + List cols = new ArrayList<>(); + cols.add("i"); + VectorIndexParams params = new VectorIndexParams.Builder(ivf).setCoveringColumns(cols).build(); + + // Mutating the caller's list after build must not affect the params. + cols.add("s"); + cols.clear(); + assertEquals( + Collections.singletonList("i"), + params.getCoveringColumns(), + "include columns must be a defensive copy, unaffected by caller mutation"); + + // The returned list must be unmodifiable. + assertThrows(UnsupportedOperationException.class, () -> params.getCoveringColumns().add("x")); + } + + /** + * {@code Index.coveringFields} must be a defensive, unmodifiable copy, matching {@code + * VectorIndexParams.coveringColumns}: mutating the caller's list (or the list the getter returns) + * must not change the value object -- equals/hashCode would silently drift, and a later JNI + * commit would read mutated covering metadata while the index files still carry the payload. + */ + @Test + public void testIndexCoveringFieldsIsDefensivelyCopied() { + List ids = new ArrayList<>(); + ids.add(3); + Index index = + Index.builder() + .uuid(UUID.randomUUID()) + .fields(Collections.singletonList(0)) + .name("covered_idx") + .datasetVersion(1L) + .indexVersion(0) + .coveringFields(ids) + .build(); + + // Mutating the caller's list after build must not affect the index. + ids.add(7); + ids.clear(); + assertEquals( + Collections.singletonList(3), + index.coveringFields(), + "covering fields must be a defensive copy, unaffected by caller mutation"); + + // The returned list must be unmodifiable. + assertThrows(UnsupportedOperationException.class, () -> index.coveringFields().add(9)); + } + + /** + * The Java builder must be able to request carried refine vectors, and the JNI must forward the + * flag: the core records them by appending the *indexed* column's own field id to the index's + * covering fields, so a build that dropped the flag still succeeds but commits an index whose + * covering declaration is empty and whose {@code refine} queries fall back to the base table. + * + *

This guards the {@code store_vectors_for_refine} read in {@code + * java/lance-jni/src/utils.rs}: restoring it to a hardcoded {@code false} leaves the index + * buildable and the field-id assertions below fail. + */ + @Test + public void testCreateIndexWithStoreVectorsForRefine(@TempDir Path tempDir) throws Exception { + try (TestVectorDataset testVectorDataset = + new TestVectorDataset(tempDir.resolve("refine_vectors"))) { + try (Dataset dataset = testVectorDataset.create()) { + IvfBuildParams ivf = new IvfBuildParams.Builder().setNumPartitions(4).build(); + // IVF_PQ, not IVF_FLAT: a flat quantizer already stores the full-precision vectors, so + // the core refuses to carry a second copy of them. + PQBuildParams pq = new PQBuildParams.Builder().setNumSubVectors(2).build(); + + VectorIndexParams vectorIndexParams = + new VectorIndexParams.Builder(ivf) + .setDistanceType(DistanceType.L2) + .setPqParams(pq) + .setStoreVectorsForRefine(true) + .build(); + assertTrue( + vectorIndexParams.isStoreVectorsForRefine(), + "the builder must round-trip the requested flag"); + IndexParams indexParams = + IndexParams.builder().setVectorIndexParams(vectorIndexParams).build(); + + dataset.createIndex( + IndexOptions.builder( + Collections.singletonList(TestVectorDataset.vectorColumnName), + IndexType.IVF_PQ, + indexParams) + .withIndexName(TestVectorDataset.indexName) + .build()); + + Index carried = + dataset.getIndexes().stream() + .filter(idx -> TestVectorDataset.indexName.equals(idx.name())) + .findFirst() + .orElse(null); + assertNotNull(carried, "Expected the index to be present"); + + int vectorFieldId = fieldId(dataset, TestVectorDataset.vectorColumnName); + assertEquals( + Collections.singletonList(vectorFieldId), + carried.coveringFields(), + "carried refine vectors are declared as the indexed column's own field id"); + // Covering fields are the trailing entries of fields; here the carried field is the + // keyed field itself, so it appears in both positions. + assertEquals( + Arrays.asList(vectorFieldId, vectorFieldId), + carried.fields(), + "the carried field must be appended after the keyed vector field"); + } + } + } + + /** + * A covered ("included") column set on {@link VectorIndexParams} must be threaded through the JNI + * create path into the built index's metadata, so the committed index reports the covered field + * id. Without the wiring the index would build but silently carry no covering columns. + * + *

This guards the {@code covering_columns} read in {@code java/lance-jni/src/utils.rs}: + * restoring it to an empty default leaves the index buildable but drops the covering declaration, + * and the field-id assertions below fail. + */ + @Test + public void testCreateIvfFlatIndexWithCoveringColumns(@TempDir Path tempDir) throws Exception { + try (TestVectorDataset testVectorDataset = + new TestVectorDataset(tempDir.resolve("ivf_flat_covering"))) { + try (Dataset dataset = testVectorDataset.create()) { + // Several partitions, so the search below really probes an IVF partition instead of + // degenerating into a scan of the one partition that holds everything. + IvfBuildParams ivf = new IvfBuildParams.Builder().setNumPartitions(4).build(); + + // Cover the non-vector "i" column so a projection of it is answered from the index. + VectorIndexParams vectorIndexParams = + new VectorIndexParams.Builder(ivf) + .setDistanceType(DistanceType.L2) + .setCoveringColumns(Collections.singletonList("i")) + .build(); + IndexParams indexParams = + IndexParams.builder().setVectorIndexParams(vectorIndexParams).build(); + + dataset.createIndex( + IndexOptions.builder( + Collections.singletonList(TestVectorDataset.vectorColumnName), + IndexType.IVF_FLAT, + indexParams) + .withIndexName(TestVectorDataset.indexName) + .build()); + + Index covered = + dataset.getIndexes().stream() + .filter(idx -> TestVectorDataset.indexName.equals(idx.name())) + .findFirst() + .orElse(null); + assertNotNull(covered, "Expected covered IVF_FLAT index to be present"); + + int coveredFieldId = fieldId(dataset, "i"); + int vectorFieldId = fieldId(dataset, TestVectorDataset.vectorColumnName); + assertEquals( + Collections.singletonList(coveredFieldId), + covered.coveringFields(), + "committed index must report the requested covering column's field id"); + // covering fields are always the trailing entries of fields, so the keyed vector field + // comes first and the covering field is appended after it. + assertEquals( + Arrays.asList(vectorFieldId, coveredFieldId), + covered.fields(), + "covering field must be appended after the keyed vector field"); + + float[] key = new float[32]; + for (int j = 0; j < key.length; j++) { + key[j] = (float) (32 + j); + } + int k = 5; + + // The assertion that actually exercises the covering payload: project the covered + // column together with an *uncovered* one and cross-check them row by row. Every row + // of the fixture satisfies s == "s-" + i, and only "i" is carried by the index, so "s" + // necessarily comes from a base-table take. A covering payload read positionally + // rather than by name, or attached to the wrong row, makes the two disagree. + Map coveredRows = searchCoveredWithBaseColumn(dataset, key, k); + assertEquals(k, coveredRows.size(), "covered search should return k distinct rows"); + coveredRows.forEach( + (i, s) -> + assertEquals( + "s-" + i, + s, + "covered 'i' must belong to the same row as the base-table 's'; a" + + " misaligned covering payload disagrees here")); + + // Recall against an exact scan. Worth keeping as a floor on search quality, but note + // what it cannot show: it does NOT establish that the *index* served the query. + // Covering is semantically transparent -- if the index were ignored and Lance fell + // back to brute-force KNN, the fallback is exact and would score 1.0 here. Java has + // no explain-plan binding, so there is no node chain to assert against as + // test_create_index_covering_columns_serve_the_query_from_the_index does on the + // Python side. That distinction is pinned instead by + // testCoveredIndexServesTheSearchNotAFullScan below, which uses fast search over a + // partially indexed dataset so the two outcomes genuinely differ. + Set exact = searchCoveredColumn(dataset, key, k, false); + Set ann = searchCoveredColumn(dataset, key, k, true); + assertEquals(k, exact.size(), "exact KNN ground truth should return k distinct rows"); + Set hits = new HashSet<>(ann); + hits.retainAll(exact); + double recall = (double) hits.size() / k; + assertTrue( + recall >= 0.5, + "recall of the covered index against exact KNN must be at least 0.5 but was " + + recall + + " (ann=" + + ann + + ", exact=" + + exact + + ")"); + } + } + } + + /** + * A covered search must actually be served by the index, not by a full scan that silently + * produces the same answer. + * + *

Recall against an exact scan cannot show this: covering is semantically transparent, so if + * the index became unusable and Lance fell back to brute-force KNN, the fallback is exact and + * scores 1.0. Java has no explain-plan binding, so the Python trick of pinning the plan's node + * chain is unavailable. + * + *

What is available is fast search, which restricts a query to indexed fragments. Index only 2 + * of the dataset's 5 fragments, and the two outcomes stop agreeing: {@code TestVectorDataset} + * writes the *same* 80 vectors into every fragment, so the query key matches one row at distance + * 0 in each of the 5 fragments ("i" = 1, 81, 161, 241, 321) with the next-nearest row ~181 away. + * Served by the index under fast search, only the two indexed fragments can contribute, so every + * returned "i" is below 160. Served by a full scan, the distance-0 rows from the three unindexed + * fragments win and "i" values above 160 appear. + */ + @Test + public void testCoveredIndexServesTheSearchNotAFullScan(@TempDir Path tempDir) throws Exception { + try (TestVectorDataset testVectorDataset = + new TestVectorDataset(tempDir.resolve("ivf_flat_covering_fast_search"))) { + try (Dataset dataset = testVectorDataset.create()) { + List fragments = dataset.getFragments(); + assertTrue(fragments.size() >= 3, "fixture must have unindexed fragments left over"); + + IvfBuildParams ivf = new IvfBuildParams.Builder().setNumPartitions(4).build(); + VectorIndexParams vectorIndexParams = + new VectorIndexParams.Builder(ivf) + .setDistanceType(DistanceType.L2) + .setCoveringColumns(Collections.singletonList("i")) + .build(); + IndexParams indexParams = + IndexParams.builder().setVectorIndexParams(vectorIndexParams).build(); + + List segments = new ArrayList<>(); + for (int f = 0; f < 2; f++) { + segments.add( + dataset.createIndex( + IndexOptions.builder( + Collections.singletonList(TestVectorDataset.vectorColumnName), + IndexType.IVF_FLAT, + indexParams) + .withIndexName(TestVectorDataset.indexName) + .withFragmentIds(Collections.singletonList(fragments.get(f).getId())) + .build())); + } + dataset.commitExistingIndexSegments( + TestVectorDataset.indexName, TestVectorDataset.vectorColumnName, segments); + + int coveredFieldId = fieldId(dataset, "i"); + for (Index segment : segments) { + assertEquals( + Collections.singletonList(coveredFieldId), + segment.coveringFields(), + "every distributed segment must carry the covering declaration"); + } + + float[] key = new float[32]; + for (int j = 0; j < key.length; j++) { + key[j] = (float) (32 + j); + } + + ScanOptions options = + new ScanOptions.Builder() + .columns(Collections.singletonList("i")) + .fastSearch(true) + .nearest( + new Query.Builder() + .setColumn(TestVectorDataset.vectorColumnName) + .setKey(key) + .setK(5) + .setUseIndex(true) + .build()) + .build(); + + Set values = new HashSet<>(); + try (Scanner scanner = dataset.newScan(options); + ArrowReader reader = scanner.scanBatches()) { + VectorSchemaRoot root = reader.getVectorSchemaRoot(); + while (reader.loadNextBatch()) { + IntVector iVector = (IntVector) root.getVector("i"); + for (int row = 0; row < root.getRowCount(); row++) { + values.add(iVector.get(row)); + } + } + } + + assertFalse(values.isEmpty(), "fast search over the covered index returned no rows"); + // Positive: the exact-match rows of both indexed fragments must be present. + assertTrue( + values.contains(1) && values.contains(81), + "fast search must return the distance-0 row of each indexed fragment, got " + values); + // Discriminating: any row from an unindexed fragment means a full scan served this. + for (int value : values) { + assertTrue( + value < 160, + "fast search must be confined to the 2 indexed fragments, but returned i=" + + value + + " (a full scan, not the covered index, served this query); got " + + values); + } + } + } + } + + /** Resolve a top-level column's Lance field id, which is what index metadata records. */ + private static int fieldId(Dataset dataset, String name) { + return dataset.getLanceSchema().fields().stream() + .filter(field -> name.equals(field.getName())) + .findFirst() + .orElseThrow(() -> new AssertionError("No field named " + name + " in the dataset schema")) + .getId(); + } + + /** + * Run a nearest-neighbor search projecting the covered "i" column together with the uncovered "s" + * column, returning i -> s for the matched rows. Only "i" is carried by the index, so "s" + * comes from a base-table take and the pair cross-checks the covering payload's row alignment. + */ + private static Map searchCoveredWithBaseColumn( + Dataset dataset, float[] key, int k) throws Exception { + ScanOptions options = + new ScanOptions.Builder() + .columns(Arrays.asList("i", "s")) + .nearest( + new Query.Builder() + .setColumn(TestVectorDataset.vectorColumnName) + .setKey(key) + .setK(k) + .setUseIndex(true) + .build()) + .build(); + + Map rows = new HashMap<>(); + try (Scanner scanner = dataset.newScan(options); + ArrowReader reader = scanner.scanBatches()) { + VectorSchemaRoot root = reader.getVectorSchemaRoot(); + while (reader.loadNextBatch()) { + IntVector iVector = (IntVector) root.getVector("i"); + VarCharVector sVector = (VarCharVector) root.getVector("s"); + for (int row = 0; row < root.getRowCount(); row++) { + rows.put(iVector.get(row), new String(sVector.get(row), StandardCharsets.UTF_8)); + } + } + } + return rows; + } + + /** + * Run a nearest-neighbor search projecting only the covered "i" column and return the matched + * values, so a covered result can be compared against exact KNN ground truth. + */ + private static Set searchCoveredColumn( + Dataset dataset, float[] key, int k, boolean useIndex) throws Exception { + ScanOptions options = + new ScanOptions.Builder() + .columns(Collections.singletonList("i")) + .nearest( + new Query.Builder() + .setColumn(TestVectorDataset.vectorColumnName) + .setKey(key) + .setK(k) + .setUseIndex(useIndex) + .build()) + .build(); + + Set values = new HashSet<>(); + try (Scanner scanner = dataset.newScan(options); + ArrowReader reader = scanner.scanBatches()) { + VectorSchemaRoot root = reader.getVectorSchemaRoot(); + while (reader.loadNextBatch()) { + IntVector iVector = (IntVector) root.getVector("i"); + for (int row = 0; row < root.getRowCount(); row++) { + values.add(iVector.get(row)); + } + } + } + return values; + } + /** * Regression test for the metric_type passthrough in the JNI VectorTrainer. * diff --git a/protos/ann.proto b/protos/ann.proto index f5de5e25e7b..fd91c3c6d27 100644 --- a/protos/ann.proto +++ b/protos/ann.proto @@ -23,6 +23,23 @@ enum VectorApproxMode { Accurate = 2; } +// The covering ("included") index columns a query needs materialized, by name. +// +// Exists as a message rather than a bare `repeated string` because proto3 has no +// presence tracking for repeated fields, and the empty list is a distinct, meaningful +// state here: +// +// * absent: no narrowing computed; materialize every covering column declared. +// * present and empty: materialize nothing, though the index does declare covering. +// * present and non-empty: materialize exactly these. +// +// Collapsing "present, empty" into "absent" costs no correctness (a covering column is +// semantically transparent -- its values otherwise arrive from a base-table read) but +// silently restores full materialization on every distributed plan. +message CoveringProjection { + repeated string columns = 1; +} + // Serialized vector query parameters. message VectorQueryProto { // Query vector as Arrow IPC bytes (supports Float16, Float32, Float64, UInt8, etc.) @@ -43,6 +60,12 @@ message VectorQueryProto { // Query-time approximation mode. Currently only affects RQ-quantized vector // indexes, such as IVF_RQ. Other index types ignore this setting. VectorApproxMode approx_mode = 14; + // Which covering columns the index must materialize for this query. Absent means + // "not computed" -- see CoveringProjection. Carried across the wire so a remote + // executor declares the same search output schema the planner did; without it the + // executor's node is wider than the plan it came from, and the surrounding nodes + // were built against the planner's narrower schema. + CoveringProjection covering_projection = 15; } // Serializable form of ANNIvfSubIndexExec — the IVF sub-index search node. diff --git a/python/python/lance/dataset.py b/python/python/lance/dataset.py index 616c5f31004..1ebb2a2a130 100644 --- a/python/python/lance/dataset.py +++ b/python/python/lance/dataset.py @@ -3848,6 +3848,8 @@ def _create_index_impl( streaming_refine_passes: Optional[int] = None, skip_transpose: bool = False, rabitq_model: Optional[str] = None, + covering_columns: Optional[List[str]] = None, + store_vectors_for_refine: bool = False, require_commit: bool = True, **kwargs, ) -> Index: @@ -4176,6 +4178,12 @@ def _create_index_impl( if rabitq_model is not None: kwargs["rabitq_model"] = rabitq_model + if covering_columns is not None: + kwargs["covering_columns"] = covering_columns + + if store_vectors_for_refine: + kwargs["store_vectors_for_refine"] = True + # Add fragment_ids and index_uuid to kwargs if provided for # distributed indexing if fragment_ids is not None: @@ -4236,6 +4244,8 @@ def create_index( streaming_coreset_rate: Optional[int] = None, streaming_refine_passes: Optional[int] = None, skip_transpose: bool = False, + covering_columns: Optional[List[str]] = None, + store_vectors_for_refine: bool = False, progress_callback: Optional[Callable[[IndexProgress], None]] = None, **kwargs, ) -> LanceDataset: @@ -4305,6 +4315,33 @@ def create_index( If True, the index will be trained on the data (e.g., compute IVF centroids, PQ codebooks). If False, an empty index structure will be created without training, which can be populated later. + covering_columns : List[str], optional + The columns to cover ("include") in the index. Their values are stored + inline in the index so a query projecting only covered columns is + answered from the index without a take from the base table. Each must + name a top-level, non-key column of the dataset; the columns are + validated by the core when the index is built. Vector indexes only. + None by default (no covering). + store_vectors_for_refine : bool, default False + Carry full-precision vectors in the index storage so a query using + ``refine_factor`` re-ranks from the index instead of taking the vectors + from the base table. Off by default: it stores a second copy of the + vector column, which is the largest column in most datasets. + + Not expressible through ``covering_columns``: the build pipeline + rewrites the indexed column in place, so a covering column under that + name would serve values that are not the vectors you wrote. Requires the + V3 index file format, and is rejected on flat quantizers (which already + store the vectors), on multivector columns, on a nested indexed column, + and with precomputed shuffle buffers. + + While such an index exists the indexed column can no longer be renamed, + nor have its nullability changed: the carried copy lives in index storage + under that column's name and nothing renames it, so ``alter_columns`` + refuses rather than leave an index whose ``optimize_indices`` and + ``compact_files`` are permanently broken. Drop the index, alter the + column, then rebuild. Dropping the column itself stays allowed -- that + removes the index along with it. fragment_ids : List[int], optional If provided, the index will be created only on the specified fragments. This enables distributed/fragment-level indexing. When provided, the @@ -4471,6 +4508,8 @@ def create_index( streaming_coreset_rate=streaming_coreset_rate, streaming_refine_passes=streaming_refine_passes, skip_transpose=skip_transpose, + covering_columns=covering_columns, + store_vectors_for_refine=store_vectors_for_refine, require_commit=True, **kwargs, ) @@ -4509,6 +4548,8 @@ def create_index_uncommitted( streaming_refine_passes: Optional[int] = None, skip_transpose: bool = False, rabitq_model: Optional[str] = None, + covering_columns: Optional[List[str]] = None, + store_vectors_for_refine: bool = False, **kwargs, ) -> Index: """ @@ -4571,6 +4612,23 @@ def create_index_uncommitted( "create_index_uncommitted requires fragment_ids " "for distributed index build" ) + # Covering columns are a vector-index feature. This branch builds a + # scalar segment and never reaches the vector params, so honouring + # covering_columns here is impossible -- reject it rather than return a + # segment with covering_fields == [] that the caller believes is covered. + if covering_columns: + raise ValueError( + "covering_columns is only supported for vector indexes, but " + f"index_type={index_type!r} builds a scalar index segment; " + f"got covering_columns={covering_columns!r}" + ) + # Same reasoning: carried refine vectors are stored by the vector build + # pipeline, which this branch never reaches. + if store_vectors_for_refine: + raise ValueError( + "store_vectors_for_refine is only supported for vector indexes, " + f"but index_type={index_type!r} builds a scalar index segment" + ) kwargs = dict(kwargs) column, rust_index_type, _ = self._prepare_scalar_index_request( @@ -4617,6 +4675,8 @@ def create_index_uncommitted( streaming_refine_passes=streaming_refine_passes, skip_transpose=skip_transpose, rabitq_model=rabitq_model, + covering_columns=covering_columns, + store_vectors_for_refine=store_vectors_for_refine, require_commit=False, **kwargs, ) diff --git a/python/python/lance/lance/indices/__init__.pyi b/python/python/lance/lance/indices/__init__.pyi index f4b0bf69592..915844774f7 100644 --- a/python/python/lance/lance/indices/__init__.pyi +++ b/python/python/lance/lance/indices/__init__.pyi @@ -92,6 +92,8 @@ class IndexDescription: num_rows_indexed: int fields: list[int] field_names: list[str] + covering_fields: list[int] + covering_field_names: list[str] segments: list[IndexSegmentDescription] details: dict total_size_bytes: Optional[int] diff --git a/python/python/tests/test_commit_index.py b/python/python/tests/test_commit_index.py index f4939434957..db8bbc16633 100644 --- a/python/python/tests/test_commit_index.py +++ b/python/python/tests/test_commit_index.py @@ -279,8 +279,23 @@ def test_index_covering_fields_roundtrip(dataset_with_index, test_table, tmp_pat # Exercises indices.rs: the IndexMetadata -> PyIndexSegmentDescription # conversion used by describe_indices(). - segment = dataset_without_index.describe_indices()[0].segments[0] - assert segment.covering_fields == [price_id] + desc = dataset_without_index.describe_indices()[0] + assert desc.segments[0].covering_fields == [price_id] + + # The index-level view lifts the same declaration out of the segments and + # resolves it against the current schema. + assert desc.covering_fields == [price_id] + assert desc.covering_field_names == ["price"] + # The covering suffix is reported apart from the key: `fields` stays the + # keyed prefix so it never claims the index can answer queries on "price". + assert desc.fields == [meta_id] + assert desc.field_names == ["meta"] + + # The repr has to carry the covering too, or `print(ds.describe_indices()[0])` + # renders a covered and an uncovered index identically. + rendered = repr(desc) + assert f"covering_fields=[{price_id}]" in rendered + assert 'covering_field_names=["price"]' in rendered def test_commit_index_rejects_invalid_covering_fields(dataset_with_index, tmp_path): @@ -316,3 +331,55 @@ def test_commit_index_rejects_invalid_covering_fields(dataset_with_index, tmp_pa create_index_op, read_version=dataset_with_index.version, ) + + +def test_commit_index_tolerates_missing_covering_fields( + dataset_with_index, test_table, tmp_path +): + """An Index pickled by a pre-covering lance has no `covering_fields` attribute. + + The dataclass default is `field(default_factory=list)`, so pickle's + `__dict__`-only restore leaves nothing to fall back to -- the attribute is + simply absent. The PyO3 conversion must treat that as "uncovered" (empty), + like it already does for `index_details`, not abort the commit with an + AttributeError mid rolling upgrade. + """ + from lance.dataset import Index + + index_id = dataset_with_index.describe_indices()[0].segments[0].uuid + + dataset_without_index = lance.write_dataset( + test_table, tmp_path / "dataset_without_index" + ) + src_index_dir = Path(dataset_with_index.uri) / "_indices" / index_id + dest_index_dir = Path(dataset_without_index.uri) / "_indices" / index_id + shutil.copytree(src_index_dir, dest_index_dir) + + field_id = _get_field_id_by_name(dataset_without_index.lance_schema, "meta") + index = Index( + uuid=index_id, + name="meta_idx", + fields=[field_id], + dataset_version=dataset_without_index.version, + fragment_ids=set( + [f.fragment_id for f in dataset_without_index.get_fragments()] + ), + index_version=0, + ) + # Simulate the old-writer object: the attribute does not exist at all, + # which is different from `covering_fields=[]`. + delattr(index, "covering_fields") + assert not hasattr(index, "covering_fields") + + create_index_op = lance.LanceOperation.CreateIndex( + new_indices=[index], + removed_indices=[], + ) + committed = lance.LanceDataset.commit( + dataset_without_index.uri, + create_index_op, + read_version=dataset_without_index.version, + ) + + described = committed.describe_indices()[0] + assert described.segments[0].covering_fields == [] diff --git a/python/python/tests/test_vector_index.py b/python/python/tests/test_vector_index.py index 253592dd35b..7013e9c7442 100644 --- a/python/python/tests/test_vector_index.py +++ b/python/python/tests/test_vector_index.py @@ -328,6 +328,322 @@ def test_ann(indexed_dataset): run(indexed_dataset) +def _field_id(dataset, name): + for field in dataset.lance_schema.fields(): + if field.name() == name: + return field.id() + raise KeyError(name) + + +def test_create_index_with_covering_columns(tmp_path): + """A covered ("included") column passed to create_index is threaded through the + PyO3 boundary into the built index's metadata, so the committed index reports the + covered field id. Without the wiring the index builds but carries no covering + columns.""" + tbl = create_table() + dataset = lance.write_dataset(tbl, tmp_path) + price_id = _field_id(dataset, "price") + + query = np.random.randn(128).astype(np.float32) + exact = dataset.to_table( + nearest={"column": "vector", "q": query, "k": 10}, columns=["id"] + )["id"].to_pylist() + + dataset = dataset.create_index( + "vector", + index_type="IVF_PQ", + num_partitions=4, + num_sub_vectors=16, + covering_columns=["price"], + ) + + # The CreateIndex transaction carries the built index's metadata, including the + # covered field ids. + created_index = dataset.get_transactions(1)[0].operation.new_indices[0] + assert created_index.covering_fields == [price_id] + + # The covering declaration also survives into the committed manifest, so it is + # visible without replaying the transaction log. + desc = dataset.describe_indices()[0] + assert [segment.covering_fields for segment in desc.segments] == [[price_id]] + + # The index-level view resolves the same declaration to schema names, so callers + # can see what an index covers without knowing field ids. + assert desc.covering_fields == [price_id] + assert desc.covering_field_names == ["price"] + # Covering columns are reported separately from the key the index answers on. + assert desc.field_names == ["vector"] + + hits = dataset.to_table( + nearest={ + "column": "vector", + "q": query, + "k": 10, + "nprobes": 4, + "refine_factor": 10, + }, + columns=["id", "price"], + ) + recall = len(set(exact) & set(hits["id"].to_pylist())) / len(exact) + assert recall >= 0.5, f"recall={recall}, exact={exact}, got={hits['id']}" + + # Declaring a covering column must not disturb the values a query projects. + prices = dict(zip(tbl["id"].to_pylist(), tbl["price"].to_pylist())) + assert hits["price"].to_pylist() == [prices[i] for i in hits["id"].to_pylist()] + + +def _covered_query_table(nvec=1000, ndim=128, k=10, id_offset=1_000_000, seed=0): + """A vector table with an unambiguous k-neighborhood and ids disjoint from row ids. + + Two fixture properties are load-bearing, and neither is cosmetic: + + `create_table` numbers ids from 0, so on a single-fragment write `id == _rowid` for + every row -- and under that fixture any defect that serves the row id where a + covered value belongs (or the reverse) satisfies a row-alignment assertion by + accident. A `UInt64` covering column substituted for `_rowid` downcasts cleanly, so + the confusion is silent. `id_offset` separates the two columns. + + The first `k` rows form one tight, well-separated cluster around the query while + every other row is orders of magnitude further away. In uniformly random 128-d data + the true top-k sit at nearly identical distances, so which ones a quantizer returns + is a coin flip and a recall gate on it flakes around its threshold. Here the top-k + is unambiguous, so the recall assertion tests the search rather than the fixture. + """ + rng = np.random.default_rng(seed) + mat = rng.standard_normal((nvec, ndim)).astype(np.float32) * 10.0 + mat[:k] = mat[0] + rng.standard_normal((k, ndim)).astype(np.float32) * 0.01 + tbl = pa.table( + { + "vector": pa.FixedSizeListArray.from_arrays( + pa.array(mat.reshape(-1), type=pa.float32()), ndim + ), + "id": pa.array([id_offset + i for i in range(nvec)], type=pa.uint64()), + "price": pa.array(rng.random(nvec) * 100), + } + ) + return tbl, mat + + +def test_create_index_with_store_vectors_for_refine(tmp_path): + """`store_vectors_for_refine` must reach the core through the PyO3 boundary. It is + not expressible as a covering column -- the core records it by appending the + *indexed* column's own field id to `covering_fields` -- so a build that dropped the + flag still succeeds and still answers queries, just from the base table. The field + ids below are what separate the two.""" + tbl = create_table() + dataset = lance.write_dataset(tbl, tmp_path) + vector_id = _field_id(dataset, "vector") + + dataset = dataset.create_index( + "vector", + index_type="IVF_PQ", + num_partitions=4, + num_sub_vectors=16, + store_vectors_for_refine=True, + ) + + desc = dataset.describe_indices()[0] + assert desc.covering_fields == [vector_id] + assert desc.covering_field_names == ["vector"] + # The carried vectors are the keyed column itself, so unlike an ordinary covering + # column the name shows up on both sides. + assert desc.field_names == ["vector"] + + # And the option is honoured end to end: a refine query is answered without the + # base-table read it exists to remove. + query = np.random.randn(128).astype(np.float32) + scanner = dataset.scanner( + nearest={ + "column": "vector", + "q": query, + "k": 10, + "nprobes": 4, + "refine_factor": 2, + }, + columns=["vector"], + ) + assert "LanceRead" not in scanner.explain_plan(True) + assert scanner.to_table().num_rows == 10 + + +def test_create_index_store_vectors_for_refine_rejected_by_core(tmp_path): + """Validation lives in the core, not the binding: the option is refused on a + multivector column, and that refusal has to surface as a Python error rather than + an index that silently carries nothing.""" + tbl = create_multivec_table() + dataset = lance.write_dataset(tbl, tmp_path) + + with pytest.raises(ValueError, match="store_vectors_for_refine"): + dataset.create_index( + "vector", + index_type="IVF_PQ", + num_partitions=4, + num_sub_vectors=16, + metric="cosine", + store_vectors_for_refine=True, + ) + + +def test_create_index_covering_columns_serve_the_query_from_the_index(tmp_path): + """The payoff of the whole create-side wiring: an index built *through the Python + binding* with covering columns must answer a query for those columns out of index + storage, with values that match an independent read of the base table. + + Creation succeeding proves nothing -- nor does a correct query result on its own, + since a plain base-table take returns the same values. The plan assertions are what + separate the two, and they are deliberately both positive and negative: pinning the + exact node chain means a take inserted between the search and the output fails the + test whatever that node is called, while a bare `"LanceRead" not in plan` would go + vacuously true the moment the node is renamed -- and then the test would pass with + covering entirely dead, since values, row alignment and recall are all satisfied by + an ordinary base-table take. + """ + id_offset = 1_000_000 + tbl, mat = _covered_query_table(id_offset=id_offset) + dataset = lance.write_dataset(tbl, tmp_path) + + dataset = dataset.create_index( + "vector", + index_type="IVF_PQ", + num_partitions=4, + num_sub_vectors=16, + covering_columns=["id", "price"], + ) + desc = dataset.describe_indices()[0] + assert desc.covering_field_names == ["id", "price"] + # Covering columns never join the keyed prefix, or a consumer would believe the + # index can be searched on them. + assert desc.field_names == ["vector"] + + query = mat[0] + scanner = dataset.scanner( + nearest={"column": "vector", "q": query, "k": 10, "nprobes": 4}, + columns=["id", "price"], + with_row_id=True, + ) + plan = scanner.explain_plan(True) + # Positive: the covered plan is exactly search -> sort -> project, with nothing in + # between fetching from the base table. Asserting the whole chain (rather than only + # the absence of a node name) is what makes a rename fail loudly instead of + # silently satisfying the test. + plan_nodes = [ + line.split(":", 1)[0].strip() for line in plan.splitlines() if line.strip() + ] + assert plan_nodes == [ + "ProjectionExec", + "SortExec", + "ANNSubIndex", + "ANNIvfPartition", + ], f"unexpected node in the covered plan; plan was:\n{plan}" + # Negative, kept for the clearer failure message when the take is what came back. + assert "LanceRead" not in plan, ( + f"a covered projection must not fall back to a take; plan was:\n{plan}" + ) + + hits = scanner.to_table() + assert hits.num_rows == 10 + + # Independent read: reopen the dataset and read the base table directly, rather + # than trusting the in-memory `tbl` the index was built from. + base = lance.dataset(tmp_path).to_table(columns=["id", "price"]) + price_by_id = dict(zip(base["id"].to_pylist(), base["price"].to_pylist())) + + ids = hits["id"].to_pylist() + row_ids = hits["_rowid"].to_pylist() + prices = hits["price"].to_pylist() + # Single-fragment, step-id dataset => a correctly covered id is the row id plus the + # offset. A covered value sourced from `_rowid`, or paired with the wrong row, + # cannot satisfy this. + assert ids == [row_id + id_offset for row_id in row_ids] + assert prices == [price_by_id[i] for i in ids] + + # Serving the columns from the index is worthless if the search returns the wrong + # neighbors, so also gate on recall against brute-force ground truth. The fixture's + # planted neighborhood makes that truth unambiguous, so this gate measures the + # search and not the concentration of random high-dimensional distances. + truth = set( + (id_offset + int(i)) + for i in np.argsort(np.linalg.norm(mat - query, axis=1))[:10] + ) + recall = len(truth & set(ids)) / len(truth) + assert recall >= 0.5, f"covered recall {recall} < 0.5 (got {ids}, truth {truth})" + + +def test_create_index_uncommitted_with_covering_columns(tmp_path): + """create_index_uncommitted is a second entry point into the same PyO3 create + path. It has to forward covering_columns as well, or a distributed build silently + produces a segment with no covering columns.""" + tbl = create_table(nvec=256) + dataset = lance.write_dataset(tbl, tmp_path) + price_id = _field_id(dataset, "price") + + segment = dataset.create_index_uncommitted( + "vector", + index_type="IVF_PQ", + num_partitions=4, + num_sub_vectors=16, + covering_columns=["price"], + fragment_ids=[fragment.fragment_id for fragment in dataset.get_fragments()], + ) + assert segment.covering_fields == [price_id] + + dataset = dataset.commit_existing_index_segments("vector_idx", "vector", [segment]) + assert dataset.describe_indices()[0].segments[0].covering_fields == [price_id] + + +@pytest.mark.parametrize( + "index_type", + [pytest.param("BTREE", id="btree"), pytest.param("BITMAP", id="bitmap")], +) +def test_create_index_uncommitted_rejects_covering_columns_for_scalar( + tmp_path, index_type +): + """The segment-native scalar branch of create_index_uncommitted early-returns + before the vector params are ever built, so it cannot honour covering_columns. It + has to say so: silently returning a segment with covering_fields == [] leaves a + distributed pipeline believing it built covered scalar segments, and the covering + turns up missing only at query time.""" + dataset = lance.write_dataset(create_table(nvec=256), tmp_path) + fragment_ids = [fragment.fragment_id for fragment in dataset.get_fragments()] + + with pytest.raises(ValueError, match="only supported for vector indexes"): + dataset.create_index_uncommitted( + "price", + index_type=index_type, + fragment_ids=fragment_ids, + covering_columns=["meta"], + ) + + # The guard must not disturb the ordinary uncovered scalar path. + segment = dataset.create_index_uncommitted( + "price", index_type=index_type, fragment_ids=fragment_ids + ) + assert segment.covering_fields == [] + + +@pytest.mark.parametrize( + "covering_columns", + [ + pytest.param(["nope"], id="unknown_column"), + pytest.param(["vector"], id="indexed_column"), + pytest.param(["price", "price"], id="duplicate_column"), + ], +) +def test_create_index_covering_columns_rejected_by_core(tmp_path, covering_columns): + """Covering columns are validated in the Rust core, never in the binding. The + binding only has to let those errors reach the caller.""" + dataset = lance.write_dataset(create_table(nvec=256), tmp_path) + + with pytest.raises(ValueError, match="covering_columns"): + dataset.create_index( + "vector", + index_type="IVF_PQ", + num_partitions=4, + num_sub_vectors=16, + covering_columns=covering_columns, + ) + + def test_create_index_progress_callback_vector(tmp_path): ds = _make_sample_dataset_base(tmp_path, "vector_progress", 1500, 128) recorder = ProgressRecorder() diff --git a/python/src/dataset.rs b/python/src/dataset.rs index 547ba91d16b..6bd1d27dde9 100644 --- a/python/src/dataset.rs +++ b/python/src/dataset.rs @@ -4968,6 +4968,8 @@ fn prepare_vector_index_params( let mut rq_params = RQBuildParams::default(); let mut index_file_version = IndexFileVersion::V3; let mut skip_transpose = false; + let mut covering_columns: Vec = Vec::new(); + let mut store_vectors_for_refine = false; if let Some(kwargs) = kwargs { // Parse metric type @@ -5129,6 +5131,24 @@ fn prepare_vector_index_params( if let Some(value) = kwargs.get_item("skip_transpose")? { skip_transpose = value.extract()?; } + + // Covering ("included") columns: names of extra dataset columns stored inline in the + // index. Empty when absent. The core validates them when the index is built. + if let Some(cols) = kwargs.get_item("covering_columns")? + && !cols.is_none() + { + covering_columns = cols.extract()?; + } + + // Carry full-precision vectors in the index so `refine` re-ranks from index + // storage rather than taking them from the base table. Deliberately separate + // from `covering_columns`: the core records it by naming the indexed column + // itself, which the covering setter rejects. Validated by the core at build. + if let Some(value) = kwargs.get_item("store_vectors_for_refine")? + && !value.is_none() + { + store_vectors_for_refine = value.extract()?; + } } let mut params = match index_type { @@ -5174,6 +5194,8 @@ fn prepare_vector_index_params( }?; params.version(index_file_version); params.skip_transpose(skip_transpose); + params.covering_columns(covering_columns); + params.store_vectors_for_refine(store_vectors_for_refine); if let Some(kwargs) = kwargs && let Some(acc) = kwargs.get_item("accelerator")? { @@ -5674,6 +5696,10 @@ impl PySearchFilter { query_parallelism, dist_q_c: 0.0, approx_mode, + // Not a user-settable search parameter: the covering projection is derived + // from what the plan reads, so it stays `None` (materialize whatever the + // index declares) at the binding boundary. + covering_projection: None, }; Ok(Self { diff --git a/python/src/indices.rs b/python/src/indices.rs index acb15491fcb..cc2f3b95fcc 100644 --- a/python/src/indices.rs +++ b/python/src/indices.rs @@ -674,15 +674,23 @@ pub struct PyIndexDescription { pub type_url: String, /// The short type of the index (may not be unique) pub index_type: String, - /// The ids of the fields that the index is built on + /// The ids of the fields the index is keyed on -- the columns it can answer queries about. + /// Covering ("included") columns are excluded; those are reported per-segment as + /// `covering_fields`. pub fields: Vec, - /// The full paths of the fields that the index is built on + /// The full paths of the fields the index is keyed on, matching `fields` /// (dotted, with backtick-quoted segments for non-identifier names) pub field_names: Vec, /// The number of rows indexed by the index pub num_rows_indexed: u64, /// The details of the index pub details: PyJson, + /// The ids of the covering ("included") fields the index co-locates alongside + /// its key fields, so covered queries can skip the base-table take. Empty for + /// indexes without covering columns. All segments of one index share the set. + pub covering_fields: Vec, + /// The resolved names of `covering_fields` in the current schema. + pub covering_field_names: Vec, /// The segments of the index pub segments: Vec, /// The total size in bytes of all files across all segments @@ -703,12 +711,28 @@ impl PyIndexDescription { }) .collect(); - let segments = index + let segments: Vec = index .metadata() .iter() .map(PyIndexSegmentDescription::from_metadata) .collect(); + // All segments of one logical index declare the same covering columns + // (enforced at every commit boundary), so the first segment is authoritative. + let covering_fields = segments + .first() + .map(|segment| segment.covering_fields.clone()) + .unwrap_or_default(); + let covering_field_names = covering_fields + .iter() + .map(|field| { + dataset + .schema() + .field_path_minimal(*field) + .unwrap_or_else(|_| "".to_string()) + }) + .collect(); + let details = index.details().unwrap_or_else(|_| "{}".to_string()); Self { @@ -716,6 +740,8 @@ impl PyIndexDescription { fields: index.field_ids().to_vec(), field_names, index_type: index.index_type().to_string(), + covering_fields, + covering_field_names, segments, type_url: index.type_url().to_string(), num_rows_indexed: index.rows_indexed(), @@ -728,13 +754,19 @@ impl PyIndexDescription { #[pymethods] impl PyIndexDescription { pub fn __repr__(&self) -> String { + // `covering_fields` is printed alongside `fields` so a covered and an uncovered + // index do not render identically -- the segment repr already prints it, and + // without it here `print(ds.describe_indices()[0])` cannot answer "what does + // this index cover?" without descending into the segments. let mut repr = format!( - "IndexDescription(name='{}', type_url='{}', num_rows_indexed={}, fields={:?}, field_names={:?}, num_segments={}", + "IndexDescription(name='{}', type_url='{}', num_rows_indexed={}, fields={:?}, field_names={:?}, covering_fields={:?}, covering_field_names={:?}, num_segments={}", self.name, self.type_url, self.num_rows_indexed, self.fields, self.field_names, + self.covering_fields, + self.covering_field_names, self.segments.len() ); if let Some(byte_size) = self.total_size_bytes { @@ -762,3 +794,173 @@ pub fn register_indices(py: Python, m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_submodule(&indices)?; Ok(()) } + +#[cfg(test)] +mod tests { + use super::*; + use arrow_array::{ArrayRef, Int32Array, RecordBatch, RecordBatchIterator, StructArray}; + use arrow_schema::{DataType, Field as ArrowField, Schema as ArrowSchema}; + + /// Enough of an `IndexDescription` to drive the read-side conversion with metadata + /// no commit boundary would accept. That is the point: every write path rejects a + /// covering field id that is not a top-level field of the schema being committed + /// (`Operation::CreateIndex` in `transaction.rs`), and every schema-shrinking path + /// rejects dropping a covered column out from under a live index + /// (`reject_covered_field_subtree_change`). So `covering_field_names`'s `""` + /// fallback is only ever reached by a manifest this writer did not produce -- a + /// future format revision or a foreign writer -- and this stub is the only way to + /// hand it one. + struct StubIndexDescription { + field_ids: Vec, + segments: Vec, + } + + impl IndexDescription for StubIndexDescription { + fn name(&self) -> &str { + "stub_idx" + } + fn metadata(&self) -> &[IndexMetadata] { + &self.segments + } + fn type_url(&self) -> &str { + "" + } + fn index_type(&self) -> &str { + "Vector" + } + fn rows_indexed(&self) -> u64 { + 0 + } + fn field_ids(&self) -> &[u32] { + &self.field_ids + } + fn details(&self) -> lance_core::Result { + Ok("{}".to_string()) + } + fn total_size_bytes(&self) -> Option { + None + } + } + + fn covered_segment(fields: Vec, covering_fields: Vec) -> IndexMetadata { + IndexMetadata { + uuid: Uuid::new_v4(), + fields, + covering_fields, + name: "stub_idx".to_string(), + dataset_version: 1, + fragment_bitmap: None, + index_details: None, + index_version: 0, + created_at: None, + base_id: None, + files: None, + } + } + + /// A dataset whose field ids deliberately do not match top-level positions: + /// `id` is field 0, the struct `nested` is field 1 (its children take 2 and 3), + /// and `price` is field 4 while sitting at top-level position 2. Resolving a + /// covered id positionally instead of by id therefore cannot accidentally + /// produce the right answer. + async fn dataset_with_offset_field_ids() -> LanceDataset { + let nested_fields = vec![ + ArrowField::new("a", DataType::Int32, false), + ArrowField::new("b", DataType::Int32, false), + ]; + let schema = Arc::new(ArrowSchema::new(vec![ + ArrowField::new("id", DataType::Int32, false), + ArrowField::new( + "nested", + DataType::Struct(nested_fields.clone().into()), + false, + ), + ArrowField::new("price", DataType::Int32, false), + ])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(Int32Array::from(vec![1, 2, 3])), + Arc::new(StructArray::new( + nested_fields.into(), + vec![ + Arc::new(Int32Array::from(vec![4, 5, 6])) as ArrayRef, + Arc::new(Int32Array::from(vec![7, 8, 9])) as ArrayRef, + ], + None, + )), + Arc::new(Int32Array::from(vec![10, 20, 30])), + ], + ) + .unwrap(); + let reader = RecordBatchIterator::new(vec![Ok(batch)], schema); + let dataset = LanceDataset::write(reader, "memory://test", None) + .await + .unwrap(); + // Pin the ids the tests below hardcode, so a change in id assignment fails + // here rather than silently making the covered id resolve to another column. + assert_eq!(dataset.schema().field("price").unwrap().id, PRICE_FIELD_ID); + dataset + } + + /// `price`'s field id -- see [`dataset_with_offset_field_ids`]. + const PRICE_FIELD_ID: i32 = 4; + + fn runtime() -> tokio::runtime::Runtime { + tokio::runtime::Builder::new_multi_thread() + .enable_all() + .build() + .unwrap() + } + + #[test] + fn covering_field_names_resolve_against_the_schema() { + let dataset = runtime().block_on(dataset_with_offset_field_ids()); + let index = StubIndexDescription { + field_ids: vec![0], + segments: vec![covered_segment( + vec![0, PRICE_FIELD_ID], + vec![PRICE_FIELD_ID], + )], + }; + + let description = PyIndexDescription::new(&index, &dataset); + + assert_eq!(description.covering_fields, vec![PRICE_FIELD_ID]); + assert_eq!(description.covering_field_names, vec!["price".to_string()]); + } + + #[test] + fn covering_field_names_fall_back_when_the_id_is_unresolvable() { + let dataset = runtime().block_on(dataset_with_offset_field_ids()); + // No field carries id 99, so `field_path_minimal` returns an error. + let index = StubIndexDescription { + field_ids: vec![0], + segments: vec![covered_segment(vec![0, 99], vec![99])], + }; + + // Introspection must still succeed -- the fallback exists so a stale or + // unrecognized covering declaration cannot make `describe_indices()` error out. + let description = PyIndexDescription::new(&index, &dataset); + + assert_eq!(description.covering_fields, vec![99]); + assert_eq!( + description.covering_field_names, + vec!["".to_string()] + ); + } + + #[test] + fn covering_fields_are_empty_without_segments() { + let dataset = runtime().block_on(dataset_with_offset_field_ids()); + let index = StubIndexDescription { + field_ids: vec![0], + segments: vec![], + }; + + let description = PyIndexDescription::new(&index, &dataset); + + assert!(description.covering_fields.is_empty()); + assert!(description.covering_field_names.is_empty()); + } +} diff --git a/rust/lance-index/src/traits.rs b/rust/lance-index/src/traits.rs index f5441bd80f0..9980530d278 100644 --- a/rust/lance-index/src/traits.rs +++ b/rust/lance-index/src/traits.rs @@ -344,7 +344,13 @@ pub trait IndexDescription: Send + Sync { /// deleted. fn rows_indexed(&self) -> u64; - /// Returns the ids of the fields that the index is built on. + /// Returns the ids of the fields that the index is built on -- the columns it can + /// answer queries for. + /// + /// This is the index's *keyed* prefix only. An index may additionally carry values + /// for columns it is not keyed on (see [`IndexMetadata::covering_fields`]); those are + /// deliberately absent here, because the index cannot be searched on them. They stay + /// reachable per segment via [`Self::metadata`]. fn field_ids(&self) -> &[u32]; /// Returns a JSON string representation of the index details diff --git a/rust/lance-index/src/vector.rs b/rust/lance-index/src/vector.rs index a41973f57c8..fd0f5639f6e 100644 --- a/rust/lance-index/src/vector.rs +++ b/rust/lance-index/src/vector.rs @@ -166,6 +166,52 @@ pub struct Query { /// This currently only affects RQ-quantized vector indexes, such as IVF_RQ. /// Other index types ignore this setting. pub approx_mode: ApproxMode, + + /// The covering ("included") columns this query needs the index to materialize, + /// by name. + /// + /// Unlike the other fields here this is not a user-facing search parameter. It is + /// derived per plan from what the query actually reads, so that an index covering a + /// wide payload does not pay to materialize that payload on searches which never + /// look at it. + /// + /// # The three states are distinct, and must stay distinct + /// + /// * `None` — no query-level narrowing or physical-capability resolution was computed. + /// A raw index search may consider every physical covering column, but an execution + /// planner must not treat this state as proof that declared values are servable. It + /// resolves storage first and converts the query to an explicit `Some` projection. + /// * `Some(&[])` — the index *does* declare covering columns, but this query needs + /// **none** of them. Consumers must do no covering work at all: not "project zero + /// columns out of a batch that was loaded anyway", but skip the covering read + /// entirely. + /// * `Some(cols)` — this query needs exactly `cols`, and the execution planner has + /// verified that every selected segment can physically serve them. + /// + /// `Some(&[])` is the state this field exists for, and the one that silently + /// degrades if it is folded into `None`: a covering column is semantically + /// transparent, so conflating the two restores full materialization while every + /// result stays byte-identical (the values simply arrive from the base table + /// instead). No result-correctness test can catch that; only a cost or plan-shape + /// assertion can. Treat `Option::unwrap_or_default()` on this field, or any + /// `is_empty()` test that does not first distinguish `None`, as a bug. + /// + /// # Physical capability is a ceiling + /// + /// A declared column left out of `cols` is simply fetched from the base table — + /// correct, just slower. A column listed here that any selected segment cannot emit + /// leaves the covered projection short a column. Planners may therefore remove an + /// unproven field, but must never add one based on the manifest declaration alone. + /// + /// # Notes + /// + /// Names rather than field ids, because consumers match against a storage batch whose + /// columns carry names. The planner separately verifies source field ids, Arrow shape, + /// and compatible physical order before producing this name projection. + /// + /// `Arc` rather than `Vec` because [`Query`] is cloned once per probed partition on + /// the search hot path. + pub covering_projection: Option>, } impl From for DistanceType { @@ -236,6 +282,17 @@ pub trait VectorIndex: Send + Sync + std::fmt::Debug + Index { /// Get the total number of partitions in the index. fn total_partitions(&self) -> usize; + /// Covering columns physically present and safe to serve from this index, paired with + /// their source dataset field ids and returned in storage order. + /// + /// [`IndexMetadata::covering_fields`](lance_table::format::IndexMetadata::covering_fields) + /// is only the logical declaration. Query planners intersect that declaration with this + /// physical capability for every segment before omitting a base-table read. Formats that + /// do not expose a verifiable covering payload use the default empty capability. + fn physical_covering_fields(&self) -> Result> { + Ok(Vec::new()) + } + /// Search a single partition for nearest neighbors. /// /// This method should return the same results as [`VectorIndex::search`] method except diff --git a/rust/lance-index/src/vector/bq/storage.rs b/rust/lance-index/src/vector/bq/storage.rs index 0f7f034b2c1..4e4d10a12f2 100644 --- a/rust/lance-index/src/vector/bq/storage.rs +++ b/rust/lance-index/src/vector/bq/storage.rs @@ -44,6 +44,7 @@ use crate::frag_reuse::{FragReuseIndex, FragReuseIndexHandle}; use crate::pb; use crate::scalar::RowIdRemapper; use crate::vector::ApproxMode; +use crate::vector::PART_ID_COLUMN; use crate::vector::bq::dist_table_quant::{ DistTableDequant, quantize_dist_table_into, quantize_dist_table_u16_into, }; @@ -78,6 +79,25 @@ pub const RABIT_EX_CODE_COLUMN: &str = "__ex_codes"; /// older versions, which fail with a missing-column error instead of /// misinterpreting the bytes. pub const RABIT_BLOCKED_EX_CODE_COLUMN: &str = "__blocked_ex_codes"; +/// RaBitQ storage's internal (non-covering) column names: the row id, the binary code +/// and extended-bit code columns, the per-row factor columns (some present only for +/// higher `num_bits`), and the legacy `__ivf_part_id` column (left in pre-#3606 +/// single-partition builds; unlike PQ, RaBitQ storage does not drop it at +/// construction). Every other column is a user-declared covering ("included") column. +/// The single authority for this set -- the distributed shard merger classifies from +/// it too, so a new factor column added here is excluded everywhere at once. +pub const RABIT_INTERNAL_COLUMNS: &[&str] = &[ + ROW_ID, + RABIT_CODE_COLUMN, + RABIT_EX_CODE_COLUMN, + RABIT_BLOCKED_EX_CODE_COLUMN, + ADD_FACTORS_COLUMN, + SCALE_FACTORS_COLUMN, + ERROR_FACTORS_COLUMN, + EX_ADD_FACTORS_COLUMN, + EX_SCALE_FACTORS_COLUMN, + PART_ID_COLUMN, +]; pub const SEGMENT_LENGTH: usize = 4; pub const SEGMENT_NUM_CODES: usize = 1 << SEGMENT_LENGTH; const RABIT_PRUNE_STATS_ENV: &str = "LANCE_RQ_PRUNE_STATS"; @@ -2040,6 +2060,9 @@ fn accumulate_filtered_distances_into_heap( impl VectorStore for RabitQuantizationStorage { type DistanceCalculator<'a> = RabitDistCalculator<'a>; + /// RaBitQ storage is `[_rowid, , ]`. + const INTERNAL_COLUMNS: &'static [&'static str] = RABIT_INTERNAL_COLUMNS; + fn as_any(&self) -> &dyn std::any::Any { self } @@ -3035,6 +3058,39 @@ mod tests { .unwrap() } + #[test] + fn test_covering_excludes_legacy_part_id_column() { + // Same legacy `__ivf_part_id` guard as the other quantizer storages: + // pre-#3606 `num_partitions=1` builds left it in the storage file, and + // it must never be classified as a covering column -- otherwise search + // emits a column the exec's declared schema (from + // `IndexMetadata.covering_fields`, empty for those indexes) lacks. + use crate::vector::PART_ID_COLUMN; + use arrow_array::UInt32Array; + use arrow_schema::{DataType, Field, Schema}; + + let code_dim = 64; + let codes = make_test_codes(10, code_dim); + let metadata = make_test_metadata(codes.value_length() as usize * 8); + let base = make_test_batch(codes); + + let mut fields: Vec<_> = base.schema().fields().iter().cloned().collect(); + fields.push(Arc::new(Field::new(PART_ID_COLUMN, DataType::UInt32, true))); + let mut columns = base.columns().to_vec(); + columns.push(Arc::new(UInt32Array::from_iter_values( + (0..base.num_rows()).map(|_| 0), + )) as ArrayRef); + let batch = RecordBatch::try_new(Arc::new(Schema::new(fields)), columns).unwrap(); + + let storage = + RabitQuantizationStorage::try_from_batch(batch, &metadata, DistanceType::L2, None) + .unwrap(); + assert!( + storage.covering_field_indices().is_empty(), + "legacy __ivf_part_id must not be classified as a covering column" + ); + } + fn make_test_ex_codes(num_vectors: usize, code_dim: usize, num_bits: u8) -> FixedSizeListArray { let ex_bits = rabit_ex_bits(num_bits).unwrap(); let ex_code_bytes = rabit_ex_code_bytes(code_dim, ex_bits).unwrap(); diff --git a/rust/lance-index/src/vector/distributed/index_merger.rs b/rust/lance-index/src/vector/distributed/index_merger.rs index 370011eb8bd..cc72ef43957 100755 --- a/rust/lance-index/src/vector/distributed/index_merger.rs +++ b/rust/lance-index/src/vector/distributed/index_merger.rs @@ -21,8 +21,8 @@ use crate::IndexMetadata as IndexMetaSchema; use crate::pb; use crate::scalar::OldIndexDataFilter; use crate::vector::bq::storage::{ - RABIT_CODE_COLUMN, RABIT_METADATA_KEY, RabitQuantizationMetadata, RabitQueryEstimator, - pack_codes, rabit_binary_code_field, rabit_ex_code_field, + RABIT_CODE_COLUMN, RABIT_INTERNAL_COLUMNS, RABIT_METADATA_KEY, RabitQuantizationMetadata, + RabitQueryEstimator, pack_codes, rabit_binary_code_field, rabit_ex_code_field, }; use crate::vector::bq::transform::{ ADD_FACTORS_FIELD, ERROR_FACTORS_FIELD, EX_ADD_FACTORS_FIELD, EX_SCALE_FACTORS_FIELD, @@ -30,14 +30,20 @@ use crate::vector::bq::transform::{ }; use crate::vector::bq::validate_rq_num_bits; use crate::vector::flat::index::FlatMetadata; +use crate::vector::flat::storage::{FLAT_COLUMN, FLAT_INTERNAL_COLUMNS}; use crate::vector::ivf::storage::{IVF_METADATA_KEY, IvfModel as IvfStorageModel}; -use crate::vector::pq::storage::{PQ_METADATA_KEY, ProductQuantizationMetadata, transpose}; +use crate::vector::pq::storage::{ + PQ_INTERNAL_COLUMNS, PQ_METADATA_KEY, ProductQuantizationMetadata, transpose, +}; use crate::vector::quantizer::QuantizerMetadata; -use crate::vector::sq::storage::{SQ_METADATA_KEY, ScalarQuantizationMetadata}; +use crate::vector::sq::storage::{ + SQ_INTERNAL_COLUMNS, SQ_METADATA_KEY, ScalarQuantizationMetadata, +}; use crate::vector::storage::STORAGE_METADATA_KEY; +use crate::vector::storage::{COVERING_FIELD_IDS_KEY, covering_field_indices_excluding}; use crate::vector::{DISTANCE_TYPE_KEY, PQ_CODE_COLUMN, SQ_CODE_COLUMN}; use crate::{INDEX_AUXILIARY_FILE_NAME, INDEX_METADATA_SCHEMA_KEY}; -use arrow_schema::{DataType, Field, Schema as ArrowSchema}; +use arrow_schema::{DataType, Field, FieldRef, Schema as ArrowSchema}; use bytes::Bytes; use lance_core::datatypes::Schema as LanceSchema; use lance_file::reader::{FileReader as V2Reader, FileReaderOptions as V2ReaderOptions}; @@ -249,6 +255,34 @@ fn init_writer_for_storage( Ok(()) } +/// The covering ("included") columns carried in a shard's auxiliary storage schema: every +/// field that is not the row id or one of the quantizer's internal code/factor columns (or the +/// legacy `__ivf_part_id`). The merger rebuilds the unified output schema from quantizer +/// metadata (it reshapes the code column via the transpose), so these fields must be +/// re-appended or the covered payload is silently dropped from the merged index. All shards +/// share the same covering set, so the first shard's schema is authoritative. +fn covering_fields_from_shard_schema( + schema: &ArrowSchema, + index_type: SupportedIvfIndexType, +) -> Vec { + use SupportedIvfIndexType::*; + // Each storage's exported internal set already excludes PART_ID_COLUMN, which + // matters here: a legacy PQ shard's raw aux schema may still physically carry + // `__ivf_part_id` even though PQ storage drops it at load, so classifying from + // the raw schema (pre-load) must exclude it or it would be wrongly classified + // as a covering column. + let internal: &[&str] = match index_type { + IvfPq | IvfHnswPq => PQ_INTERNAL_COLUMNS, + IvfSq | IvfHnswSq => SQ_INTERNAL_COLUMNS, + IvfFlat | IvfHnswFlat => FLAT_INTERNAL_COLUMNS, + IvfRq => RABIT_INTERNAL_COLUMNS, + }; + covering_field_indices_excluding(schema, internal) + .into_iter() + .map(|i| schema.fields()[i].clone()) + .collect() +} + /// Create and initialize a unified writer for FLAT storage. pub async fn init_writer_for_flat( object_store: &lance_io::object_store::ObjectStore, @@ -257,18 +291,21 @@ pub async fn init_writer_for_flat( item_type: &DataType, dt: DistanceType, format_version: ConcreteFileVersion, + covering_fields: &[FieldRef], ) -> Result { - let arrow_schema = ArrowSchema::new(vec![ + let mut fields = vec![ (*ROW_ID_FIELD).clone(), Field::new( - crate::vector::flat::storage::FLAT_COLUMN, + FLAT_COLUMN, DataType::FixedSizeList( Arc::new(Field::new("item", item_type.clone(), true)), d0 as i32, ), true, ), - ]); + ]; + fields.extend(covering_fields.iter().map(|f| f.as_ref().clone())); + let arrow_schema = ArrowSchema::new(fields); let writer = object_store.create(aux_out).await?; let mut w = versions::create_writer( format_version, @@ -291,13 +328,14 @@ pub async fn init_writer_for_pq( dt: DistanceType, pm: &ProductQuantizationMetadata, format_version: ConcreteFileVersion, + covering_fields: &[FieldRef], ) -> Result { let num_bytes = if pm.nbits == 4 { pm.num_sub_vectors / 2 } else { pm.num_sub_vectors }; - let arrow_schema = ArrowSchema::new(vec![ + let mut fields = vec![ (*ROW_ID_FIELD).clone(), Field::new( PQ_CODE_COLUMN, @@ -307,7 +345,9 @@ pub async fn init_writer_for_pq( ), true, ), - ]); + ]; + fields.extend(covering_fields.iter().map(|f| f.as_ref().clone())); + let arrow_schema = ArrowSchema::new(fields); let writer = object_store.create(aux_out).await?; let mut w = versions::create_writer( format_version, @@ -336,9 +376,10 @@ pub async fn init_writer_for_sq( dt: DistanceType, sq_meta: &ScalarQuantizationMetadata, format_version: ConcreteFileVersion, + covering_fields: &[FieldRef], ) -> Result { let d0 = sq_meta.dim; - let arrow_schema = ArrowSchema::new(vec![ + let mut fields = vec![ (*ROW_ID_FIELD).clone(), Field::new( SQ_CODE_COLUMN, @@ -348,7 +389,9 @@ pub async fn init_writer_for_sq( ), true, ), - ]); + ]; + fields.extend(covering_fields.iter().map(|f| f.as_ref().clone())); + let arrow_schema = ArrowSchema::new(fields); let writer = object_store.create(aux_out).await?; let mut w = versions::create_writer( format_version, @@ -368,6 +411,7 @@ pub async fn init_writer_for_rq( dt: DistanceType, rq_meta: &RabitQuantizationMetadata, format_version: ConcreteFileVersion, + covering_fields: &[FieldRef], ) -> Result { let mut fields = vec![ (*ROW_ID_FIELD).clone(), @@ -383,6 +427,7 @@ pub async fn init_writer_for_rq( fields.push(EX_ADD_FACTORS_FIELD.clone()); fields.push(EX_SCALE_FACTORS_FIELD.clone()); } + fields.extend(covering_fields.iter().map(|f| f.as_ref().clone())); let arrow_schema = ArrowSchema::new(fields); let writer = object_store.create(aux_out).await?; let mut w = versions::create_writer( @@ -898,6 +943,11 @@ async fn merge_partial_vector_auxiliary_files_inner( let mut nlist_opt: Option = None; let mut accumulated_lengths: Vec = Vec::new(); let mut first_centroids: Option = None; + // Covering ("included") columns of the first shard, used as the authoritative set every + // later shard must match (see the per-shard check below). + // (covering fields, source dataset field ids) of the first shard; every later + // shard must match both. + let mut first_covering_fields: Option<(Vec, Option)> = None; // Track per-shard readers, IVF lengths, and precomputed partition offsets. // This avoids reopening each shard file for every partition during merge. @@ -1044,6 +1094,68 @@ async fn merge_partial_vector_auxiliary_files_inner( // Preserve the historical fallback while keeping the writer boundary exact. let fv = format_version.unwrap_or(ConcreteFileVersion::V2_0); + // Covering ("included") columns carried in this shard's storage. The unified writer + // schemas below are rebuilt from quantizer metadata, so these must be re-appended or + // the covered payload is dropped from the merged index. The output schema is built + // from the *first* shard's covering fields and `concat_batches` combines columns + // positionally, so a later shard whose covering columns differ in name, type, or order + // would be silently stored under the first shard's names -> wrong covered values. + // Reject any such mismatch rather than corrupt the merged index. + let shard_arrow: ArrowSchema = reader.schema().as_ref().into(); + let covering_fields = covering_fields_from_shard_schema(&shard_arrow, idx_type); + // Name and type alone cannot establish that two shards cover the *same logical + // column*: Arrow fields carry no Lance field id, and the ids in a shard's own + // Lance schema are file-local. A column dropped and re-added between two shard + // builds yields the same name and type under a different field id, and + // `concat_batches` would stack the old values under the new field. The build + // stamps the source field ids for exactly this comparison. + let covering_source_ids = reader + .metadata() + .file_schema + .metadata + .get(COVERING_FIELD_IDS_KEY) + .cloned(); + if !covering_fields.is_empty() && covering_source_ids.is_none() { + return Err(Error::index(format!( + "Distributed merge: shard {idx} carries covering (included) columns but records \ + no source field ids, so it cannot be proven to cover the same fields as the \ + other shards. Rebuild the shard." + ))); + } + match first_covering_fields.as_ref() { + None => { + first_covering_fields = Some((covering_fields.clone(), covering_source_ids.clone())) + } + Some((first, first_source_ids)) => { + let matches = first.len() == covering_fields.len() + && first + .iter() + .zip(&covering_fields) + .all(|(a, b)| a.name() == b.name() && a.data_type() == b.data_type()) + && *first_source_ids == covering_source_ids; + if !matches { + let describe = |fields: &[FieldRef], source_ids: &Option| { + let names = fields + .iter() + .map(|f| format!("{}: {}", f.name(), f.data_type())) + .collect::>() + .join(", "); + match source_ids { + Some(ids) => format!("{names} (source field ids {ids})"), + None => names, + } + }; + return Err(Error::index(format!( + "Distributed merge: covering (included) columns differ across shards; \ + shard 0 has [{}] but shard {idx} has [{}]. All shards must declare the \ + same covering columns (same source fields, names, types, and order).", + describe(first, first_source_ids), + describe(&covering_fields, &covering_source_ids), + ))); + } + } + } + match idx_type { SupportedIvfIndexType::IvfSq => { // Handle Scalar Quantization (SQ) storage for IVF_SQ @@ -1097,8 +1209,15 @@ async fn merge_partial_vector_auxiliary_files_inner( sq_meta = Some(sq_meta_parsed.clone()); } if v2w_opt.is_none() { - let w = - init_writer_for_sq(object_store, &aux_out, dt, &sq_meta_parsed, fv).await?; + let w = init_writer_for_sq( + object_store, + &aux_out, + dt, + &sq_meta_parsed, + fv, + &covering_fields, + ) + .await?; v2w_opt = Some(w); } } @@ -1205,8 +1324,15 @@ async fn merge_partial_vector_auxiliary_files_inner( rq_meta = Some(rq_meta_parsed.clone()); } if v2w_opt.is_none() { - let w = - init_writer_for_rq(object_store, &aux_out, dt, &rq_meta_parsed, fv).await?; + let w = init_writer_for_rq( + object_store, + &aux_out, + dt, + &rq_meta_parsed, + fv, + &covering_fields, + ) + .await?; v2w_opt = Some(w); } } @@ -1305,8 +1431,15 @@ async fn merge_partial_vector_auxiliary_files_inner( if v2w_opt.is_none() { let mut pm_for_unified = pm.clone(); pm_for_unified.transposed = true; - let w = - init_writer_for_pq(object_store, &aux_out, dt, &pm_for_unified, fv).await?; + let w = init_writer_for_pq( + object_store, + &aux_out, + dt, + &pm_for_unified, + fv, + &covering_fields, + ) + .await?; v2w_opt = Some(w); } } @@ -1334,8 +1467,16 @@ async fn merge_partial_vector_auxiliary_files_inner( return Err(Error::index("Dimension mismatch across shards".to_string())); } if v2w_opt.is_none() { - let w = init_writer_for_flat(object_store, &aux_out, d0, &item_type, dt, fv) - .await?; + let w = init_writer_for_flat( + object_store, + &aux_out, + d0, + &item_type, + dt, + fv, + &covering_fields, + ) + .await?; v2w_opt = Some(w); } } @@ -1367,8 +1508,16 @@ async fn merge_partial_vector_auxiliary_files_inner( return Err(Error::index("Dimension mismatch across shards".to_string())); } if v2w_opt.is_none() { - let w = init_writer_for_flat(object_store, &aux_out, d0, &item_type, dt, fv) - .await?; + let w = init_writer_for_flat( + object_store, + &aux_out, + d0, + &item_type, + dt, + fv, + &covering_fields, + ) + .await?; v2w_opt = Some(w); } } @@ -1463,8 +1612,15 @@ async fn merge_partial_vector_auxiliary_files_inner( if v2w_opt.is_none() { let mut pm_for_unified = pm.clone(); pm_for_unified.transposed = true; - let w = - init_writer_for_pq(object_store, &aux_out, dt, &pm_for_unified, fv).await?; + let w = init_writer_for_pq( + object_store, + &aux_out, + dt, + &pm_for_unified, + fv, + &covering_fields, + ) + .await?; v2w_opt = Some(w); } } @@ -1515,8 +1671,15 @@ async fn merge_partial_vector_auxiliary_files_inner( sq_meta = Some(sq_meta_parsed.clone()); } if v2w_opt.is_none() { - let w = - init_writer_for_sq(object_store, &aux_out, dt, &sq_meta_parsed, fv).await?; + let w = init_writer_for_sq( + object_store, + &aux_out, + dt, + &sq_meta_parsed, + fv, + &covering_fields, + ) + .await?; v2w_opt = Some(w); } } @@ -1716,6 +1879,15 @@ async fn merge_partial_vector_auxiliary_files_inner( for len in merged_lengths.iter() { ivf_model.add_partition(*len); } + // Carry the covering provenance onto the merged output. Every shard has already + // been checked to agree on this stamp, so the first shard's value IS the merged + // file's value. Without it the merged segment carries covering columns but no + // source ids, and feeding it back in as a shard of a later hierarchical merge + // would be rejected by the very check that verified its inputs -- with an error + // telling the user to rebuild a shard the merger itself produced. + if let Some((_, Some(covering_source_ids))) = first_covering_fields.as_ref() { + w.add_schema_metadata(COVERING_FIELD_IDS_KEY, covering_source_ids.clone()); + } let dt2 = distance_type.ok_or_else(|| Error::index("Distance type missing".to_string()))?; write_unified_ivf_and_index_metadata(w, &ivf_model, dt2, idx_type_final).await?; let summary = w.finish().await?; diff --git a/rust/lance-index/src/vector/flat/storage.rs b/rust/lance-index/src/vector/flat/storage.rs index 1de4c4387fb..de84897e92e 100644 --- a/rust/lance-index/src/vector/flat/storage.rs +++ b/rust/lance-index/src/vector/flat/storage.rs @@ -6,8 +6,9 @@ use std::{borrow::Cow, sync::Arc}; use super::index::FlatMetadata; use crate::frag_reuse::{FragReuseIndex, FragReuseIndexHandle}; use crate::scalar::RowIdRemapper; +use crate::vector::PART_ID_COLUMN; use crate::vector::quantizer::QuantizerStorage; -use crate::vector::storage::{DistCalculator, VectorStore}; +use crate::vector::storage::{DistCalculator, VectorStore, remap_row_ids_by_name}; use crate::vector::utils::do_prefetch; use arrow::array::AsArray; use arrow::compute::concat_batches; @@ -26,6 +27,14 @@ use lance_linalg::distance::{Cosine, DistanceType, Dot, L2, Normalize, norm_l2_f pub const FLAT_COLUMN: &str = "flat"; +/// Flat storage's internal (non-covering) column names: the row id, the flat vector +/// column, and the legacy `__ivf_part_id` column (left in pre-#3606 single-partition +/// builds; unlike PQ, flat storage does not drop it at construction). Every other +/// column in a flat storage schema is a user-declared covering ("included") column. +/// The single authority for this set -- the distributed shard merger classifies from +/// it too, so a new internal column added here is excluded everywhere at once. +pub const FLAT_INTERNAL_COLUMNS: &[&str] = &[ROW_ID, FLAT_COLUMN, PART_ID_COLUMN]; + /// Per-vector L2 norms cached for Cosine distance, so `cosine_with_norms` can /// skip recomputing each stored vector's norm per comparison. `None` for other /// metrics and for value types without a norm kernel. @@ -85,7 +94,7 @@ impl QuantizerStorage for FlatFloatStorage { frag_reuse_index: Option>, ) -> Result { let batch = if let Some(frag_reuse_index_ref) = frag_reuse_index.as_ref() { - frag_reuse_index_ref.remap_row_ids_record_batch(batch, 0)? + remap_row_ids_by_name(batch, frag_reuse_index_ref.as_ref())? } else { batch }; @@ -163,6 +172,9 @@ impl FlatFloatStorage { impl VectorStore for FlatFloatStorage { type DistanceCalculator<'a> = FlatFloatDistanceCalc<'a>; + /// Flat storage is `[_rowid, flat, ]`. + const INTERNAL_COLUMNS: &'static [&'static str] = FLAT_INTERNAL_COLUMNS; + fn to_batches(&self) -> Result> { Ok([self.batch.clone()].into_iter()) } @@ -265,7 +277,7 @@ impl QuantizerStorage for FlatBinStorage { frag_reuse_index: Option>, ) -> Result { let batch = if let Some(frag_reuse_index_ref) = frag_reuse_index.as_ref() { - frag_reuse_index_ref.remap_row_ids_record_batch(batch, 0)? + remap_row_ids_by_name(batch, frag_reuse_index_ref.as_ref())? } else { batch }; @@ -339,6 +351,9 @@ impl FlatBinStorage { impl VectorStore for FlatBinStorage { type DistanceCalculator<'a> = FlatDistanceCal<'a, UInt8Type>; + /// Flat storage is `[_rowid, flat, ]`. + const INTERNAL_COLUMNS: &'static [&'static str] = FLAT_INTERNAL_COLUMNS; + fn to_batches(&self) -> Result> { Ok([self.batch.clone()].into_iter()) } @@ -685,6 +700,70 @@ mod tests { FlatFloatStorage::new(vectors, DistanceType::L2) } + // Pre-#3606 `num_partitions=1` builds left `__ivf_part_id` in the flat + // storage file. It must never be classified as a covering column -- + // otherwise search emits a column the exec's declared schema (from + // `IndexMetadata.covering_fields`, empty for those indexes) lacks, + // desyncing the two and failing every query on that legacy index. + #[test] + fn test_flat_covering_excludes_legacy_part_id_column() { + use crate::vector::PART_ID_COLUMN; + const DIM: i32 = 8; + const N: usize = 8; + + let part_id_batch = |vectors: ArrayRef| { + RecordBatch::try_from_iter_with_nullable(vec![ + ( + ROW_ID, + Arc::new(UInt64Array::from_iter_values(0..N as u64)) as ArrayRef, + true, + ), + (FLAT_COLUMN, vectors, true), + ( + PART_ID_COLUMN, + Arc::new(arrow_array::UInt32Array::from_iter_values( + (0..N).map(|_| 0), + )) as ArrayRef, + true, + ), + ]) + .unwrap() + }; + let meta = FlatMetadata { dim: DIM as usize }; + + let floats = + arrow_array::Float32Array::from_iter_values((0..N * DIM as usize).map(|v| v as f32)); + let float_storage = FlatFloatStorage::try_from_batch( + part_id_batch(Arc::new( + FixedSizeListArray::try_new_from_values(floats, DIM).unwrap(), + )), + &meta, + DistanceType::L2, + None, + ) + .unwrap(); + assert!( + float_storage.covering_field_indices().is_empty(), + "legacy __ivf_part_id must not be classified as covering (float storage)" + ); + + let codes = + arrow_array::UInt8Array::from_iter_values((0..N * DIM as usize).map(|v| v as u8)); + let bin_storage = FlatBinStorage::try_from_batch( + part_id_batch(Arc::new( + FixedSizeListArray::try_new_from_values(codes, DIM).unwrap(), + )), + &meta, + DistanceType::L2, + None, + ) + .unwrap(); + assert!( + bin_storage.covering_field_indices().is_empty(), + "legacy __ivf_part_id must not be classified as covering (binary storage)" + ); + } + #[test] fn test_flat_float_storage_distance_f16() { let storage = make_f16_storage(); diff --git a/rust/lance-index/src/vector/pq/storage.rs b/rust/lance-index/src/vector/pq/storage.rs index a5ee3496571..c2d7c0f340d 100644 --- a/rust/lance-index/src/vector/pq/storage.rs +++ b/rust/lance-index/src/vector/pq/storage.rs @@ -14,7 +14,7 @@ use std::{ use arrow::datatypes::{self, UInt8Type}; use arrow_array::{ArrayRef, ArrowPrimitiveType, PrimitiveArray}; use arrow_array::{ - FixedSizeListArray, RecordBatch, UInt8Array, UInt64Array, + FixedSizeListArray, RecordBatch, UInt8Array, UInt32Array, UInt64Array, cast::AsArray, types::{Float32Type, UInt64Type}, }; @@ -43,7 +43,7 @@ use crate::vector::graph::{OrderedFloat, OrderedNode}; use crate::{ INDEX_METADATA_SCHEMA_KEY, IndexMetadata, pb, vector::{ - PQ_CODE_COLUMN, + PART_ID_COLUMN, PQ_CODE_COLUMN, pq::transform::PQTransformer, quantizer::{QuantizerMetadata, QuantizerStorage}, storage::{DistCalculator, VectorStore}, @@ -53,6 +53,14 @@ use crate::{ pub const PQ_METADATA_KEY: &str = "lance:pq"; +/// PQ storage's internal (non-covering) column names: the row id, the PQ code column, +/// and the legacy `__ivf_part_id` column (written by pre-#3606 single-partition +/// builds; PQ storage drops it at construction, so post-load it is never present and +/// excluding it is a no-op -- but pre-load classification, e.g. from a raw shard +/// schema in the distributed merger, needs it). Every other column is a user-declared +/// covering ("included") column. The single authority for this set. +pub const PQ_INTERNAL_COLUMNS: &[&str] = &[ROW_ID, PQ_CODE_COLUMN, PART_ID_COLUMN]; + #[derive(Debug, Clone, Serialize, Deserialize)] pub struct ProductQuantizationMetadata { pub codebook_position: usize, @@ -230,16 +238,23 @@ impl ProductQuantizationStorage { transposed: bool, frag_reuse_index: Option>, ) -> Result { - if batch.num_columns() != 2 { - log::warn!( - "PQ storage should have 2 columns, but got {} columns: {}", - batch.num_columns(), - batch.schema(), - ); - batch = batch.project(&[ - batch.schema().index_of(ROW_ID)?, - batch.schema().index_of(PQ_CODE_COLUMN)?, - ])?; + // Require `_rowid` and the PQ code column, but preserve any additional + // ("included"/covering) columns row-aligned. + // Normalize the order to `[_rowid, __pq_code, ]` so the schema + // is stable across builds. `__ivf_part_id` is never a covering column: + // legacy `num_partitions=1` builds (pre-#3606) wrote it into the index + // file, so keep dropping it on load as before. + let row_id_idx = batch.schema().index_of(ROW_ID)?; + let pq_idx = batch.schema().index_of(PQ_CODE_COLUMN)?; + if row_id_idx != 0 || pq_idx != 1 || batch.num_columns() > 2 { + let schema = batch.schema(); + let mut order = Vec::with_capacity(batch.num_columns()); + order.push(row_id_idx); + order.push(pq_idx); + order.extend((0..batch.num_columns()).filter(|&i| { + i != row_id_idx && i != pq_idx && schema.field(i).name() != PART_ID_COLUMN + })); + batch = batch.project(&order)?; } let Some(row_ids) = batch.column_by_name(ROW_ID) else { @@ -285,6 +300,7 @@ impl ProductQuantizationStorage { let transposed_codes = pq_code.values(); let mut new_row_ids = Vec::with_capacity(row_ids.len()); let mut new_codes = Vec::with_capacity(row_ids.len() * num_sub_vectors); + let mut kept_indices: Vec = Vec::with_capacity(row_ids.len()); let row_ids_values = row_ids.values(); for (i, row_id) in row_ids_values.iter().enumerate() { @@ -296,6 +312,7 @@ impl ProductQuantizationStorage { num_sub_vectors, i as u32, )); + kept_indices.push(i as u32); } } @@ -311,7 +328,14 @@ impl ProductQuantizationStorage { new_transposed_codes, num_bytes_in_code as i32, )?); - RecordBatch::try_new(batch.schema(), vec![new_row_ids, codes_fsl])? + // Preserve any included/covering columns row-aligned to the kept rows. + rebuild_storage_batch( + batch.schema(), + &batch, + new_row_ids, + codes_fsl, + &UInt32Array::from(kept_indices), + )? }; pq_code = batch[PQ_CODE_COLUMN] .as_fixed_size_list() @@ -529,6 +553,40 @@ where transposed_codes.into() } +/// Rebuild a PQ-storage batch after a row selection (fragment-reuse remap in +/// [`ProductQuantizationStorage::new`] or [`ProductQuantizationStorage::remap`]), +/// preserving any extra ("included"/covering) columns beyond `_rowid` and the PQ +/// code. The already-remapped `row_ids` and transposed `pq_codes` are supplied +/// directly; every other column is gathered from `original` at the surviving +/// positions `kept` so it stays row-aligned. `schema` carries the column order +/// `[_rowid, __pq_code, ]`. +fn rebuild_storage_batch( + schema: SchemaRef, + original: &RecordBatch, + row_ids: ArrayRef, + pq_codes: ArrayRef, + kept: &UInt32Array, +) -> Result { + let columns = schema + .fields() + .iter() + .map(|field| { + let name = field.name().as_str(); + if name == ROW_ID { + Ok(row_ids.clone()) + } else if name == PQ_CODE_COLUMN { + Ok(pq_codes.clone()) + } else { + let col = original.column_by_name(name).ok_or_else(|| { + Error::index(format!("column '{name}' missing from PQ storage batch")) + })?; + Ok(arrow::compute::take(col.as_ref(), kept, None)?) + } + }) + .collect::>>()?; + Ok(RecordBatch::try_new(schema, columns)?) +} + #[async_trait] impl QuantizerStorage for ProductQuantizationStorage { type Metadata = ProductQuantizationMetadata; @@ -610,6 +668,7 @@ impl QuantizerStorage for ProductQuantizationStorage { let transposed_codes = self.pq_code.values(); let mut new_row_ids = Vec::with_capacity(self.len()); let mut new_codes = Vec::with_capacity(self.len() * self.metadata.num_sub_vectors); + let mut kept_indices: Vec = Vec::with_capacity(self.len()); let row_ids = self.row_ids.values(); for (i, row_id) in row_ids.iter().enumerate() { @@ -622,6 +681,7 @@ impl QuantizerStorage for ProductQuantizationStorage { self.metadata.num_sub_vectors, i as u32, )); + kept_indices.push(i as u32); } Some(None) => {} None => { @@ -632,6 +692,7 @@ impl QuantizerStorage for ProductQuantizationStorage { self.metadata.num_sub_vectors, i as u32, )); + kept_indices.push(i as u32); } } } @@ -647,7 +708,14 @@ impl QuantizerStorage for ProductQuantizationStorage { new_transposed_codes, num_bytes_in_code as i32, )?); - RecordBatch::try_new(self.schema(), vec![new_row_ids.clone(), codes_fsl])? + // Preserve any included/covering columns row-aligned to the kept rows. + rebuild_storage_batch( + self.schema(), + &self.batch, + new_row_ids.clone(), + codes_fsl, + &UInt32Array::from(kept_indices), + )? }; let transposed_codes = batch[PQ_CODE_COLUMN] .as_fixed_size_list() @@ -710,6 +778,9 @@ impl QuantizerStorage for ProductQuantizationStorage { impl VectorStore for ProductQuantizationStorage { type DistanceCalculator<'a> = PQDistCalculator; + /// PQ storage is `[_rowid, __pq_code, ]`. + const INTERNAL_COLUMNS: &'static [&'static str] = PQ_INTERNAL_COLUMNS; + fn to_batches(&self) -> Result> { Ok(std::iter::once(self.batch.clone())) } @@ -1290,6 +1361,132 @@ mod tests { assert_eq!(storage.row_ids.len(), TOTAL); } + #[tokio::test] + async fn test_build_preserves_extra_column() { + // A covering ("included") column beyond _rowid + __pq_code must be kept + // in the storage batch, row-aligned to the rows, not dropped. + let storage = create_pq_storage_with_extra_column().await; + let batch = storage.batch(); + assert_eq!(batch.num_columns(), 3); + assert!(batch.column_by_name(ROW_ID).is_some()); + assert!(batch.column_by_name(PQ_CODE_COLUMN).is_some()); + let extra = batch + .column_by_name("extra") + .expect("extra column should be preserved") + .as_primitive::(); + // In the fixture extra[i] == row_ids[i] == i, so it must stay aligned. + for (i, row_id) in storage.row_ids().enumerate() { + assert_eq!(extra.value(i) as u64, *row_id); + } + } + + #[tokio::test] + async fn test_build_rejects_mismatched_batch_columns() { + // Batches merged into one storage must carry the same column set. A + // batch with a column the FIRST batch lacks must be rejected: aligning + // to the first batch's schema would silently drop it, losing covering + // data while the index metadata still advertises it. + let codebook = Float32Array::from_iter_values((0..256 * DIM).map(|_| rand::random())); + let codebook = FixedSizeListArray::try_new_from_values(codebook, DIM as i32).unwrap(); + let pq = ProductQuantizer::new(NUM_SUB_VECTORS, 8, DIM, codebook, DistanceType::Dot); + + let vec_field = Field::new( + "vec", + DataType::FixedSizeList( + Field::new_list_field(DataType::Float32, true).into(), + DIM as i32, + ), + true, + ); + let make_fsl = || { + let vectors = Float32Array::from_iter_values((0..TOTAL * DIM).map(|_| rand::random())); + FixedSizeListArray::try_new_from_values(vectors, DIM as i32).unwrap() + }; + let row_ids = || UInt64Array::from_iter_values((0..TOTAL).map(|v| v as u64)); + + let plain_schema = ArrowSchema::new(vec![vec_field.clone(), ROW_ID_FIELD.clone()]); + let plain = RecordBatch::try_new( + plain_schema.into(), + vec![Arc::new(make_fsl()), Arc::new(row_ids())], + ) + .unwrap(); + + let covered_schema = ArrowSchema::new(vec![ + vec_field, + ROW_ID_FIELD.clone(), + Field::new("extra", DataType::UInt32, true), + ]); + let covered = RecordBatch::try_new( + covered_schema.into(), + vec![ + Arc::new(make_fsl()), + Arc::new(row_ids()), + Arc::new(UInt32Array::from_iter_values((0..TOTAL).map(|v| v as u32))), + ], + ) + .unwrap(); + + let err = StorageBuilder::new("vec".to_owned(), pq.distance_type, pq, None) + .unwrap() + .build(vec![plain, covered]) + .expect_err("a batch with extra columns must be rejected, not silently projected"); + assert!(err.to_string().contains("extra"), "got: {err}"); + } + + #[tokio::test] + async fn test_new_drops_legacy_part_id_column() { + // Pre-#3606 `num_partitions=1` builds wrote `__ivf_part_id` into the + // index file. It must still be dropped on load (the legacy read shim), + // NOT preserved and misclassified as a covering column -- that would + // make search emit an extra column the exec's declared schema (from + // `IndexMetadata.covering_fields`, empty for those indexes) lacks. + use crate::vector::PART_ID_COLUMN; + let codebook = Float32Array::from_iter_values((0..256 * DIM).map(|_| rand::random())); + let codebook = FixedSizeListArray::try_new_from_values(codebook, DIM as i32).unwrap(); + + let schema = ArrowSchema::new(vec![ + ROW_ID_FIELD.clone(), + Field::new( + PQ_CODE_COLUMN, + DataType::FixedSizeList( + Field::new_list_field(DataType::UInt8, true).into(), + NUM_SUB_VECTORS as i32, + ), + true, + ), + Field::new(PART_ID_COLUMN, DataType::UInt32, true), + ]); + let row_ids = UInt64Array::from_iter_values((0..TOTAL).map(|v| v as u64)); + let codes = UInt8Array::from_iter_values((0..TOTAL * NUM_SUB_VECTORS).map(|v| v as u8)); + let codes = FixedSizeListArray::try_new_from_values(codes, NUM_SUB_VECTORS as i32).unwrap(); + let part_ids = UInt32Array::from_iter_values((0..TOTAL).map(|_| 0)); + let batch = RecordBatch::try_new( + schema.into(), + vec![Arc::new(row_ids), Arc::new(codes), Arc::new(part_ids)], + ) + .unwrap(); + + let storage = ProductQuantizationStorage::new( + codebook, + batch, + 8, + NUM_SUB_VECTORS, + DIM, + DistanceType::L2, + true, // codes already transposed + None, + ) + .unwrap(); + assert!( + storage.batch().column_by_name(PART_ID_COLUMN).is_none(), + "legacy __ivf_part_id must be dropped on load, not kept" + ); + assert!( + storage.covering_field_indices().is_empty(), + "legacy __ivf_part_id must not be classified as a covering column" + ); + } + #[tokio::test] async fn test_distance_all() { let storage = create_pq_storage().await; @@ -1392,8 +1589,19 @@ mod tests { // Rewritten row i lands at offset i of frag 1. assert_eq!(*row_id, (1u64 << 32) | i as u64); } - assert_eq!(new_storage.batch.num_columns(), 2); + // The covering ("extra") column must survive remap, row-aligned to the + // surviving rows — otherwise PK/covering data is lost on compaction. + assert_eq!(new_storage.batch.num_columns(), 3); assert!(new_storage.batch.column_by_name(ROW_ID).is_some()); assert!(new_storage.batch.column_by_name(PQ_CODE_COLUMN).is_some()); + let extra = new_storage + .batch + .column_by_name("extra") + .expect("extra column should survive remap") + .as_primitive::(); + // Surviving rows are original indices 0..TOTAL/2, whose extra values are 0..TOTAL/2. + for i in 0..TOTAL / 2 { + assert_eq!(extra.value(i), i as u32); + } } } diff --git a/rust/lance-index/src/vector/quantizer.rs b/rust/lance-index/src/vector/quantizer.rs index 433fa5031d4..c8b43636e24 100644 --- a/rust/lance-index/src/vector/quantizer.rs +++ b/rust/lance-index/src/vector/quantizer.rs @@ -290,7 +290,15 @@ pub trait QuantizerStorage: Clone + Sized + DeepSizeOf + VectorStore { let mut indices = Vec::with_capacity(b.num_rows()); let mut new_row_ids = Vec::with_capacity(b.num_rows()); - let row_ids = b.column(0).as_primitive::().values(); + // Look up by name, not position: a covering storage batch is + // `[_rowid, code, ]` for some quantizers but the + // covering columns can precede `_rowid` if the writer schema was + // inferred from a batch rather than declared, so position 0 is + // not reliably the row id column. + let row_id_col = b + .column_by_name(ROW_ID) + .ok_or_else(|| Error::index(format!("column {ROW_ID} not found in batch")))?; + let row_ids = row_id_col.as_primitive::().values(); for (i, row_id) in row_ids.iter().enumerate() { match mapping.get(*row_id) { Some(Some(new_id)) => { diff --git a/rust/lance-index/src/vector/sq/storage.rs b/rust/lance-index/src/vector/sq/storage.rs index 6062c22242e..35fd8a64acf 100644 --- a/rust/lance-index/src/vector/sq/storage.rs +++ b/rust/lance-index/src/vector/sq/storage.rs @@ -30,15 +30,25 @@ use crate::scalar::RowIdRemapper; use crate::{ INDEX_METADATA_SCHEMA_KEY, IndexMetadata, vector::{ - SQ_CODE_COLUMN, + PART_ID_COLUMN, SQ_CODE_COLUMN, quantizer::{QuantizerMetadata, QuantizerStorage}, - storage::{DistCalculator, DistanceCalculatorOptions, QueryResidual, VectorStore}, + storage::{ + DistCalculator, DistanceCalculatorOptions, QueryResidual, VectorStore, + remap_row_ids_by_name, + }, transform::Transformer, }, }; pub const SQ_METADATA_KEY: &str = "lance:sq"; +/// SQ storage's internal (non-covering) column names: the row id, the SQ code column, +/// and the legacy `__ivf_part_id` column (left in pre-#3606 single-partition builds; +/// unlike PQ, SQ does not drop it at construction). Every other column in an SQ +/// storage schema is a user-declared covering ("included") column. The single +/// authority for this set -- the distributed shard merger classifies from it too. +pub const SQ_INTERNAL_COLUMNS: &[&str] = &[ROW_ID, SQ_CODE_COLUMN, PART_ID_COLUMN]; + #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] pub struct ScalarQuantizationMetadata { pub dim: usize, @@ -193,7 +203,7 @@ impl ScalarQuantizationStorage { offsets.push(0); for mut batch in batches.into_iter() { if let Some(frag_reuse_index_ref) = frag_reuse_index.as_ref() { - batch = frag_reuse_index_ref.remap_row_ids_record_batch(batch, 0)? + batch = remap_row_ids_by_name(batch, frag_reuse_index_ref.as_ref())? } offsets.push(offsets.last().unwrap() + batch.num_rows() as u32); let chunk = SQStorageChunk::new(batch)?; @@ -342,6 +352,9 @@ impl QuantizerStorage for ScalarQuantizationStorage { impl VectorStore for ScalarQuantizationStorage { type DistanceCalculator<'a> = SQDistCalculator<'a>; + /// SQ storage is `[_rowid, __sq_code, ]`. + const INTERNAL_COLUMNS: &'static [&'static str] = SQ_INTERNAL_COLUMNS; + fn to_batches(&self) -> Result> { Ok(self.chunks.iter().map(|c| c.batch.clone())) } @@ -785,6 +798,50 @@ mod tests { RecordBatch::try_new(schema, vec![Arc::new(row_ids), Arc::new(code_arr)]).unwrap() } + #[test] + fn test_covering_excludes_legacy_part_id_column() { + // Pre-#3606 `num_partitions=1` builds left `__ivf_part_id` in the SQ + // storage file. It must never be classified as a covering column -- + // otherwise search emits a column the exec's declared schema (from + // `IndexMetadata.covering_fields`, empty for those indexes) lacks, + // desyncing the two and failing every query on that legacy index. + use crate::vector::PART_ID_COLUMN; + use arrow_array::UInt32Array; + const DIM: usize = 64; + const N: usize = 10; + + let row_ids = UInt64Array::from_iter_values(0..N as u64); + let sq_code = UInt8Array::from_iter_values((0..N * DIM).map(|v| v as u8)); + let code_arr = FixedSizeListArray::try_new_from_values(sq_code, DIM as i32).unwrap(); + let part_ids = UInt32Array::from_iter_values((0..N).map(|_| 0)); + + let schema = Arc::new(Schema::new(vec![ + Field::new(ROW_ID, DataType::UInt64, false), + Field::new( + SQ_CODE_COLUMN, + DataType::FixedSizeList( + Arc::new(Field::new("item", DataType::UInt8, true)), + DIM as i32, + ), + false, + ), + Field::new(PART_ID_COLUMN, DataType::UInt32, true), + ])); + let batch = RecordBatch::try_new( + schema, + vec![Arc::new(row_ids), Arc::new(code_arr), Arc::new(part_ids)], + ) + .unwrap(); + + let storage = + ScalarQuantizationStorage::try_new(8, DistanceType::L2, -0.7..0.7, [batch], None) + .unwrap(); + assert!( + storage.covering_field_indices().is_empty(), + "legacy __ivf_part_id must not be classified as a covering column" + ); + } + #[test] fn test_get_chunks() { const DIM: usize = 64; diff --git a/rust/lance-index/src/vector/storage.rs b/rust/lance-index/src/vector/storage.rs index 151be945741..8313c1a1cf9 100644 --- a/rust/lance-index/src/vector/storage.rs +++ b/rust/lance-index/src/vector/storage.rs @@ -5,14 +5,14 @@ use crate::vector::quantizer::QuantizerStorage; use arrow::compute::concat_batches; -use arrow_array::{ArrayRef, RecordBatch}; -use arrow_schema::SchemaRef; +use arrow_array::{ArrayRef, RecordBatch, UInt32Array}; +use arrow_schema::{Field, Schema, SchemaRef}; use futures::prelude::stream::TryStreamExt; use lance_arrow::RecordBatchExt; use lance_core::deepsize::DeepSizeOf; use lance_core::{Error, ROW_ID, Result}; use lance_encoding::decoder::FilterExpression; -use lance_file::reader::FileReader; +use lance_file::reader::{FileReader, ReaderProjection}; use lance_io::ReadBatchParams; use lance_io::scheduler::IoStats; use lance_linalg::distance::DistanceType; @@ -20,7 +20,7 @@ use prost::Message; use std::{ any::Any, borrow::Cow, - collections::BinaryHeap, + collections::{BinaryHeap, HashSet}, mem::size_of, ops::{Deref, DerefMut}, sync::Arc, @@ -155,6 +155,33 @@ pub trait DistCalculator { pub const STORAGE_METADATA_KEY: &str = "storage_metadata"; +/// Schema-metadata key recording the *source dataset* field ids of a storage file's +/// covering ("included") columns, comma separated in physical schema order. +/// +/// Arrow fields carry no Lance field id, so a storage file's own schema cannot say +/// which logical column a covering column came from -- and the ids in the file's Lance +/// schema are file-local, assigned when the writer was built. Without this, a +/// distributed merge comparing shards by name and type would accept two shards whose +/// covering columns share a name and type but belong to different fields (a column +/// dropped and re-added between shard builds) and concatenate them as one. +/// +/// Absent when the storage has no covering columns. +pub const COVERING_FIELD_IDS_KEY: &str = "covering_field_ids"; + +/// Full-precision vectors carried in index storage so `refine` can re-rank without +/// taking them from the base table. +pub const REFINE_VECTOR_COLUMN: &str = "__refine_vector"; + +/// Names the build pipeline uses in flight that must never be a user column. +/// +/// [`REFINE_VECTOR_COLUMN`] carries the copy of the indexed column past the transform +/// chain, which consumes a column of that column's own name. The copy takes the real name +/// back before it is written, so this never appears in a stored schema -- but a *user* +/// covering column of the same name would collide with the copy mid-flight, so the name is +/// reserved. Reservation is all this is for: it is chained into `RESERVED_STORAGE_COLUMNS` +/// at index-create time and rejected there. +pub const DEFERRED_INTERNAL_COLUMNS: &[&str] = &[REFINE_VECTOR_COLUMN]; + #[derive(Debug)] pub struct QueryScratch { pub distances: Vec, @@ -373,6 +400,108 @@ impl DeepSizeOf for QueryScratchPool { /// /// It abstracts away the logic to compute the distance between vectors. /// +/// Indices of the covering ("included") columns in a vector-storage schema: every +/// field whose name is not one of `internal`. Each storage lists its own non-covering +/// column names (the row id, its quantization code columns, and any legacy bookkeeping +/// column it carries) and shares this filter, so covering detection stays a per-storage +/// data decision rather than duplicated iteration logic. See +/// [`VectorStore::covering_field_indices`] for why the set cannot be inferred generically. +pub(crate) fn covering_field_indices_excluding( + schema: &arrow_schema::Schema, + internal: &[&str], +) -> Vec { + schema + .fields() + .iter() + .enumerate() + .filter(|(_, f)| !internal.contains(&f.name().as_str())) + .map(|(i, _)| i) + .collect() +} + +/// Resolve the covering columns a physical storage schema can safely serve. +/// +/// Arrow fields do not carry Lance field ids, so the physical columns are usable only +/// when the storage also stamps one source id for each column. Invalid or incomplete +/// capability metadata deliberately resolves to no columns: the manifest declaration is +/// still valid, but readers must fetch those values from the base table instead. +fn physical_covering_fields_from_schema(schema: &Schema, internal: &[&str]) -> Vec<(i32, Field)> { + let covering_indices = covering_field_indices_excluding(schema, internal); + if covering_indices.is_empty() { + return Vec::new(); + } + + let Some(encoded_ids) = schema.metadata().get(COVERING_FIELD_IDS_KEY) else { + // Every return of `Vec::new()` below withdraws the whole covering payload, and the + // query then reads those columns from the base table with correct results and no + // error -- so without a line here the operator cannot tell a covered index that + // stopped eliding takes from one that never did. + log::debug!( + "Index storage carries {} covering column(s) but no `{}` stamp, so none can be \ + served; their values will come from a base-table read.", + covering_indices.len(), + COVERING_FIELD_IDS_KEY + ); + return Vec::new(); + }; + let Ok(source_ids) = encoded_ids + .split(',') + .map(str::trim) + .map(str::parse::) + .collect::, _>>() + else { + return Vec::new(); + }; + if source_ids.len() != covering_indices.len() + || source_ids.iter().collect::>().len() != source_ids.len() + { + log::debug!( + "Index storage's `{}` stamp ({:?}) does not pair one distinct source id with each \ + of its {} covering column(s), so none can be served; their values will come from \ + a base-table read.", + COVERING_FIELD_IDS_KEY, + source_ids, + covering_indices.len() + ); + return Vec::new(); + } + + let covering_names = covering_indices + .iter() + .map(|index| schema.field(*index).name()) + .collect::>(); + if covering_names.len() != covering_indices.len() { + log::debug!( + "Index storage has duplicate covering column names, so a value cannot be bound to \ + one source id; none can be served and their values will come from a base-table \ + read." + ); + return Vec::new(); + } + + source_ids + .into_iter() + .zip(covering_indices) + .map(|(source_id, index)| (source_id, schema.field(index).clone())) + .collect() +} + +/// Remap `batch`'s row ids through `frag_reuse_index`, locating the row id column +/// by name rather than assuming it is at a fixed position. A covering storage batch +/// is built from a scan projection of `[, vector]` with `_rowid` +/// appended by the scan, so `_rowid` lands wherever the covering columns end -- +/// never reliably at a fixed index once any are configured. +pub(crate) fn remap_row_ids_by_name( + batch: RecordBatch, + frag_reuse_index: &dyn RowIdRemapper, +) -> Result { + let row_id_idx = batch + .schema() + .index_of(ROW_ID) + .map_err(|_| Error::schema(format!("column {} not found", ROW_ID)))?; + frag_reuse_index.remap_row_ids_record_batch(batch, row_id_idx) +} + /// TODO: should we rename this to "VectorDistance"?; /// ///

@@ -385,6 +514,17 @@ pub trait VectorStore: Send + Sync + Sized + Clone { where Self: 'a; + /// This storage's own column names: the row id, its quantization code + /// columns, and any legacy bookkeeping column it carries. Everything else in + /// the schema is a covering ("included") column. + /// + /// Required rather than defaulted on purpose: only the storage knows which of + /// its columns are code columns, and a wrong answer here silently mislabels + /// covering data (e.g. it would mistake RaBitQ's code columns for covering + /// ones). Making it a required associated const means a new storage cannot + /// forget to answer. + const INTERNAL_COLUMNS: &'static [&'static str]; + fn as_any(&self) -> &dyn Any; fn schema(&self) -> &SchemaRef; @@ -410,6 +550,40 @@ pub trait VectorStore: Send + Sync + Sized + Clone { /// The storage implement will perform quantization if necessary. fn append_batch(&self, batch: RecordBatch, vector_column: &str) -> Result; + /// Field indices of the "included"/covering columns: the extra columns + /// stored alongside the row id and quantization code so a covered query can + /// skip the take from the base table. Derived from [`Self::INTERNAL_COLUMNS`]. + fn covering_field_indices(&self) -> Vec { + covering_field_indices_excluding(self.schema().as_ref(), Self::INTERNAL_COLUMNS) + } + + /// `[_rowid, ]` for the whole storage. Captured while a + /// partition is loaded during search so covering columns can be emitted with + /// the result without a separate take. Returns `None` if there are no + /// included columns (ordinary index — nothing to cover). + fn covering_batch(&self) -> Result> { + let included = self.covering_field_indices(); + if included.is_empty() { + return Ok(None); + } + let schema = self.schema().clone(); + let row_id_idx = schema.index_of(ROW_ID)?; + let mut indices = Vec::with_capacity(included.len() + 1); + indices.push(row_id_idx); + indices.extend(included); + // Project each batch BEFORE concatenating. Concatenating the full partition first + // would copy every quantization-code column only to discard it on the next line -- + // and those columns dominate the partition (HNSW partitions target 1 << 20 rows), so + // on a multi-probe query that copy is hundreds of MB of pure waste on the ANN hot + // path. Only `[_rowid, ]` is ever read from the result. + let projected: Vec = self + .to_batches()? + .map(|batch| batch.project(&indices)) + .collect::>()?; + let projected_schema = Arc::new(schema.project(&indices)?); + Ok(Some(concat_batches(&projected_schema, projected.iter())?)) + } + /// Create a [DistCalculator] to compute the distance between the query. /// /// Using dist calculator can be more efficient as it can pre-compute some @@ -480,7 +654,43 @@ impl StorageBuilder { } pub fn build(&self, batches: Vec) -> Result { - let mut batch = concat_batches(batches[0].schema_ref(), batches.iter())?; + // Batches can come from different sources (existing partitions loaded + // from disk vs freshly shuffled/split/reassigned data) that carry the + // same columns in a different order -- notably when the index has + // included/covering columns. `concat_batches` matches columns by + // position, so align every batch to the first batch's column order (by + // name). No-op when all batches already share a schema. + let schema = batches[0].schema(); + let aligned: Vec = batches + .iter() + .map(|b| { + if b.schema() == schema { + return Ok(b.clone()); + } + // `project_by_schema` reorders by name and errors on a missing + // column, but it cannot see extras -- and a batch with MORE + // columns must error rather than have them silently dropped + // (e.g. covering columns the index metadata still advertises). + // Equal count + every expected column present = same set. + if b.num_columns() != schema.fields().len() { + let names = |s: &arrow_schema::Schema| { + s.fields() + .iter() + .map(|f| f.name().as_str()) + .collect::>() + .join(", ") + }; + return Err(Error::index(format!( + "mismatched columns while merging vector storage batches: \ + expected [{}], got [{}]", + names(&schema), + names(&b.schema()), + ))); + } + Ok(b.project_by_schema(schema.as_ref())?) + }) + .collect::>>()?; + let mut batch = concat_batches(&schema, aligned.iter())?; if batch.column_by_name(self.quantizer.column()).is_none() { let vectors = batch @@ -508,6 +718,61 @@ impl StorageBuilder { } } +/// How [`IvfQuantizationStorage::take_covering`] reads a partition's covering values. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum CoveringRead { + /// [`ReadBatchParams::Indices`] -- read only the survivors' rows. + Scattered, + /// [`ReadBatchParams::Range`] -- read the partition's whole row range. + Sequential, +} + +/// Survivors at or above this percentage of a partition's rows make the scattered read +/// degenerate: `Indices` over most of a range costs more than the sequential read of that +/// range it replaced, because the scattered form gives up contiguity while still touching +/// nearly every page. +/// +/// # Derivation +/// +/// Measured, not guessed. At production partition sizes (~8,192 rows) and `k <= 100` the +/// survivors of one probe are **<= 1.2%** of a partition; on the 2,048-row fixture +/// partitions with `k = 10` it is 0.49%. Ten percent therefore sits roughly an order of +/// magnitude above the worst observed common case, so the fallback costs nothing in +/// practice while still bounding the degenerate case at "no worse than the read this +/// replaced". The design's §1 warns explicitly against assuming random access always +/// wins; this is that guard, deliberately conservative in the direction that cannot +/// regress. +pub const COVERING_SCATTERED_READ_MAX_PERCENT: usize = 10; + +/// Which read shape to use for `survivors` covering rows out of `partition_rows`. +/// +/// A zero-row partition has nothing to read either way; it is reported as sequential so +/// callers have a single empty-range path. +pub fn covering_read_for(survivors: usize, partition_rows: usize) -> CoveringRead { + if partition_rows == 0 + || survivors.saturating_mul(100) + >= partition_rows.saturating_mul(COVERING_SCATTERED_READ_MAX_PERCENT) + { + CoveringRead::Sequential + } else { + CoveringRead::Scattered + } +} + +/// Which columns a partition load reads out of the storage file. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum PartitionColumns { + /// Only the storage's own columns ([`VectorStore::INTERNAL_COLUMNS`]): the row id, + /// the quantization codes and any per-row factor columns. Covering columns are left + /// unread, which is what makes the loaded partition independent of the query that + /// loaded it -- and therefore sound to share under a partition-id-only cache key. + Internal, + /// Every column the file carries, covering columns included. Required by the paths + /// that rewrite a partition into a new index file: a covering column that was never + /// read cannot be written back. + All, +} + /// Loader to load partitioned PQ storage from disk. #[derive(Debug)] pub struct IvfQuantizationStorage { @@ -518,6 +783,17 @@ pub struct IvfQuantizationStorage { ivf: IvfModel, frag_reuse_index: Option>, + /// Lazily-computed covering schema (see [`Self::covering_schema`]). The schema is + /// fixed for the storage's lifetime and is consulted on every search, so it is + /// resolved once rather than per query. Constructing it builds an empty `Q::Storage`: + /// cheap for a current-format index (zero rows, and the codebook is cloned, not + /// decoded), but the legacy `codebook_tensor` path does decode, so this is not free + /// for every index. + covering_schema: std::sync::OnceLock>, + /// Lazily-computed internal-column projection (see [`Self::internal_projection`]). + /// Fixed for the storage's lifetime and consulted on every partition load, so it is + /// resolved once rather than per probe. + internal_projection: std::sync::OnceLock, } impl DeepSizeOf for IvfQuantizationStorage { @@ -589,6 +865,8 @@ impl IvfQuantizationStorage { metadata, ivf, frag_reuse_index, + covering_schema: std::sync::OnceLock::new(), + internal_projection: std::sync::OnceLock::new(), }) } @@ -620,6 +898,8 @@ impl IvfQuantizationStorage { metadata, ivf, frag_reuse_index, + covering_schema: std::sync::OnceLock::new(), + internal_projection: std::sync::OnceLock::new(), } } @@ -656,13 +936,80 @@ impl IvfQuantizationStorage { Arc::new(self.reader.schema().as_ref().into()) } + /// The `[_rowid, ]` schema for this index's covering + /// columns, or `None` if the index has no covering columns. Derived from an + /// empty storage instance (no I/O) so callers can emit the correct covered + /// output schema even when zero partitions are searched. Cached because it is queried + /// per search. The cache is best-effort under concurrency: racing cold callers may each + /// construct one and all but the first are discarded. That is deliberate -- the empty + /// construction is cheap on the current format, and `OnceLock::get_or_init` cannot + /// carry the `Result` this returns. + pub fn covering_schema(&self) -> Result> { + if let Some(cached) = self.covering_schema.get() { + return Ok(cached.clone()); + } + let arrow_schema = arrow_schema::Schema::from(self.reader.schema().as_ref()); + let empty = RecordBatch::new_empty(Arc::new(arrow_schema)); + // No remapper: the batch is empty, so there are no row ids to remap and the + // remapper cannot influence the schema this derives. Passing one would also be + // wrong now that the field holds a `dyn RowIdRemapper` -- the default + // `try_from_batch_with_remapper` rejects a remapper outright for every storage + // that does not override it, which would fail this call on an SQ/RQ/FLAT index + // whose dataset happens to carry a fragment-reuse index. + let storage = Q::Storage::try_from_batch(empty, self.metadata(), self.distance_type, None)?; + let computed = storage.covering_batch()?.map(|b| b.schema()); + Ok(self.covering_schema.get_or_init(|| computed).clone()) + } + + /// Physical covering columns this storage can safely serve, paired with their source + /// dataset field ids and returned in storage order. + /// + /// A missing or malformed source-id stamp makes the physical covering payload + /// unavailable rather than making the index unreadable. Query planning will then use + /// the ordinary base-table take for the declared columns. + pub fn physical_covering_fields(&self) -> Vec<(i32, Field)> { + let schema = Schema::from(self.reader.schema().as_ref()); + physical_covering_fields_from_schema(&schema, Q::Storage::INTERNAL_COLUMNS) + } + /// Get the number of partitions in the storage. pub fn num_partitions(&self) -> usize { self.ivf.num_partitions() } - /// Load a partition's quantization storage, optionally measuring the exact - /// I/O it performs into `io_stats`. + /// The read projection selecting only this storage's own columns. + /// + /// [`VectorStore::INTERNAL_COLUMNS`] is the union over every build a storage type + /// has ever produced -- the legacy `__ivf_part_id`, RaBitQ's factor columns that + /// exist only above a `num_bits` threshold -- so it is intersected with this file's + /// schema first. Projecting the const directly would reject any file that legitimately + /// omits one of those columns, which is most of them. + fn internal_projection(&self) -> Result { + if let Some(cached) = self.internal_projection.get() { + return Ok(cached.clone()); + } + let schema = self.reader.schema(); + let names: Vec<&str> = Q::Storage::INTERNAL_COLUMNS + .iter() + .copied() + .filter(|name| schema.field(name).is_some()) + .collect(); + let projection = lance_file::versions::reader_projection_from_column_names( + self.reader.metadata().version(), + schema, + &names, + )?; + Ok(self.internal_projection.get_or_init(|| projection).clone()) + } + + /// Load a partition's quantization storage, reading the columns named by + /// `columns` and optionally measuring the exact I/O it performs into `io_stats`. + /// + /// [`PartitionColumns::Internal`] leaves the covering columns on disk. Search + /// wants that: covering is needed for at most `k` survivors while the codes are + /// scanned for every row on every probe, and an entry that holds no covering is + /// the same entry for every query -- which is what makes caching it under a + /// partition id alone correct rather than accidentally correct. /// /// When `io_stats` is `Some`, the partition is read through a reader whose /// scheduler also records into the sink (a cheap clone that shares all @@ -671,29 +1018,49 @@ impl IvfQuantizationStorage { pub async fn load_partition( &self, part_id: usize, + columns: PartitionColumns, io_stats: Option, ) -> Result { + let projection = match columns { + PartitionColumns::Internal => Some(self.internal_projection()?), + PartitionColumns::All => None, + }; + // The batch schema must describe exactly what was read: `concat_batches` + // indexes each batch by this schema's field positions without comparing the + // two, so the full file schema over a projected read would index past the + // last column. + let schema: SchemaRef = Arc::new(match &projection { + Some(projection) => projection.schema.as_ref().into(), + None => self.reader.schema().as_ref().into(), + }); let range = self.ivf.row_range(part_id); let batch = if range.is_empty() { - let schema = self.reader.schema(); - let arrow_schema = arrow_schema::Schema::from(schema.as_ref()); - RecordBatch::new_empty(Arc::new(arrow_schema)) + RecordBatch::new_empty(schema) } else { let reader = match &io_stats { Some(io_stats) => Cow::Owned(self.reader.with_io_stats(io_stats.recorder())), None => Cow::Borrowed(&self.reader), }; - let batches = reader - .read_stream( - ReadBatchParams::Range(range), - u32::MAX, - 1, - FilterExpression::no_filter(), - ) - .await? - .try_collect::>() - .await?; - let schema = Arc::new(self.reader.schema().as_ref().into()); + let params = ReadBatchParams::Range(range); + let stream = match projection { + Some(projection) => { + reader + .read_stream_projected( + params, + u32::MAX, + 1, + projection, + FilterExpression::no_filter(), + ) + .await? + } + None => { + reader + .read_stream(params, u32::MAX, 1, FilterExpression::no_filter()) + .await? + } + }; + let batches = stream.try_collect::>().await?; concat_batches(&schema, batches.iter())? }; Q::Storage::try_from_batch_with_remapper( @@ -703,12 +1070,242 @@ impl IvfQuantizationStorage { self.frag_reuse_index.clone(), ) } + + /// This storage's covering ("included") column names in storage order, narrowed to + /// `wanted`. + /// + /// `wanted` is [`Query::covering_projection`], whose three states are all meaningful: + /// `None` keeps every physical column, `Some(&[])` keeps none, and `Some(cols)` keeps + /// the intersection. Execution planners pass an explicit, storage-verified `Some`; + /// `None` is only the raw-index behavior before capability resolution. The result is + /// empty for an ordinary index and for a query that needs no covering column -- in + /// both cases the caller skips the covering read entirely. + /// + /// Order follows storage, never `wanted`: that is the order the gathered batch comes + /// back in. The planner verifies that this physical order is compatible with its + /// declaration-ordered output schema before enabling covering. + /// + /// [`Query::covering_projection`]: crate::vector::Query::covering_projection + pub fn covering_columns(&self, wanted: Option<&[String]>) -> Result> { + let Some(schema) = self.covering_schema()? else { + return Ok(Vec::new()); + }; + Ok(schema + .fields() + .iter() + .map(|field| field.name()) + .filter(|name| name.as_str() != ROW_ID) + .filter(|name| wanted.is_none_or(|wanted| wanted.contains(name))) + .cloned() + .collect()) + } + + /// The read projection [`Self::take_covering`] uses: `[_rowid, ]`. + fn covering_projection(&self, columns: &[String]) -> Result { + let mut names: Vec<&str> = Vec::with_capacity(columns.len() + 1); + names.push(ROW_ID); + names.extend(columns.iter().map(String::as_str)); + lance_file::versions::reader_projection_from_column_names( + self.reader.metadata().version(), + self.reader.schema(), + &names, + ) + } + + /// The `[_rowid, ]` schema [`Self::take_covering`] returns. + /// + /// Callers declare their covered output schema from this rather than from + /// [`Self::covering_schema`] so that the schema they promise and the batches they + /// emit come from a single derivation. The two agree today, but they are computed + /// from different inputs (a projected file schema against an empty `Q::Storage`) and + /// a divergence would surface as a stream that mixes batch schemas. + pub fn covering_read_schema(&self, columns: &[String]) -> Result { + Ok(Arc::new( + self.covering_projection(columns)?.schema.as_ref().into(), + )) + } + + /// Read `[_rowid, ]` for one partition's search survivors. + /// + /// This is the covering read the search path pays after scoring has settled the heap, + /// so it is proportional to `k` rather than to the partition size -- which is why the + /// partition entry the search caches carries no covering at all (see + /// [`PartitionColumns::Internal`]) and why this result is deliberately **not** cached: + /// it is a different set of rows for every query. + /// + /// `positions` are offsets into the partition's row range, strictly ascending and + /// deduplicated (the reader's take path requires that). `None` means the caller could + /// not derive them -- a deferred fragment-reuse remap drops rows from the loaded + /// partition, so its row order no longer matches the file's -- and the whole range is + /// read instead, with the caller matching by row id. + /// + /// Row ids come back with the values, so the caller can align covering values to its + /// survivors **by row id** rather than by position, and a survivor whose row the + /// gather did not return is an error rather than a silent null. + pub async fn take_covering( + &self, + part_id: usize, + positions: Option<&[u32]>, + columns: &[String], + io_stats: Option, + ) -> Result { + let projection = self.covering_projection(columns)?; + // As in `load_partition`: `concat_batches` indexes each batch by this schema's + // field positions without comparing the two, so it must describe the projected + // read rather than the whole file. + let schema: SchemaRef = Arc::new(projection.schema.as_ref().into()); + let range = self.ivf.row_range(part_id); + if range.is_empty() { + return Ok(RecordBatch::new_empty(schema)); + } + // `ReadBatchParams::Indices` addresses rows with a `u32`, so a partition ending + // beyond that cannot be expressed as a scattered read at all. Falling back keeps + // such a file readable (just slower) instead of failing the query. + let addressable = range.end <= u32::MAX as usize; + let read = match positions { + Some(positions) if addressable => covering_read_for(positions.len(), range.len()), + _ => CoveringRead::Sequential, + }; + let params = match read { + CoveringRead::Sequential => ReadBatchParams::Range(range.clone()), + CoveringRead::Scattered => ReadBatchParams::Indices(UInt32Array::from_iter_values( + // `addressable` bounds the sum, but saturate anyway: a caller passing a + // position past the partition must reach the reader's own bounds check as + // an error, not overflow here. + positions + .unwrap_or_default() + .iter() + .map(|position| (range.start as u32).saturating_add(*position)), + )), + }; + let reader = match &io_stats { + Some(io_stats) => Cow::Owned(self.reader.with_io_stats(io_stats.recorder())), + None => Cow::Borrowed(&self.reader), + }; + let batches = reader + .read_stream_projected( + params, + u32::MAX, + 1, + projection, + FilterExpression::no_filter(), + ) + .await? + .try_collect::>() + .await?; + let batch = concat_batches(&schema, batches.iter())?; + match &self.frag_reuse_index { + // The loaded partition's row ids were remapped through the same index when it + // was built from this file, so the gathered ids must be too or nothing matches. + Some(frag_reuse_index) => remap_row_ids_by_name(batch, frag_reuse_index.as_ref()), + None => Ok(batch), + } + } } #[cfg(test)] mod tests { - use super::{QueryScratchCapacity, QueryScratchPool}; + use super::{ + COVERING_FIELD_IDS_KEY, COVERING_SCATTERED_READ_MAX_PERCENT, CoveringRead, + QueryScratchCapacity, QueryScratchPool, covering_read_for, + physical_covering_fields_from_schema, + }; + use arrow_schema::{DataType, Field, Schema}; use lance_core::deepsize::DeepSizeOf; + use std::collections::HashMap; + + const TEST_INTERNAL_COLUMNS: &[&str] = &["_rowid", "__pq_code"]; + + fn storage_schema(covering_fields: Vec, source_ids: Option<&str>) -> Schema { + let mut fields = vec![ + Field::new("_rowid", DataType::UInt64, false), + Field::new("__pq_code", DataType::Binary, false), + ]; + fields.extend(covering_fields); + let metadata = source_ids + .map(|ids| HashMap::from([(COVERING_FIELD_IDS_KEY.to_string(), ids.to_string())])) + .unwrap_or_default(); + Schema::new_with_metadata(fields, metadata) + } + + #[test] + fn test_physical_covering_fields_require_valid_source_ids() { + let covering = vec![ + Field::new("price", DataType::Int64, true), + Field::new("payload", DataType::Utf8, false), + ]; + let schema = storage_schema(covering.clone(), Some("17, 23")); + + let resolved = physical_covering_fields_from_schema(&schema, TEST_INTERNAL_COLUMNS); + + assert_eq!( + resolved, + vec![(17, covering[0].clone()), (23, covering[1].clone())] + ); + } + + #[test] + fn test_physical_covering_fields_treat_invalid_capability_as_absent() { + let covering = vec![ + Field::new("price", DataType::Int64, true), + Field::new("payload", DataType::Utf8, false), + ]; + for source_ids in [None, Some(""), Some("17"), Some("17,nope"), Some("17,17")] { + let schema = storage_schema(covering.clone(), source_ids); + assert!( + physical_covering_fields_from_schema(&schema, TEST_INTERNAL_COLUMNS).is_empty(), + "source ids {source_ids:?} must not prove a physical covering capability" + ); + } + + let duplicate_names = storage_schema( + vec![ + Field::new("price", DataType::Int64, true), + Field::new("price", DataType::Int64, true), + ], + Some("17,23"), + ); + assert!( + physical_covering_fields_from_schema(&duplicate_names, TEST_INTERNAL_COLUMNS) + .is_empty() + ); + } + + #[test] + fn test_physical_covering_fields_ignore_orphaned_source_id_metadata() { + let schema = storage_schema(Vec::new(), Some("17")); + assert!(physical_covering_fields_from_schema(&schema, TEST_INTERNAL_COLUMNS).is_empty()); + } + + /// The scattered read is the point of the survivor gather, so it must be chosen across + /// the whole range of survivor fractions a real query produces -- and abandoned once + /// `Indices` would scatter over most of the partition anyway, where the sequential read + /// it replaced is cheaper. + /// + /// The inputs are the measured ones: at production partition sizes (~8,192 rows) the + /// survivors of one probe are <= 1.2% at `k <= 100`, and 0.49% on the 2,048-row fixture + /// partitions at `k = 10`. Both must land on `Scattered` with room to spare, or the + /// threshold is not conservative -- it is in the way. + /// + /// The two rows either side of 10% are what make this a threshold test rather than a + /// "scattered usually wins" test: moving `COVERING_SCATTERED_READ_MAX_PERCENT` in + /// either direction fails one of them. + #[test] + fn test_covering_read_falls_back_to_a_sequential_read_near_the_threshold() { + assert_eq!(COVERING_SCATTERED_READ_MAX_PERCENT, 10); + + // k = 10 on a 2,048-row fixture partition: 0.49%. + assert_eq!(covering_read_for(10, 2048), CoveringRead::Scattered); + // k = 100 on an 8,192-row production partition: 1.2%, the measured worst case. + assert_eq!(covering_read_for(100, 8192), CoveringRead::Scattered); + // Either side of the threshold: 819/8192 = 9.998%, 820/8192 = 10.01%. + assert_eq!(covering_read_for(819, 8192), CoveringRead::Scattered); + assert_eq!(covering_read_for(820, 8192), CoveringRead::Sequential); + // The degenerate case the threshold exists for. + assert_eq!(covering_read_for(8192, 8192), CoveringRead::Sequential); + // An empty partition has nothing to read either way; one path, not two. + assert_eq!(covering_read_for(0, 0), CoveringRead::Sequential); + } #[test] fn test_query_scratch_pool_reuses_buffers() { diff --git a/rust/lance-index/src/vector/transform.rs b/rust/lance-index/src/vector/transform.rs index 2faae7b9ce6..1ad89afd237 100644 --- a/rust/lance-index/src/vector/transform.rs +++ b/rust/lance-index/src/vector/transform.rs @@ -7,18 +7,18 @@ use std::fmt::Debug; use std::sync::Arc; -use arrow::datatypes::UInt64Type; -use arrow_array::UInt64Array; use arrow_array::types::{Float16Type, Float32Type, Float64Type}; use arrow_array::{Array, ArrowPrimitiveType, RecordBatch, UInt32Array, cast::AsArray}; use arrow_schema::{DataType, Field, Schema}; use lance_arrow::RecordBatchExt; use num_traits::Float; -use lance_core::{Error, ROW_ID, ROW_ID_FIELD, Result}; +use lance_core::{Error, Result}; use lance_linalg::kernels::normalize_fsl; use tracing::instrument; +use super::{CENTROID_DIST_COLUMN, PART_ID_COLUMN}; + /// Transform of a Vector Matrix. /// /// @@ -194,25 +194,48 @@ impl Transformer for Flatten { match arr.data_type() { DataType::FixedSizeList(_, _) => Ok(batch.clone()), DataType::List(_) => { - let row_ids = batch[ROW_ID].as_primitive::(); let vectors = arr.as_list::(); - - let row_ids = row_ids.values().iter().zip(vectors.iter()).flat_map( - |(row_id, multivector)| { - std::iter::repeat_n( - *row_id, - multivector.map(|multivec| multivec.len()).unwrap_or(0), - ) - }, - ); - let row_ids = UInt64Array::from_iter_values(row_ids); - let vectors = vectors.values().as_fixed_size_list().clone(); - let schema = Arc::new(Schema::new(vec![ - ROW_ID_FIELD.clone(), - Field::new(self.column.as_str(), vectors.data_type().clone(), true), - ])); - let batch = - RecordBatch::try_new(schema, vec![Arc::new(row_ids), Arc::new(vectors)])?; + // Each source row expands into one output row per vector in its list. + // Replicate the row id AND every other column (e.g. covering / "included" + // columns) across the expansion via a gather; dropping them here would + // silently lose covering data on the multivector build/split path. + // + // The IVF partition-assignment columns are the exception and must be + // DROPPED, not replicated: they describe the source *row*, and after the + // expansion each sub-vector needs its own nearest centroid. Carrying them + // through makes `PartitionTransformer` see its output column already + // present and early-return, silently collapsing every sub-vector of a row + // into that row's single partition -- no error, just degraded recall. + let take_indices = + UInt32Array::from_iter_values((0..vectors.len()).flat_map(|i| { + let n = if vectors.is_valid(i) { + vectors.value(i).len() + } else { + 0 + }; + std::iter::repeat_n(i as u32, n) + })); + let flat_vectors = vectors.values().as_fixed_size_list().clone(); + let mut fields: Vec> = Vec::with_capacity(batch.num_columns()); + let mut columns: Vec = + Vec::with_capacity(batch.num_columns()); + for (i, field) in batch.schema().fields().iter().enumerate() { + if field.name() == &self.column { + fields.push(Arc::new(Field::new( + self.column.as_str(), + flat_vectors.data_type().clone(), + true, + ))); + columns.push(Arc::new(flat_vectors.clone())); + } else if field.name() == PART_ID_COLUMN || field.name() == CENTROID_DIST_COLUMN + { + continue; + } else { + fields.push(field.clone()); + columns.push(arrow::compute::take(batch.column(i), &take_indices, None)?); + } + } + let batch = RecordBatch::try_new(Arc::new(Schema::new(fields)), columns)?; Ok(batch) } _ => Err(Error::index(format!( @@ -236,6 +259,114 @@ mod tests { use lance_arrow::*; use lance_linalg::distance::L2; + #[test] + fn test_flatten_preserves_extra_columns_for_multivector() { + use arrow::buffer::OffsetBuffer; + use arrow::datatypes::UInt64Type; + use arrow_array::{ListArray, UInt64Array}; + use lance_core::{ROW_ID, ROW_ID_FIELD}; + + // row 10 has 2 vectors, row 20 has 1 vector; a covering column "meta" must be + // replicated across each row's expanded vectors, not dropped. + let values = Float32Array::from(vec![1.0, 1.0, 2.0, 2.0, 3.0, 3.0]); + let fsl = FixedSizeListArray::try_new_from_values(values, 2).unwrap(); + let offsets = OffsetBuffer::new(vec![0i32, 2, 3].into()); + let item = Arc::new(Field::new("item", fsl.data_type().clone(), true)); + let list = ListArray::new(item, offsets, Arc::new(fsl), None); + + let batch = RecordBatch::try_new( + Arc::new(Schema::new(vec![ + ROW_ID_FIELD.clone(), + Field::new("vector", list.data_type().clone(), true), + Field::new("meta", DataType::Int32, true), + ])), + vec![ + Arc::new(UInt64Array::from(vec![10u64, 20])), + Arc::new(list), + Arc::new(Int32Array::from(vec![100, 200])), + ], + ) + .unwrap(); + + let out = Flatten::new("vector").transform(&batch).unwrap(); + + assert_eq!(out.num_rows(), 3, "2 + 1 vectors expand to 3 rows"); + assert_eq!( + out[ROW_ID].as_primitive::().values(), + &[10, 10, 20] + ); + let meta = out + .column_by_name("meta") + .expect("covering column must survive Flatten") + .as_primitive::(); + assert_eq!( + meta.values(), + &[100, 100, 200], + "covering column replicated per expanded vector" + ); + assert!(matches!( + out.column_by_name("vector").unwrap().data_type(), + DataType::FixedSizeList(_, 2) + )); + } + + /// The IVF partition-assignment columns describe the source *row*. Replicating them + /// across the multivector expansion makes `PartitionTransformer` see its output column + /// already present and early-return, so every sub-vector inherits its row's single + /// partition instead of being assigned to its own nearest centroid -- silent recall + /// loss with no error. They must be dropped even though covering columns are kept. + #[test] + fn test_flatten_drops_partition_columns_for_multivector() { + use arrow::buffer::OffsetBuffer; + use arrow_array::{ListArray, UInt32Array as U32, UInt64Array}; + use lance_core::ROW_ID_FIELD; + + let values = Float32Array::from(vec![1.0, 1.0, 2.0, 2.0, 3.0, 3.0]); + let fsl = FixedSizeListArray::try_new_from_values(values, 2).unwrap(); + let offsets = OffsetBuffer::new(vec![0i32, 2, 3].into()); + let item = Arc::new(Field::new("item", fsl.data_type().clone(), true)); + let list = ListArray::new(item, offsets, Arc::new(fsl), None); + + let batch = RecordBatch::try_new( + Arc::new(Schema::new(vec![ + ROW_ID_FIELD.clone(), + Field::new("vector", list.data_type().clone(), true), + Field::new("meta", DataType::Int32, true), + Field::new(PART_ID_COLUMN, DataType::UInt32, true), + Field::new(CENTROID_DIST_COLUMN, DataType::Float32, true), + ])), + vec![ + Arc::new(UInt64Array::from(vec![10u64, 20])), + Arc::new(list), + Arc::new(Int32Array::from(vec![100, 200])), + Arc::new(U32::from(vec![7u32, 9])), + Arc::new(Float32Array::from(vec![0.5f32, 0.25])), + ], + ) + .unwrap(); + + let out = Flatten::new("vector").transform(&batch).unwrap(); + + assert_eq!(out.num_rows(), 3); + assert!( + out.column_by_name(PART_ID_COLUMN).is_none(), + "a row's partition assignment must not survive the multivector expansion, or \ + PartitionTransformer early-returns and every sub-vector shares one partition" + ); + assert!( + out.column_by_name(CENTROID_DIST_COLUMN).is_none(), + "the row's centroid distance is likewise meaningless per sub-vector" + ); + // ...while genuine covering columns are still replicated. + assert_eq!( + out.column_by_name("meta") + .expect("covering column must survive") + .as_primitive::() + .values(), + &[100, 100, 200] + ); + } + #[tokio::test] async fn test_normalize_transformer_f32() { let data = Float32Array::from_iter_values([1.0, 1.0, 2.0, 2.0].into_iter()); diff --git a/rust/lance-table/src/format/overlay/staleness.rs b/rust/lance-table/src/format/overlay/staleness.rs index a6eebe61254..4e096b25ce4 100644 --- a/rust/lance-table/src/format/overlay/staleness.rs +++ b/rust/lance-table/src/format/overlay/staleness.rs @@ -108,6 +108,12 @@ pub fn overlaid_fragments(fragments: &[Fragment]) -> HashMap { /// stale because an overlay committed after the segment was built touches a field the segment /// indexes. Field-aware and version-gated via [`overlay_exclusion_offsets`]. /// +/// Walking `segment.fields` alone covers both the key fields and any covering +/// (`covering_fields`) columns materialized in the index storage: `fields` always lists +/// `covering_fields` as its trailing entries (see [`IndexMetadata::covering_fields`]). An +/// overlay touching either makes the segment's stored copy stale, so both must gate overlay +/// exclusion. +/// /// `overlaid_frags` holds only the fragments that actually carry overlays (rare), so the loop is /// `O(overlaid_frags)` rather than `O(fragments the segment covers)`. pub fn collect_overlay_stale_frags( @@ -412,4 +418,48 @@ mod tests { "fragment absent from bitmap contributes no rows" ); } + + #[test] + fn test_collect_rows_flags_overlay_on_covered_column() { + let schema = flat_test_schema(); + // Segment indexes field 3 (key) and COVERS field 2 (included). An overlay touching + // ONLY the covered field must still flag its rows stale -- otherwise a covered query + // serves the stale index-stored copy of that column instead of the fresh overlay value. + let fragment = fragment_with_overlay(3, dense_overlay(vec![2], [1, 2], 9)); + let overlaid: HashMap = HashMap::from([(3u32, &fragment)]); + // `covering_fields` is always the trailing entries of `fields`. + let seg = IndexMetadata { + covering_fields: vec![2], + ..segment(vec![3, 2], 1, Some(bitmap([3]))) + }; + + let mut stale = HashMap::new(); + collect_overlay_stale_rows_for_segment(&seg, &overlaid, &mut stale, &schema).unwrap(); + assert_eq!( + stale.get(&3), + Some(&bitmap([1, 2])), + "overlay on a covered (included) column must flag its rows stale" + ); + } + + #[test] + fn test_collect_frags_flags_overlay_on_covered_column() { + let schema = flat_test_schema(); + // Same at fragment granularity: an overlay on a covered column marks the fragment stale. + let fragment = fragment_with_overlay(3, dense_overlay(vec![2], [1, 2], 9)); + let overlaid: HashMap = HashMap::from([(3u32, &fragment)]); + // `covering_fields` is always the trailing entries of `fields`. + let seg = IndexMetadata { + covering_fields: vec![2], + ..segment(vec![3, 2], 1, Some(bitmap([3]))) + }; + + let mut stale = RoaringBitmap::new(); + collect_overlay_stale_frags(&seg, &overlaid, &mut stale, &schema).unwrap(); + assert_eq!( + stale, + bitmap([3]), + "overlay on a covered (included) column must flag the fragment stale" + ); + } } diff --git a/rust/lance-table/src/system_index/frag_reuse.rs b/rust/lance-table/src/system_index/frag_reuse.rs index 401cbde7a87..20b0f99fbed 100644 --- a/rust/lance-table/src/system_index/frag_reuse.rs +++ b/rust/lance-table/src/system_index/frag_reuse.rs @@ -397,11 +397,11 @@ impl CompactFragReuseIndex { RoaringTreemap::from_iter(row_ids.iter().filter_map(|addr| self.remap_row_id(addr))) } - /// Remap a record batch that contains a row_id column at index `row_id_idx` - /// Currently this assumes there are only 2 columns in the schema, - /// which is the case for all indexes. - /// For example, for btree, the schema is (value, row_id). - /// For vector index storage, the schema is (row_id, vector). + /// Remap a record batch that contains a row_id column at index `row_id_idx`. + /// Every other column (there may be one, as for scalar indexes -- `(value, + /// row_id)` -- or several, as for a covered vector index's storage -- + /// `(row_id, code, )`) is row-aligned to the surviving, + /// remapped row ids by the same take. pub fn remap_row_ids_record_batch( &self, batch: RecordBatch, @@ -426,8 +426,9 @@ fn remap_row_ids_record_batch( row_id_idx: usize, remap: impl FnOnce(&mut [Option]), ) -> Result { - assert_eq!(batch.schema().fields().len(), 2); - let other_column_idx = 1 - row_id_idx; + // Every column but `_rowid` is carried through row-aligned, however many there are. + // A covered index's storage batch holds its covering ("included") columns next to the + // codes, so this cannot assume the two-column `[_rowid, ]` shape. let row_ids = batch.column(row_id_idx).as_primitive::(); let mut remapped_row_ids = row_ids .values() @@ -436,26 +437,27 @@ fn remap_row_ids_record_batch( .map(Some) .collect::>(); remap(&mut remapped_row_ids); - let (val_indices, new_row_ids): (Vec, Vec) = remapped_row_ids + let (keep_indices, new_row_ids): (Vec, Vec) = remapped_row_ids .iter() .enumerate() .filter_map(|(idx, new_id)| new_id.map(|new_id| (idx as u64, new_id))) .unzip(); - let new_val_indices = UInt64Array::from_iter_values(val_indices); - let new_vals = arrow::compute::take(batch.column(other_column_idx), &new_val_indices, None)?; - - let mut batch_data: Vec<(usize, ArrayRef)> = vec![ - ( - row_id_idx, - Arc::new(UInt64Array::from_iter_values(new_row_ids)) as ArrayRef, - ), - (other_column_idx, Arc::new(new_vals)), - ]; - batch_data.sort_by_key(|(i, _)| *i); - Ok(RecordBatch::try_new( - batch.schema(), - batch_data.into_iter().map(|(_, item)| item).collect(), - )?) + let keep_indices = UInt64Array::from_iter_values(keep_indices); + let new_row_ids: ArrayRef = Arc::new(UInt64Array::from_iter_values(new_row_ids)); + + let columns = batch + .columns() + .iter() + .enumerate() + .map(|(idx, column)| { + if idx == row_id_idx { + Ok(new_row_ids.clone()) + } else { + Ok(arrow::compute::take(column, &keep_indices, None)?) + } + }) + .collect::>>()?; + Ok(RecordBatch::try_new(batch.schema(), columns)?) } fn remap_row_ids_array( diff --git a/rust/lance-table/src/transaction/index_maintenance.rs b/rust/lance-table/src/transaction/index_maintenance.rs index 06cf9648a88..e76b0ca8db2 100644 --- a/rust/lance-table/src/transaction/index_maintenance.rs +++ b/rust/lance-table/src/transaction/index_maintenance.rs @@ -19,6 +19,18 @@ use lance_core::{Error, Result}; use roaring::RoaringBitmap; use std::collections::{HashMap, HashSet}; +/// Insert `field`'s id and every descendant's id into `field_ids`. +/// +/// A private copy of `lance::index::collect_subtree_field_ids`: the walk is four +/// lines and the alternative is exporting a lance-table helper solely so the lance +/// crate's copy could delegate to it. +fn collect_subtree_field_ids(field: &lance_core::datatypes::Field, field_ids: &mut HashSet) { + field_ids.insert(field.id); + for child in &field.children { + collect_subtree_field_ids(child, field_ids); + } +} + impl Transaction { pub(super) fn register_pure_rewrite_rows_update_frags_in_indices( indices: &mut [IndexMetadata], @@ -32,6 +44,13 @@ impl Transaction { return Ok(()); } + // Read the parameter name as the inverse of what it does. Despite + // "for_preserving_frag_bitmap", this list is the set of fields whose VALUES were + // updated by the rewrite, and an index depending on any of them is *excluded* from + // the bitmap extension below -- i.e. passing a field here withholds preservation + // rather than granting it. A caller that cannot prove any field survived the rewrite + // unchanged therefore passes *every* field id, which correctly extends no bitmap at + // all. The name is kept only because renaming it would churn shared callers. let value_updated_field_set = fields_for_preserving_frag_bitmap .iter() .collect::>(); @@ -42,9 +61,21 @@ impl Transaction { if index.results_are_row_addrs() { continue; } - let index_covers_modified_field = index.fields.iter().any(|field_id| { - value_updated_field_set.contains(&u32::try_from(*field_id).unwrap()) - }); + // Covering (`covering_fields`) values are materialized in the index + // storage just like the indexed `fields`, so a rewrite of either + // makes this fragment's entries stale. Reusing them would serve the + // pre-update value, so treat both as "modified field" here. + // + // Expand to the leaf subtree, exactly as every sibling prune does: + // `covering_fields` can only ever hold TOP-LEVEL ids (index creation + // resolves whole columns and rejects dotted names), while the caller + // may report the individual leaves it rewrote. Comparing raw ids would + // miss a covered struct whose child was rewritten and wrongly extend + // this index's bitmap onto the moved fragment, serving the pre-update + // subfield value from index storage. + let index_covers_modified_field = Self::index_dependent_field_ids(index, schema) + .iter() + .any(|field_id| value_updated_field_set.contains(field_id)); if index_covers_modified_field { continue; } @@ -86,27 +117,92 @@ impl Transaction { Ok(()) } - /// If an operation modifies one or more fields in a fragment then we need to remove - /// that fragment from any indices that cover one of the modified fields. + /// The full set of leaf field ids an index depends on: its indexed key fields plus + /// every covered (included) field, each expanded to its whole subtree. `fields` + /// already lists `covering_fields` as its trailing entries (see + /// [`IndexMetadata::covering_fields`]), so walking `fields` alone covers both; a + /// modified/overlaid data file lists leaf ids, so a covered struct's parent id must + /// still be expanded to recognize a change to one of its subfields. Returned as raw + /// `i32` (the field id space of `DataFile.fields` / overlay fields). Shared by the + /// freshness prunes and the conflict-rebase checks. + pub fn index_dependent_leaf_ids(index: &IndexMetadata, schema: &Schema) -> HashSet { + let mut ids: HashSet = HashSet::new(); + for &id in index.fields.iter() { + match schema.field_by_id(id) { + Some(field) => collect_subtree_field_ids(field, &mut ids), + None => { + ids.insert(id); + } + } + } + ids + } + + fn index_dependent_field_ids(index: &IndexMetadata, schema: &Schema) -> HashSet { + Self::index_dependent_leaf_ids(index, schema) + .into_iter() + .filter_map(|id| u32::try_from(id).ok()) + .collect() + } + + /// Drop `updated_fragments` from the coverage of any index that depends on a field + /// listed in `fields_modified`. Covered struct fields are expanded to their leaf + /// subtree: `covering_fields` records the parent struct id while a modified data + /// file lists leaf ids, so exact-id matching would miss an update to a covered + /// struct's subfield and serve its stale value from the index. Every caller must + /// therefore supply the schema (the conflict-rebase path captures it at the read + /// version). pub fn prune_updated_fields_from_indices( indices: &mut [IndexMetadata], updated_fragments: &[Fragment], fields_modified: &[u32], + schema: &Schema, + ) { + if fields_modified.is_empty() { + return; + } + + let deps = Self::index_dependent_field_ids_for_each(indices, schema); + Self::prune_updated_fields_with_deps(indices, updated_fragments, fields_modified, &deps); + } + + /// The dependent-leaf-id set for each index, positionally parallel to `indices`. + /// + /// Resolving these walks the schema and allocates per index, so callers that prune + /// fragment-by-fragment must compute this once and reuse it rather than paying it per + /// fragment -- on a compaction rewriting N fragments across M indices the per-fragment + /// form is N*M schema walks on the commit path. + fn index_dependent_field_ids_for_each( + indices: &[IndexMetadata], + schema: &Schema, + ) -> Vec> { + indices + .iter() + .map(|index| Self::index_dependent_field_ids(index, schema)) + .collect() + } + + /// [`Self::prune_updated_fields_from_indices`] with the per-index dependent-id sets + /// already resolved. `deps` must be positionally parallel to `indices`. + fn prune_updated_fields_with_deps( + indices: &mut [IndexMetadata], + updated_fragments: &[Fragment], + fields_modified: &[u32], + deps: &[HashSet], ) { if fields_modified.is_empty() { return; } + debug_assert_eq!(indices.len(), deps.len()); // If we modified any fields in the fragments then we need to remove those fragments // from the index if the index covers one of those modified fields. - let fields_modified_set = fields_modified.iter().collect::>(); - for index in indices.iter_mut() { - if index - .fields + let fields_modified_set = fields_modified.iter().copied().collect::>(); + for (index, dependent_ids) in indices.iter_mut().zip(deps) { + let touches_index = dependent_ids .iter() - .any(|field_id| fields_modified_set.contains(&u32::try_from(*field_id).unwrap())) - && let Some(fragment_bitmap) = &mut index.fragment_bitmap - { + .any(|field_id| fields_modified_set.contains(field_id)); + if touches_index && let Some(fragment_bitmap) = &mut index.fragment_bitmap { for fragment_id in updated_fragments.iter().map(|f| f.id as u32) { fragment_bitmap.remove(fragment_id); } @@ -116,7 +212,9 @@ impl Transaction { /// Map each (non-tombstoned) field id in a fragment to the path of the data /// file that backs it. - fn fragment_field_paths(frag: &Fragment) -> HashMap { + // `pub`: the conflict-rebase path in the `lance` crate resolves per-fragment + // field ownership with this too. + pub fn fragment_field_paths(frag: &Fragment) -> HashMap { let mut map = HashMap::new(); for file in &frag.files { for &field_id in file.fields.iter() { @@ -138,9 +236,13 @@ impl Transaction { indices: &mut [IndexMetadata], prev_fragments: &[Fragment], new_fragments: &[Fragment], + schema: &Schema, ) { let prev_by_id: HashMap = prev_fragments.iter().map(|f| (f.id, f)).collect(); + // Resolved once for all fragments: this is the commit path, and the per-fragment + // form costs a schema walk and two HashSet allocations per index per fragment. + let deps = Self::index_dependent_field_ids_for_each(indices, schema); for new_frag in new_fragments { let Some(prev) = prev_by_id.get(&new_frag.id) else { continue; // brand-new fragment: nothing stale to prune @@ -160,10 +262,11 @@ impl Transaction { if changed.is_empty() { continue; } - Self::prune_updated_fields_from_indices( + Self::prune_updated_fields_with_deps( indices, std::slice::from_ref(new_frag), &changed, + &deps, ); } } @@ -178,17 +281,22 @@ impl Transaction { pub(super) fn prune_overlay_stale_fields_from_indices( indices: &mut [IndexMetadata], groups: &[RewriteGroup], + schema: &Schema, ) { + // Resolved once for all groups, as in `prune_merge_rewritten_fields_from_indices`: + // this is the commit path, and the per-index form costs a schema walk and two + // HashSet allocations for every (group, index) pair. + let deps = Self::index_dependent_field_ids_for_each(indices, schema); for group in groups { // field id -> newest overlay committed_version supplying that field - let mut overlaid_field_versions: HashMap = HashMap::new(); + let mut overlaid_field_versions: HashMap = HashMap::new(); for old_frag in &group.old_fragments { for overlay in &old_frag.overlays { for &field_id in overlay.data_file.fields.iter() { - if field_id < 0 { - // Tombstoned (obsolete) overlay field: supplies nothing. + // Tombstoned (obsolete) overlay fields (< 0) supply nothing. + let Ok(field_id) = u32::try_from(field_id) else { continue; - } + }; let entry = overlaid_field_versions.entry(field_id).or_insert(0); *entry = (*entry).max(overlay.committed_version); } @@ -203,8 +311,12 @@ impl Transaction { .iter() .map(|f| f.id as u32) .collect::>(); - for index in indices.iter_mut() { - let is_stale = index.fields.iter().any(|field_id| { + for (index, dependent_field_ids) in indices.iter_mut().zip(deps.iter()) { + // A covered (included) field an overlay supplied makes the index just as + // stale as an indexed field would -- the index storage still holds the + // pre-overlay copy of that column. `deps` expands covered structs to their + // leaf subtree, since the overlay lists leaf ids. + let is_stale = dependent_field_ids.iter().any(|field_id| { overlaid_field_versions .get(field_id) .is_some_and(|&overlay_version| overlay_version > index.dataset_version) @@ -218,6 +330,103 @@ impl Transaction { } } + /// Reject a commit that would drop, rename, retype, or otherwise change the subtree of + /// a field an index *covers* (an "included" column). A covered field's physical payload + /// schema is fixed at index build time, so an index whose key survives but whose covered + /// column's subtree changed would keep emitting the old payload schema while covered ANN + /// queries declare the new one. This fires for two commit operations: + /// - `Project`, which can drop/rename/retype a covered field (public drop/alter APIs + /// preflight this, but a raw `Operation::Project` reaches `build_manifest` directly); + /// - `Merge` (how `add_columns` commits), which can grow a covered *struct* by adding a + /// child -- even an AllNulls, metadata-only child that writes no data file, so the + /// file-path-keyed coverage prune never sees it. + /// + /// `old_schema` is the pre-commit schema (needed to detect renames/retypes/child-adds, + /// which keep the covered field's id). `op` names the operation for the error message. + pub(super) fn reject_covered_field_subtree_change( + indices: &[IndexMetadata], + old_schema: Option<&Schema>, + new_schema: &Schema, + op: &str, + ) -> Result<()> { + let new_ids: HashSet = new_schema.fields_pre_order().map(|f| f.id).collect(); + for index in indices { + // Skip only if a KEYED field left the schema: `retain_relevant_indices` drops + // the whole index once any of `fields` is gone, so a dropped key field means + // the index is on its way out anyway and there is no dangling covering to + // protect. Checking the *keyed* prefix specifically (not all of `fields`) + // matters because `fields` also lists the covering suffix (see + // `IndexMetadata::covering_fields`): a commit that drops *only* a covering + // field, leaving the key intact, must NOT be skipped here. If it were, this + // guard would let the drop through uncontested, and `retain_relevant_indices` + // -- which keys its own retention on the very same full `fields` list -- would + // then silently delete the entire index (not just desync its covering) with no + // error. Below, that same dropped covering id is caught by + // `covered_field_subtree_changed`'s "removed" case and rejected instead. + // (`FLAG_COVERED_INDEX_METADATA` fences pre-covering builds off the whole + // dataset, so a declaration erased by an old writer -- once this guard's + // other concern -- cannot arise.) + if !index.keyed_fields().iter().all(|id| new_ids.contains(id)) { + continue; + } + // A field the index is also KEYED on gets NO exemption: carried refine vectors + // (`store_vectors_for_refine`) declare the indexed column in `covering_fields`, + // and index storage holds that copy under the column's build-time name. A rename + // preserves the field id but not that name, leaving storage the rebuild paths + // cannot resolve, so it is refused here exactly as it is in the `schema_evolution` + // preflight and in `prune_stale_segment_coverage`. A change that *removes* the id + // (a drop, or a cast, which reassigns a fresh one) never reaches this loop: the + // keyed-field check above skipped the index and `retain_relevant_indices` drops it. + for &covered_id in index.covering_fields.iter() { + let changed = match old_schema { + Some(old) => Self::covered_field_subtree_changed(old, new_schema, covered_id), + // Without the pre-commit schema we can only detect a drop. + None => new_schema.field_by_id(covered_id).is_none(), + }; + if changed { + return Err(Error::invalid_input(format!( + "{op} would drop or alter covered (included) field id {covered_id} \ + (its subtree changed), still used by index '{}'. Drop the index with \ + drop_index() before changing the column.", + index.name + ))); + } + } + } + Ok(()) + } + + /// Whether a covered ("included") field's subtree differs between `old_schema` and + /// `new_schema` -- dropped, renamed, retyped, a nullability flip, or any child change. + /// Such a change means the index's stored physical payload no longer matches the + /// schema a query declares. `data_type()` is recursive for structs, so it also covers + /// a child's name/type/nullability change. + pub fn covered_field_subtree_changed( + old_schema: &Schema, + new_schema: &Schema, + covered_id: i32, + ) -> bool { + let Some(old) = old_schema.field_by_id(covered_id) else { + return false; + }; + match new_schema.field_by_id(covered_id) { + None => true, + Some(new) => { + // Metadata is compared for the same reason name, type and nullability are: + // this guard exists to turn a covering desync into an explicit error, and + // the read path's acceptance test (`covering_fields_match` in + // `io::exec::knn`) compares metadata too. A difference this guard ignores + // but that one does not is not a caught error -- it is a covered index + // that silently stops serving its columns and falls back to a base-table + // read, with no way for the user to learn why. + old.name != new.name + || old.nullable != new.nullable + || old.data_type() != new.data_type() + || old.metadata != new.metadata + } + } + } + pub(crate) fn retain_relevant_indices( indices: &mut Vec, schema: &Schema, @@ -451,6 +660,8 @@ impl Transaction { mod tests { use super::*; use crate::transaction::test_support::overlay_with_field; + use arrow_schema::{DataType, Field as ArrowField, Schema as ArrowSchema}; + use lance_core::datatypes::Schema as LanceSchema; use uuid::Uuid; #[test] @@ -929,6 +1140,60 @@ mod tests { assert!(indices.is_empty()); } + /// A covered column whose *only* change is field metadata must be rejected here. + /// + /// The read path intersects a declaration against each segment's physical columns with + /// `covering_fields_match`, which compares name, type, nullability AND metadata. If + /// this guard compared any less, a metadata-only edit would commit cleanly and then + /// silently withdraw covering at query time -- exactly the silent degradation the + /// guard is here to convert into an error. + #[test] + fn test_covered_field_subtree_change_notices_metadata_only_edit() { + let with_metadata = |value: Option<&str>| { + let mut field = ArrowField::new("carried", DataType::Int32, false); + if let Some(value) = value { + field = field.with_metadata( + [("unit".to_string(), value.to_string())] + .into_iter() + .collect(), + ); + } + LanceSchema::try_from(&ArrowSchema::new(vec![ + ArrowField::new("key", DataType::Int32, false), + field, + ])) + .unwrap() + }; + + let old_schema = with_metadata(Some("metres")); + let carried_id = old_schema.field("carried").unwrap().id; + + assert!( + Transaction::covered_field_subtree_changed( + &old_schema, + &with_metadata(Some("feet")), + carried_id + ), + "a changed metadata value must count as a covering change" + ); + assert!( + Transaction::covered_field_subtree_changed( + &old_schema, + &with_metadata(None), + carried_id + ), + "dropping metadata entirely must count as a covering change" + ); + assert!( + !Transaction::covered_field_subtree_changed( + &old_schema, + &with_metadata(Some("metres")), + carried_id + ), + "an unchanged field must not be reported as changed" + ); + } + #[test] fn test_prune_overlay_stale_fields_from_indices() { // Fragment 0 carried an overlay on field 1 committed at v5, and was @@ -942,6 +1207,13 @@ mod tests { // Post-remap state: every index already covers the new fragment (7). let covering = || Some(RoaringBitmap::from_iter([7u32])); + // Indexes an un-overlaid field (3) but *covers* the overlaid field 1 as an + // included column; built (v2) before the overlay, so its stored copy of + // field 1 is stale and the rewritten fragment must be dropped. + let mut covered_stale = create_test_index("covered_stale", 3, 2, covering(), false); + // `covering_fields` is always the trailing entries of `fields`. + covered_stale.fields = vec![3, 1]; + covered_stale.covering_fields = vec![1]; let mut indices = vec![ // Stale: covers the overlaid field 1, built (v2) before the overlay. create_test_index("stale", 1, 2, covering(), false), @@ -950,9 +1222,16 @@ mod tests { create_test_index("fresh", 1, 5, covering(), false), // Unrelated: covers field 2, which the overlay never touched. create_test_index("unrelated", 2, 2, covering(), false), + covered_stale, ]; - Transaction::prune_overlay_stale_fields_from_indices(&mut indices, &groups); + // Flat field ids with no struct nesting: an empty schema makes subtree + // expansion a no-op (each id maps to itself), exercising the id-level logic. + Transaction::prune_overlay_stale_fields_from_indices( + &mut indices, + &groups, + &LanceSchema::default(), + ); assert!( !indices[0].fragment_bitmap.as_ref().unwrap().contains(7), @@ -966,6 +1245,207 @@ mod tests { indices[2].fragment_bitmap.as_ref().unwrap().contains(7), "an index on an un-overlaid field is unaffected" ); + assert!( + !indices[3].fragment_bitmap.as_ref().unwrap().contains(7), + "an index whose *covered* field was overlaid is stale too" + ); + } + + /// `covering_fields` records a covered struct's parent id, but a modified data file + /// lists leaf ids. Schema-aware pruning must expand the covered struct to its subtree + /// so an update to a subfield still invalidates the fragment's coverage. + #[test] + fn test_prune_updated_fields_expands_covered_struct_subtree() { + use arrow_schema::Fields as ArrowFields; + + let arrow = ArrowSchema::new(vec![ + ArrowField::new( + "s", + DataType::Struct(ArrowFields::from(vec![ + ArrowField::new("a", DataType::Int32, true), + ArrowField::new("b", DataType::Int32, true), + ])), + true, + ), + ArrowField::new("x", DataType::Int32, true), + ]); + let schema = LanceSchema::try_from(&arrow).unwrap(); + let s_id = schema.field("s").unwrap().id; + let a_id = u32::try_from(schema.field("s.a").unwrap().id).unwrap(); + let x_id = schema.field("x").unwrap().id; + + // Index keyed on scalar `x`, covering the whole struct `s`, over fragment 0. + let mut idx = create_test_index( + "cov", + x_id, + 1, + Some(RoaringBitmap::from_iter([0u32])), + false, + ); + // `covering_fields` is always the trailing entries of `fields`. + idx.fields = vec![x_id, s_id]; + idx.covering_fields = vec![s_id]; + + // Updating leaf `s.a` prunes the covered struct's fragment: the subtree + // expansion bridges the parent-id-vs-leaf-id mismatch that exact-id matching + // would miss. + let mut indices = vec![idx]; + Transaction::prune_updated_fields_from_indices( + &mut indices, + &[Fragment::new(0)], + &[a_id], + &schema, + ); + assert!( + !indices[0].fragment_bitmap.as_ref().unwrap().contains(0), + "updating a covered struct's subfield must drop the fragment from coverage" + ); + } + + /// A Project that drops or renames a covered (included) column must be rejected at + /// the commit boundary, even when the index's key column survives. + #[test] + fn test_reject_projected_covered_fields() { + let arrow = ArrowSchema::new(vec![ + ArrowField::new("vec", DataType::Int32, false), + ArrowField::new("meta", DataType::Utf8, true), + ArrowField::new("other", DataType::Int32, true), + ]); + let old_schema = LanceSchema::try_from(&arrow).unwrap(); + let vec_id = old_schema.field("vec").unwrap().id; + let meta_id = old_schema.field("meta").unwrap().id; + + // Index keyed on `vec`, covering `meta`. + let mut idx = create_test_index( + "cov", + vec_id, + 1, + Some(RoaringBitmap::from_iter([0u32])), + true, + ); + // `covering_fields` is always the trailing entries of `fields`. + idx.fields = vec![vec_id, meta_id]; + idx.covering_fields = vec![meta_id]; + let indices = vec![idx]; + + // Keeping every field -> Ok. + assert!( + Transaction::reject_covered_field_subtree_change( + &indices, + Some(&old_schema), + &old_schema, + "Project" + ) + .is_ok() + ); + + // Dropping the covered `meta` while keeping the key -> rejected. + let dropped = old_schema.project(&["vec", "other"]).unwrap(); + let err = Transaction::reject_covered_field_subtree_change( + &indices, + Some(&old_schema), + &dropped, + "Project", + ) + .unwrap_err(); + assert!( + err.to_string().contains("drop or alter covered"), + "expected a covered-subtree rejection, got: {err}" + ); + + // Dropping the key too (the index would be dropped) -> not our concern. + let drop_all = old_schema.project(&["other"]).unwrap(); + assert!( + Transaction::reject_covered_field_subtree_change( + &indices, + Some(&old_schema), + &drop_all, + "Project" + ) + .is_ok() + ); + + // Renaming the covered field (id stays, name changes) -> rejected. + let mut renamed = old_schema.clone(); + renamed.field_by_id_mut(meta_id).unwrap().name = "meta_renamed".to_string(); + assert!( + Transaction::reject_covered_field_subtree_change( + &indices, + Some(&old_schema), + &renamed, + "Project" + ) + .is_err(), + "renaming a covered field must be rejected" + ); + + // Flipping the covered field's nullability (id + name stay) -> rejected. + let mut retyped_null = old_schema.clone(); + retyped_null.field_by_id_mut(meta_id).unwrap().nullable = false; + assert!( + Transaction::reject_covered_field_subtree_change( + &indices, + Some(&old_schema), + &retyped_null, + "Project" + ) + .is_err(), + "a covered field nullability change must be rejected" + ); + } + + /// Growing a covered *struct*'s subtree (adding a child) is what `add_columns`/Merge + /// does; the covered field id is unchanged but its `data_type()` differs, so the commit + /// boundary must reject it just like the Project drop/alter cases above. + #[test] + fn test_reject_covered_struct_child_add() { + let child = ArrowField::new("a", DataType::Int32, false); + let arrow = ArrowSchema::new(vec![ + ArrowField::new("vec", DataType::Int32, false), + ArrowField::new("meta", DataType::Struct(vec![child.clone()].into()), true), + ]); + let old_schema = LanceSchema::try_from(&arrow).unwrap(); + let vec_id = old_schema.field("vec").unwrap().id; + let meta_id = old_schema.field("meta").unwrap().id; + + let mut idx = create_test_index( + "cov", + vec_id, + 1, + Some(RoaringBitmap::from_iter([0u32])), + true, + ); + // `covering_fields` is always the trailing entries of `fields`. + idx.fields = vec![vec_id, meta_id]; + idx.covering_fields = vec![meta_id]; // covers the struct's parent id + let indices = vec![idx]; + + // New schema grows `meta` with a second child -> the struct's data_type changes. + let grown_arrow = ArrowSchema::new(vec![ + ArrowField::new("vec", DataType::Int32, false), + ArrowField::new( + "meta", + DataType::Struct(vec![child, ArrowField::new("b", DataType::Int32, true)].into()), + true, + ), + ]); + // Preorder id assignment gives `meta` the same id (1) in both schemas, so the + // covered id resolves to the struct on each side and only its data_type differs. + let grown = LanceSchema::try_from(&grown_arrow).unwrap(); + assert_eq!(grown.field("meta").unwrap().id, meta_id); + + let err = Transaction::reject_covered_field_subtree_change( + &indices, + Some(&old_schema), + &grown, + "add_columns (Merge)", + ) + .unwrap_err(); + assert!( + err.to_string().contains("drop or alter covered") + && err.to_string().contains("add_columns"), + "expected a covered-struct-growth rejection, got: {err}" + ); } // Helper functions for retain_relevant_indices tests diff --git a/rust/lance-table/src/transaction/manifest_build.rs b/rust/lance-table/src/transaction/manifest_build.rs index 6f216f78c06..bbc16c9cd8e 100644 --- a/rust/lance-table/src/transaction/manifest_build.rs +++ b/rust/lance-table/src/transaction/manifest_build.rs @@ -697,6 +697,7 @@ impl Transaction { &mut final_indices, updated_fragments, fields_modified, + &schema, ); let mut new_fragments = @@ -841,7 +842,7 @@ impl Transaction { // base data. Any index older than one of those overlays was built on // the pre-overlay values, so drop the rewritten fragment from its // coverage to keep it from serving stale values. - Self::prune_overlay_stale_fields_from_indices(&mut final_indices, groups); + Self::prune_overlay_stale_fields_from_indices(&mut final_indices, groups, &schema); if let Some(frag_reuse_index) = frag_reuse_index { final_indices.retain(|idx| idx.name != frag_reuse_index.name); @@ -870,6 +871,45 @@ impl Transaction { new_index.validate_covering_fields()?; } final_indices.extend(new_indices.clone()); + // All delta segments of one logical index (same name) must declare + // the same covering columns: the covered read path derives its + // output schema from one delta but executes all of them, so + // committing a mixed set would fail every query on the index. + // Scoped to names this operation adds, so a commit untouched by + // covering cannot be blocked by pre-existing state. + for new_index in new_indices { + // A declared covering field must resolve to a TOP-LEVEL field of + // the schema being committed: the read path resolves the id to its + // bare name, so a nested id (e.g. a struct child) would silently + // bind to a same-named top-level column, and a nonexistent id + // fails every later query at planning time ("index metadata and + // schema are inconsistent"). Reject at the commit boundary. + for id in &new_index.covering_fields { + if !schema.fields.iter().any(|field| field.id == *id) { + return Err(Error::invalid_input(format!( + "CreateIndex: index '{}' (segment {}) declares covering \ + field id {}, which is not a top-level field of the \ + schema (covering columns must be top-level)", + new_index.name, new_index.uuid, id + ))); + } + } + if let Some(conflict) = final_indices.iter().find(|idx| { + idx.name == new_index.name + && idx.covering_fields != new_index.covering_fields + }) { + return Err(Error::invalid_input(format!( + "CreateIndex: delta segment {} of index '{}' declares covering \ + fields {:?}, but delta segment {} declares {:?}; all segments of \ + one index must declare the same covering columns", + new_index.uuid, + new_index.name, + new_index.covering_fields, + conflict.uuid, + conflict.covering_fields + ))); + } + } } Operation::ReserveFragments { .. } | Operation::UpdateConfig { .. } => { final_fragments.extend(maybe_existing_fragments?.clone()); @@ -906,6 +946,18 @@ impl Transaction { } final_fragments.extend(merged_fragments); + // `add_columns` commits as a Merge; growing a *covered* field's subtree + // (e.g. an AllNulls child under a covered struct, which writes no data file + // and so is invisible to the file-path-keyed prune below) would leave the + // index emitting the old payload schema while covered queries declare the + // new one. Reject at the commit boundary, mirroring the Project guard. + Self::reject_covered_field_subtree_change( + &final_indices, + current_manifest.map(|m| &m.schema), + &schema, + "add_columns (Merge)", + )?; + // A Merge can rewrite a column's data file in place; the field stays // in the schema, so the index is retained -- prune its now-stale // entries for the rewritten fragments. @@ -913,6 +965,7 @@ impl Transaction { &mut final_indices, existing_fragments, fragments, + &schema, ); // Some fields that have indices may have been removed, so we should @@ -936,6 +989,16 @@ impl Transaction { }); } + // A Project that drops or renames a *covered* field would leave dangling + // index metadata whose storage still emits the column -- reject it before + // dropping/retaining any indices. + Self::reject_covered_field_subtree_change( + &final_indices, + current_manifest.map(|m| &m.schema), + &schema, + "Project", + )?; + // Some fields that have indices may have been removed, so we should // remove those indices as well. Self::retain_relevant_indices(&mut final_indices, &schema, &final_fragments) @@ -1162,6 +1225,7 @@ impl Transaction { &mut final_indices, &modified_fragments, &replaced_fields, + &schema, ); } Operation::DataOverlay { groups } => { @@ -1573,6 +1637,95 @@ mod tests { ) } + /// A stable-row-id `RewriteRows` update that modifies ONLY a covering + /// (included) column must not re-extend a covered index's fragment + /// coverage onto the moved fragment: the index still holds the pre-update + /// materialized value, so reusing it would serve a stale covered result. + /// The modified field is in `covering_fields`, not `fields`. + /// + /// Both spellings of "the covering column changed" must be caught. `covering_fields` + /// only ever holds TOP-LEVEL ids (index creation resolves whole columns and rejects + /// dotted names), but the caller may report either the covered struct's own id or the + /// leaf ids it actually rewrote -- `Operation::Update` is public, so + /// `fields_for_preserving_frag_bitmap` is caller-supplied and unvalidated. Comparing raw + /// ids catches only the first spelling and silently serves stale subfield values for the + /// second. + #[test] + fn test_pure_rewrite_preserves_covering_index_coverage() { + use arrow_schema::Fields as ArrowFields; + + // meta{a} => ids 0 (struct) and 1 (leaf); vector => id 2. + let arrow = ArrowSchema::new(vec![ + ArrowField::new( + "meta", + DataType::Struct(ArrowFields::from(vec![ArrowField::new( + "a", + DataType::Int32, + true, + )])), + true, + ), + ArrowField::new("vector", DataType::Int32, true), + ]); + let schema = LanceSchema::try_from(&arrow).unwrap(); + let meta_id = schema.field("meta").unwrap().id; + let leaf_id = schema.field("meta.a").unwrap().id; + let vector_id = schema.field("vector").unwrap().id; + + // Returns (covered index re-extended?, plain index re-extended?). + let run = |modified: &[u32]| { + let mut covering = sample_index_metadata("covering"); + // `covering_fields` is always the trailing entries of `fields`. + covering.fields = vec![vector_id, meta_id]; // top-level, as creation guarantees + covering.covering_fields = vec![meta_id]; + covering.fragment_bitmap = Some([0].into_iter().collect()); + + let mut plain = sample_index_metadata("plain"); + plain.fields = vec![vector_id]; + plain.covering_fields = vec![]; + plain.fragment_bitmap = Some([0].into_iter().collect()); + + let mut indices = vec![covering, plain]; + let overlaid: HashMap = HashMap::new(); + Transaction::register_pure_rewrite_rows_update_frags_in_indices( + &mut indices, + &[1], // new fragment holding the moved rows + &[0], // original fragment (covered by both indices) + modified, + &overlaid, + &schema, + ) + .unwrap(); + ( + indices[0].fragment_bitmap.as_ref().unwrap().contains(1), + indices[1].fragment_bitmap.as_ref().unwrap().contains(1), + ) + }; + + for (label, modified) in [ + ( + "covered struct reported by its own id", + vec![meta_id as u32], + ), + ( + "covered struct reported by its leaf id", + vec![leaf_id as u32], + ), + ] { + let (covered_extended, plain_extended) = run(&modified); + assert!( + !covered_extended, + "covered index must not re-extend coverage onto a fragment whose covering \ + column was rewritten ({label})" + ); + assert!( + plain_extended, + "index that neither indexes nor covers the modified field still reuses its \ + entries ({label})" + ); + } + } + #[test] fn test_create_index_build_manifest_keeps_unremoved_same_name_indices() { let manifest = sample_manifest(); @@ -1649,6 +1802,251 @@ mod tests { ); } + /// A raw CreateIndex commit whose new index declares a covering field id missing + /// from the schema must be rejected at the commit boundary -- otherwise every later + /// query on the index fails at planning time with an inconsistent-metadata error. + /// + /// `covering_fields` must also be a suffix of `fields` (`fields = [0, 9999]` here), + /// or `validate_covering_fields`'s cardinality check rejects the index earlier, for + /// an unrelated reason, before this test's intended "unresolvable in schema" check + /// ever runs. + #[test] + fn test_create_index_build_manifest_rejects_unresolvable_covered_field() { + let manifest = sample_manifest(); + let mut covered = sample_index_metadata("vector_idx"); + covered.fields = vec![0, 9999]; + covered.covering_fields = vec![9999]; + let transaction = Transaction::new( + manifest.version, + Operation::CreateIndex { + new_indices: vec![covered], + removed_indices: vec![], + }, + None, + ); + let err = transaction + .build_manifest(Some(&manifest), vec![], "txn", &default_build_config()) + .unwrap_err(); + assert!( + err.to_string().contains("9999") && err.to_string().contains("top-level"), + "an unresolvable covered field id must be rejected at commit: {err}" + ); + } + + /// A covering declaration must reference a TOP-LEVEL field: `covering_fields` + /// resolution on the read path goes id -> name -> arrow field by bare name, so a + /// nested field id (e.g. a struct child) would silently bind to a same-named + /// top-level column instead of erroring. The create path rejects dotted names; + /// raw commits must reject nested ids for the same reason. + #[test] + fn test_create_index_build_manifest_rejects_nested_covered_field() { + use arrow_schema::Fields as ArrowFields; + + let arrow = ArrowSchema::new(vec![ + ArrowField::new("id", DataType::Int32, false), + ArrowField::new( + "s", + DataType::Struct(ArrowFields::from(vec![ArrowField::new( + "id", + DataType::Int64, + true, + )])), + true, + ), + ]); + let schema = LanceSchema::try_from(&arrow).unwrap(); + let nested_id = schema.field("s.id").unwrap().id; + let manifest = Manifest::new( + schema, + Arc::new(vec![Fragment::new(0)]), + DataStorageFormat::new(ConcreteFileVersion::V2_0), + HashMap::new(), + ); + + let mut covered = sample_index_metadata("vector_idx"); + covered.fields = vec![0, nested_id]; + covered.covering_fields = vec![nested_id]; + let transaction = Transaction::new( + manifest.version, + Operation::CreateIndex { + new_indices: vec![covered], + removed_indices: vec![], + }, + None, + ); + let err = transaction + .build_manifest(Some(&manifest), vec![], "txn", &default_build_config()) + .unwrap_err(); + assert!( + err.to_string().contains("top-level"), + "a nested covered field id must be rejected at commit: {err}" + ); + } + + /// A raw `Operation::Project` that drops ONLY a covering (included) field, keeping the + /// index's key field intact, must be rejected at the commit boundary. `fields` lists the + /// covering suffix alongside the key (see `IndexMetadata::covering_fields`), so + /// `retain_relevant_indices` -- which requires every entry of `fields` to survive -- + /// would otherwise auto-prune the WHOLE index the moment the commit went through, with + /// no error: a working vector index silently disappearing rather than merely desyncing + /// its covering payload. + #[test] + fn test_project_build_manifest_rejects_dropped_covering_field() { + let arrow = ArrowSchema::new(vec![ + ArrowField::new("vec", DataType::Int32, false), + ArrowField::new("meta", DataType::Utf8, true), + ]); + let schema = LanceSchema::try_from(&arrow).unwrap(); + let vec_id = schema.field("vec").unwrap().id; + let meta_id = schema.field("meta").unwrap().id; + let manifest = Manifest::new( + schema, + Arc::new(vec![Fragment::new(0)]), + DataStorageFormat::new(ConcreteFileVersion::V2_0), + HashMap::new(), + ); + + let mut covered = sample_index_metadata("vector_idx"); + covered.fields = vec![vec_id, meta_id]; + covered.covering_fields = vec![meta_id]; + + let projected_schema = manifest.schema.project(&["vec"]).unwrap(); + let transaction = Transaction::new( + manifest.version, + Operation::Project { + schema: projected_schema, + preserves_nullability: true, + }, + None, + ); + let err = transaction + .build_manifest( + Some(&manifest), + vec![covered], + "txn", + &default_build_config(), + ) + .unwrap_err(); + assert!( + err.to_string().contains("drop or alter covered"), + "dropping a covering field while its index's key survives must be rejected \ + instead of silently letting the whole index be auto-pruned: {err}" + ); + } + + /// All delta segments of one logical index must declare the same covering columns: + /// the covered read path derives its output schema from one delta but executes all + /// of them, so a mixed set fails every query on the index at execution time. + #[test] + fn test_create_index_build_manifest_rejects_mixed_covering_deltas() { + // A second top-level field is needed so the covered index can declare it as a + // trailing, schema-resolvable `covering_fields` entry (`sample_manifest`'s + // single-field schema has no field left over once field 0 is keyed on). + let arrow = ArrowSchema::new(vec![ + ArrowField::new("id", DataType::Int32, false), + ArrowField::new("vector", DataType::Int32, false), + ]); + let schema = LanceSchema::try_from(&arrow).unwrap(); + let covered_field_id = schema.field("vector").unwrap().id; + let manifest = Manifest::new( + schema, + Arc::new(vec![Fragment::new(0)]), + DataStorageFormat::new(ConcreteFileVersion::V2_0), + HashMap::new(), + ); + let mut covered = sample_index_metadata("vector_idx"); + covered.fields = vec![0, covered_field_id]; + covered.covering_fields = vec![covered_field_id]; + let plain = sample_index_metadata("vector_idx"); + + // Adding a non-covered delta to an existing covered index is rejected. + let transaction = Transaction::new( + manifest.version, + Operation::CreateIndex { + new_indices: vec![plain.clone()], + removed_indices: vec![], + }, + None, + ); + let err = transaction + .build_manifest( + Some(&manifest), + vec![covered.clone()], + "txn", + &default_build_config(), + ) + .unwrap_err(); + assert!( + err.to_string().contains("covering"), + "mixed covering across deltas must be rejected: {err}" + ); + + // A mixed covering set within one commit's own new segments is rejected too. + let transaction = Transaction::new( + manifest.version, + Operation::CreateIndex { + new_indices: vec![covered.clone(), plain.clone()], + removed_indices: vec![], + }, + None, + ); + let err = transaction + .build_manifest(Some(&manifest), vec![], "txn", &default_build_config()) + .unwrap_err(); + assert!( + err.to_string().contains("covering"), + "mixed covering within one commit must be rejected: {err}" + ); + + // Removing the disagreeing delta in the same commit leaves a homogeneous + // final state, which is fine. + let mut covered2 = sample_index_metadata("vector_idx"); + covered2.fields = vec![0, covered_field_id]; + covered2.covering_fields = vec![covered_field_id]; + let transaction = Transaction::new( + manifest.version, + Operation::CreateIndex { + new_indices: vec![covered2], + removed_indices: vec![plain.clone()], + }, + None, + ); + let (_, final_indices) = transaction + .build_manifest( + Some(&manifest), + vec![covered.clone(), plain], + "txn", + &default_build_config(), + ) + .unwrap(); + assert_eq!(final_indices.len(), 2); + assert!( + final_indices + .iter() + .all(|idx| idx.covering_fields == vec![covered_field_id]), + "the homogeneous replacement set should commit" + ); + + // An index with a different name is unaffected. + let other = sample_index_metadata("other_idx"); + let transaction = Transaction::new( + manifest.version, + Operation::CreateIndex { + new_indices: vec![other], + removed_indices: vec![], + }, + None, + ); + transaction + .build_manifest( + Some(&manifest), + vec![covered], + "txn", + &default_build_config(), + ) + .unwrap(); + } + #[test] fn test_update_build_manifest_replaces_and_removes_fragments() { let manifest = sample_manifest_with_fragments(0..5); diff --git a/rust/lance/benches/mem_wal/vector/hnsw/mem_wal_recall_hnsw.rs b/rust/lance/benches/mem_wal/vector/hnsw/mem_wal_recall_hnsw.rs index 625900cf834..add484e8d07 100644 --- a/rust/lance/benches/mem_wal/vector/hnsw/mem_wal_recall_hnsw.rs +++ b/rust/lance/benches/mem_wal/vector/hnsw/mem_wal_recall_hnsw.rs @@ -596,6 +596,7 @@ async fn run_checkpoint( query_parallelism: 1, dist_q_c: 0.0, approx_mode: Default::default(), + covering_projection: None, }; // IVFIndex::search is intentionally unimplemented (top-level does // partition-aware search); replicate the ANN exec node: pick the diff --git a/rust/lance/src/dataset/index.rs b/rust/lance/src/dataset/index.rs index a3e0917d432..ff46d27adae 100644 --- a/rust/lance/src/dataset/index.rs +++ b/rust/lance/src/dataset/index.rs @@ -421,79 +421,6 @@ mod tests { assert_ne!(remapped[0].new_id, unaffected_segment_id); } - /// A covered index must be withdrawn from remapping rather than remapped. - /// No index type carries the declared payload through a remap, so a - /// replacement would republish a covering claim its storage does not back. - /// Withdrawal also means the legacy `index_details: None` migration path is - /// never reached for such an index, so it cannot panic there either. - #[tokio::test] - async fn test_remapper_migration_path_withdraws_covered_index() { - let reader = lance_datagen::gen_batch() - .col("a", array::step::()) - .col("b", array::step::()) - .into_reader_rows(RowCount::from(20), BatchCount::from(1)); - let mut dataset = Dataset::write(reader, "memory://", None).await.unwrap(); - dataset - .create_index( - &["a"], - IndexType::BTree, - None, - &ScalarIndexParams::for_builtin(BuiltinIndexType::BTree), - false, - ) - .await - .unwrap(); - - let a_id = dataset.schema().field("a").unwrap().id; - let b_id = dataset.schema().field("b").unwrap().id; - let current = dataset.load_indices().await.unwrap(); - let mut legacy_covered = current[0].clone(); - legacy_covered.fields = vec![a_id, b_id]; - legacy_covered.covering_fields = vec![b_id]; - // Force the legacy migration path this fix touches. - legacy_covered.index_details = None; - - let transaction = Transaction::new( - dataset.manifest.version, - Operation::CreateIndex { - new_indices: vec![legacy_covered], - removed_indices: current.to_vec(), - }, - None, - ); - dataset - .apply_commit(transaction, &Default::default(), &Default::default()) - .await - .unwrap(); - - let index_uuid = dataset.load_indices().await.unwrap()[0].uuid; - - // Fully delete every row so `remap_index` returns `RemapResult::Keep`, - // landing in the `index_details: None` migration branch this fix touches. - let remap_to_empty = (0..dataset.count_all_rows().await.unwrap()) - .map(|i| (i as u64, None)) - .collect::>(); - let remapper = DatasetIndexRemapperOptions::default() - .create_remapper(&dataset) - .await - .unwrap() - .expect("a real index should require a remapper"); - let remapped = remapper - .remap_indices(RowAddrRemap::direct(remap_to_empty), &[0]) - .await - .unwrap(); - - // Withdrawn, not remapped: no index type carries the declared payload - // through a remap, so producing a replacement would republish a covering - // claim its storage does not back. The original entry stays in the - // manifest and simply stops covering the rewritten fragments. - assert!( - remapped.is_empty(), - "a covered index must be withdrawn from remapping, got {remapped:?}" - ); - let _ = index_uuid; - } - /// The same migration path must reject -- cleanly, not with a panic -- a /// malformed `covering_fields` longer than `fields`. /// diff --git a/rust/lance/src/dataset/optimize.rs b/rust/lance/src/dataset/optimize.rs index 81f01882e55..e65a4a64315 100644 --- a/rust/lance/src/dataset/optimize.rs +++ b/rust/lance/src/dataset/optimize.rs @@ -6994,6 +6994,8 @@ mod tests { version: crate::index::vector::IndexFileVersion::V3, skip_transpose: false, runtime_hints: Default::default(), + covering_columns: Default::default(), + store_vectors_for_refine: false, }, false, ) @@ -7129,6 +7131,8 @@ mod tests { version: crate::index::vector::IndexFileVersion::V3, skip_transpose: false, runtime_hints: Default::default(), + covering_columns: Default::default(), + store_vectors_for_refine: false, }, false, ) diff --git a/rust/lance/src/dataset/optimize/remapping.rs b/rust/lance/src/dataset/optimize/remapping.rs index 8a9c8898cb6..065612ed862 100644 --- a/rust/lance/src/dataset/optimize/remapping.rs +++ b/rust/lance/src/dataset/optimize/remapping.rs @@ -307,16 +307,23 @@ async fn remap_index(dataset: &mut Dataset, index_id: &Uuid) -> Result<()> { // does not incorporate overlays committed after the source index was built. // Exclude those fragments so queries scan their current values instead. if let Some(fragment_bitmap) = &mut bitmap_after_remap { + // Fields whose value this index materializes: its indexed key fields plus every + // covered (included) column expanded to its leaf subtree (an overlay lists leaf + // ids, while `covering_fields` records a parent struct id). + let relevant_field_ids = + Transaction::index_dependent_leaf_ids(&curr_index_meta, dataset.schema()); for fragment in dataset.manifest.fragments.iter() { - let has_newer_indexed_overlay = fragment.overlays.iter().any(|overlay| { + // An overlay refreshing an indexed or covered field makes the index storage + // stale for that fragment: it still holds the pre-overlay copy of the value. + let has_newer_relevant_overlay = fragment.overlays.iter().any(|overlay| { overlay.committed_version > curr_index_meta.dataset_version && overlay .data_file .fields .iter() - .any(|field_id| curr_index_meta.fields.contains(field_id)) + .any(|field_id| relevant_field_ids.contains(field_id)) }); - if has_newer_indexed_overlay { + if has_newer_relevant_overlay { fragment_bitmap.remove(fragment.id as u32); } } diff --git a/rust/lance/src/dataset/scanner.rs b/rust/lance/src/dataset/scanner.rs index 87f2d83a281..af8064cd187 100644 --- a/rust/lance/src/dataset/scanner.rs +++ b/rust/lance/src/dataset/scanner.rs @@ -100,6 +100,7 @@ use crate::dataset::row_offsets_to_row_addresses; use crate::dataset::rowids::{live_row_addrs_to_row_ids, translate_addr_treemap_to_row_ids}; use crate::dataset::utils::SchemaAdapter; use crate::index::DatasetIndexInternalExt; +use crate::index::covering::effective_covering; use crate::index::scalar::fetch_index_details; use crate::index::scalar::inverted::{ fts_index_fragment_bitmap, load_segment_details, load_segments, normalize_inverted_details, @@ -645,6 +646,33 @@ pub struct FilterPlan { expr_filter_plan: ExprFilterPlan, } +/// Whether an FTS query scores from the columns of the plan it is handed, rather than +/// building its own FTS plan and joining the result on `_rowid`. +/// +/// Only a `Match` does: it is the shape `FlatMatchQueryExec` / `FlatMatchFilterExec` +/// support, and those read the document column straight out of their input. +/// +/// This is the single authority for that rule. Two very different callers depend on it and +/// must not drift apart: +/// +/// * [`FilterPlan::refine_columns`], which decides what the refine pass has to *take* from +/// the base table, and +/// * [`Scanner::fts_scored_columns`], which decides what a covered vector index has to +/// *materialize* so that take is not needed. +/// +/// They are two halves of one decision. If a new FTS shape learns to score from its input +/// and only the first is updated, the take appears and the covered index stops declaring +/// the column it serves -- silently, because the results are identical either way. +fn fts_scores_from_input(query: &FtsQuery) -> bool { + match query { + FtsQuery::Match(_) => true, + FtsQuery::Phrase(_) + | FtsQuery::Boost(_) + | FtsQuery::MultiMatch(_) + | FtsQuery::Boolean(_) => false, + } +} + impl FilterPlan { pub fn new(query_filter: Option, expr_filter_plan: ExprFilterPlan) -> Self { Self { @@ -700,9 +728,11 @@ impl FilterPlan { fts_query.columns() }; - // Add refine column for match query since it supports `FlatMatchQueryExec`. - // Other fts query use join so we don't need to add refine column. - if let FtsQuery::Match(_) = &fts_query.query { + // A shape that scores from its input needs that column taken for it; + // the others join on `_rowid` and need nothing. Same rule, same + // authority, as the covering declaration -- see + // `fts_scores_from_input`. + if fts_scores_from_input(&fts_query.query) { columns.extend(cols.iter().cloned().collect::>()); } } @@ -2021,6 +2051,7 @@ impl Scanner { query_parallelism: DEFAULT_QUERY_PARALLELISM, dist_q_c: 0.0, approx_mode: Default::default(), + covering_projection: None, }); self.nearest_query_count = query_count; self.is_batch_nearest = is_batch_nearest; @@ -3779,9 +3810,17 @@ impl Scanner { if self.prefilter { let source: Arc = match &filter_plan.vector_filter() { Some(vector_query) => { - // Perform vector search first then rerank according to BM25 scores + // Perform vector search first then rerank according to BM25 scores. + // The rerank reads the FTS document column from this plan, so the + // vector search must be told to materialize it when the index covers + // it -- the scan's own projection need not mention it at all. + let vector_query = self.with_covering_projection( + vector_query, + filter_plan, + &Self::fts_scored_columns(query), + ); let vector_plan = self - .vector_search(&filter_plan.expr_filter_plan, vector_query) + .vector_search(&filter_plan.expr_filter_plan, &vector_query) .await?; self.fts_rerank(vector_plan, query).await? } @@ -3810,6 +3849,15 @@ impl Scanner { let Some(query) = self.nearest.as_ref() else { return Err(Error::invalid_input("No nearest query".to_string())); }; + // A postfiltered FTS *query filter* is scored by the refine pass on top of this + // search (`FilterPlan::refine_columns` takes its document column), so name it here + // too. Free in the prefilter branch below: that branch plans the FTS index as the + // source and scores with a flat KNN, so it builds no ANN node to declare anything. + let fts_scored = filter_plan + .fts_filter() + .map(|fts| Self::fts_scored_columns(&fts)) + .unwrap_or_default(); + let query = self.with_covering_projection(query, filter_plan, &fts_scored); if self.prefilter { log::trace!("source is a vector search (prefilter)"); @@ -3826,10 +3874,10 @@ impl Scanner { .union_column(&query.column, OnMissing::Error)?; let plan = self.take(fts_plan, projection)?; - self.flat_knn(plan, query)? + self.flat_knn(plan, &query)? } None => { - self.vector_search(&filter_plan.expr_filter_plan, query) + self.vector_search(&filter_plan.expr_filter_plan, &query) .await? } }; @@ -3841,7 +3889,7 @@ impl Scanner { // If we are postfiltering then we can't use scalar indices for the filter // and will need to run the postfilter in memory filter_plan.make_refine_only(); - self.vector_search(&ExprFilterPlan::default(), query).await + self.vector_search(&ExprFilterPlan::default(), &query).await } } @@ -5258,7 +5306,195 @@ impl Scanner { Ok(flat_match_plan) } + /// The dataset columns this scan still has to read once the vector search has + /// produced its `[_distance, _rowid]` rows. + /// + /// This is the input to [`Query::covering_projection`]. A covering column only pays + /// for itself if it saves a base-table read, so the set an index should materialize + /// is this set intersected with what the index declares. The intersection (and the + /// emission order, which is the index's, not this one's) is done by + /// [`effective_covering`]; this method only answers "what does the rest of the plan + /// read". + /// + /// Every column named here is one a later [`Scanner::take`] would otherwise fetch: + /// the output projection (`create_plan`'s final take), the sort keys and the + /// aggregate inputs, plus whatever the caller passes in `also_read`. Returning a + /// *superset* is harmless -- a name the index does not declare is dropped by the + /// intersection -- so this errs wide on purpose. + /// + /// Names are top-level only, matching the covering declaration: `covering_columns` + /// rejects dotted paths, so a projection of `nested.x` contributes `nested` and a + /// filter on `nested.x` contributes nothing. Both merely under-narrow. + /// + /// `also_read` carries the columns a node *between* the search and the final + /// projection consumes, which this method cannot see from the scan options alone: + /// + /// * on a postfiltered scan, every column the filter mentions -- not just the refine + /// ones, because an index-satisfied conjunct can still be rechecked, and the recheck + /// reads from whatever the search node emitted. A *prefiltered* scan contributes + /// none of them; see [`Self::with_covering_projection`] for why; + /// * the document column an FTS scorer stacked on the vector search reads from it + /// (see [`Self::fts_scored_columns`]); it takes the column from the base table when + /// the search did not emit it. + /// + /// # An empty result is meaningful + /// + /// A query selecting only `_distance` / `_rowid` reads no dataset column at all, and + /// must yield an empty vec, which becomes `Some(&[])` -- "this query needs no + /// covering column". That is the state the whole projection pushdown exists for; + /// see [`Query::covering_projection`] for why it must not collapse into `None`. + /// + /// [`Query::covering_projection`]: lance_index::vector::Query::covering_projection + /// [`effective_covering`]: crate::index::covering::effective_covering + fn columns_read_after_vector_search(&self, also_read: &[String]) -> Vec { + let mut columns: Vec = Vec::new(); + let push = |columns: &mut Vec, name: &str| { + if !columns.iter().any(|existing| existing == name) { + columns.push(name.to_string()); + } + }; + + let projection = &self.projection_plan.physical_projection; + for field in &self.dataset.schema().fields { + if projection.contains_field_id(field.id) { + push(&mut columns, &field.name); + } + } + + if let Some(ordering) = &self.ordering { + for column in ordering { + push(&mut columns, &column.column_name); + } + } + + if let Some(aggregate) = &self.aggregate { + for column in aggregate.required_columns() { + push(&mut columns, &column); + } + } + + for column in also_read { + push(&mut columns, column); + } + + columns + } + + /// The columns an FTS scorer reads out of the plan it is handed, rather than fetching + /// for itself. + /// + /// Two nodes do this, and both sit *between* a vector search and the final + /// projection, so neither is visible to [`Self::columns_read_after_vector_search`]: + /// [`Scanner::fts_rerank`] (`FlatMatchQueryExec`, when FTS is the search and a vector + /// query is the filter) and the refine pass over an FTS query filter + /// (`FlatMatchFilter`, via `FilterPlan::refine_columns`). + /// + /// Which shapes score from their input is decided by [`fts_scores_from_input`], which + /// `FilterPlan::refine_columns` consults for the matching take. Narrowing must not be + /// looser than that rule: declaring a covering column no one reads is precisely the + /// cost this pushdown exists to remove. + /// + /// Names are passed through as the query states them. A dotted FTS field path names + /// no top-level column and so contributes nothing, which under-narrows exactly like a + /// filter on `nested.x` -- the scorer then takes the column from the base table, as it + /// does today. + /// Both callers pass queries `create_plan` has already resolved through + /// [`Self::resolve_full_text_search_query`], so the default unnamed `Match` arrives + /// with its column filled from the FTS-indexed columns -- no fill is needed (or + /// wanted) here. A caller handing this an UNRESOLVED query would silently narrow + /// the filled column away and un-engage covering for the most common query shape, + /// with results staying byte-identical; the unnamed-Match cases in + /// `test_covered_ann_keeps_the_column_an_fts_scorer_reads` pin against that. + fn fts_scored_columns(query: &FullTextSearchQuery) -> Vec { + if !fts_scores_from_input(&query.query) { + return Vec::new(); + } + let mut columns: Vec = query.columns().into_iter().collect(); + // `columns()` is a set; sort so the resolved projection is deterministic across + // runs. The emission order is the index's, but a stable request order keeps plans + // reproducible. + columns.sort_unstable(); + columns + } + + /// `query` with [`Query::covering_projection`] resolved against this scan: the one + /// place a covered index is told what to materialize. + /// + /// Called from the plan sources rather than from `vector_search`, because they still + /// hold the complete [`FilterPlan`]. By the time `vector_search` runs the plan has + /// been split into a prefilter and a refine half and it sees only one of them, so + /// resolving there would forget the filter columns of a postfiltered scan -- which + /// is the default. Deliberately not *also* resolved there as a fallback: two sites + /// computing the same policy from different inputs means the poorer one silently + /// covers for a mistake in the richer one. + /// + /// Always `Some`. The empty case is the point (see + /// [`Self::columns_read_after_vector_search`]); a `vector_search` reached without + /// passing through here keeps `None` and materializes everything, which is the safe + /// direction. + /// + /// `also_read` is for consumers the scan options do not describe -- today the FTS + /// scorers stacked on a vector search (see [`Self::fts_scored_columns`]). Callers with + /// no such consumer pass `&[]`. + /// + /// # The expression filter counts only when it is rechecked above the search + /// + /// A postfiltered scan re-evaluates the whole filter on the search's output: both plan + /// sources call `FilterPlan::make_refine_only`, so `create_plan` takes + /// `FilterPlan::refine_columns` and runs `refine_filter` above the search. Those + /// columns are read from whatever the search emitted, so a covered index should serve + /// them. + /// + /// A *prefiltered* scan is the opposite: the filter is consumed entirely below the + /// search (as a `ScalarIndexQuery` or a filtered read feeding the prefilter), and both + /// sources then call `FilterPlan::disable_refine`, which empties the expression filter + /// so `create_plan`'s `if filter_plan.has_refine()` take never runs. Nothing above the + /// search reads those columns, so naming them here would make the index materialize a + /// payload that the final projection discards. Storage honors this narrowed, + /// physically verified projection, so naming those columns would also cause wasted + /// index reads. + /// + /// The unindexed-fragment and stale-row paths in `knn_combined` do apply the filter, + /// but they read their columns from the base table, not from the search, so they do not + /// change this. + /// + /// [`Query::covering_projection`]: lance_index::vector::Query::covering_projection + fn with_covering_projection( + &self, + query: &Query, + filter_plan: &FilterPlan, + also_read: &[String], + ) -> Query { + let mut query = query.clone(); + let mut columns = if self.prefilter { + Vec::new() + } else { + filter_plan.expr_filter_plan.all_columns() + }; + columns.extend_from_slice(also_read); + // `refine` re-ranks against the indexed column, reading it out of the search output + // like the FTS scorers above. The narrowing below is computed from the projection, + // ordering and aggregates -- none of which name the vector column for the common + // `SELECT payload ORDER BY vector` shape -- so without this it would be narrowed + // away and refine would silently fall back to a base-table take. + if query.refine_factor.is_some() { + columns.push(query.column.clone()); + } + // `Some`, unconditionally: an empty list means "materialize no covering column", + // and folding it into `None` restores full materialization with every result + // still correct. See `Query::covering_projection`. + query.covering_projection = + Some(Arc::from(self.columns_read_after_vector_search(&columns))); + query + } + // ANN/KNN search execution node with optional prefilter + // + // `q.covering_projection` is expected to be resolved already (see + // `with_covering_projection`): the ANN node and every flat/fallback path in + // `knn_combined` derive their covering columns from it, and they are unioned against + // a single schema, so a flat path emitting fewer covering columns than the ANN node + // declares fails the projection with a missing column. #[async_recursion] async fn vector_search( &self, @@ -5452,6 +5688,11 @@ impl Scanner { }; let mut knn_node = if q.refine_factor.is_some() { + // No guard here on whether the index already carries the vectors: `take` + // subtracts its input's schema and returns the input untouched when nothing + // new is needed, so an index that emitted them makes this a no-op. A second + // check would be a redundant mechanism, and a regression in `take`'s + // elision would then go unnoticed. let vector_projection = self .dataset .empty_projection() @@ -5640,6 +5881,12 @@ impl Scanner { let q = q.clone(); debug_assert!(q.metric_type.is_some()); + // The ANN schema has already been narrowed against every selected segment's + // physical storage. Keep it as the authority for which declared covering columns + // the flat paths must mirror; the vector take below can widen `knn_node` with + // unrelated base-table columns. + let ann_schema = knn_node.schema(); + // Ensure the vector column is present for distance computation. if knn_node.schema().column_with_name(&q.column).is_none() { let vector_projection = self @@ -5649,11 +5896,32 @@ impl Scanner { knn_node = self.take(knn_node, vector_projection)?; } + // Every flat path below must produce the covering columns the ANN node actually + // serves before it is unioned with `knn_node`. The manifest/query pair supplies the + // declaration order, while `ann_schema` supplies the storage-authoritative subset. + let covering_columns: Vec = effective_covering( + indexed_segments + .first() + .map(|s| s.covering_fields.as_slice()) + .unwrap_or(&[]), + q.covering_projection.as_deref(), + self.dataset.schema(), + )? + .into_iter() + .filter(|field| ann_schema.column_with_name(field.name()).is_some()) + .map(|field| field.name().clone()) + .collect(); + let mut columns = vec![q.column.clone()]; if let Some(expr) = filter_plan.full_expr.as_ref() { let filter_columns = Planner::column_names_in_expr(expr); columns.extend(filter_columns); } + for name in &covering_columns { + if !columns.contains(name) { + columns.push(name.clone()); + } + } // Collect flat-path plans; union order matches original (flat before ANN) so test snapshots // and downstream plan analyses remain stable. @@ -5708,6 +5976,9 @@ impl Scanner { let filter_columns = Planner::column_names_in_expr(expr); take_proj = take_proj.union_columns(filter_columns, OnMissing::Error)?; } + if !covering_columns.is_empty() { + take_proj = take_proj.union_columns(covering_columns.clone(), OnMissing::Error)?; + } let mut stale_node = self.stale_rows_take(stale_rows, take_proj).await?; if let Some(expr) = filter_plan.full_expr.as_ref() { stale_node = Arc::new(LanceFilterExec::try_new(expr.clone(), stale_node)?); @@ -6644,7 +6915,8 @@ impl Scanner { prefilter_source, overlay_block, self.external_row_mask.clone(), - )?; + ) + .await?; let sort_expr = PhysicalSortExpr { expr: expressions::col(DIST_COL, inner_fanout_search.schema().as_ref())?, options: SortOptions { @@ -6706,7 +6978,8 @@ impl Scanner { prefilter_source.clone(), overlay_block.clone(), self.external_row_mask.clone(), - )?; + ) + .await?; let sort_expr = PhysicalSortExpr { expr: expressions::col(DIST_COL, ann_node.schema().as_ref())?, options: SortOptions { @@ -15032,6 +15305,8 @@ full_filter=name LIKE Utf8(\"test%2\"), refine_filter=name LIKE Utf8(\"test%2\") version: crate::index::vector::IndexFileVersion::Legacy, skip_transpose: false, runtime_hints: Default::default(), + covering_columns: Default::default(), + store_vectors_for_refine: false, }, false, ) @@ -16604,4 +16879,426 @@ full_filter=name LIKE Utf8(\"test%2\"), refine_filter=name LIKE Utf8(\"test%2\") assert_eq!(i_array.values(), &[expected_i]); } } + + /// Number of rows written before the index build, and again appended after it. + /// + /// Bounded from below on two sides, so this is close to the smallest fixture these + /// measurements can be made on: + /// + /// - PQ refuses to train an 8-bit codebook on fewer than 256 rows. + /// - below ~500 rows the covering column's whole partition fits in one data page, so + /// reading one survivor's covering reads every row's. At 300 rows the streaming + /// branch of [`test_covered_projected_query_serves_covering_from_the_index`] + /// measures **4.51x** its plain arm with the survivors-only gather *working* and + /// **4.52x** with it disabled -- 0.1% apart, and both far past its bound. The signal + /// is not weakened at that size, it is gone. + const COVERED_FIXTURE_ROWS: usize = 500; + + /// PQ codebook training iterations for every covering fixture below. + /// + /// No assertion in these tests depends on codebook quality -- none measures recall, and + /// both arms of every comparison train with the same parameters -- so this is set to + /// the smallest value that trains. + const COVERED_PQ_ITERS: usize = 2; + + /// The rows every covering fixture below is built from. + /// + /// Shared rather than repeated because the byte counts these fixtures produce are only + /// comparable to each other while the rows are identically shaped: `payload` is + /// ~1 KB/row, so one copy of it (~0.5 MB) dominates everything else these queries read + /// (~0.15 MB) and no measurement is swamped by fixed costs. + /// + /// The width is load-bearing in *both* directions. Narrower rows stop dominating; + /// wider ones break the streaming branch, whose covering read is per probed partition + /// rather than per survivor -- at 4 KB/row it measures 2.1x its plain arm with the + /// gather working, over its own bound. Trade rows for width in neither direction. + fn covered_fixture_rows() -> impl arrow_array::RecordBatchReader + Send + 'static { + gen_batch() + .col( + "vector", + array::rand_vec::(Dimension::from(32)), + ) + .col("payload", array::rand_utf8(ByteCount::from(1024), false)) + .col("extra", array::step::()) + .into_reader_rows( + RowCount::from(COVERED_FIXTURE_ROWS as u64), + BatchCount::from(1), + ) + } + + /// A dataset whose `payload` column is wide (~1 KB/row), half of it indexed and half + /// of it not, for the read-byte measurements below. + /// + /// Every property here is load-bearing for + /// [`test_covered_unprojected_query_pays_for_no_covering_copy`]: + /// + /// - The rows are appended *after* the index build, so a vector query unions an ANN + /// branch (which reads the index) with a flat scan of the unindexed fragment + /// (which reads the base table). A covering column is therefore reachable on + /// *two* paths, and scanner-side narrowing removes exactly the flat-scan one. + /// - The two halves are the same size and come from the same generator, so the + /// unindexed half's `payload` cost is a faithful unit for the indexed half's too. + /// - Four IVF partitions rather than one, so the query reaches the real probe path. + async fn covered_scan_fixture(uri: &str, cover_payload: bool) -> Dataset { + let mut dataset = Dataset::write(covered_fixture_rows(), uri, None) + .await + .unwrap(); + + let mut params = VectorIndexParams::ivf_pq(4, 8, 4, MetricType::L2, COVERED_PQ_ITERS); + if cover_payload { + params.covering_columns(vec!["payload".to_string()]); + } + dataset + .create_index(&["vector"], IndexType::Vector, None, ¶ms, true) + .await + .unwrap(); + + dataset.append(covered_fixture_rows(), None).await.unwrap(); + dataset + } + + /// Bytes read while running an ANN query that projects no dataset column at all. + /// + /// The counter is reset first, so the fixture build is excluded and only the query is + /// measured. + async fn unprojected_ann_query_bytes(dataset: &Dataset) -> u64 { + let query = Float32Array::from(vec![0.3f32; 32]); + let _ = dataset.object_store.as_ref().io_stats_incremental(); + let mut scan = dataset.scan(); + scan.nearest("vector", &query, 10).unwrap(); + scan.project::<&str>(&[]).unwrap(); + let _ = scan.try_into_batch().await.unwrap(); + dataset + .object_store + .as_ref() + .io_stats_incremental() + .read_bytes + } + + /// Bytes read by scanning `payload` over the fragment appended after the index build: + /// exactly the read that scanner-side covering narrowing removes, measured on the same + /// dataset and the same column so it tracks any encoding change the bound must survive. + async fn unindexed_payload_bytes(dataset: &Dataset) -> u64 { + let fragments = dataset.get_fragments(); + assert_eq!( + fragments.len(), + 2, + "fixture must be one indexed fragment plus one appended unindexed fragment" + ); + let unindexed = fragments.last().unwrap(); + assert_eq!( + unindexed.count_rows(None).await.unwrap(), + COVERED_FIXTURE_ROWS, + "the appended fragment must hold the whole unindexed half" + ); + let unindexed = unindexed.metadata().clone(); + + let _ = dataset.object_store.as_ref().io_stats_incremental(); + let mut scan = dataset.scan(); + scan.with_fragments(vec![unindexed]); + scan.project(&["payload"]).unwrap(); + let _ = scan.try_into_batch().await.unwrap(); + dataset + .object_store + .as_ref() + .io_stats_incremental() + .read_bytes + } + + /// The covering columns a dataset's sole index declares, by field id. + async fn declared_covering_field_ids(dataset: &Dataset) -> Vec { + let indices = dataset.load_indices().await.unwrap(); + assert_eq!(indices.len(), 1, "fixture builds exactly one index"); + indices[0].covering_fields.clone() + } + + /// A covered index must make an unrelated query pay for the covering column *not at + /// all*. + /// + /// Without the projection narrowing, a query that reads no covering column still + /// materializes `payload` on both paths it can reach: the index storage serves it for + /// the indexed half, and the flat scan over the unindexed fragment reads it from the + /// base table. The narrowing removes the flat-scan copy; the survivors-only gather + /// removes the index one, so a covered query now reads neither. + /// + /// # The bound is half a copy, because that is what the reachable regression costs + /// + /// This bound was one and a half copies, calibrated when losing the narrowing cost two + /// copies. The survivors-only gather then removed the index-side copy independently, so + /// losing the narrowing costs **one** copy -- under the old bound, which therefore + /// passed with its own subject disabled. Measured on this fixture, each row taken with + /// `covering_projection` forced to `None` / the gather forced onto its whole-range + /// fallback: + /// + /// | | read bytes | copies of `payload` above plain | + /// |---|---|---| + /// | plain index | 140,930 | -- | + /// | covered, narrowed + gathered (today) | 145,034 | 0.008 | + /// | covered, narrowing disabled | 674,978 | 1.03 | + /// | covered, narrowing and gather both disabled | 1,163,967 | 2.01 | + /// + /// Each row's copy count is against its own run's `plain` and `one_copy`; the plain arm + /// itself varies by up to ~10% run to run, which is why the assertion is a ratio. + /// + /// Half a copy sits an order of magnitude above today's cost and half an order below + /// the cheapest regression, so it discriminates without pinning the exact number. + /// A recorded measurement is a claim with an expiry date: any later change that + /// *reduces* what a regression here costs has to re-derive this bound, not just the + /// baseline row. + /// + /// Expressing it in measured copies rather than as a byte constant keeps its margins if + /// unrelated encoding work moves the absolute totals. The subtraction of `plain` is + /// what makes this a covering measurement rather than a query-cost measurement: it + /// removes the vectors, PQ codes and metadata that both indexes read identically. + #[tokio::test] + async fn test_covered_unprojected_query_pays_for_no_covering_copy() { + let plain = covered_scan_fixture("memory://cov_io_one_copy_plain", false).await; + let covered = covered_scan_fixture("memory://cov_io_one_copy_covered", true).await; + + // Guard the fixture itself. An `covering_columns` that silently did nothing would + // make `covered` cost the same as `plain`, and an upper bound on their difference + // would then pass for the wrong reason. + let payload_id = covered.schema().field("payload").unwrap().id; + assert_eq!( + declared_covering_field_ids(&covered).await, + vec![payload_id], + "the covered fixture's index must actually declare `payload` as covering" + ); + assert!( + declared_covering_field_ids(&plain).await.is_empty(), + "the plain fixture's index must declare no covering column" + ); + + let plain_bytes = unprojected_ann_query_bytes(&plain).await; + let covered_bytes = unprojected_ann_query_bytes(&covered).await; + let one_copy = unindexed_payload_bytes(&covered).await; + + let covering_overhead = covered_bytes.saturating_sub(plain_bytes); + let bound = one_copy / 2; + assert!( + covering_overhead < bound, + "a query projecting no covering column must not pay for `payload` at all: \ + covered read {covered_bytes} bytes vs plain {plain_bytes}, an overhead of \ + {covering_overhead} against a bound of {bound} (half the {one_copy}-byte \ + covering copy). Half a copy or more means the scanner stopped narrowing what \ + the covered index materializes." + ); + + // The same two measurements as a plain ratio. Redundant with the bound above + // against today's code -- both fail on a lost narrowing -- but it is the one that + // stays meaningful if `unindexed_payload_bytes` ever stops being a faithful unit + // for the indexed half, and it costs no extra fixture. + assert!( + covered_bytes < plain_bytes * 2, + "a covered index must not materially inflate a query that projects no covering \ + column: covered read {covered_bytes} bytes vs plain {plain_bytes}" + ); + } + + /// The same fixture with every row indexed, so a covering column has exactly one + /// source -- the index -- and a measurement over it is the index's own read rather + /// than a mix of index and base table. + async fn covered_index_only_fixture( + uri: &str, + mut params: VectorIndexParams, + cover_payload: bool, + ) -> Dataset { + let mut dataset = Dataset::write(covered_fixture_rows(), uri, None) + .await + .unwrap(); + if cover_payload { + params.covering_columns(vec!["payload".to_string()]); + } + dataset + .create_index(&["vector"], IndexType::Vector, None, ¶ms, true) + .await + .unwrap(); + dataset + } + + /// Bytes read while running an ANN query that projects the covering column, with the + /// rows it returned. + /// + /// All four partitions are probed, so the covering read spans several partitions + /// rather than the single one the default `nprobes` would reach. `nprobes` and + /// `query_parallelism` are set after `nearest` because they apply to the query that + /// call installs. + /// + /// `query_parallelism(1)` routes the search through + /// [`VectorIndex::search_partitions`], whose branch is then chosen by the sub-index: + /// the flat sub-index (IVF_PQ) takes the global top-k heap, HNSW has no such heap and + /// takes the streaming send loop. That is how one query shape measures both covering + /// emit paths. + async fn projected_ann_query(dataset: &Dataset) -> (u64, RecordBatch) { + let query = Float32Array::from(vec![0.3f32; 32]); + let _ = dataset.object_store.as_ref().io_stats_incremental(); + let mut scan = dataset.scan(); + scan.nearest("vector", &query, 10).unwrap(); + scan.nprobes(4); + scan.query_parallelism(1); + scan.project(&["payload"]).unwrap(); + let batch = scan.try_into_batch().await.unwrap(); + let bytes = dataset + .object_store + .as_ref() + .io_stats_incremental() + .read_bytes; + (bytes, batch) + } + + /// The plan for the query [`projected_ann_query`] runs, so the two arms' plans are + /// derived the same way as their byte counts. + async fn projected_ann_plan(dataset: &Dataset) -> String { + let query = Float32Array::from(vec![0.3f32; 32]); + let mut scan = dataset.scan(); + scan.nearest("vector", &query, 10).unwrap(); + scan.nprobes(4); + scan.query_parallelism(1); + scan.project(&["payload"]).unwrap(); + scan.explain_plan(true).await.unwrap() + } + + /// The query a covered index exists for: it projects the covering column, and the + /// index serves it without touching the base table. + /// + /// A plain index answers this with a base-table take of the `k` winners, which renders + /// as `LanceRead`; a covered index answers it out of index storage and takes nothing. + /// Both plans are asserted, because a covered plan that quietly fell back to the take + /// would return byte-identical rows -- the covering column is semantically transparent, + /// so results alone cannot tell the two apart. + /// + /// # The two emit paths get different bounds, because they bound the read differently + /// + /// The global top-k heap reads once per partition that owns a *survivor* (at most + /// `min(nprobe, k)` reads); the streaming loop every `IVF_HNSW_*` index takes reads once + /// per *probed* partition, so its covering cost grows with `nprobe` and, at a fixture + /// this size, with the data page a partition's covering lands in. One shared bound would + /// have to be the looser of the two, which would stop discriminating for the heap. Each + /// case therefore carries its own multiple of its own plain arm, each measured against + /// the gather forced onto its whole-range fallback -- the regression the survivors-only + /// gather replaced, and what a pending fragment-reuse remap still falls back to: + /// + /// | branch | plain | covered | ratio | bound | gather disabled | + /// |---|---|---|---|---|---| + /// | global heap (`pq`) | 64,197 | 67,725 | 1.06x | 1.5x | 4.74x | + /// | streaming (`hnsw_pq`) | 95,512 | 147,009 | 1.54x | 3x | 6.42x | + /// + /// Both bounds are order-of-magnitude claims, set so unrelated encoding work moving the + /// absolute totals cannot make them flaky. Neither is a floor on the fixture: shrinking + /// it further collapses the streaming row's ratio into its own regression -- see + /// [`COVERED_FIXTURE_ROWS`]. + /// + /// Higher up the sweep the two separate sharply -- on a 32-partition, 32,768-row index + /// with the same ~1 KB payload, warm, `nprobes = 32`: global heap 15,202 B / 5 requests + /// (flat in `nprobe`), streaming 454,558 B / 106 requests, 2,069 vs 737 us against its + /// own plain arm. That is recorded, not pinned here: this test is the shape check that + /// each branch serves covering from the index at all, on a fixture small enough to + /// stay a unit test. + /// + /// HNSW is built with a small graph for the same reason the PQ codebook trains for two + /// iterations: nothing here measures recall, and both arms use identical parameters. + #[rstest] + #[case::global_heap_branch( + "pq", + VectorIndexParams::ivf_pq(4, 8, 4, MetricType::L2, COVERED_PQ_ITERS), + 3 + )] + #[case::streaming_branch( + "hnsw_pq", + VectorIndexParams::with_ivf_hnsw_pq_params( + MetricType::L2, + IvfBuildParams::new(4), + HnswBuildParams { + max_level: 3, + m: 8, + ef_construction: 30, + ..HnswBuildParams::default() + }, + PQBuildParams::new(4, 8), + ), + 6 + )] + #[tokio::test] + async fn test_covered_projected_query_serves_covering_from_the_index( + #[case] family: &str, + #[case] params: VectorIndexParams, + // Halves: the global heap's bound is 1.5x, the streaming loop's 3x. An integer + // multiple cannot express the former, and 2x would not discriminate for the latter. + #[case] bound_halves: u64, + ) { + let plain = covered_index_only_fixture( + &format!("memory://cov_projected_plain_{family}"), + params.clone(), + false, + ) + .await; + let covered = covered_index_only_fixture( + &format!("memory://cov_projected_covered_{family}"), + params, + true, + ) + .await; + + // The two cases are the two covering emit paths, not two spellings of one: + // `search_partitions` takes the global top-k heap only when the sub-index has + // one, and only IVF_PQ's flat sub-index does. + use lance_index::vector::{flat::index::FlatIndex, hnsw::HNSW, v3::subindex::IvfSubIndex}; + assert!( + FlatIndex::supports_global_topk_heap(), + "the `pq` case must exercise the global-heap branch" + ); + assert!( + !HNSW::supports_global_topk_heap(), + "the `hnsw_pq` case must exercise the streaming branch" + ); + + // Guard the fixture itself: an `covering_columns` that silently did nothing would + // make the covered arm cost what the plain one costs, and the bound below would + // then pass for the wrong reason. + let payload_id = covered.schema().field("payload").unwrap().id; + assert_eq!( + declared_covering_field_ids(&covered).await, + vec![payload_id], + "the covered fixture's index must actually declare `payload` as covering" + ); + + // Each dataset serves exactly one measured query, so both are measured against a + // cold partition cache. A second query on the same dataset would find the + // partitions already resident and read far less. + let (plain_bytes, plain_batch) = projected_ann_query(&plain).await; + let (covered_bytes, covered_batch) = projected_ann_query(&covered).await; + assert_eq!(plain_batch.num_rows(), 10); + assert_eq!(covered_batch.num_rows(), 10); + assert_eq!( + covered_batch + .column_by_name("payload") + .expect("the covered result must carry the projected covering column") + .null_count(), + 0, + "the covering values must come back from the index, not as a null fill" + ); + let bound = plain_bytes * bound_halves / 2; + assert!( + covered_bytes < bound, + "a covered index must serve the column it covers for about what a plain index \ + pays to take it: the {family} branch read {covered_bytes} bytes covered vs \ + {plain_bytes} plain, against a bound of {bound} \ + ({bound_halves} halves of plain). A covering read proportional to the \ + partition's size rather than to its survivors overshoots this by several times." + ); + + let plain_plan = projected_ann_plan(&plain).await; + assert!( + plain_plan.contains("LanceRead"), + "the plain arm must take `payload` from the base table, which is the work the \ + covered arm is expected to avoid; without it the assertion below is vacuous. \ + Plan:\n{plain_plan}" + ); + let covered_plan = projected_ann_plan(&covered).await; + assert!( + !covered_plan.contains("LanceRead"), + "a covered index must serve `payload` from the index, with no base-table take. \ + Plan:\n{covered_plan}" + ); + } } diff --git a/rust/lance/src/dataset/schema_evolution.rs b/rust/lance/src/dataset/schema_evolution.rs index af0e1ab2622..42babda0b1b 100644 --- a/rust/lance/src/dataset/schema_evolution.rs +++ b/rust/lance/src/dataset/schema_evolution.rs @@ -793,37 +793,120 @@ pub(super) async fn alter_columns( new_schema.validate()?; new_schema.verify_primary_key()?; - // If any column being cast has an attached index, fail fast. Cast operations - // rewrite the underlying column data and silently invalidate any index on the - // affected column(s). The current behavior is to drop such indices without - // warning, which has caused production incidents where vector search silently - // regressed to brute-force scan. We require users to explicitly drop the - // index before altering the column type, so the action is never silent. That - // includes an index this build has no reader for: the cast reassigns the - // field id, so carrying it forward is impossible and staying quiet about it - // is the silent drop this guard exists to abolish. - if !cast_fields.is_empty() { + // Fail fast on alterations that would invalidate an index or desync its + // covering data: + // * CAST of an indexed *key* column -- cast rewrites the column data and + // silently invalidates the index. This has caused production incidents + // where vector search regressed to a brute-force scan. (Indices are + // keyed by field id, so a plain rename or nullability change of a key + // column is safe and stays allowed.) + // * ANY alter (rename, cast, or nullability) of a covering ("included") + // column, including one the index is also keyed on. `fields` lists the + // covering suffix right alongside the key (see + // `IndexMetadata::covering_fields`), so the auto-prune + // (`retain_relevant_indices`) DOES see a covering column's id. + // + // For a CAST this guard is load-bearing: a cast reassigns the column a + // fresh field id, so the old covering id disappears from the schema and + // auto-prune -- which requires every id in `fields` to survive -- would + // delete the WHOLE index, not just leave its covering stale. + // + // For a plain rename or nullability flip the id survives, so auto-prune + // never fires and the index is left declaring covering that its storage + // no longer matches. That is NOT silent corruption: the read path + // intersects the declaration with each segment's proven physical + // capability (`common_physical_covering`, which compares name, type, + // nullability and metadata), so a mismatch withdraws covering and the + // query falls back to a base-table take -- correct results, minus the + // elision. The guard is conservative here rather than load-bearing, and + // is kept so a lost elision surfaces as an explicit error instead of a + // silent slowdown. Relaxing it to degrade instead of reject is a tracked + // follow-up. + // + // Both cases include an index this build has no reader for -- hence + // `load_all_indices`. A cast reassigns the field id, so carrying such an index + // forward is impossible, and staying quiet about it is the silent drop these + // guards exist to abolish. + // Reject rather than let the index quietly change behaviour; the cast case + // above is the one that would otherwise destroy the index outright. + let cast_ids: Vec = alterations + .iter() + .filter(|a| a.data_type.is_some()) + .filter_map(|a| dataset.schema().field(&a.path).map(|f| f.id)) + .collect(); + let altered: Vec<(&str, i32)> = alterations + .iter() + .filter(|a| a.rename.is_some() || a.data_type.is_some() || a.nullable.is_some()) + .filter_map(|a| { + dataset + .schema() + .field(&a.path) + .map(|f| (a.path.as_str(), f.id)) + }) + .collect(); + if !altered.is_empty() { let indices = load_all_indices(dataset).await?; - let affected: Vec<&lance_table::format::IndexMetadata> = indices - .iter() - .filter(|idx| { - cast_fields - .iter() - .any(|(old, _)| idx.fields.contains(&old.id)) + // A covered column is stored as a whole subtree, so altering a *descendant* of a + // covered field (e.g. a subfield of a covered struct) invalidates the covering + // too; expand each altered field to its ancestors so the covered ancestor's id is + // caught. The expansion depends only on the schema, so do it once per altered + // field here rather than per (index x altered field) below. + let altered: Vec<(&str, i32, Vec)> = altered + .into_iter() + .map(|(name, id)| { + let ancestors = + field_ids_with_ancestors(dataset.schema(), std::slice::from_ref(&id)); + (name, id, ancestors) }) .collect(); + // An index is broken if a covering column of it is altered in any way, or if one + // of its indexed key columns is being cast. The cast check stays on the field's + // own id (only a cast of the indexed key column rewrites index data). + // + // A column the index is also KEYED on gets NO exemption here, even though a rename + // or nullability flip preserves its field id. Carried refine vectors + // (`store_vectors_for_refine`) declare the indexed column in `covering_fields`, and + // index storage holds that carried copy under the column's name *as it was at build + // time* -- nothing renames it. Letting the rename through leaves an index that still + // answers queries while both maintenance paths are permanently broken: + // `optimize_indices` fails with "field does not exist in the RecordBatch" + // and `compact_files` with "cannot remap index: its storage carries covering column + // ''". Recovering needs a drop and rebuild, so refuse up front and say so. + // + // Dropping the column is a different case and stays allowed -- see the `drop_columns` + // guard below: a drop removes the whole index via `retain_relevant_indices`, so no + // storage is stranded. `prune_stale_segment_coverage` (`index.rs`) refuses this the + // same way for uncommitted distributed segments; all three sites must agree. + let breaks_covering = |idx: &lance_table::format::IndexMetadata, ancestors: &[i32]| { + ancestors + .iter() + .any(|ancestor| idx.covering_fields.contains(ancestor)) + }; + let is_broken = |idx: &lance_table::format::IndexMetadata| { + altered + .iter() + .any(|(_, _, ancestors)| breaks_covering(idx, ancestors)) + || cast_ids.iter().any(|id| idx.fields.contains(id)) + }; + let affected: Vec<&lance_table::format::IndexMetadata> = + indices.iter().filter(|idx| is_broken(idx)).collect(); if !affected.is_empty() { - let affected_cols: Vec = cast_fields + let affected_cols: Vec = altered .iter() - .filter(|(old, _)| affected.iter().any(|i| i.fields.contains(&old.id))) - .map(|(old, _)| old.name.clone()) + .filter(|(_, id, ancestors)| { + affected.iter().any(|i| { + breaks_covering(i, ancestors) + || (cast_ids.contains(id) && i.fields.contains(id)) + }) + }) + .map(|(name, _, _)| name.to_string()) .collect(); let affected_idx_names: Vec = affected.iter().map(|i| i.name.clone()).collect(); return Err(Error::invalid_input(format!( - "Cannot cast column(s) [{}] to a new type: they have {} index(es) \ - attached: [{}]. Cast rewrites column data and invalidates any index \ - on the affected column(s). Drop the index(es) with drop_index() \ - before altering, then recreate them after the cast completes.", + "Cannot alter column(s) [{}]: they are used by {} index(es) [{}] -- as a \ + covering (included) column (any change), or as an indexed key column being \ + cast (which rewrites the data and invalidates the index). Drop the index(es) \ + with drop_index() before altering, then recreate them.", affected_cols.join(", "), affected.len(), affected_idx_names.join(", "), @@ -1017,6 +1100,23 @@ pub(super) async fn alter_columns( Ok(()) } +/// Expand `field_ids` to also include every ancestor field id (root -> field). +/// A covering column is stored as a whole subtree, so touching a *descendant* of a +/// covered field (e.g. a subfield of a covered struct) still invalidates the covering +/// even though `covering_fields` records only the ancestor's id. Adding each targeted +/// field's ancestors makes the ancestor's id appear in the set so the exact-membership +/// checks catch it. +fn field_ids_with_ancestors(schema: &Schema, field_ids: &[i32]) -> Vec { + let mut expanded = Vec::with_capacity(field_ids.len()); + for &id in field_ids { + match schema.field_ancestry_by_id(id) { + Some(ancestry) => expanded.extend(ancestry.iter().map(|field| field.id)), + None => expanded.push(id), + } + } + expanded +} + /// Remove columns from the dataset. /// /// This is a metadata-only operation and does not remove the data from the @@ -1045,6 +1145,63 @@ pub(super) async fn drop_columns(dataset: &mut Dataset, columns: &[&str]) -> Res )); } + // Fail fast if any column being dropped is a covering ("included") column of + // an index, while the index's key survives. Dropping the key too is fine -- + // `retain_relevant_indices` auto-prunes the whole index once any id in `fields` + // leaves the schema, and `fields` lists the covering suffix right alongside the + // key (see `IndexMetadata::covering_fields`), so a key+covering drop cleanly + // removes the whole thing. But that is exactly why dropping *only* the covering + // column, without this guard, would not merely leave the index's storage + // referencing a missing column -- the same auto-prune check would see the + // covering id vanish from `fields` and silently delete the ENTIRE index (the key, + // its ANN structure, everything), not just desync its covering. Require the user + // to drop the index first so the action is never silent. + let drop_field_ids: Vec = columns + .iter() + .filter_map(|c| dataset.schema().field(c).map(|f| f.id)) + .collect(); + let indices = load_all_indices(dataset).await?; + // Treat a descendant of a covered (e.g. struct) column as covered too: the covering + // declaration names the struct while the drop names a leaf inside it. + let drop_field_ids = field_ids_with_ancestors(dataset.schema(), &drop_field_ids); + // An index whose *key* is also being dropped is auto-pruned by + // `retain_relevant_indices`, so its covering metadata cannot dangle -- the whole index + // goes away. Rejecting here would block the perfectly safe "drop the vector column and + // its covered payload together" case. The commit-boundary guard + // (`reject_covered_field_subtree_change`) skips these indices for the same reason, and + // this preflight must agree with it or it forbids what the commit would accept. + // + // Agreement requires testing what that guard tests: whether a KEYED field survives the + // post-drop schema. Neither half of that can be read off `drop_field_ids`. It is + // expanded *upwards* to ancestors, so dropping `s.a` lists the struct `s` that in fact + // survives; and the keyed prefix is not all of `fields`, which also carries the + // covering suffix. Testing `fields` against the expanded set got both wrong in + // opposite directions -- it rejected dropping a key alongside a strict subset of its + // covering columns (the commit accepts that: the key is gone, the index goes with it), + // and it skipped an index keyed on a struct when a leaf under that struct was dropped, + // deferring a case this preflight is meant to catch to a far less legible error. + let surviving: HashSet = new_schema.fields_pre_order().map(|f| f.id).collect(); + let affected: Vec<&lance_table::format::IndexMetadata> = indices + .iter() + .filter(|idx| idx.keyed_fields().iter().all(|id| surviving.contains(id))) + .filter(|idx| { + drop_field_ids + .iter() + .any(|id| idx.covering_fields.contains(id)) + }) + .collect(); + if !affected.is_empty() { + let affected_idx_names: Vec = affected.iter().map(|i| i.name.clone()).collect(); + return Err(Error::invalid_input(format!( + "Cannot drop column(s) [{}]: they are a covering (included) column of {} \ + index(es) [{}]. Drop the index(es) with drop_index() before dropping the \ + column(s), then recreate them.", + columns.join(", "), + affected.len(), + affected_idx_names.join(", "), + ))); + } + let transaction = Transaction::new( dataset.manifest.version, Operation::Project { @@ -3589,6 +3746,485 @@ mod test { Ok(()) } + /// The covered-column schema guards must treat a *descendant* of a covered field as + /// covered too: a covered column is stored as a whole subtree, so `covering_fields` + /// records only the ancestor's id. `field_ids_with_ancestors` expands a targeted + /// subfield to include that ancestor so the exact-membership guard catches it. + #[test] + fn test_field_ids_with_ancestors_expands_struct_subtree() { + let arrow_schema = ArrowSchema::new(vec![ + ArrowField::new( + "s", + DataType::Struct(ArrowFields::from(vec![ + ArrowField::new("a", DataType::Int32, true), + ArrowField::new("b", DataType::Int32, true), + ])), + true, + ), + ArrowField::new("x", DataType::Int32, true), + ]); + let schema = Schema::try_from(&arrow_schema).unwrap(); + let s_id = schema.field("s").unwrap().id; + let a_id = schema.field("s.a").unwrap().id; + + let expanded = field_ids_with_ancestors(&schema, &[a_id]); + assert!( + expanded.contains(&s_id), + "expanding a subfield must include its covered struct ancestor's id" + ); + assert!( + expanded.contains(&a_id), + "the field's own id must be included" + ); + } + + /// Dropping a covered column TOGETHER WITH its index's key column is safe and must be + /// allowed: `retain_relevant_indices` drops the whole index once its key leaves the + /// schema, so no covering metadata can dangle. The commit-boundary guard + /// (`reject_covered_field_subtree_change`) already skips such indices; this preflight + /// must agree with it, or it forbids an operation the commit would accept. + /// + /// The `key_and_covering_subset` case is what makes the two agree in general. The + /// preflight used to skip an index only when *every* id in `fields` was being dropped, + /// while the commit guard skips as soon as a *keyed* id is. With two covering columns + /// and only one of them dropped alongside the key, those disagree: the commit accepts + /// (the key is gone, so the whole index goes) and the preflight rejected. + #[rstest] + #[case::key_and_all_covering(&["vec", "id", "extra"], vec!["keep"])] + #[case::key_and_covering_subset(&["vec", "id"], vec!["extra", "keep"])] + #[tokio::test] + async fn test_drop_covered_column_with_its_index_key_is_allowed( + #[case] to_drop: &[&str], + #[case] expected_remaining: Vec<&str>, + ) -> Result<()> { + use lance_arrow::FixedSizeListArrayExt; + use lance_index::IndexType; + use lance_linalg::distance::MetricType; + use lance_testing::datagen::generate_random_array; + + use crate::index::vector::VectorIndexParams; + + let schema = Arc::new(ArrowSchema::new(vec![ + ArrowField::new( + "vec", + DataType::FixedSizeList( + Arc::new(ArrowField::new("item", DataType::Float32, true)), + 64, + ), + false, + ), + ArrowField::new("id", DataType::Int32, false), + ArrowField::new("extra", DataType::Int32, false), + ArrowField::new("keep", DataType::Int32, false), + ])); + let nrows: i32 = 256; + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new( + ::try_new_from_values( + generate_random_array(64 * nrows as usize), + 64, + ) + .unwrap(), + ), + Arc::new(Int32Array::from_iter_values(0..nrows)), + Arc::new(Int32Array::from_iter_values(0..nrows)), + Arc::new(Int32Array::from_iter_values(0..nrows)), + ], + )?; + + let test_dir = TempStrDir::default(); + let mut dataset = Dataset::write( + RecordBatchIterator::new(vec![Ok(batch)], schema.clone()), + &test_dir, + None, + ) + .await?; + + let mut params = VectorIndexParams::ivf_pq(4, 8, 8, MetricType::L2, 50); + params.covering_columns(vec!["id".to_string(), "extra".to_string()]); + dataset + .create_index(&["vec"], IndexType::Vector, None, ¶ms, false) + .await?; + + // Drop the index key together with its covering column(s) in one call. + dataset.drop_columns(to_drop).await?; + + assert!( + dataset.load_indices().await?.is_empty(), + "the index must be pruned once its key column is gone" + ); + let remaining: Vec<&str> = dataset + .schema() + .fields + .iter() + .map(|f| f.name.as_str()) + .collect(); + assert_eq!(remaining, expected_remaining); + Ok(()) + } + + /// Dropping a column that an index *covers* (an "included"/covering column, + /// not the key column) must fail fast. The covering id lives in `fields`, so + /// letting the drop through would make `retain_relevant_indices` delete the + /// ENTIRE index -- key, ANN structure and all -- not merely desync its + /// covering. The user must drop the index first. + #[tokio::test] + async fn test_drop_columns_fails_on_covered_column() -> Result<()> { + use lance_arrow::FixedSizeListArrayExt; + use lance_index::IndexType; + use lance_linalg::distance::MetricType; + use lance_testing::datagen::generate_random_array; + + use crate::index::vector::VectorIndexParams; + + let schema = Arc::new(ArrowSchema::new(vec![ + ArrowField::new( + "vec", + DataType::FixedSizeList( + Arc::new(ArrowField::new("item", DataType::Float32, true)), + 64, + ), + false, + ), + ArrowField::new("id", DataType::Int32, false), + ])); + let nrows: i32 = 256; + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new( + ::try_new_from_values( + generate_random_array(64 * nrows as usize), + 64, + ) + .unwrap(), + ), + Arc::new(Int32Array::from_iter_values(0..nrows)), + ], + )?; + + let test_dir = TempStrDir::default(); + let mut dataset = Dataset::write( + RecordBatchIterator::new(vec![Ok(batch)], schema.clone()), + &test_dir, + None, + ) + .await?; + + // IVF_PQ index on `vec` that COVERS `id`. + let mut params = VectorIndexParams::ivf_pq(4, 8, 8, MetricType::L2, 50); + params.covering_columns(vec!["id".to_string()]); + dataset + .create_index(&["vec"], IndexType::Vector, None, ¶ms, false) + .await?; + let index_name = dataset.load_indices().await?[0].name.clone(); + + // Dropping the covered column must fail, naming the column, the index, + // and the remediation. + let err = dataset + .drop_columns(&["id"]) + .await + .expect_err("dropping a covered column should fail"); + assert!( + matches!(err, Error::InvalidInput { .. }), + "expected Error::InvalidInput, got: {err:?}" + ); + let msg = err.to_string(); + // `[id]` (the rendered column list), not a bare `id`: the auto-generated index name + // is `vec_idx`, which contains `id` as a substring, so a bare `contains("id")` would + // pass even if the message never named the column. + assert!( + msg.contains("[id]") && msg.contains(&index_name) && msg.contains("drop_index"), + "error should mention the column, index, and remediation, got: {msg}" + ); + + // Unchanged: column still present, index still there. + assert!(dataset.schema().field("id").is_some()); + assert_eq!(dataset.load_indices().await?.len(), 1); + + // After dropping the index, the same drop succeeds. + dataset.drop_index(&index_name).await?; + dataset.drop_columns(&["id"]).await?; + assert!(dataset.schema().field("id").is_none()); + + Ok(()) + } + + /// Dropping a column that is only an INDEXED key column (not a covering + /// column) must remain allowed: lance auto-drops the dependent index. Only + /// a covering ("included") column dropped *while its index's key survives* + /// is protected by the drop guard: dropping the key removes the whole + /// index either way (via `retain_relevant_indices`), so there is nothing to + /// protect, whereas dropping just the covering column would otherwise + /// silently take the whole index down with it (see `drop_columns`). + /// Regression guard for the Python `test_drop_columns` behavior. + #[tokio::test] + async fn test_drop_indexed_noncovered_column_drops_index() -> Result<()> { + use crate::index::vector::VectorIndexParams; + use lance_arrow::FixedSizeListArrayExt; + use lance_index::IndexType; + use lance_linalg::distance::MetricType; + use lance_testing::datagen::generate_random_array; + + let schema = Arc::new(ArrowSchema::new(vec![ + ArrowField::new( + "vec", + DataType::FixedSizeList( + Arc::new(ArrowField::new("item", DataType::Float32, true)), + 64, + ), + false, + ), + ArrowField::new("id", DataType::Int32, false), + ])); + let nrows: i32 = 256; + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new( + ::try_new_from_values( + generate_random_array(64 * nrows as usize), + 64, + ) + .unwrap(), + ), + Arc::new(Int32Array::from_iter_values(0..nrows)), + ], + )?; + let test_dir = TempStrDir::default(); + let mut dataset = Dataset::write( + RecordBatchIterator::new(vec![Ok(batch)], schema.clone()), + &test_dir, + None, + ) + .await?; + + // Plain IVF_PQ index on `vec` -- NO covering columns. + let params = VectorIndexParams::ivf_pq(4, 8, 8, MetricType::L2, 50); + dataset + .create_index(&["vec"], IndexType::Vector, None, ¶ms, false) + .await?; + assert_eq!(dataset.load_indices().await?.len(), 1); + + // Dropping the indexed key column is allowed and auto-drops the index. + dataset.drop_columns(&["vec"]).await?; + assert!(dataset.schema().field("vec").is_none()); + assert_eq!( + dataset.load_indices().await?.len(), + 0, + "dropping the indexed column should have dropped its index" + ); + + Ok(()) + } + + /// Fix 1: renaming or changing the nullability of a covered ("included") + /// column must fail fast (alongside cast) -- it would leave the index storage + /// referencing a stale name/schema and break covered queries. + #[tokio::test] + async fn test_alter_columns_rename_nullable_fail_on_covered_column() -> Result<()> { + use crate::index::vector::VectorIndexParams; + use lance_arrow::FixedSizeListArrayExt; + use lance_index::IndexType; + use lance_linalg::distance::MetricType; + use lance_testing::datagen::generate_random_array; + + let schema = Arc::new(ArrowSchema::new(vec![ + ArrowField::new( + "vec", + DataType::FixedSizeList( + Arc::new(ArrowField::new("item", DataType::Float32, true)), + 64, + ), + false, + ), + ArrowField::new("id", DataType::Int32, false), + ])); + let nrows: i32 = 256; + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new( + ::try_new_from_values( + generate_random_array(64 * nrows as usize), + 64, + ) + .unwrap(), + ), + Arc::new(Int32Array::from_iter_values(0..nrows)), + ], + )?; + let test_dir = TempStrDir::default(); + let mut dataset = Dataset::write( + RecordBatchIterator::new(vec![Ok(batch)], schema.clone()), + &test_dir, + None, + ) + .await?; + + let mut params = VectorIndexParams::ivf_pq(4, 8, 8, MetricType::L2, 50); + params.covering_columns(vec!["id".to_string()]); + dataset + .create_index(&["vec"], IndexType::Vector, None, ¶ms, false) + .await?; + + // Rename of the covered column must fail with the remediation hint. + let err = dataset + .alter_columns(&[ColumnAlteration::new("id".into()).rename("uid".into())]) + .await + .expect_err("rename of covered column should fail"); + assert!( + matches!(err, Error::InvalidInput { .. }), + "expected Error::InvalidInput, got: {err:?}" + ); + assert!( + err.to_string().contains("drop_index"), + "rename error should suggest drop_index, got: {err}" + ); + assert!( + dataset.schema().field("id").is_some(), + "rename should not have applied" + ); + + // Nullability change of the covered column must also fail. + let err = dataset + .alter_columns(&[ColumnAlteration::new("id".into()).set_nullable(true)]) + .await + .expect_err("nullability change of covered column should fail"); + assert!( + matches!(err, Error::InvalidInput { .. }), + "expected Error::InvalidInput, got: {err:?}" + ); + assert!( + err.to_string().contains("drop_index"), + "nullable error should suggest drop_index, got: {err}" + ); + + Ok(()) + } + + /// Carried refine vectors put the *indexed* column's own id into `covering_fields` + /// (`CreateIndexBuilder` appends it), so an alter of that column is an alter of a + /// covering column and must be refused -- exactly like any other covering column. + /// + /// A rename looks harmless (the field id survives, so the declaration still resolves) + /// but is NOT: index storage holds the carried copy under the column's *old* name, and + /// nothing renames it. Allowing the rename leaves an index that answers queries while + /// both maintenance paths are permanently broken -- `optimize_indices` dies with + /// "field does not exist in the RecordBatch" and `compact_files` with "cannot + /// remap index: its storage carries covering column ''". Refusing up front costs + /// the user one `drop_index()`; allowing it costs them the index. + /// + /// Dropping the same column IS allowed, and the asymmetry is the point: a drop takes + /// the index with it (`retain_relevant_indices` prunes any index whose keyed field left + /// the schema), so no storage is stranded. A rename keeps the index alive. + #[tokio::test] + async fn test_alter_of_indexed_column_refused_with_carried_refine() -> Result<()> { + use crate::index::vector::VectorIndexParams; + use lance_arrow::FixedSizeListArrayExt; + use lance_index::IndexType; + use lance_linalg::distance::MetricType; + use lance_testing::datagen::generate_random_array; + + let schema = Arc::new(ArrowSchema::new(vec![ + ArrowField::new( + "vec", + DataType::FixedSizeList( + Arc::new(ArrowField::new("item", DataType::Float32, true)), + 64, + ), + false, + ), + ArrowField::new("id", DataType::Int32, false), + ])); + let nrows: i32 = 256; + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new( + ::try_new_from_values( + generate_random_array(64 * nrows as usize), + 64, + ) + .unwrap(), + ), + Arc::new(Int32Array::from_iter_values(0..nrows)), + ], + )?; + let test_dir = TempStrDir::default(); + let mut dataset = Dataset::write( + RecordBatchIterator::new(vec![Ok(batch)], schema.clone()), + &test_dir, + None, + ) + .await?; + + let mut params = VectorIndexParams::ivf_pq(4, 8, 8, MetricType::L2, 50); + params.store_vectors_for_refine(true); + dataset + .create_index(&["vec"], IndexType::Vector, None, ¶ms, false) + .await?; + let indexed_field_id = dataset.schema().field("vec").unwrap().id; + assert!( + dataset.load_indices().await?[0] + .covering_fields + .contains(&indexed_field_id), + "precondition: carried refine vectors must declare the indexed column" + ); + + // Rename: refused, with the remedy named. + let err = dataset + .alter_columns(&[ColumnAlteration::new("vec".into()).rename("embedding".into())]) + .await + .expect_err("rename of a refine-carrying indexed column must be refused"); + assert!( + matches!(err, Error::InvalidInput { .. }), + "expected Error::InvalidInput, got: {err:?}" + ); + assert!( + err.to_string().contains("drop_index"), + "the refusal must name the remedy, got: {err}" + ); + assert!( + dataset.schema().field("vec").is_some(), + "the rename must not have applied" + ); + + // Nullability flip: refused for the same reason. + assert!( + dataset + .alter_columns(&[ColumnAlteration::new("vec".into()).set_nullable(true)]) + .await + .is_err(), + "nullability change of a refine-carrying indexed column must be refused" + ); + + // Cast: refused, as it always was. + assert!( + dataset + .alter_columns(&[ColumnAlteration::new("vec".into()).cast_to( + DataType::FixedSizeList( + Arc::new(ArrowField::new("item", DataType::Float16, true)), + 64, + ), + )]) + .await + .is_err(), + "cast of the indexed column must be refused" + ); + + // Dropping it is allowed: the index goes with the column, so nothing is stranded. + dataset.drop_columns(&["vec"]).await?; + assert!( + dataset.load_indices().await?.is_empty(), + "dropping the indexed column must prune the index rather than strand it" + ); + + Ok(()) + } + #[test] fn test_is_upcast_downcast_dictionary() { use DataType::*; diff --git a/rust/lance/src/dataset/take.rs b/rust/lance/src/dataset/take.rs index ddb76a0f720..971d1d189f4 100644 --- a/rust/lance/src/dataset/take.rs +++ b/rust/lance/src/dataset/take.rs @@ -530,6 +530,19 @@ impl TakeBuilder { self } + /// Add the row id to the output *without* requesting it as a projected column. + /// + /// Naming `_rowid` in a [`ProjectionRequest`] also sets `must_add_row_offset` + /// (see `ProjectionPlan::from_schema`), which arms the "must not target deleted + /// rows" check in [`do_take_rows`] and turns a take that legitimately returns + /// fewer rows than requested into an error. It also computes a row-offset column + /// per batch. A caller that wants the ids only to realign its own output against + /// the rows that actually came back needs neither, so it uses this instead. + pub(crate) fn include_row_id(mut self) -> Self { + Arc::make_mut(&mut self.projection).include_row_id(); + self + } + pub(super) fn with_missing_row_policy(mut self, policy: MissingRowPolicy) -> Self { self.missing_row_policy = policy; self diff --git a/rust/lance/src/dataset/tests/dataset_index.rs b/rust/lance/src/dataset/tests/dataset_index.rs index 19673e83ae1..dc3f1ce7bf2 100644 --- a/rust/lance/src/dataset/tests/dataset_index.rs +++ b/rust/lance/src/dataset/tests/dataset_index.rs @@ -13,7 +13,7 @@ use crate::dataset::WriteDestination; use crate::dataset::builder::DatasetBuilder; use crate::dataset::tests::dataset_migrations::scan_dataset; use crate::dataset::tests::dataset_transactions::{assert_results, execute_sql}; -use crate::dataset::transaction::{DataReplacementGroup, Operation, Transaction}; +use crate::dataset::transaction::{DataReplacementGroup, Operation}; use crate::index::vector::VectorIndexParams; use crate::session::Session; use crate::utils::test::covering; @@ -21,7 +21,7 @@ use crate::{Dataset, Error, Result}; use lance_arrow::FixedSizeListArrayExt; use crate::dataset::write::{WriteMode, WriteParams}; -use crate::index::DatasetIndexExt; +use crate::index::{CreateIndexBuilder, DatasetIndexExt}; use arrow::array::{AsArray, GenericListBuilder, GenericStringBuilder}; use arrow::datatypes::UInt64Type; use arrow_array::RecordBatch; @@ -56,6 +56,8 @@ use lance_index::{IndexType, scalar::ScalarIndexParams, vector::DIST_COL}; use lance_io::scheduler::{ScanScheduler, SchedulerConfig}; use lance_io::utils::CachedFileSize; use lance_linalg::distance::MetricType; +use lance_table::format::DataFile; +use object_store::path::Path; use datafusion::common::{assert_contains, assert_not_contains}; use futures::{StreamExt, TryStreamExt}; @@ -64,7 +66,7 @@ use lance_arrow::json::ARROW_JSON_EXT_NAME; use lance_index::scalar::inverted::query::{FtsQuery, MultiMatchQuery}; use lance_table::format::BasePath; use lance_testing::datagen::generate_random_array; -use rand::Rng; +use rand::{Rng, SeedableRng, rngs::StdRng}; use rstest::rstest; #[rstest] @@ -225,6 +227,1037 @@ async fn test_covered_vector_column_is_not_selected_for_ann() { ); } +/// The entire point of a covered IVF_PQ index: a query whose projection is +/// satisfied by the keyed vector column's own `_distance`/`_rowid` plus the +/// index's declared `covering_fields` must be answered from the index alone, +/// with no take against the base table. A query that also needs a column the +/// index does not carry must still take. +/// +/// Two traps established empirically on this project make a broken +/// take-elision path look correct anyway: +/// - `num_partitions = 1` never reaches the probe path at all. +/// - Even 4 partitions is not enough with uniform-random vectors: the +/// early-pruning heuristic (`early_pruning` in `io/exec/knn.rs`) still +/// searches every partition, so `late_search`'s split is never +/// exercised. The data must have well-separated clusters, and +/// `early_pruning` must stay off. +/// So this test uses 4 well-separated clusters and asserts on the plan the +/// scanner actually built (`explain_plan`), never on query results -- a take +/// can run and simply return the same values. +#[tokio::test] +async fn test_covered_ann_query_elides_base_table_take() { + const DIMS: usize = 16; + const NUM_CLUSTERS: usize = 4; + const ROWS_PER_CLUSTER: usize = 64; + const TOTAL: usize = NUM_CLUSTERS * ROWS_PER_CLUSTER; + + let test_uri = TempStrDir::default(); + + // Well-separated clusters (centers 1000 apart along dim 0), not + // uniform-random vectors -- see the module doc above. + let mut rng = StdRng::seed_from_u64(42); + let mut payload = Vec::with_capacity(TOTAL); + let mut extra = Vec::with_capacity(TOTAL); + let mut values = Vec::with_capacity(TOTAL * DIMS); + for cluster in 0..NUM_CLUSTERS { + let center = (cluster * 1000) as f32; + for row in 0..ROWS_PER_CLUSTER { + let row_id = (cluster * ROWS_PER_CLUSTER + row) as i32; + payload.push(row_id); + extra.push(row_id * 2); + for dim in 0..DIMS { + let base = if dim == 0 { center } else { 0.0 }; + values.push(base + (rng.random::() - 0.5) * 0.02); + } + } + } + + let vectors: ArrayRef = Arc::new( + ::try_new_from_values( + Float32Array::from(values), + DIMS as i32, + ) + .unwrap(), + ); + let schema = Arc::new(ArrowSchema::new(vec![ + ArrowField::new("vector", vectors.data_type().clone(), false), + ArrowField::new("payload", DataType::Int32, false), + ArrowField::new("extra", DataType::Int32, false), + ])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + vectors, + Arc::new(Int32Array::from(payload)), + Arc::new(Int32Array::from(extra)), + ], + ) + .unwrap(); + let reader = RecordBatchIterator::new(vec![batch].into_iter().map(Ok), schema); + let mut dataset = Dataset::write(reader, &test_uri, None).await.unwrap(); + + // 4 partitions, matching the cluster count: the covering columns' + // reachability doesn't hinge on which internal probe path runs, but a + // covered-ANN test with too few effectively-reachable partitions has + // already once passed against broken code on this project (see doc + // above), so this follows the same defensive recipe regardless. + let mut params = VectorIndexParams::ivf_pq(NUM_CLUSTERS, 8, 4, MetricType::L2, 2); + params.covering_columns(vec!["payload".to_string()]); + dataset + .create_index(&["vector"], IndexType::Vector, None, ¶ms, true) + .await + .unwrap(); + + // Query near the last cluster's center. + let mut q_values = vec![0.0f32; DIMS]; + q_values[0] = ((NUM_CLUSTERS - 1) * 1000) as f32; + let query = Float32Array::from(q_values); + + // Covered case: the projection (`payload`) is entirely satisfied by the + // index's own output columns -- the plan must contain no base-table take. + // A base-table take on this (stable/v2) storage format is a + // `FilteredReadExec`, displayed as `LanceRead: ...`; checking for the + // literal substring "Take" would be a no-op here, since that name is + // only ever used by the legacy (v1) `TakeExec`. + let mut scan = dataset.scan(); + scan.nearest("vector", &query, 10).unwrap(); + scan.project(&["payload"]).unwrap(); + let plan = scan.explain_plan(true).await.unwrap(); + assert!( + !plan.contains("LanceRead"), + "a covered projection ['payload'] must skip the base-table take entirely:\n{plan}" + ); + let batch = scan.try_into_batch().await.unwrap(); + assert!(batch.column_by_name("payload").is_some()); + + // Uncovered case: `extra` is not a declared covering column, so the + // scanner must take it from the base table. `source=stream` is the + // marker for a row-id-driven take (as opposed to a full-table scan), + // which is what a base-table take against ANN results looks like. + let mut scan = dataset.scan(); + scan.nearest("vector", &query, 10).unwrap(); + scan.project(&["extra"]).unwrap(); + let plan = scan.explain_plan(true).await.unwrap(); + assert!( + plan.contains("LanceRead") && plan.contains("source=stream"), + "a projection needing the uncovered column 'extra' must take from the base table:\n{plan}" + ); + let batch = scan.try_into_batch().await.unwrap(); + assert!(batch.column_by_name("extra").is_some()); +} + +/// A dataset with a covered IVF_PQ index over two covering columns, for the +/// projection-pushdown tests below. +/// +/// Deliberate properties, each load-bearing for at least one assertion: +/// - `covering_columns(["payload", "price"])` is the *reverse* of schema order, so an +/// implementation that walks the schema and filters by membership produces +/// `["price", "payload"]` and fails. +/// - `price` and `payload` have different types and values that are neither the row id +/// nor each other (`price` is negative, `payload` is a string), so a positional read +/// of the storage batch cannot pass for a by-name one. +/// - `extra` is a real column the index does not carry, so "narrowed away" and +/// "never declared" stay distinguishable. +/// - Four well-separated clusters and four partitions, so the runtime probe path is +/// genuinely reached rather than short-circuited (see the module doc above). +async fn covered_two_column_dataset(uri: &str) -> Dataset { + const DIMS: usize = 16; + const NUM_CLUSTERS: usize = 4; + const ROWS_PER_CLUSTER: usize = 64; + const TOTAL: usize = NUM_CLUSTERS * ROWS_PER_CLUSTER; + + let mut rng = StdRng::seed_from_u64(7); + let mut price = Vec::with_capacity(TOTAL); + let mut payload = Vec::with_capacity(TOTAL); + let mut extra = Vec::with_capacity(TOTAL); + let mut values = Vec::with_capacity(TOTAL * DIMS); + for cluster in 0..NUM_CLUSTERS { + let center = (cluster * 1000) as f32; + for row in 0..ROWS_PER_CLUSTER { + let index = (cluster * ROWS_PER_CLUSTER + row) as i32; + price.push(-index - 7); + payload.push(format!("p{}", index * 3 + 11)); + extra.push(index * 2); + for dim in 0..DIMS { + let base = if dim == 0 { center } else { 0.0 }; + values.push(base + (rng.random::() - 0.5) * 0.02); + } + } + } + + let vectors: ArrayRef = Arc::new( + ::try_new_from_values( + Float32Array::from(values), + DIMS as i32, + ) + .unwrap(), + ); + let schema = Arc::new(ArrowSchema::new(vec![ + ArrowField::new("vector", vectors.data_type().clone(), false), + ArrowField::new("price", DataType::Int32, false), + ArrowField::new("payload", DataType::Utf8, false), + ArrowField::new("extra", DataType::Int32, false), + ])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + vectors, + Arc::new(Int32Array::from(price)), + Arc::new(StringArray::from(payload)), + Arc::new(Int32Array::from(extra)), + ], + ) + .unwrap(); + let reader = RecordBatchIterator::new(vec![batch].into_iter().map(Ok), schema); + let mut dataset = Dataset::write(reader, uri, None).await.unwrap(); + + let mut params = VectorIndexParams::ivf_pq(NUM_CLUSTERS, 8, 4, MetricType::L2, 2); + params.covering_columns(vec!["payload".to_string(), "price".to_string()]); + dataset + .create_index(&["vector"], IndexType::Vector, None, ¶ms, true) + .await + .unwrap(); + dataset +} + +/// A query vector near the last cluster's centre of [`covered_two_column_dataset`]. +fn covered_query_vector() -> Float32Array { + let mut values = vec![0.0f32; 16]; + values[0] = 3000.0; + Float32Array::from(values) +} + +/// The ANN search node's declared output schema: `[_distance, _rowid]` plus exactly the +/// covering columns this scan asks the index to materialize. +fn ann_declared_schema(plan: &Arc) -> Vec { + fn find( + plan: &Arc, + ) -> Option> { + if plan.name() == "ANNSubIndexExec" { + return Some(plan.clone()); + } + plan.children().into_iter().find_map(find) + } + let ann = find(plan).expect("the scan must plan an ANN sub-index search"); + ann.schema() + .fields() + .iter() + .map(|field| field.name().clone()) + .collect() +} + +/// The `projection=[...]` list of the (single) flat `LanceScan` in an explained plan. +fn flat_scan_projection(plan: &str) -> Vec { + let line = plan + .lines() + .find(|line| line.contains("LanceScan:")) + .unwrap_or_else(|| panic!("plan has no flat LanceScan node:\n{plan}")); + let start = line + .find("projection=[") + .expect("LanceScan prints a projection") + + "projection=[".len(); + let rest = &line[start..]; + let end = rest.find(']').expect("projection list is bracketed"); + rest[..end] + .split(", ") + .filter(|name| !name.is_empty()) + .map(str::to_string) + .collect() +} + +/// The covering set is per query, not per index: an index carrying `payload` and `price` +/// must be asked for exactly the subset a given scan would otherwise take from the base +/// table, and for *nothing* when the scan reads no dataset column at all. +/// +/// Asserted on the ANN node's declared schema rather than on results, because a covering +/// column is semantically transparent: with the narrowing disabled every case below +/// still returns identical, correct values -- they just arrive via a base-table read. +/// The `Some(vec![])` case (no covering column projected) is the one the whole pushdown +/// exists for, and the one that silently disappears if it is folded into "not computed". +/// +/// This is a plan-shape assertion, decided by a static schema comparison at plan +/// construction, so it needs no clustering and no minimum partition count. +#[tokio::test] +async fn test_covered_ann_declares_only_the_covering_columns_the_query_projects() { + let test_uri = TempStrDir::default(); + let dataset = covered_two_column_dataset(&test_uri).await; + let query = covered_query_vector(); + + let declared_with = + |columns: Vec<&'static str>, filter: Option<&'static str>, prefilter: bool| { + let dataset = dataset.clone(); + let query = query.clone(); + async move { + let mut scan = dataset.scan(); + scan.nearest("vector", &query, 10).unwrap(); + scan.project(&columns).unwrap(); + if let Some(filter) = filter { + scan.prefilter(prefilter); + scan.filter(filter).unwrap(); + } + ann_declared_schema(&scan.create_plan().await.unwrap()) + } + }; + // Postfiltering is the scanner default; the prefilter cases say so explicitly. + let declared = |columns: Vec<&'static str>, filter: Option<&'static str>| { + declared_with(columns, filter, false) + }; + + // Nothing but the search's own output columns: the index declares covering columns, + // but this query needs none of them, so it must ask for none. Collapsing this into + // "no narrowing computed" restores full materialization invisibly. + assert_eq!( + declared(vec![ROW_ID], None).await, + vec![DIST_COL, ROW_ID], + "a query reading no dataset column must ask the index for no covering column" + ); + + // One of two. + assert_eq!( + declared(vec!["price"], None).await, + vec![DIST_COL, ROW_ID, "price"], + "only the covering column this query reads may be materialized" + ); + assert_eq!( + declared(vec!["payload"], None).await, + vec![DIST_COL, ROW_ID, "payload"], + ); + + // Both, in the index's declaration order -- which is the reverse of schema order, so + // this also pins that the emitted order follows the index and not the projection. + assert_eq!( + declared(vec!["price", "payload"], None).await, + vec![DIST_COL, ROW_ID, "payload", "price"], + "covering columns are emitted in the index's declared order" + ); + + // A column the index does not carry narrows to nothing, and is not confused with a + // covering column that was narrowed away. + assert_eq!( + declared(vec!["extra"], None).await, + vec![DIST_COL, ROW_ID], + "an uncovered projection cannot be served by the index" + ); + + // A filter column is read by the refine pass, so it counts even though it is not in + // the output projection. Postfiltering is the default, so this also pins that the + // resolution happens where the whole filter is still visible. + assert_eq!( + declared(vec!["extra"], Some("price < 0")).await, + vec![DIST_COL, ROW_ID, "price"], + "a covering column used only by the filter still saves a base-table read" + ); + + // ...but only when the filter is rechecked above the search. A *prefiltered* scan + // consumes the filter below it and then calls `disable_refine`, so nothing above the + // search reads `price` and materializing it would be pure waste -- free today, a real + // read once storage honours the declaration. The paired invariant test + // (`..._refetches_a_covering_column`, `prefilter_on_covering`) is what proves this + // narrowing does not cost a base-table take; this assertion is what proves it happens. + assert_eq!( + declared_with(vec!["extra"], Some("price < 0"), true).await, + vec![DIST_COL, ROW_ID], + "a prefiltered scan's filter columns are consumed below the search, so the index \ + must not be asked to materialize them" + ); +} + +/// An FTS scorer sitting on top of a vector search reads one more column than the scan +/// options describe: it scores from whatever the search emitted, so the FTS document +/// column is consumed *between* the search and the final projection. Narrowing that is +/// blind to it makes the scorer take the column from the base table -- correct, but +/// exactly the read a covered index exists to avoid. +/// +/// Both directions of the pairing are covered, because they are separate code paths: +/// +/// * FTS is the search and the vector query is the filter -- `Scanner::fts_rerank`; +/// * the vector query is the search and FTS is the (postfiltered) query filter -- +/// the refine pass, via `FilterPlan::refine_columns`. +/// +/// Each is paired with a non-`Match` shape that discriminates rather than merely +/// differing: only `Match` scores from its input, every other shape builds its own plan +/// and joins on `_rowid`. So the `Match` cases must keep `payload` and the others must +/// still narrow it away -- a "fix" that simply stopped narrowing on FTS plans passes the +/// first of each pair and fails the second. +#[tokio::test] +async fn test_covered_ann_keeps_the_column_an_fts_scorer_reads() { + use crate::dataset::scanner::QueryFilter; + + let test_uri = TempStrDir::default(); + let mut dataset = covered_two_column_dataset(&test_uri).await; + dataset + .create_index( + &["payload"], + IndexType::Inverted, + None, + &InvertedIndexParams::default(), + true, + ) + .await + .unwrap(); + + let vector_query = lance_index::vector::Query { + column: "vector".to_string(), + key: Arc::new(covered_query_vector()), + // Every row passes the vector filter, so the rerank below is exercised on a + // populated input regardless of where the jitter puts any single row. + k: 256, + lower_bound: None, + upper_bound: None, + minimum_nprobes: 4, + maximum_nprobes: None, + ef: None, + refine_factor: None, + metric_type: Some(MetricType::L2), + use_index: true, + query_parallelism: lance_index::vector::DEFAULT_QUERY_PARALLELISM, + dist_q_c: 0.0, + approx_mode: Default::default(), + covering_projection: None, + }; + + // Neither scan projects a dataset column, so `payload` can only ever be declared + // because an FTS scorer is known to read it. + + // FTS is the search, the vector query is the filter: `Scanner::fts_rerank`. + let reranked = |fts: FullTextSearchQuery| { + let dataset = dataset.clone(); + let vector_query = vector_query.clone(); + async move { + let mut scan = dataset.scan(); + scan.full_text_search(fts) + .unwrap() + .prefilter(true) + .filter_query(QueryFilter::Vector(vector_query)) + .unwrap(); + scan.project::<&str>(&[]).unwrap(); + scan + } + }; + + // The vector query is the search, FTS is the query filter. Postfiltering is the + // default, and it is the postfilter path whose refine pass scores from the search. + let postfiltered = |fts: FullTextSearchQuery| { + let dataset = dataset.clone(); + async move { + let mut scan = dataset.scan(); + scan.nearest("vector", &covered_query_vector(), 256) + .unwrap(); + scan.minimum_nprobes(4); + scan.filter_query(QueryFilter::Fts(fts)).unwrap(); + scan.project::<&str>(&[]).unwrap(); + scan + } + }; + + // `p11` is the `payload` value of row 0 (see `covered_two_column_dataset`). + let match_query = || { + FullTextSearchQuery::new_query(FtsQuery::Match( + MatchQuery::new("p11".to_string()).with_column(Some("payload".to_string())), + )) + }; + // Same terms, but a shape that joins on `_rowid` instead of scoring from the input. + let joined_query = || { + FullTextSearchQuery::new_query(FtsQuery::Boolean(BooleanQuery::new(vec![( + Occur::Should, + FtsQuery::Match( + MatchQuery::new("p11".to_string()).with_column(Some("payload".to_string())), + ), + )]))) + }; + + let scan = reranked(match_query()).await; + assert_eq!( + ann_declared_schema(&scan.create_plan().await.unwrap()), + vec![DIST_COL, ROW_ID, "payload"], + "the covered document column the FTS rerank scores from must survive the narrowing" + ); + // The plan is not merely well-shaped: it runs and scores. + let batch = scan.try_into_batch().await.unwrap(); + assert!( + batch.num_rows() > 0, + "the reranked search must return the matching row" + ); + assert!( + batch.column_by_name(SCORE_COL).is_some(), + "an FTS rerank emits a score column" + ); + + let scan = reranked(joined_query()).await; + assert_eq!( + ann_declared_schema(&scan.create_plan().await.unwrap()), + vec![DIST_COL, ROW_ID], + "a non-Match rerank joins on _rowid and reads nothing from the search, so the \ + covering column must still be narrowed away" + ); + + let scan = postfiltered(match_query()).await; + assert_eq!( + ann_declared_schema(&scan.create_plan().await.unwrap()), + vec![DIST_COL, ROW_ID, "payload"], + "the refine pass over an FTS query filter scores from the search too" + ); + let batch = scan.try_into_batch().await.unwrap(); + assert!( + batch.num_rows() > 0, + "the postfiltered search must return the matching row" + ); + + // No non-Match counterpart for this direction: `Scanner::flat_fts_filter` rejects + // every other shape outright, so there is no plan to assert a schema on. The + // Match-only rule is pinned by the rerank pair above, which shares the same resolver + // (`fts_scores_from_input`). + // + // Asserted on the error's *content*, not merely on `is_err()`: this case exists to + // fire when the restriction is lifted, and any unrelated planning failure -- a fixture + // that stopped building the inverted index, a change to `filter_query` validation -- + // would otherwise keep it green while the claim it protects went untested. + let err = postfiltered(joined_query()) + .await + .create_plan() + .await + .expect_err("a non-Match FTS post-filter is not plannable today"); + assert!( + err.to_string().contains("Only Match queries are supported"), + "if a non-Match FTS post-filter ever becomes plannable it needs its own \ + declared-schema case here; got a different error instead: {err}" + ); + + // The DEFAULT shape: a Match that names no column at all. `create_plan` resolves + // both the FTS search query and the FTS query filter through + // `resolve_full_text_search_query`, which fills the column from the FTS-indexed + // columns BEFORE any covering resolution sees the query -- so the filled column + // must reach the covering declaration on both directions. If that resolution were + // ever skipped (or `fts_scored_columns` handed an unresolved query), the most + // common query shape would silently stop engaging covering, taking `payload` from + // the base table on every query with byte-identical results. + let unnamed_match_query = + || FullTextSearchQuery::new_query(FtsQuery::Match(MatchQuery::new("p11".to_string()))); + + let scan = reranked(unnamed_match_query()).await; + assert_eq!( + ann_declared_schema(&scan.create_plan().await.unwrap()), + vec![DIST_COL, ROW_ID, "payload"], + "an unnamed Match resolves to the FTS-indexed column, which the covering \ + declaration must include" + ); + let batch = scan.try_into_batch().await.unwrap(); + assert!( + batch.num_rows() > 0, + "the unnamed-Match rerank must return the matching row" + ); + + let scan = postfiltered(unnamed_match_query()).await; + assert_eq!( + ann_declared_schema(&scan.create_plan().await.unwrap()), + vec![DIST_COL, ROW_ID, "payload"], + "the refine pass fills the unnamed Match's column, so the covering must too" + ); + let batch = scan.try_into_batch().await.unwrap(); + assert!( + batch.num_rows() > 0, + "the unnamed-Match postfilter must return the matching row" + ); +} + +// --------------------------------------------------------------------------------------- +// The covering-elision invariant +// +// `Scanner::columns_read_after_vector_search` is a hand-written enumeration of the columns +// something above the vector search reads. Nothing ties it to the nodes `create_plan` +// actually builds, and under-declaration is silent: results stay byte-identical because +// `Scanner::take` simply fetches the column from the base table instead. Phase 7 shipped +// two readers missing from that enumeration (`FlatMatchQueryExec` and `FlatMatchFilter`), +// both found by reading plan text, not by a failing test. +// +// The oracle is that plan text. For an index that covers a column, a take above the search +// fetching that column *is* the lost elision -- there is nothing else to look for. So the +// tests below plan a matrix of scan shapes over an index covering **every** non-vector +// column and assert, structurally, that no take above the search fetches any of them. +// Any newly inserted reader trips this in whatever shapes the matrix exercises. +// +// This is a test-level guard, not a self-maintaining derivation; a genuinely new plan shape +// still needs its own case. See the task report for the two-pass planning alternative that +// would remove the enumeration outright. +// --------------------------------------------------------------------------------------- + +/// The covering columns that some take *above* the vector search fetches from the base +/// table -- i.e. the elisions this scan shape lost. +/// +/// Returns `None` when the plan contains no ANN search at all, so a caller can tell +/// "nothing was taken" apart from "there was nothing to take above". +/// +/// Only the ancestors of the ANN node are inspected. The prefilter hangs *below* the +/// search (it is a child of `ANNSubIndexExec`) and the flat/unindexed fallback is a +/// sibling branch of the union; both legitimately read from the base table, and neither is +/// on the path this walks. Of the ancestors, only the two nodes that can fetch from the +/// dataset are counted -- a `ProjectionExec` can rename or extract a struct child, which +/// would otherwise look like a newly appeared column. +fn covering_taken_above_search( + plan: &Arc, + covering: &HashSet, +) -> Option> { + if plan.name() == "ANNSubIndexExec" { + return Some(Vec::new()); + } + for child in plan.children() { + let Some(mut taken) = covering_taken_above_search(child, covering) else { + continue; + }; + if matches!(plan.name(), "FilteredReadExec" | "TakeExec") { + let child_schema = child.schema(); + let below: HashSet<&str> = child_schema + .fields() + .iter() + .map(|field| field.name().as_str()) + .collect(); + for field in plan.schema().fields() { + if covering.contains(field.name()) && !below.contains(field.name().as_str()) { + taken.push(field.name().clone()); + } + } + } + return Some(taken); + } + None +} + +/// The scan shapes the invariant is checked against. Each is a distinct route through +/// `create_plan`'s post-source pipeline, which is where a reader can hide. +#[derive(Debug, Clone, Copy)] +enum CoveredShape { + /// The state the pushdown exists for: nothing read from the dataset at all. + ProjectNothing, + ProjectOneCovering, + ProjectEveryCovering, + /// Postfiltering is the default; the filter is rechecked above the search. + PostfilterOnCovering, + /// Prefiltering consumes the filter below the search, so nothing above reads it. + PrefilterOnCovering, + /// `create_plan` takes the sort keys above the search, twice. + OrderByCovering, + /// Aggregate inputs are taken above the search on their own path. + AggregateOverCovering, + /// Inserts a `take` of the vector column above the ANN node before the flat re-score. + RefineFactor, + LimitOffset, + WithRowAddress, + /// Skips `knn_combined` entirely. + FastSearch, + /// FTS is the search, the vector query is the filter: `FlatMatchQueryExec` scores from + /// the search's own columns. + FtsRerank, + /// The vector query is the search, FTS is the query filter: the refine pass scores + /// from the search's own columns via `FlatMatchFilter`. + FtsPostfilterRefine, +} + +/// A dataset whose covered index carries **every** non-vector column, so that any column +/// any scan shape reads above the search is a covering column and therefore a detectable +/// elision loss. `payload` also carries an inverted index, for the two FTS shapes. +/// +/// Four well-separated clusters and four partitions, so the runtime probe path is genuinely +/// reached rather than short-circuited by a single partition. +async fn fully_covered_dataset(uri: &str, append_unindexed: bool) -> Dataset { + const DIMS: usize = 16; + const NUM_CLUSTERS: usize = 4; + const ROWS_PER_CLUSTER: usize = 64; + const TOTAL: usize = NUM_CLUSTERS * ROWS_PER_CLUSTER; + + let mut rng = StdRng::seed_from_u64(11); + let mut make_batch = |tag: &str| { + let mut price = Vec::with_capacity(TOTAL); + let mut payload = Vec::with_capacity(TOTAL); + let mut extra = Vec::with_capacity(TOTAL); + let mut values = Vec::with_capacity(TOTAL * DIMS); + for cluster in 0..NUM_CLUSTERS { + let center = (cluster * 1000) as f32; + for row in 0..ROWS_PER_CLUSTER { + let index = (cluster * ROWS_PER_CLUSTER + row) as i32; + price.push(-index - 7); + payload.push(format!("{tag}{}", index * 3 + 11)); + extra.push(index * 2); + for dim in 0..DIMS { + let base = if dim == 0 { center } else { 0.0 }; + values.push(base + (rng.random::() - 0.5) * 0.02); + } + } + } + let vectors: ArrayRef = Arc::new( + ::try_new_from_values( + Float32Array::from(values), + DIMS as i32, + ) + .unwrap(), + ); + let schema = Arc::new(ArrowSchema::new(vec![ + ArrowField::new("vector", vectors.data_type().clone(), false), + ArrowField::new("price", DataType::Int32, false), + ArrowField::new("payload", DataType::Utf8, false), + ArrowField::new("extra", DataType::Int32, false), + ])); + RecordBatch::try_new( + schema, + vec![ + vectors, + Arc::new(Int32Array::from(price)), + Arc::new(StringArray::from(payload)), + Arc::new(Int32Array::from(extra)), + ], + ) + .unwrap() + }; + + let batch = make_batch("p"); + let schema = batch.schema(); + let reader = RecordBatchIterator::new(vec![batch].into_iter().map(Ok), schema.clone()); + let mut dataset = Dataset::write(reader, uri, None).await.unwrap(); + + let mut params = VectorIndexParams::ivf_pq(NUM_CLUSTERS, 8, 4, MetricType::L2, 2); + // Every non-vector column. `covering_columns` rejects the indexed vector column itself, + // so this is the widest legal declaration for this schema. + params.covering_columns(vec![ + "payload".to_string(), + "price".to_string(), + "extra".to_string(), + ]); + dataset + .create_index(&["vector"], IndexType::Vector, None, ¶ms, true) + .await + .unwrap(); + dataset + .create_index( + &["payload"], + IndexType::Inverted, + None, + &InvertedIndexParams::default(), + true, + ) + .await + .unwrap(); + + if append_unindexed { + // After the build, so `knn_combined` plans its flat fallback branch and the extra + // `take` of the vector column that sits above the ANN node. + let more = make_batch("q"); + let reader = RecordBatchIterator::new(vec![more].into_iter().map(Ok), schema); + dataset.append(reader, None).await.unwrap(); + } + dataset +} + +/// Build the scan for one shape. Kept in one place so the matrix stays a list of names. +async fn plan_covered_shape( + dataset: &Dataset, + shape: CoveredShape, +) -> Arc { + use crate::dataset::scanner::{AggregateExpr, ColumnOrdering, QueryFilter}; + + let query = covered_query_vector(); + let vector_filter = lance_index::vector::Query { + column: "vector".to_string(), + key: Arc::new(query.clone()), + k: 32, + lower_bound: None, + upper_bound: None, + minimum_nprobes: 4, + maximum_nprobes: None, + ef: None, + refine_factor: None, + metric_type: Some(MetricType::L2), + use_index: true, + query_parallelism: lance_index::vector::DEFAULT_QUERY_PARALLELISM, + dist_q_c: 0.0, + approx_mode: Default::default(), + covering_projection: None, + }; + // `p11` is row 0's payload; any term would do, the shape is what is under test. + let match_query = || { + FullTextSearchQuery::new_query(FtsQuery::Match( + MatchQuery::new("p11".to_string()).with_column(Some("payload".to_string())), + )) + }; + + let mut scan = dataset.scan(); + // Every shape except the two FTS ones is a plain vector search. + if !matches!( + shape, + CoveredShape::FtsRerank | CoveredShape::FtsPostfilterRefine + ) { + scan.nearest("vector", &query, 10).unwrap(); + scan.minimum_nprobes(4); + } + scan.with_row_id(); + + match shape { + CoveredShape::ProjectNothing => { + scan.project::<&str>(&[]).unwrap(); + } + CoveredShape::ProjectOneCovering => { + scan.project(&["payload"]).unwrap(); + } + CoveredShape::ProjectEveryCovering => { + scan.project(&["payload", "price", "extra"]).unwrap(); + } + CoveredShape::PostfilterOnCovering => { + scan.prefilter(false); + scan.filter("price < 0 AND extra >= 0").unwrap(); + scan.project::<&str>(&[]).unwrap(); + } + CoveredShape::PrefilterOnCovering => { + scan.prefilter(true); + scan.filter("price < 0 AND extra >= 0").unwrap(); + scan.project::<&str>(&[]).unwrap(); + } + CoveredShape::OrderByCovering => { + scan.project::<&str>(&[]).unwrap(); + scan.order_by(Some(vec![ColumnOrdering::asc_nulls_first( + "price".to_string(), + )])) + .unwrap(); + } + CoveredShape::AggregateOverCovering => { + scan.project::<&str>(&[]).unwrap(); + scan.aggregate(AggregateExpr::builder().sum("extra").build()) + .unwrap(); + } + CoveredShape::RefineFactor => { + scan.refine(2); + scan.project(&["payload"]).unwrap(); + } + CoveredShape::LimitOffset => { + scan.limit(Some(5), Some(2)).unwrap(); + scan.project(&["payload"]).unwrap(); + } + CoveredShape::WithRowAddress => { + scan.with_row_address(); + scan.project(&["price"]).unwrap(); + } + CoveredShape::FastSearch => { + scan.fast_search(); + scan.project(&["payload"]).unwrap(); + } + CoveredShape::FtsRerank => { + scan.full_text_search(match_query()) + .unwrap() + .prefilter(true) + .filter_query(QueryFilter::Vector(vector_filter)) + .unwrap(); + scan.project::<&str>(&[]).unwrap(); + } + CoveredShape::FtsPostfilterRefine => { + scan.nearest("vector", &query, 32).unwrap(); + scan.minimum_nprobes(4); + scan.prefilter(false); + scan.filter_query(QueryFilter::Fts(match_query())).unwrap(); + scan.project::<&str>(&[]).unwrap(); + } + } + scan.create_plan().await.unwrap() +} + +/// No take above a covered vector search may fetch a column the index already carries. +/// +/// This is the structural counterpart to +/// `test_covered_ann_declares_only_the_covering_columns_the_query_projects`: that test +/// pins what the *declaration* says for a handful of shapes, this one pins the +/// *consequence* -- that the declaration was complete enough that nothing had to be +/// re-fetched -- across every shape in the matrix, without naming a single column. +/// +/// The assertion is paired, not bare: each case first asserts the shape really planned an +/// ANN search (otherwise "no take above it" is vacuous), then that the set of covering +/// columns taken above it is empty. +/// +/// Run with and without an unindexed fragment, because the fallback rewrites the whole +/// region above the search into a union with an extra take of the vector column. +#[rstest] +#[case::project_nothing(CoveredShape::ProjectNothing)] +#[case::project_one_covering(CoveredShape::ProjectOneCovering)] +#[case::project_every_covering(CoveredShape::ProjectEveryCovering)] +#[case::postfilter_on_covering(CoveredShape::PostfilterOnCovering)] +#[case::prefilter_on_covering(CoveredShape::PrefilterOnCovering)] +#[case::order_by_covering(CoveredShape::OrderByCovering)] +#[case::aggregate_over_covering(CoveredShape::AggregateOverCovering)] +#[case::refine_factor(CoveredShape::RefineFactor)] +#[case::limit_offset(CoveredShape::LimitOffset)] +#[case::with_row_address(CoveredShape::WithRowAddress)] +#[case::fast_search(CoveredShape::FastSearch)] +#[case::fts_rerank(CoveredShape::FtsRerank)] +#[case::fts_postfilter_refine(CoveredShape::FtsPostfilterRefine)] +#[tokio::test] +async fn test_no_take_above_a_covered_search_refetches_a_covering_column( + #[case] shape: CoveredShape, + #[values(false, true)] unindexed_fragment: bool, +) { + let test_uri = TempStrDir::default(); + let dataset = fully_covered_dataset(&test_uri, unindexed_fragment).await; + let covering: HashSet = ["payload", "price", "extra"] + .into_iter() + .map(str::to_string) + .collect(); + + let plan = plan_covered_shape(&dataset, shape).await; + let taken = covering_taken_above_search(&plan, &covering).unwrap_or_else(|| { + panic!( + "{shape:?} (unindexed_fragment={unindexed_fragment}) planned no ANN search, so \ + the invariant below would hold vacuously; either the shape stopped using the \ + index or it does not belong in this matrix. Plan:\n{}", + datafusion::physical_plan::displayable(plan.as_ref()).indent(true) + ) + }); + assert!( + taken.is_empty(), + "{shape:?} (unindexed_fragment={unindexed_fragment}) re-fetches covering column(s) \ + {taken:?} from the base table above the search. The index already carries them, so \ + some node above the search reads a column \ + `Scanner::columns_read_after_vector_search` does not name. Plan:\n{}", + datafusion::physical_plan::displayable(plan.as_ref()).indent(true) + ); +} + +/// The ANN node and the flat fallback are unioned against a single schema, so they must +/// narrow together. A flat path that emits *fewer* covering columns than the ANN node +/// declares is not a slow plan, it is a failed one -- which is why the pre-pushdown code +/// simply emitted every declared covering column on both sides. +/// +/// Asserts the flat scan's projection positively (exact list), so it fails both when the +/// flat path keeps reading a narrowed-away column and when it stops reading a needed one. +#[tokio::test] +async fn test_covered_ann_flat_fallback_narrows_with_the_index() { + let test_uri = TempStrDir::default(); + let mut dataset = covered_two_column_dataset(&test_uri).await; + + // Append after indexing so `knn_combined` has an unindexed fragment to scan flat. + let schema = Arc::new(ArrowSchema::from(dataset.schema())); + let appended: ArrayRef = Arc::new( + ::try_new_from_values( + Float32Array::from(vec![3000.0f32; 16]), + 16, + ) + .unwrap(), + ); + let append = RecordBatch::try_new( + schema.clone(), + vec![ + appended, + Arc::new(Int32Array::from(vec![-9999])), + Arc::new(StringArray::from(vec!["appended"])), + Arc::new(Int32Array::from(vec![12345])), + ], + ) + .unwrap(); + dataset + .append( + RecordBatchIterator::new(vec![Ok(append)], schema.clone()), + None, + ) + .await + .unwrap(); + + let query = covered_query_vector(); + let mut scan = dataset.scan(); + scan.nearest("vector", &query, 10).unwrap(); + scan.project(&["price"]).unwrap(); + + // The index really is used -- without this the assertions below would hold just as + // well for a plan that fell back to brute force entirely. + let plan = scan.explain_plan(true).await.unwrap(); + assert!( + plan.contains("ANNSubIndex"), + "the scan must use the covered index:\n{plan}" + ); + assert_eq!( + ann_declared_schema(&scan.create_plan().await.unwrap()), + vec![DIST_COL, ROW_ID, "price"], + ); + assert_eq!( + flat_scan_projection(&plan), + vec!["vector", "price"], + "the flat fallback must read the vector it scores plus exactly the covering \ + columns the ANN node declares -- no more (wasted IO) and no fewer (broken \ + union):\n{plan}" + ); + + // And the plan actually runs, with the narrowed covering values arriving by name out + // of a storage batch that still physically carries both covering columns. + let batch = scan.try_into_batch().await.unwrap(); + assert!( + batch.num_rows() > 0, + "the covered search must return the appended row's cluster" + ); + let prices = batch + .column_by_name("price") + .expect("price is projected") + .as_primitive::(); + assert!( + prices.iter().all(|price| price.unwrap() < 0), + "price values must be the negative fixture values, not row ids or payload \ + positions: {prices:?}" + ); +} + +/// Every covering value must survive a narrowed declaration, by name. The storage batch +/// still carries both covering columns at this point (narrowing storage is a later +/// phase), so a positional read would hand back `payload` where `price` was asked for -- +/// which the fixture's disjoint types and values make impossible to miss. +/// +/// Uses the four-cluster fixture and asserts against `take_rows` ground truth rather +/// than recall: recall cannot discriminate an index read from a brute-force fallback, +/// since a covering column is semantically transparent and a fallback scores 1.0. +#[tokio::test] +async fn test_covered_ann_narrowed_values_match_the_base_table() { + let test_uri = TempStrDir::default(); + let dataset = covered_two_column_dataset(&test_uri).await; + let query = covered_query_vector(); + + let mut scan = dataset.scan(); + scan.nearest("vector", &query, 10).unwrap(); + scan.project(&["payload"]).unwrap(); + scan.with_row_id(); + + let plan = scan.explain_plan(true).await.unwrap(); + assert!( + plan.contains("ANNSubIndex"), + "the values below must come from the index, not a fallback scan:\n{plan}" + ); + assert!( + !plan.contains("LanceRead"), + "a covered projection must not take from the base table:\n{plan}" + ); + assert_eq!( + ann_declared_schema(&scan.create_plan().await.unwrap()), + vec![DIST_COL, ROW_ID, "payload"], + ); + + let batch = scan.try_into_batch().await.unwrap(); + assert_eq!(batch.num_rows(), 10); + let row_ids = batch + .column_by_name(ROW_ID) + .expect("row ids requested by name") + .as_primitive::() + .values() + .to_vec(); + let from_index = batch + .column_by_name("payload") + .expect("payload is projected") + .as_string::() + .clone(); + + let ground_truth = dataset + .take_rows(&row_ids, dataset.schema().project(&["payload"]).unwrap()) + .await + .unwrap(); + let expected = ground_truth + .column_by_name("payload") + .expect("take returns the requested column") + .as_string::() + .clone(); + assert_eq!( + from_index, expected, + "the narrowed covering column must carry the same values the base table would" + ); +} + /// A filtered `describe_indices` must still find a covered index by its keyed /// column. The matcher compares the caller's resolved field slice against the one /// column named by `for_column`, so a caller that passes all of `index.fields` -- @@ -273,85 +1306,23 @@ async fn test_describe_indices_filters_a_covered_index_by_its_keyed_field() { ); } -/// An unfiltered `optimize_indices()` is a table-wide request, so a covered index -/// must not abort it -- that would block optimization of every other index on the -/// table over one this build merely cannot rebuild. It is skipped with a warning. -/// Only a caller that names the covered index gets an error (see -/// `test_optimize_indices_rejects_a_covered_index`). +/// The scalar counterpart to the covered-vector optimize case above. Before a producer +/// existed optimize skipped a covered group outright; the producer makes covered rebuilds +/// real, so what needs pinning is that rebuilding a covered *scalar* group preserves its +/// covering declaration rather than silently dropping it. A carried column is inert to +/// every scalar read path, so a lost declaration would stay invisible until a covered +/// planner consulted it. /// -/// Both the current and the stale case are covered: erroring on the stale one -/// aborts the loop before the replacements accumulated for the other groups are -/// committed, leaving an unrelated index stale too. -#[rstest] -#[case::current(false)] -#[case::stale(true)] -#[tokio::test] -async fn test_optimize_skips_a_covered_index_without_blocking_others(#[case] stale: bool) { - let test_uri = TempStrDir::default(); - let mut dataset = covering::write_vector_payload_dataset(&test_uri).await; - - // A covered vector index, and a plain scalar index that optimize may touch. - covering::create_ivf_pq_index(&mut dataset, "vec").await; - covering::create_btree_index(&mut dataset, "payload", Some("payload_idx")).await; - - let (_, payload_id) = covering::declare_covering(&mut dataset, "vec", "payload").await; - let covered_uuid = dataset - .load_indices() - .await - .unwrap() - .iter() - .find(|idx| !idx.covering_fields.is_empty()) - .expect("the covered index should exist") - .uuid; - - if stale { - // Now *both* groups have an unindexed fragment, so the covered group - // would genuinely be rebuilt -- this is where the refusal used to fire. - covering::append_vector_payload_rows(&mut dataset, 256).await; - } - - dataset - .optimize_indices(&OptimizeOptions::default()) - .await - .expect("a covered index must not abort an unfiltered optimize"); - - let after = dataset.load_indices().await.unwrap(); - assert!( - after - .iter() - .any(|idx| idx.uuid == covered_uuid && idx.covering_fields == vec![payload_id]), - "the covered index must be left exactly as it was" - ); - - if stale { - let payload_idx = after - .iter() - .filter(|idx| idx.name == "payload_idx") - .filter_map(|idx| idx.fragment_bitmap.as_ref()) - .fold(roaring::RoaringBitmap::new(), |mut acc, bitmap| { - acc |= bitmap; - acc - }); - assert!( - payload_idx.contains(1), - "the unrelated index must still have been optimized onto the new fragment, got {payload_idx:?}" - ); - } else { - assert_eq!(after.len(), 2, "both indices must survive"); - } -} - -/// The same skip, for a *scalar* covered index. The rule does not branch on index -/// type, but scalar groups take their own no-work path a few lines below the -/// covering gate, so a covered scalar index reaching that gate first is worth -/// pinning separately from the vector case above. +/// The append is what gives this test teeth: with no unindexed rows the covered group has +/// no work, and the assertions would hold with the rebuild never happening. /// -/// The append is what gives this test teeth. Without it the covered group has no -/// work either way, so the scalar no-work path below the gate produces the same -/// observable outcome as the gate itself and the test passes with the gate -/// removed entirely. +/// The reload right before `optimize_indices` is load-bearing. Without it optimize +/// intermittently commits a manifest missing one of the two groups. That race needs no +/// covering to reproduce -- two plain btree indexes over an append are enough -- so it is +/// not this branch's to fix, and the reload keeps this test measuring covering rather than +/// that race. #[tokio::test] -async fn test_optimize_skips_a_stale_covered_scalar_index() { +async fn test_optimize_rebuild_preserves_a_covered_scalar_declaration() { let test_uri = TempStrDir::default(); let mut dataset = covering::write_three_int_column_dataset(&test_uri).await; covering::create_btree_index(&mut dataset, "a", None).await; @@ -367,9 +1338,10 @@ async fn test_optimize_skips_a_stale_covered_scalar_index() { .expect("the covered index should exist") .uuid; - // Both scalar groups now have an unindexed fragment, so the covered one - // would genuinely be rebuilt if the gate did not skip it first. + // Both scalar groups now have an unindexed fragment, so the covered one is genuinely + // rebuilt rather than reaching a no-work path. covering::append_three_int_column_rows(&mut dataset, 64).await; + dataset.load_indices().await.unwrap(); dataset .optimize_indices(&OptimizeOptions::default()) @@ -377,14 +1349,36 @@ async fn test_optimize_skips_a_stale_covered_scalar_index() { .expect("a stale covered scalar index must not abort optimize"); let after = dataset.load_indices().await.unwrap(); + let covered = after + .iter() + .filter(|idx| !idx.covering_fields.is_empty()) + .collect::>(); + assert_eq!( + covered.len(), + 1, + "exactly one index must still declare covering, got {covered:?}" + ); + assert_eq!( + covered[0].covering_fields, + vec![carried_id], + "the rebuild must carry the covering declaration forward, not drop it" + ); + assert_ne!( + covered[0].uuid, covered_uuid, + "precondition: the covered group must actually have been rebuilt, or this test \ + would pass with the rebuild path never entered" + ); + let covered_coverage = covered[0] + .fragment_bitmap + .as_ref() + .expect("a rebuilt index records its coverage"); assert!( - after - .iter() - .any(|idx| idx.uuid == covered_uuid && idx.covering_fields == vec![carried_id]), - "the covered scalar index must be left exactly as it was, not rebuilt" + covered_coverage.contains(1), + "the rebuilt covered index must span the appended fragment, got {covered_coverage:?}" ); - // The unrelated scalar index was still maintained, so the skip is scoped to - // the covered group rather than aborting the loop. + + // The unrelated scalar index was still maintained, so the rebuild is scoped to the + // covered group rather than aborting the loop. let b_id = dataset.schema().field_id("b").unwrap(); let b_coverage = after .iter() @@ -400,60 +1394,6 @@ async fn test_optimize_skips_a_stale_covered_scalar_index() { ); } -/// A caller that names the covered index asked for it specifically, so the -/// refusal is loud rather than a skip. -#[tokio::test] -async fn test_optimize_indices_rejects_a_covered_index() { - let test_uri = TempStrDir::default(); - let mut dataset = covering::write_vector_payload_dataset(&test_uri).await; - covering::create_ivf_pq_index(&mut dataset, "vec").await; - - // Nothing writes carried values, so the storage does not contain `payload` - // -- which is exactly why optimize must refuse: it rebuilds from a scan - // projecting the keyed field and `_rowid` only, and would republish the - // declaration on a segment that still has no payload. - let (vec_id, payload_id) = covering::declare_covering(&mut dataset, "vec", "payload").await; - - // Append AFTER declaring covering, so the group really would be rebuilt. - // Without it this would assert the refusal against an index optimize had no - // work for, and would keep passing if the refusal moved behind a no-work - // check. - covering::append_vector_payload_rows(&mut dataset, 256).await; - - let before = dataset.load_indices().await.unwrap(); - let before_uuid = before[0].uuid; - let covered_name = before[0].name.clone(); - - // Name the covered index: the refusal is reserved for a caller that targeted - // it. An unfiltered call skips it instead, which - // `test_optimize_skips_a_covered_index_without_blocking_others` covers. - let err = dataset - .optimize_indices(&OptimizeOptions::default().index_names(vec![covered_name])) - .await - .expect_err("optimizing a targeted covered index must be refused"); - assert!( - err.to_string().contains("declares covering fields"), - "unexpected message: {err}" - ); - - // Refused, not partially applied: the index is exactly as it was. - let after = dataset.load_indices().await.unwrap(); - assert_eq!(after.len(), 1); - assert_eq!( - after[0].uuid, before_uuid, - "a refused optimize must not replace the index" - ); - assert_eq!(after[0].covering_fields, vec![payload_id]); - assert_eq!(after[0].fields, vec![vec_id, payload_id]); - - // The appended data above is unindexed, so this really is a case optimize - // would otherwise have merged -- the refusal is not the no-op path. - assert!( - after.iter().all(|idx| !idx.covering_fields.is_empty()), - "precondition: the only index is still the covered one" - ); -} - #[rstest] #[tokio::test] async fn test_create_scalar_index( @@ -7448,17 +8388,91 @@ async fn test_load_segment_params_full_fidelity() { assert_eq!(&read, opened.params()); } -/// Compaction of a covered index must succeed. `remap_index` rejected any -/// index with more than one field, so this failed outright -- a covered -/// dataset could be created but never compacted. -#[tokio::test] -async fn test_compaction_withdraws_a_covered_index_without_failing() { - use crate::dataset::optimize::{CompactionOptions, compact_files}; +/// Write `values` as a standalone single-field `payload` data file (field `payload_id`) +/// and return the resulting [`DataFile`], ready to use in an `Operation::DataReplacement` +/// group. Kept in its own file from `vec` throughout this test so replacing it is a +/// straightforward single-field swap, never a split of a shared file. +async fn write_payload_data_file( + dataset: &Dataset, + name: &str, + payload_id: i32, + values: Int32Array, +) -> DataFile { + let payload_schema = Arc::new(ArrowSchema::new(vec![ArrowField::new( + "payload", + DataType::Int32, + true, + )])); + let path = format!("{name}.lance"); + let object_writer = dataset + .object_store + .create(&Path::from(format!("data/{path}"))) + .await + .unwrap(); + let mut writer = lance_file::versions::v2_1::create_writer( + object_writer, + payload_schema.as_ref().try_into().unwrap(), + Default::default(), + ) + .unwrap(); + let write_batch = RecordBatch::try_new(payload_schema, vec![Arc::new(values)]).unwrap(); + writer.write_batch(&write_batch).await.unwrap(); + writer.finish().await.unwrap(); - let test_uri = TempStrDir::default(); + let (major, minor) = LanceFileVersion::Stable.resolve().to_data_file_numbers(); + DataFile { + path, + fields: Arc::from([payload_id]), + column_indices: Arc::from([0]), + file_major_version: major, + file_minor_version: minor, + file_size_bytes: CachedFileSize::unknown(), + base_id: None, + } +} + +/// `prune_stale_segment_coverage` must judge staleness from every field a segment +/// declares -- keyed and carried alike -- not just the keyed field's subtree. +/// +/// Before the fix, only the keyed (`vec`) field fed the staleness check: a fragment +/// whose *carried* (`payload`) column was physically rewritten between when the +/// segment was built and when it is committed stayed marked "covered" regardless, so +/// the committed index would go on serving that fragment's now-stale carried values. +/// `dataset/overlay.rs`'s per-read staleness check already considers every declared +/// field; this is the build-time counterpart disagreeing with it. +#[tokio::test] +async fn test_commit_existing_index_segments_prunes_fragment_with_rewritten_carried_column() { + // An in-memory store, not `TempStrDir`: `write_payload_data_file` writes bare + // relative paths (matching this crate's other `DataReplacement` tests), which for + // a `file://` store would resolve against the filesystem root rather than the + // dataset's own directory. + let test_uri = "memory://"; let dimension = 16; - let schema = Arc::new(ArrowSchema::new(vec![ + let vec_schema = Arc::new(ArrowSchema::new(vec![ArrowField::new( + "vec", + DataType::FixedSizeList( + Arc::new(ArrowField::new("item", DataType::Float32, true)), + dimension, + ), + false, + )])); + let vectors = Arc::new( + ::try_new_from_values( + generate_random_array(512 * dimension as usize), + dimension, + ) + .unwrap(), + ); + let batch = RecordBatch::try_new(vec_schema.clone(), vec![vectors]).unwrap(); + let reader = RecordBatchIterator::new(vec![Ok(batch)], vec_schema.clone()); + let dataset = Dataset::write(reader, test_uri, None).await.unwrap(); + + // Add `payload` as an all-null column. It stays fileless until populated below, so + // it never shares a data file with `vec` -- a later replacement of just its file is + // then a straightforward single-field swap, not a split of a multi-field file (which + // `DataReplacement` does not support). + let extended_schema = Arc::new(ArrowSchema::new(vec![ ArrowField::new( "vec", DataType::FixedSizeList( @@ -7467,81 +8481,105 @@ async fn test_compaction_withdraws_a_covered_index_without_failing() { ), false, ), - ArrowField::new("payload", DataType::Int32, false), + ArrowField::new("payload", DataType::Int32, true), ])); + let fragment = dataset.get_fragments().pop().unwrap().metadata; + let read_version = dataset.manifest.version; + let dataset = Dataset::commit( + WriteDestination::Dataset(Arc::new(dataset)), + Operation::Merge { + fragments: vec![fragment], + schema: extended_schema.as_ref().try_into().unwrap(), + preserves_nullability: true, + }, + Some(read_version), + None, + None, + Arc::new(Default::default()), + false, + ) + .await + .unwrap(); + let payload_id = dataset.schema().field_id("payload").unwrap(); - let make_batch = |offset: i32| { - let vectors = Arc::new( - ::try_new_from_values( - generate_random_array(256 * dimension as usize), - dimension, - ) - .unwrap(), - ); - let payload = Arc::new(Int32Array::from_iter_values(offset..offset + 256)); - RecordBatch::try_new(schema.clone(), vec![vectors, payload]).unwrap() - }; - - // Two fragments, so compaction has something to compact. - let reader = RecordBatchIterator::new(vec![Ok(make_batch(0))], schema.clone()); - let mut dataset = Dataset::write(reader, &test_uri, None).await.unwrap(); - let reader = RecordBatchIterator::new(vec![Ok(make_batch(256))], schema.clone()); - dataset.append(reader, None).await.unwrap(); + // Give `payload` its first real values, in a data file of its own. + let payload_v1 = write_payload_data_file( + &dataset, + "payload_v1", + payload_id, + Int32Array::from_iter_values(0..512), + ) + .await; + let read_version = dataset.manifest.version; + let mut dataset = Dataset::commit( + WriteDestination::Dataset(Arc::new(dataset)), + Operation::DataReplacement { + replacements: vec![DataReplacementGroup(0, payload_v1)], + }, + Some(read_version), + None, + None, + Arc::new(Default::default()), + false, + ) + .await + .unwrap(); - let params = VectorIndexParams::ivf_pq(4, 8, 2, MetricType::L2, 50); - dataset - .create_index(&["vec"], IndexType::Vector, None, ¶ms, true) + // Build a covered segment (Task 1's `covering_columns`) while `payload` still holds + // its first values, but do not commit it yet. This mirrors a distributed build, + // whose segment metadata records the dataset version it was built from. + let mut params = VectorIndexParams::ivf_pq(4, 8, 2, MetricType::L2, 50); + params.covering_columns(vec!["payload".to_string()]); + let built = CreateIndexBuilder::new(&mut dataset, &["vec"], IndexType::Vector, ¶ms) + .name("vec_idx".to_string()) + .execute_uncommitted() .await .unwrap(); let vec_id = dataset.schema().field_id("vec").unwrap(); - let payload_id = dataset.schema().field_id("payload").unwrap(); - let current = dataset.load_indices().await.unwrap(); - let mut covered = current[0].clone(); - covered.fields = vec![vec_id, payload_id]; - covered.covering_fields = vec![payload_id]; - - let transaction = Transaction::new( - dataset.manifest.version, - Operation::CreateIndex { - new_indices: vec![covered], - removed_indices: current.to_vec(), + assert_eq!(built.fields, vec![vec_id, payload_id]); + assert_eq!(built.covering_fields, vec![payload_id]); + assert!( + built.fragment_bitmap.as_ref().unwrap().contains(0), + "sanity: the segment must cover fragment 0 before any rewrite" + ); + + // Rewrite the *carried* `payload` column's data file -- `vec`'s file is untouched -- + // advancing the dataset past the version the segment above was built at. + let payload_v2 = write_payload_data_file( + &dataset, + "payload_v2", + payload_id, + Int32Array::from_iter_values(1000..1512), + ) + .await; + let read_version = dataset.manifest.version; + let mut dataset = Dataset::commit( + WriteDestination::Dataset(Arc::new(dataset)), + Operation::DataReplacement { + replacements: vec![DataReplacementGroup(0, payload_v2)], }, + Some(read_version), None, - ); + None, + Arc::new(Default::default()), + false, + ) + .await + .unwrap(); + dataset - .apply_commit(transaction, &Default::default(), &Default::default()) + .commit_existing_index_segments("vec_idx", "vec", vec![built]) .await .unwrap(); - let fragments_before: Vec = dataset.fragments().iter().map(|f| f.id).collect(); - assert!( - fragments_before.len() > 1, - "precondition: there must be something to compact" - ); - - // Compaction of the table must not be blocked by an index it cannot remap. - compact_files(&mut dataset, CompactionOptions::default(), None) - .await - .expect("compaction of a covered index must succeed"); - - let fragments_after: Vec = dataset.fragments().iter().map(|f| f.id).collect(); - assert_ne!( - fragments_after, fragments_before, - "compaction rewrote nothing, so the remap path never ran" - ); - - // The entry survives untouched -- withdrawal skips remapping rather than - // deleting metadata -- but it now covers none of the rewritten fragments, so - // no query can be answered from a payload the storage never held. - let after = dataset.load_indices().await.unwrap(); - assert_eq!(after.len(), 1); - assert_eq!(after[0].covering_fields, vec![payload_id]); - let live: roaring::RoaringBitmap = dataset.fragments().iter().map(|f| f.id as u32).collect(); - let effective = after[0].effective_fragment_bitmap(&live); + let committed = dataset.load_indices_by_name("vec_idx").await.unwrap(); + assert_eq!(committed.len(), 1); + let coverage = committed[0].fragment_bitmap.as_ref().unwrap(); assert!( - effective.is_none_or(|bitmap| bitmap.is_empty()), - "a withdrawn covered index must stop covering fragments, got {:?}", - after[0].fragment_bitmap + !coverage.contains(0), + "fragment 0's carried `payload` column was rewritten after the segment was built \ + at dataset version {}; its coverage must be pruned, got {coverage:?}", + committed[0].dataset_version, ); } diff --git a/rust/lance/src/dataset/tests/dataset_scanner.rs b/rust/lance/src/dataset/tests/dataset_scanner.rs index 97cecaa6245..22b9e8e69d2 100644 --- a/rust/lance/src/dataset/tests/dataset_scanner.rs +++ b/rust/lance/src/dataset/tests/dataset_scanner.rs @@ -282,6 +282,7 @@ async fn test_vector_filter_fts_search() { query_parallelism: DEFAULT_QUERY_PARALLELISM, dist_q_c: 0.0, approx_mode: Default::default(), + covering_projection: None, }; // Case 1: search with prefilter=true, query_filter=vector([300,300,300,300]) diff --git a/rust/lance/src/dataset/write/merge_insert.rs b/rust/lance/src/dataset/write/merge_insert.rs index 7234e8baa0b..f02f1f24a3f 100644 --- a/rust/lance/src/dataset/write/merge_insert.rs +++ b/rust/lance/src/dataset/write/merge_insert.rs @@ -50,6 +50,7 @@ use super::{ use crate::dataset::rowids::get_row_id_index; use crate::dataset::transaction::UpdateMode::{RewriteColumns, RewriteRows}; use crate::dataset::utils::CapturedRowIds; +use crate::dataset::{ProjectionRequest, TakeBuilder}; use crate::index::DatasetIndexExt; use crate::{ Dataset, @@ -1520,6 +1521,337 @@ impl MergeInsertJob { self.create_full_table_joined_stream(source).await } + /// Names of `source_schema` columns this merge would update that are covered + /// ("included") by some index. Patching such a column in place makes the index's + /// materialized copy of it stale, which forces the commit to drop the whole updated + /// fragment from every index covering it; the partial-schema path therefore rewrites + /// those updates as row-moves, where it can, so that invalidation is confined to the + /// rows actually updated. See the branch in `execute_uncommitted_impl` for what that + /// does and does not buy, and `covered_move_blocker` for the cases that fall back to + /// the in-place patch. Join keys are matched, not updated, so they are excluded. + async fn covered_columns_updated(&self, source_schema: &Schema) -> Result> { + // Only the update-matched modes actually patch matched rows into fragments; + // DoNothing / Fail leave matched rows untouched (Delete is handled separately + // and never reaches this path). Without this gate, a partial source carrying a + // covered payload under the default DoNothing + InsertAll would be wrongly + // flagged and rejected even though no covered column is ever patched in place. + if !matches!( + self.params.when_matched, + WhenMatched::UpdateAll | WhenMatched::UpdateIf(_) | WhenMatched::UpdateIfExpr(_) + ) { + return Ok(Vec::new()); + } + let indices = self.dataset.load_indices().await?; + let covered: HashSet = indices + .iter() + .flat_map(|idx| idx.covering_fields.iter().copied()) + .collect(); + if covered.is_empty() { + return Ok(Vec::new()); + } + let target = self.dataset.schema(); + let keys: HashSet<&str> = self.params.on.iter().map(|s| s.as_str()).collect(); + let mut updated = Vec::new(); + for field in source_schema.fields() { + let name = field.name(); + if keys.contains(name.as_str()) { + continue; + } + if let Some(tf) = target.field(name) + && covered.contains(&tf.id) + { + updated.push(name.clone()); + } + } + Ok(updated) + } + + /// Why `rewrite_covered_update_as_move` cannot handle this merge, if it cannot. + /// + /// Every case below is a limitation of the *move*, never of the operation: the in-place + /// patch handles all of them, and it is correct for covered columns (see the caller -- + /// the suffix-subset shape already drops the fragment from the index, so no stale copy + /// is reachable). The move is an optimisation, and an optimisation must not be able to + /// fail an operation that would otherwise succeed, so the caller logs this reason and + /// patches in place, forfeiting only confinement. Returns the reason so the log can name + /// it; the phrasing is a "because ..." clause. + fn covered_move_blocker(&self, source_schema: &Schema) -> Option { + // The move would corrupt row identity. Moving a row to a new fragment must carry its + // stable row id along, and the full-schema `RewriteRows` path does exactly that -- it + // rechunks the merger's captured `row_id_sequence` onto the new fragments and stamps + // `row_id_meta` (see the `updating_row_ids` block in `execute_uncommitted_impl`). + // `rewrite_covered_update_as_move` has no such step, so the commit would reach + // `Transaction::assign_row_ids` and mint *fresh* ids for every moved row, silently + // breaking the one guarantee stable row ids exist to provide. The in-place patch moves + // no rows at all, so it preserves those ids by construction -- strictly safer here than + // the move, not less safe. + if self.dataset.manifest.uses_stable_row_ids() { + return Some( + "the move cannot carry stable row ids onto the new fragments (it has no \ + `row_id_sequence` rechunk step), while the in-place patch preserves them by \ + moving no rows" + .into(), + ); + } + // The move would delete rows the user never touched. With `insert_not_matched`, + // `Merger::execute_batch` pushes a *second* batch holding the unmatched (`left_only`) + // rows, projected with the row-address column appended. Those rows come from the left + // side of an outer join, so their `_rowaddr` is NULL. The move extracts addresses with + // `batch[ROW_ADDR].as_primitive::().values()`, which reads the raw value + // buffer and *ignores the null mask* -- so every inserted row would contribute a bogus + // address (in practice 0), used both to take the wrong row and, via `moved_addrs`, to + // tombstone it. Never entering the move prevents that as completely as refusing the + // merge did, and `update_fragments` already routes the NULL-`_rowaddr` group to + // `handle_new_fragments`. + if self.params.insert_not_matched { + return Some( + "the move reads row addresses through the null mask and would tombstone \ + unrelated rows for the unmatched (inserted) group, which arrives with a NULL \ + `_rowaddr`" + .into(), + ); + } + // The move re-reads full rows via a take, which materializes legacy (v1) blob columns + // as description structs, not binary -- writing those back would corrupt the blob. The + // full-schema path reads them as binary via the scan provider (`SomeBlobsBinary`); the + // take path has no equivalent. The in-place patch never reads the blob column at all. + // + // Scanning the *whole dataset schema* is deliberate, not sloppy. Do not narrow this to + // the covered or updated columns: the take reads every column, so an unrelated v1 blob + // elsewhere in the schema is corrupted just the same. In fact a covered column can + // never itself be a blob -- index creation rejects blob covering columns outright + // (`crate::index::vector`, the `covering_columns` validation) -- so this guard is + // entirely about the *collateral* columns the full-row take drags in. Narrowing it + // would therefore disable it completely and reintroduce the corruption. + if self + .dataset + .schema() + .fields_pre_order() + .any(|f| f.is_blob() && !f.is_blob_v2()) + { + return Some( + "the move's full-row take reads legacy (v1) blob columns as description \ + structs rather than binary, which would corrupt them on write back" + .into(), + ); + } + // The move overlays each source column onto a full-row take by REPLACING the whole + // column, which needs the source column's Arrow type to equal the target's exactly: + // `replace_column_by_name` keeps the target schema and revalidates through + // `RecordBatch::try_new`. The partial-schema contract does not promise that. A source + // is admitted as a *subschema*, and `compare_fields` recurses with the same options, + // so a column may legally differ in its struct children, in nested nullability, or in + // nested field metadata -- all three are part of `DataType` and none is compared by + // the check that admitted the source. Note a field's OWN nullability is not part of + // its `DataType`, so the ordinary "nullable source, non-null target" partial update + // does not land here. + let target_arrow = Schema::from(self.dataset.schema()); + let type_mismatched = source_schema + .fields() + .iter() + .filter(|f| { + target_arrow + .field_with_name(f.name()) + .is_ok_and(|tf| tf.data_type() != f.data_type()) + }) + .map(|f| f.name().clone()) + .collect::>(); + if !type_mismatched.is_empty() { + return Some(format!( + "the move replaces whole columns and source column(s) {type_mismatched:?} do \ + not have exactly the target's type (struct children, nested nullability or \ + nested metadata may differ)" + )); + } + None + } + + /// Rewrite a partial-schema covered-column update as a row-move: read the full + /// existing rows for the touched addresses, overlay the updated columns, write + /// them as new fragments, and tombstone the originals. Only the moved rows leave + /// the index's reach; the rest of the fragment keeps its coverage, which is the + /// point of the move (a recall/cost property, not a correctness one -- see the + /// caller for why). + /// + /// `source` is the merger's partial output stream (`_rowaddr` + updated cols); it is consumed + /// in bounded batches so peak memory is one batch of full rows rather than the whole update. + /// Mirrors `dataset.update()`'s streamed `RewriteRows` move (`update.rs::execute_impl`). + /// Returns the resulting `RewriteRows` operation together with the moved rows' addresses, so + /// the caller can publish them as `affected_rows` and let a concurrent delete/update rebase + /// instead of hard-conflicting. + /// + /// Non-stable-row-id datasets only. The caller **falls back to the in-place path** for + /// stable-row-id datasets, insert-carrying combinations, legacy v1 blob columns, and + /// partial struct-subschema sources — it does not reject them. That fallback branch is + /// live and load-bearing, not dead code: the in-place path is correct for covered columns + /// (it drops the updated fragment from the index via `prune_updated_fields_from_indices`), + /// so this move is an optimisation that confines the bitmap damage, and an optimisation + /// must not be able to fail the operation. Restoring a hard rejection here would + /// re-break merge shapes that work today. + async fn rewrite_covered_update_as_move( + &self, + source: SendableRecordBatchStream, + target_bases_info: Option>, + ) -> Result<(Operation, RoaringTreemap)> { + let full_schema = self.dataset.schema().clone(); + // EVERY field id, leaves included (`fields_pre_order`), never just the + // top-level ids: the consumer + // (`register_pure_rewrite_rows_update_frags_in_indices`) treats this list + // as "fields whose values were updated" and compares it against each + // index's leaf-expanded dependency set -- its contract for a caller that + // cannot prove any field survived unchanged is to pass *every* field id. + // An index keyed on a nested field stores a leaf id that a top-level-only + // set does not contain, so it would escape the exclusion and have its + // bitmap wrongly extended onto the moved fragment. Both sibling row-move + // producers (`update.rs`, the insert exec) already expand to leaves for + // this exact reason. + let preserving: Vec = full_schema + .fields_pre_order() + .map(|f| f.id as u32) + .collect(); + let empty_op = |fields_bitmap: Vec| Operation::Update { + removed_fragment_ids: Vec::new(), + updated_fragments: Vec::new(), + new_fragments: Vec::new(), + fields_modified: Vec::new(), + compacted_sstables: self.params.compacted_sstables.clone(), + fields_for_preserving_frag_bitmap: fields_bitmap, + update_mode: Some(RewriteRows), + inserted_rows_filter: None, + updated_fragment_offsets: None, + }; + + // Per merger batch: take the full existing rows for its addresses and overlay the updated + // columns, streaming the result so peak memory is one batch of full rows rather than the + // whole update. On the partial-schema path the merger appends `_rowaddr` to its output + // (`Merger::output_schema`), so every batch names the originals to tombstone; the take's + // result is built from the dataset schema and never carries that column onward. + let dataset = self.dataset.clone(); + let plan = Arc::new( + ProjectionRequest::from_schema(full_schema.clone()) + .into_projection_plan(dataset.clone())?, + ); + let out_arrow = Arc::new(Schema::from(&full_schema)); + + // Record which originals to tombstone as the overlay runs. `make_rowid_capture_stream` is + // deliberately not reused here: its `AddressStyle` capture bulk-loads a `RoaringTreemap` + // with `append`, which rejects any value <= the running maximum, so it is only safe on a + // stream whose addresses are strictly ascending. That holds today, but only as an accident + // of routing: this function is reachable only via the indexed-scan join (every other + // partial source goes to v2 -- see the caller), and that scan reads in fragment order, so + // addresses arrive ascending however the user ordered the source. Swapping this `extend` + // for `append` accordingly passes the covered suite, reversed source included. `extend` is + // kept anyway: it costs nothing, and it means a future routing or planner change cannot + // turn an ordering assumption into an abort that fires *after* fragments are written. + let moved_addrs = Arc::new(Mutex::new(RoaringTreemap::new())); + let moved_addrs_for_stream = moved_addrs.clone(); + let out_arrow_for_stream = out_arrow.clone(); + // `.then` processes one batch at a time (fully bounded); a small `.buffered(k)` could later + // pipeline the take of the next batch with the write of the current one. + let overlaid = source.then(move |batch| { + let dataset = dataset.clone(); + let plan = plan.clone(); + let moved_addrs = moved_addrs_for_stream.clone(); + async move { + let batch = batch?; + let addresses: Vec = batch[ROW_ADDR] + .as_primitive::() + .values() + .to_vec(); + let full_rows = + TakeBuilder::try_new_from_addresses(dataset, addresses.clone(), plan)? + .execute() + .await?; + let mut overlaid = full_rows; + for field in batch.schema().fields() { + let name = field.name(); + if name == ROW_ADDR || name == ROW_ID { + continue; + } + let new_col = batch + .column_by_name(name) + .expect("updated column present in merger output") + .clone(); + overlaid = overlaid.replace_column_by_name(name, new_col)?; + } + // Only after the take succeeded: a batch that never reaches the writer must not + // tombstone its originals. + moved_addrs + .lock() + .expect("covered move address set poisoned") + .extend(addresses); + Ok::<_, DataFusionError>(overlaid) + } + }); + let write_stream = Box::pin(RecordBatchStreamAdapter::new( + out_arrow_for_stream, + overlaid, + )); + + // Write the moved rows as new fragments. + // Retain the target-base info: `write_fragments_internal` consumes it, but a later cleanup + // must still resolve fragments routed to external bases (otherwise `cleanup_data_fragments` + // skips every file with a `base_id`, orphaning it). + let bases_for_cleanup = target_bases_info.clone(); + let (new_fragments, _) = write_fragments_internal( + self.dataset + .manifest() + .data_storage_format + .lance_file_format(), + Some(&self.dataset), + self.dataset.object_store.clone(), + &self.dataset.base, + full_schema.clone(), + write_stream, + WriteParams::default(), + target_bases_info, + ) + .await?; + + // The overlay recorded every address it handed to the writer; the write above has now + // drained the stream, so the set is complete. + let removed_row_addrs = std::mem::take( + &mut *moved_addrs + .lock() + .expect("covered move address set poisoned"), + ); + if removed_row_addrs.is_empty() { + // Nothing matched -> nothing written; return a no-op update. + return Ok((empty_op(preserving), RoaringTreemap::new())); + } + + // Tombstone the originals. + let deletions = Self::apply_deletions(&self.dataset, &removed_row_addrs).await; + let (updated_fragments, removed_fragment_ids) = match deletions { + Ok(v) => v, + Err(e) => { + cleanup_data_fragments( + &self.dataset.object_store, + &self.dataset.base, + bases_for_cleanup.as_deref(), + &new_fragments, + ) + .await; + return Err(e); + } + }; + + Ok(( + Operation::Update { + removed_fragment_ids, + updated_fragments, + new_fragments, + fields_modified: Vec::new(), + compacted_sstables: self.params.compacted_sstables.clone(), + fields_for_preserving_frag_bitmap: preserving, + update_mode: Some(RewriteRows), + inserted_rows_filter: None, + updated_fragment_offsets: None, + }, + removed_row_addrs, + )) + } + /// Patches the columns carried by `source` into the fragments that hold the /// rows it names, and writes the rows with no target address as new /// fragments. @@ -2736,7 +3068,7 @@ impl MergeInsertJob { let joined = self.create_joined_stream(source).await?; let merger = Merger::try_new( self.params.clone(), - source_schema, + source_schema.clone(), !is_full_schema, self.dataset.manifest.uses_stable_row_ids(), )?; @@ -2818,38 +3150,108 @@ impl MergeInsertJob { return Err(Error::not_supported_source("Deleting rows from the target table when there is no match in the source table is not supported when the source data has a different schema than the target data".into())); } - // We will have a different commit path here too, as we are modifying - // fragments rather than writing new ones - let PatchedFragments { - updated_fragments, - new_fragments, - fields_modified, - matched_offsets, - } = Self::update_fragments( - self.dataset.clone(), - Box::pin(stream), - self.dataset.manifest.version + 1, - target_bases_info, - ) - .await?; - - let operation = Operation::Update { - removed_fragment_ids: Vec::new(), - updated_fragments, - new_fragments, - fields_modified, - compacted_sstables: self.params.compacted_sstables.clone(), - fields_for_preserving_frag_bitmap: vec![], // in-place update do not affect preserving frag bitmap - update_mode: Some(RewriteColumns), - inserted_rows_filter: None, // not implemented for v1 - // The version stamped above is a guess; carry the patched offsets - // so `build_manifest` can re-stamp them at the real commit - // version after a rebase. - updated_fragment_offsets: Some(matched_offsets), + // A covered column's values are materialized in index storage, so patching one in + // place leaves that copy behind. Correctness is NOT what the row-move buys, and the + // comment here used to claim otherwise: `covering_fields` is stored as a suffix of + // `IndexMetadata::fields`, so `Transaction::index_dependent_leaf_ids` already counts + // a covered column as index-dependent and `prune_updated_fields_from_indices` drops + // the whole updated fragment from every index covering it. Disabling the move + // entirely leaves every covered correctness test in this module passing, both + // optimize variants included -- the stale copy is already unreachable, and no + // append-optimize failure mode could be reproduced. + // + // What the move buys is CONFINEMENT. That prune is all-or-nothing per fragment: one + // partial update evicts the ENTIRE fragment from the vector index, including rows the + // merge never matched, and the vectors are the expensive part to rebuild. The move + // tombstones only the matched rows and re-appends them, so the fragment keeps its + // coverage for everything else. Measured on the fixture in + // `test_covered_update_is_fresh_and_confines_index_invalidation`: with the move the + // index bitmap stays `{0, 1}` over fragments `[0, 1, 2]`; patched in place it + // collapses to `{0}`. A recall/cost property rather than a correctness one, then -- + // and that test is the only thing pinning it. + // + // It also makes the two merge_insert paths AGREE, which is the stronger argument. A + // partial source only reaches this branch when every join key has a scalar index + // (`can_use_create_plan` routes the rest to v2), and v2's `create_plan` fills every + // missing dataset column from the target side and writes FULL rows -- already a + // row-move, already confining. Without this branch the same merge would invalidate a + // whole fragment or only the matched rows depending on whether an unrelated scalar + // index happened to exist on the join key. + // + // Known cost, deliberately not addressed here: there is no match-ratio heuristic, so + // when the matched rows are (nearly) a whole fragment the move rewrites every column + // including the vectors for zero confinement gain, where the in-place patch writes one + // column file for the same coverage outcome. Adding one is not cheap -- the strategy + // must be chosen before the stream is consumed, and the per-fragment match + // distribution is not known until the merger has run, so it would take either a + // pre-pass scan or buffering the whole merge, and the latter breaks the bounded-memory + // property this path is built around. + let covered = self.covered_columns_updated(source_schema.as_ref()).await?; + // The move cannot handle every merge (see `covered_move_blocker`). Because it is an + // optimisation over an already-correct in-place patch, a case it cannot handle + // degrades to that patch instead of failing the merge -- otherwise adding + // `covering_columns` to a vector index would silently delete merge_insert shapes that + // worked before, and only on the half of the routing that reaches this code at all. + let move_blocked = if covered.is_empty() { + None + } else { + self.covered_move_blocker(source_schema.as_ref()) }; - // We have rewritten the fragments, not just the deletion files, so - // we can't use affected rows here. - (operation, None) + if !covered.is_empty() && move_blocked.is_none() { + let (operation, removed_row_addrs) = self + .rewrite_covered_update_as_move( + Box::pin(stream) as SendableRecordBatchStream, + target_bases_info, + ) + .await?; + // Same shape as the full-schema `RewriteRows` paths below -- deletions on the + // originals plus new fragments -- so publish the moved addresses. Without them + // `check_update_txn` turns any concurrent delete/update touching these fragments + // into a hard conflict, costing covered datasets concurrency that non-covered + // ones keep. + (operation, Some(RowAddrTreeMap::from(removed_row_addrs))) + } else { + if let Some(reason) = &move_blocked { + info!( + "merge_insert is updating covered column(s) {covered:?} in place rather \ + than moving the matched rows, because {reason}. The updated fragment(s) \ + lose their coverage in every index covering those columns until the next \ + optimize; the merge itself is unaffected." + ); + } + // We will have a different commit path here too, as we are modifying + // fragments rather than writing new ones + let PatchedFragments { + updated_fragments, + new_fragments, + fields_modified, + matched_offsets, + } = Self::update_fragments( + self.dataset.clone(), + Box::pin(stream), + self.dataset.manifest.version + 1, + target_bases_info, + ) + .await?; + + let operation = Operation::Update { + removed_fragment_ids: Vec::new(), + updated_fragments, + new_fragments, + fields_modified, + compacted_sstables: self.params.compacted_sstables.clone(), + fields_for_preserving_frag_bitmap: vec![], // in-place update do not affect preserving frag bitmap + update_mode: Some(RewriteColumns), + inserted_rows_filter: None, // not implemented for v1 + // The version stamped above is a guess; carry the patched offsets + // so `build_manifest` can re-stamp them at the real commit + // version after a rebase. + updated_fragment_offsets: Some(matched_offsets), + }; + // We have rewritten the fragments, not just the deletion files, so + // we can't use affected rows here. + (operation, None) + } } else { let cleanup_bases = target_bases_info.clone(); let (mut new_fragments, _) = write_fragments_internal( @@ -13248,7 +13650,7 @@ MergeInsert: on=[id], when_matched=DoNothing, when_not_matched=InsertAll, when_n mod external_error { use super::*; - use arrow_schema::{ArrowError, Field as ArrowField, Schema as ArrowSchema}; + use arrow_schema::{ArrowError, Field as ArrowField, Schema}; use std::fmt; #[derive(Debug)] @@ -13267,7 +13669,7 @@ MergeInsert: on=[id], when_matched=DoNothing, when_not_matched=InsertAll, when_n #[tokio::test] async fn test_merge_insert_execute_reader_preserves_error_message() { - let schema = Arc::new(ArrowSchema::new(vec![ + let schema = Arc::new(Schema::new(vec![ ArrowField::new("key", DataType::Int32, false), ArrowField::new("value", DataType::Int32, false), ])); @@ -13699,16 +14101,14 @@ MergeInsert: on=[id], when_matched=DoNothing, when_not_matched=InsertAll, when_n assert_eq!(combined.num_rows(), 300); } - // Regression test: after a partial-schema merge_insert drops a fragment from the vector - // index bitmap, a vector search should not return duplicate rows. The stale vector index - // data still references the dropped fragment, and the scanner also flat-scans unindexed - // fragments, causing the same rows to appear from both paths. - #[tokio::test] - async fn test_partial_merge_insert_stale_vector_index_duplicates() { + /// Build a dataset with a plain (non-covered) IVF_FLAT vector index whose fragment 1 + /// has been pruned from the index bitmap by a partial-schema merge_insert that patched + /// the indexed vector column in place -- the pruned fragment's stale rows still live in + /// the old index segment's storage. Returns `(dataset, total_rows, rows_per_frag, dim)`. + async fn build_pruned_plain_vector_dataset(uri: &str) -> (Arc, usize, usize, i32) { let dim = 4i32; let rows_per_frag = 10usize; let num_frags = 3usize; - let total_rows = rows_per_frag * num_frags; let schema = Arc::new(Schema::new(vec![ Field::new("id", DataType::Utf8, false), @@ -13745,9 +14145,7 @@ MergeInsert: on=[id], when_matched=DoNothing, when_not_matched=InsertAll, when_n // Write 3 fragments let batch0 = make_batch(0, 0.0); let reader = Box::new(RecordBatchIterator::new([Ok(batch0)], schema.clone())); - let mut ds = Dataset::write(reader, "memory://vector_stale_test", None) - .await - .unwrap(); + let mut ds = Dataset::write(reader, uri, None).await.unwrap(); for frag_idx in 1..num_frags { let batch = make_batch(frag_idx, 0.0); let reader = Box::new(RecordBatchIterator::new([Ok(batch)], schema.clone())); @@ -13797,7 +14195,18 @@ MergeInsert: on=[id], when_matched=DoNothing, when_not_matched=InsertAll, when_n .await .unwrap(); - // KNN search with k = total_rows to retrieve all rows + (ds, rows_per_frag * num_frags, rows_per_frag, dim) + } + + /// A KNN query over every row must return each id exactly once -- no stale duplicate + /// served from old index storage alongside the fresh copy. + async fn assert_knn_no_duplicate_ids( + ds: &Dataset, + total_rows: usize, + rows_per_frag: usize, + dim: i32, + ) { + let frag1_start = rows_per_frag; let query: Float32Array = (0..dim) .map(|i| (frag1_start * dim as usize + i as usize) as f32 + 0.5) .collect(); @@ -13809,7 +14218,6 @@ MergeInsert: on=[id], when_matched=DoNothing, when_not_matched=InsertAll, when_n .await .unwrap(); - // Check no duplicate ids let ids = results .column_by_name("id") .unwrap() @@ -13827,85 +14235,191 @@ MergeInsert: on=[id], when_matched=DoNothing, when_not_matched=InsertAll, when_n ); } - // Regression test: after a partial-schema merge_insert drops a fragment from the FTS - // index bitmap, a full text search should not return duplicate rows. The stale inverted - // index data still references the dropped fragment, and the scanner also flat-scans - // unindexed fragments, causing the same rows to appear from both paths. + // Regression test: after a partial-schema merge_insert drops a fragment from the vector + // index bitmap, a vector search should not return duplicate rows. The stale vector index + // data still references the dropped fragment, and the scanner also flat-scans unindexed + // fragments, causing the same rows to appear from both paths. #[tokio::test] - async fn test_partial_merge_insert_stale_fts_index_duplicates() { - let rows_per_frag = 10usize; - let num_frags = 3usize; + async fn test_partial_merge_insert_stale_vector_index_duplicates() { + let (ds, total_rows, rows_per_frag, dim) = + build_pruned_plain_vector_dataset("memory://vector_stale_test").await; + assert_knn_no_duplicate_ids(&ds, total_rows, rows_per_frag, dim).await; + } + + /// Same pruned-fragment state, but merged back into the index: + /// `optimize_indices(merge)` re-scans the pruned fragment as unindexed data AND + /// carries the fragment's stale rows from the old segment's storage -- the merge + /// must drop the stale copies for plain (non-covered) indexes just as it does for + /// covered ones, or ANN queries return duplicate row ids from the merged index. + #[tokio::test] + async fn test_optimize_merge_after_partial_update_no_duplicates_non_covered() { + use lance_index::optimize::OptimizeOptions; + let (ds, total_rows, rows_per_frag, dim) = + build_pruned_plain_vector_dataset("memory://vector_stale_opt_merge").await; + let mut ds = Arc::try_unwrap(ds).unwrap_or_else(|arc| (*arc).clone()); + ds.optimize_indices(&OptimizeOptions::merge(1)) + .await + .unwrap(); + assert_knn_no_duplicate_ids(&ds, total_rows, rows_per_frag, dim).await; + } + + /// A partial merge_insert that rewrites only a covered column is a row-MOVE: the + /// touched rows are tombstoned at their old addresses and reinserted into a new + /// fragment, so the index's stale copy is hidden by the deletion mask rather than + /// patched in place. Fragment 1 is updated in full, so it disappears entirely, and a + /// covered query must observe the fresh values for every row. + #[tokio::test] + async fn test_partial_merge_insert_moves_covered_column_rows() { + let (ds, total_rows, rows_per_frag, dim) = + build_moved_covering_dataset("memory://covering_move_test").await; + assert!( + !ds.get_fragments().iter().any(|f| f.id() == 1), + "fragment 1's rows should have been moved out (tombstoned + reinserted), \ + but fragment 1 is still present" + ); + + assert_covered_query_no_duplicates_fresh(&ds, total_rows, rows_per_frag, dim).await; + } + + /// Pins the two things a partial merge_insert on a covered column must deliver: covered reads + /// are FRESH with no intervening optimize, and the index invalidation is confined to the rows + /// actually updated. + /// + /// The second half is what the row-move uniquely buys, and it is the half that discriminates. + /// Freshness alone does not: `covering_fields` is stored as a suffix of `IndexMetadata::fields`, + /// so `prune_updated_fields_from_indices` already counts a covered column as index-dependent + /// and drops the whole updated fragment from the index whenever a `RewriteColumns` update + /// touches one. That backstop hides the stale copy on its own -- verified by disabling the + /// row-move entirely, which leaves every covered correctness assertion in this module passing. + /// What it costs is coverage: an in-place update of one payload column evicts the ENTIRE + /// fragment from the vector index, including rows the merge never touched, and the vectors are + /// the expensive part to rebuild. The row-move instead tombstones only the matched rows and + /// re-appends them, so the fragment keeps its coverage for everything else. Measured here: + /// with the move the bitmap stays `{0, 1}` over fragments `[0, 1, 2]`; patched in place it + /// collapses to `{0}`. + /// + /// Freshness is still checked, against an INDEPENDENT `take_rows` from the base table rather + /// than against a formula this test also wrote -- asserting that a column came back, or that + /// the query returned something, passes just as happily against a stale copy. `take_rows` + /// re-orders its output to match the requested row ids (see the remapping branch in + /// `dataset::take`), so the two reads line up row for row. + /// + /// Shape follows the recipe in `dataset::tests::dataset_index`: four well-separated clusters, + /// since a covered-ANN test with one partition never reaches the probe path at all and passes + /// against a broken covered read. + #[tokio::test] + async fn test_covered_update_is_fresh_and_confines_index_invalidation() { + use arrow_array::cast::AsArray; + use arrow_array::types::UInt64Type; + use lance_core::ROW_ID; + use std::collections::HashSet; + + const DIMS: i32 = 16; + const NUM_CLUSTERS: usize = 4; + const ROWS_PER_CLUSTER: usize = 64; + const TOTAL: usize = NUM_CLUSTERS * ROWS_PER_CLUSTER; + // Only the last cluster is rewritten, so one query spans both moved rows and untouched + // rows still served from the original index storage. + const UPDATED_FROM: i32 = ((NUM_CLUSTERS - 1) * ROWS_PER_CLUSTER) as i32; let schema = Arc::new(Schema::new(vec![ - Field::new("id", DataType::Utf8, false), - Field::new("category", DataType::Utf8, false), - Field::new("text", DataType::Utf8, false), + Field::new("id", DataType::Int32, false), + Field::new("payload", DataType::Int32, false), + Field::new( + "vector", + DataType::FixedSizeList( + Arc::new(Field::new("item", DataType::Float32, true)), + DIMS, + ), + false, + ), ])); - let make_batch = |frag_idx: usize| { - let start = frag_idx * rows_per_frag; - let ids: Vec = (start..start + rows_per_frag) - .map(|j| format!("id-{j:04}")) - .collect(); - let cats: Vec<&str> = vec!["A"; rows_per_frag]; - // Every row contains "common" so we can search for it and expect all rows - let texts: Vec = (start..start + rows_per_frag) - .map(|j| format!("common unique{j:04}")) - .collect(); + // Cluster centers 1000 apart on dim 0 with deterministic jitter: uniform-random vectors + // leave the early-pruning heuristic searching every partition. + let make_batch = |clusters: std::ops::Range| { + let mut ids = Vec::new(); + let mut payloads = Vec::new(); + let mut values = Vec::new(); + for cluster in clusters { + let center = (cluster * 1000) as f32; + for row in 0..ROWS_PER_CLUSTER { + let n = cluster * ROWS_PER_CLUSTER + row; + ids.push(n as i32); + payloads.push(n as i32 * 10); + for dim in 0..DIMS as usize { + let base = if dim == 0 { center } else { 0.0 }; + values.push(base + ((n * 7919 + dim) % 97) as f32 * 0.0001); + } + } + } + let vectors = + FixedSizeListArray::try_new_from_values(Float32Array::from(values), DIMS).unwrap(); RecordBatch::try_new( schema.clone(), vec![ - Arc::new(StringArray::from(ids)), - Arc::new(StringArray::from(cats)), - Arc::new(StringArray::from(texts)), + Arc::new(Int32Array::from(ids)), + Arc::new(Int32Array::from(payloads)), + Arc::new(vectors), ], ) .unwrap() }; - // Write 3 fragments - let batch0 = make_batch(0); - let reader = Box::new(RecordBatchIterator::new([Ok(batch0)], schema.clone())); - let mut ds = Dataset::write(reader, "memory://fts_stale_test", None) + // Two fragments, so the move has to reach across more than one. + let reader = Box::new(RecordBatchIterator::new( + [Ok(make_batch(0..2))], + schema.clone(), + )); + let mut ds = Dataset::write(reader, "memory://covered_fresh_no_optimize", None) .await .unwrap(); - for frag_idx in 1..num_frags { - let batch = make_batch(frag_idx); - let reader = Box::new(RecordBatchIterator::new([Ok(batch)], schema.clone())); - ds.append(reader, None).await.unwrap(); - } + let reader = Box::new(RecordBatchIterator::new( + [Ok(make_batch(2..NUM_CLUSTERS))], + schema.clone(), + )); + ds.append(reader, None).await.unwrap(); - // Create inverted index on text - let params = InvertedIndexParams::default(); - ds.create_index(&["text"], IndexType::Inverted, None, ¶ms, true) + ds.create_index( + &["id"], + IndexType::BTree, + None, + &ScalarIndexParams::default(), + false, + ) + .await + .unwrap(); + let mut params = VectorIndexParams::ivf_pq(NUM_CLUSTERS, 8, 4, MetricType::L2, 2); + params.covering_columns(vec!["payload".to_string()]); + ds.create_index(&["vector"], IndexType::Vector, None, ¶ms, false) .await .unwrap(); - let ds = Arc::new(ds); - // Partial merge_insert with (id, text) on fragment 1 rows. - // Text still contains "common" so FTS will find them via both paths. - // This drops fragment 1 from the inverted index bitmap. - let frag1_start = rows_per_frag; - let ids: Vec = (frag1_start..frag1_start + rows_per_frag) - .map(|j| format!("id-{j:04}")) - .collect(); - let texts: Vec = (frag1_start..frag1_start + rows_per_frag) - .map(|j| format!("common updated{j:04}")) - .collect(); + // Precondition the assertion depends on: the index really does carry `payload`, so a + // projection of just `payload` is answered from index storage. + let payload_field_id = ds.schema().field("payload").unwrap().id; + let indices = ds.load_indices().await.unwrap(); + let vec_idx = indices.iter().find(|i| i.name == "vector_idx").unwrap(); + assert_eq!(vec_idx.covering_fields, vec![payload_field_id]); + + // Partial-schema merge_insert rewriting ONLY the covered column, for the last cluster. let sub_schema = Arc::new(Schema::new(vec![ - Field::new("id", DataType::Utf8, false), - Field::new("text", DataType::Utf8, false), + Field::new("id", DataType::Int32, false), + Field::new("payload", DataType::Int32, false), ])); + let updated_ids: Vec = (UPDATED_FROM..TOTAL as i32).collect(); + let updated_payloads: Vec = updated_ids.iter().map(|id| id * 10 + 7).collect(); let update_batch = RecordBatch::try_new( sub_schema.clone(), vec![ - Arc::new(StringArray::from(ids)), - Arc::new(StringArray::from(texts)), + Arc::new(Int32Array::from(updated_ids)), + Arc::new(Int32Array::from(updated_payloads)), ], ) .unwrap(); let reader = Box::new(RecordBatchIterator::new([Ok(update_batch)], sub_schema)); + let frags_before: HashSet = ds.get_fragments().iter().map(|f| f.id()).collect(); let (ds, _) = MergeInsertBuilder::try_new(ds, vec!["id".to_string()]) .unwrap() .when_matched(WhenMatched::UpdateAll) @@ -13915,17 +14429,1653 @@ MergeInsert: on=[id], when_matched=DoNothing, when_not_matched=InsertAll, when_n .execute_reader(reader) .await .unwrap(); + // Deliberately no `optimize_indices` here: the point is that the stale copy is hidden + // without one. - // FTS search for "common" — every row should match exactly once - let query = FullTextSearchQuery::new("common".to_string()); - let results = ds + // The move half: the matched rows left fragment 1 for a brand-new fragment, and + // fragment 1 -- which also holds 64 rows the merge never matched -- keeps its place in + // the vector index. An in-place patch evicts it instead. + let frags_after: HashSet = ds.get_fragments().iter().map(|f| f.id()).collect(); + assert!( + frags_after.difference(&frags_before).count() == 1, + "the matched rows should have been re-appended to one new fragment; \ + fragments went from {frags_before:?} to {frags_after:?}" + ); + let updated_frag = *frags_before.iter().max().unwrap(); + let indices = ds.load_indices().await.unwrap(); + let vec_idx = indices.iter().find(|i| i.name == "vector_idx").unwrap(); + let bitmap = vec_idx.fragment_bitmap.as_ref().unwrap(); + assert!( + bitmap.contains(updated_frag as u32), + "a covered update must not evict the whole fragment from the vector index: \ + fragment {updated_frag} still holds {ROWS_PER_CLUSTER} untouched rows, but the \ + index now covers only {bitmap:?}" + ); + + let mut q_values = vec![0.0f32; DIMS as usize]; + q_values[0] = ((NUM_CLUSTERS - 1) * 1000) as f32; + let query = Float32Array::from(q_values); + + // Project ONLY the covered column. Adding an uncovered column here would send the whole + // projection to the base table and the covered storage would never be read, so the test + // would pass against a stale copy. + let mut scan = ds.scan(); + scan.nearest("vector", &query, TOTAL).unwrap(); + scan.with_row_id(); + scan.project(&["payload"]).unwrap(); + let plan = scan.explain_plan(true).await.unwrap(); + assert!( + plan.contains("ANNIvfPartition"), + "the covered values must come from the vector index, not a flat rescan:\n{plan}" + ); + let results = scan.try_into_batch().await.unwrap(); + + assert_eq!( + results.num_rows(), + TOTAL, + "covered query returned {} rows, expected {TOTAL}", + results.num_rows() + ); + let row_ids = results[ROW_ID].as_primitive::(); + let unique: HashSet = row_ids.values().iter().copied().collect(); + assert_eq!( + unique.len(), + TOTAL, + "a stale copy survived at its old address: duplicate row ids in covered results" + ); + + // The independent read: same rows, fetched from the base table by row id. + let queried_row_ids: Vec = row_ids.values().to_vec(); + let base = ds + .take_rows( + &queried_row_ids, + ProjectionRequest::from_schema(ds.schema().project(&["payload"]).unwrap()), + ) + .await + .unwrap(); + assert_eq!(base.num_rows(), TOTAL); + + let covered_payload = results["payload"].as_primitive::(); + let base_payload = base["payload"].as_primitive::(); + let mut updated_seen = 0usize; + for i in 0..TOTAL { + assert_eq!( + covered_payload.value(i), + base_payload.value(i), + "covered read disagrees with the base table for row id {}: index says {}, \ + table says {}", + row_ids.value(i), + covered_payload.value(i), + base_payload.value(i) + ); + if base_payload.value(i) % 10 == 7 { + updated_seen += 1; + } + } + // Guards the assertion above from passing vacuously if the merge silently no-ops. + assert_eq!( + updated_seen, ROWS_PER_CLUSTER, + "expected {ROWS_PER_CLUSTER} rewritten rows in the covered results, saw {updated_seen}" + ); + + // The comparison above would still agree if the move had taken the WRONG addresses -- + // the base table would carry the same corruption, so both reads would match. Row + // integrity is what catches that: the move rebuilds each row from a full-row take at an + // address, so a wrongly addressed take pairs one row's `id` with another's payload and + // duplicates or drops ids. + let full = ds .scan() - .full_text_search(query) + .project(&["id", "payload"]) .unwrap() .try_into_batch() .await .unwrap(); - + assert_eq!(full.num_rows(), TOTAL, "the move changed the row count"); + let full_ids = full["id"].as_primitive::(); + let full_payload = full["payload"].as_primitive::(); + let mut seen_ids = HashSet::new(); + for i in 0..TOTAL { + let id = full_ids.value(i); + assert!(seen_ids.insert(id), "id {id} duplicated by the move"); + let expected = if id >= UPDATED_FROM { + id * 10 + 7 + } else { + id * 10 + }; + assert_eq!( + full_payload.value(i), + expected, + "row {id} came back paired with another row's payload" + ); + } + assert_eq!(seen_ids.len(), TOTAL, "the move dropped rows"); + } + + /// Bounded-streaming regression for the covered row-move: a covered-column update spanning + /// MORE than one scan batch (> `BATCH_SIZE_FALLBACK` = 8192 rows) must move every row exactly + /// once with its fresh value. `rewrite_covered_update_as_move` streams the take/overlay/write + /// per batch and accumulates the tombstoned addresses across batches, so this exercises the + /// multi-batch accumulation path the small covered tests never reach. + #[tokio::test] + async fn test_covered_update_move_streams_multiple_batches() { + use arrow_array::cast::AsArray; + use std::collections::HashSet; + + let dim = 2i32; + let rows_per_frag = 5000usize; // 2 frags => 10_000 rows > BATCH_SIZE_FALLBACK (8192) + let num_frags = 2usize; + let total_rows = rows_per_frag * num_frags; + + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Utf8, false), + Field::new("category", DataType::Utf8, false), + Field::new( + "vec", + DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Float32, true)), dim), + false, + ), + ])); + let make_batch = |frag_idx: usize| { + let start = frag_idx * rows_per_frag; + let ids: Vec = (start..start + rows_per_frag) + .map(|j| format!("id-{j:05}")) + .collect(); + let cats: Vec = (start..start + rows_per_frag) + .map(|j| format!("cat-{j:05}")) + .collect(); + let values: Vec = (0..rows_per_frag * dim as usize) + .map(|i| (start * dim as usize + i) as f32) + .collect(); + let vectors = + FixedSizeListArray::try_new_from_values(Float32Array::from(values), dim).unwrap(); + RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(StringArray::from(ids)), + Arc::new(StringArray::from(cats)), + Arc::new(vectors), + ], + ) + .unwrap() + }; + + let reader = Box::new(RecordBatchIterator::new( + [Ok(make_batch(0))], + schema.clone(), + )); + let mut ds = Dataset::write(reader, "memory://covered_move_multibatch", None) + .await + .unwrap(); + for frag_idx in 1..num_frags { + let reader = Box::new(RecordBatchIterator::new( + [Ok(make_batch(frag_idx))], + schema.clone(), + )); + ds.append(reader, None).await.unwrap(); + } + + // BTree on the join key + a vector index covering `category`. + ds.create_index( + &["id"], + IndexType::BTree, + None, + &ScalarIndexParams::default(), + false, + ) + .await + .unwrap(); + let mut params = VectorIndexParams::ivf_flat(2, MetricType::L2); + params.covering_columns(vec!["category".to_string()]); + ds.create_index(&["vec"], IndexType::Vector, None, ¶ms, false) + .await + .unwrap(); + let ds = Arc::new(ds); + + // Partial merge_insert rewriting `category` for EVERY row -> a >1-batch covered move. + let ids: Vec = (0..total_rows).map(|j| format!("id-{j:05}")).collect(); + let new_cats: Vec = (0..total_rows).map(|j| format!("updated-{j:05}")).collect(); + let sub_schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Utf8, false), + Field::new("category", DataType::Utf8, false), + ])); + let update_batch = RecordBatch::try_new( + sub_schema.clone(), + vec![ + Arc::new(StringArray::from(ids)), + Arc::new(StringArray::from(new_cats)), + ], + ) + .unwrap(); + let reader = Box::new(RecordBatchIterator::new([Ok(update_batch)], sub_schema)); + let (ds, _) = MergeInsertBuilder::try_new(ds, vec!["id".to_string()]) + .unwrap() + .when_matched(WhenMatched::UpdateAll) + .when_not_matched(WhenNotMatched::DoNothing) + .try_build() + .unwrap() + .execute_reader(reader) + .await + .unwrap(); + + // Every row must survive exactly once with its fresh category -- nothing lost or + // duplicated across the batch boundary the streaming move crosses. + let batch = ds + .scan() + .project(&["id", "category"]) + .unwrap() + .try_into_batch() + .await + .unwrap(); + assert_eq!( + batch.num_rows(), + total_rows, + "row count changed after a multi-batch covered move" + ); + let ids_col = batch["id"].as_string::(); + let cats_col = batch["category"].as_string::(); + let mut seen = HashSet::new(); + for i in 0..batch.num_rows() { + let id = ids_col.value(i); + assert!( + seen.insert(id.to_string()), + "duplicate id {id} after covered move" + ); + let n: usize = id[3..].parse().unwrap(); + assert_eq!( + cats_col.value(i), + format!("updated-{n:05}"), + "stale or missing covered value for {id}" + ); + } + assert_eq!(seen.len(), total_rows); + } + + /// Runs a partial-schema merge_insert whose source carries a STRUCT with only one of the + /// target's two children, alongside `category`. `cover` decides whether the vector index + /// includes `category`, i.e. whether the covered row-move branch is taken at all. + async fn partial_struct_source_merge(uri: &str, cover: bool) -> Result> { + const ROWS: usize = 256; + const DIM: i32 = 2; + let struct_fields: arrow_schema::Fields = vec![ + Field::new("a", DataType::Int32, false), + Field::new("b", DataType::Int32, false), + ] + .into(); + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Utf8, false), + Field::new("category", DataType::Utf8, false), + Field::new("s", DataType::Struct(struct_fields.clone()), false), + Field::new( + "vec", + DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Float32, true)), DIM), + false, + ), + ])); + let ids: Vec = (0..ROWS).map(|j| format!("id-{j:03}")).collect(); + let cats: Vec = (0..ROWS).map(|j| format!("cat-{j:03}")).collect(); + let a = Int32Array::from((0..ROWS as i32).collect::>()); + let b = Int32Array::from((0..ROWS as i32).map(|v| v + 1000).collect::>()); + let s = StructArray::new(struct_fields, vec![Arc::new(a), Arc::new(b)], None); + let values: Vec = (0..ROWS * DIM as usize).map(|i| i as f32).collect(); + let vectors = + FixedSizeListArray::try_new_from_values(Float32Array::from(values), DIM).unwrap(); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(StringArray::from(ids.clone())), + Arc::new(StringArray::from(cats)), + Arc::new(s), + Arc::new(vectors), + ], + ) + .unwrap(); + let reader = Box::new(RecordBatchIterator::new([Ok(batch)], schema.clone())); + let mut ds = Dataset::write(reader, uri, None).await.unwrap(); + ds.create_index( + &["id"], + IndexType::BTree, + None, + &ScalarIndexParams::default(), + false, + ) + .await + .unwrap(); + let mut params = VectorIndexParams::ivf_flat(2, MetricType::L2); + if cover { + params.covering_columns(vec!["category".to_string()]); + } + ds.create_index(&["vec"], IndexType::Vector, None, ¶ms, false) + .await + .unwrap(); + let ds = Arc::new(ds); + + // Partial source: the join key, the (maybe-covered) `category`, and a STRUCT with + // only one of its two children. + let partial_struct_fields: arrow_schema::Fields = + vec![Field::new("a", DataType::Int32, false)].into(); + let sub_schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Utf8, false), + Field::new("category", DataType::Utf8, false), + Field::new("s", DataType::Struct(partial_struct_fields.clone()), false), + ])); + let new_cats: Vec = (0..ROWS).map(|j| format!("upd-{j:03}")).collect(); + let new_a = Int32Array::from((0..ROWS as i32).map(|v| v + 500).collect::>()); + let new_s = StructArray::new(partial_struct_fields, vec![Arc::new(new_a)], None); + let update_batch = RecordBatch::try_new( + sub_schema.clone(), + vec![ + Arc::new(StringArray::from(ids)), + Arc::new(StringArray::from(new_cats)), + Arc::new(new_s), + ], + ) + .unwrap(); + let reader = Box::new(RecordBatchIterator::new([Ok(update_batch)], sub_schema)); + MergeInsertBuilder::try_new(ds, vec!["id".to_string()]) + .unwrap() + .when_matched(WhenMatched::UpdateAll) + .when_not_matched(WhenNotMatched::DoNothing) + .try_build() + .unwrap() + .execute_reader(reader) + .await + .map(|(ds, _)| ds) + } + + /// `id -> (s.a, s.b, category)` for the whole dataset, so the covered and uncovered runs of + /// `partial_struct_source_merge` can be compared row for row. + async fn partial_struct_merge_contents(ds: &Dataset) -> HashMap { + let batch = ds + .scan() + .project(&["id", "category", "s"]) + .unwrap() + .try_into_batch() + .await + .unwrap(); + let ids = batch["id"].as_string::(); + let cats = batch["category"].as_string::(); + let s = batch["s"].as_struct(); + let a = s.column_by_name("a").unwrap().as_primitive::(); + let b = s.column_by_name("b").unwrap().as_primitive::(); + (0..batch.num_rows()) + .map(|i| { + ( + ids.value(i).to_string(), + (a.value(i), b.value(i), cats.value(i).to_string()), + ) + }) + .collect() + } + + /// The covered row-move overlays each source column onto a full-row take by REPLACING the + /// whole column, so it needs the source column's Arrow type to equal the target's exactly. A + /// partial-schema source is only required to be a *subschema*, and `compare_fields` recurses + /// with the same options, so a struct may legally arrive carrying only some of its children. + /// Left unguarded this fails with an opaque Arrow "column types must match schema types" from + /// `replace_column_by_name`; guarded, it falls back to the in-place patch. + /// + /// The control is the whole point: the same merge on a dataset whose index does NOT cover + /// `category` takes the in-place branch, and both runs must produce **identical data**. That + /// pins the fallback as behaviour-preserving rather than merely non-failing -- adding + /// `covering_columns` to a vector index must not change what a merge does, only how much index + /// coverage survives it. + #[tokio::test] + async fn test_merge_insert_covered_update_falls_back_with_partial_struct_source() { + let plain = partial_struct_source_merge("memory://partial_struct_plain", false) + .await + .expect( + "control: a partial struct source is handled by the in-place branch, so if this \ + fails the comparison below is no longer evidence of anything", + ); + let covered = partial_struct_source_merge("memory://partial_struct_covered", true) + .await + .expect("a partial struct source must fall back to in-place, not fail"); + + let plain_rows = partial_struct_merge_contents(&plain).await; + let covered_rows = partial_struct_merge_contents(&covered).await; + assert_eq!(plain_rows.len(), 256); + assert_eq!( + plain_rows, covered_rows, + "covering a column must not change what the merge writes" + ); + // Spot-check against the expected values too, so an identical-but-wrong pair of runs + // cannot pass: `s.a` was updated, `s.b` was never in the source and must survive. + assert_eq!(covered_rows["id-007"], (507, 1007, "upd-007".to_string())); + + // Confinement is what covering forfeits here, and only that. + let indices = covered.load_indices().await.unwrap(); + let vec_idx = indices.iter().find(|i| i.name == "vec_idx").unwrap(); + let bitmap = vec_idx.fragment_bitmap.as_ref().unwrap(); + assert!( + bitmap.is_empty(), + "the in-place fallback should have evicted the only fragment from the vector index; \ + still covered means the fallback was not taken: {bitmap:?}" + ); + } + + /// The covered row-move must not depend on the order its source arrives in: it accumulates + /// tombstoned addresses with `extend` rather than the bulk `RoaringTreemap::append`, which + /// rejects any value <= the running maximum and would abort *after* fragments were written, + /// orphaning them. + /// + /// Note what this can and cannot reach. The BTree is required, not incidental: without a + /// scalar index on every join key, `can_use_create_plan` routes the merge to v2 and this + /// function never runs (the earlier version of this test omitted the index specifically to + /// get the hash join, and so tested nothing here). With it, the merge takes the indexed-scan + /// join -- the only route to the move -- which reads in fragment order, so descending + /// addresses cannot be produced through the public API at all. What is pinned is that a + /// reversed *source* is handled correctly; `extend` remains a deliberate defence against a + /// future routing change rather than something reachable today. + #[tokio::test] + async fn test_covered_update_move_accepts_unsorted_source_order() { + use arrow_array::cast::AsArray; + use std::collections::HashSet; + + let dim = 2i32; + let total_rows = 128usize; + + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Utf8, false), + Field::new("category", DataType::Utf8, false), + Field::new( + "vec", + DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Float32, true)), dim), + false, + ), + ])); + let ids: Vec = (0..total_rows).map(|j| format!("id-{j:04}")).collect(); + let cats: Vec = (0..total_rows).map(|j| format!("cat-{j:04}")).collect(); + let values: Vec = (0..total_rows * dim as usize).map(|i| i as f32).collect(); + let vectors = + FixedSizeListArray::try_new_from_values(Float32Array::from(values), dim).unwrap(); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(StringArray::from(ids.clone())), + Arc::new(StringArray::from(cats)), + Arc::new(vectors), + ], + ) + .unwrap(); + + let reader = Box::new(RecordBatchIterator::new([Ok(batch)], schema.clone())); + let mut ds = Dataset::write(reader, "memory://covered_move_unsorted", None) + .await + .unwrap(); + // The BTree is what makes this test reach `rewrite_covered_update_as_move` at all -- + // see the doc comment. + ds.create_index( + &["id"], + IndexType::BTree, + None, + &ScalarIndexParams::default(), + false, + ) + .await + .unwrap(); + let mut params = VectorIndexParams::ivf_flat(2, MetricType::L2); + params.covering_columns(vec!["category".to_string()]); + ds.create_index(&["vec"], IndexType::Vector, None, ¶ms, false) + .await + .unwrap(); + let ds = Arc::new(ds); + + // Feed the source fully reversed. + let mut update_ids: Vec = ids.clone(); + update_ids.reverse(); + let new_cats: Vec = update_ids + .iter() + .map(|id| format!("updated-{}", &id[3..])) + .collect(); + let sub_schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Utf8, false), + Field::new("category", DataType::Utf8, false), + ])); + let update_batch = RecordBatch::try_new( + sub_schema.clone(), + vec![ + Arc::new(StringArray::from(update_ids)), + Arc::new(StringArray::from(new_cats)), + ], + ) + .unwrap(); + let reader = Box::new(RecordBatchIterator::new([Ok(update_batch)], sub_schema)); + let frags_before: HashSet = ds.get_fragments().iter().map(|f| f.id()).collect(); + let (ds, _) = MergeInsertBuilder::try_new(ds, vec!["id".to_string()]) + .unwrap() + .when_matched(WhenMatched::UpdateAll) + .when_not_matched(WhenNotMatched::DoNothing) + .try_build() + .unwrap() + .execute_reader(reader) + .await + .unwrap(); + + // Proof this test reached the move rather than the in-place patch: the move re-appends the + // matched rows to a new fragment, where the in-place patch leaves the fragment ids alone. + // Without this the value checks below pass on either path. + // + // This does NOT by itself exclude the v2 path, which also writes full rows and would also + // append. v2 is excluded structurally instead: the scalar index created above means + // `can_use_create_plan`'s `!would_use_scalar_index` is false. That coupling is load-bearing + // -- it is exactly what this test lacked when it silently never ran the move at all. + let frags_after: HashSet = ds.get_fragments().iter().map(|f| f.id()).collect(); + assert!( + !frags_after + .difference(&frags_before) + .collect::>() + .is_empty(), + "expected the covered row-move to append a new fragment; fragments went from \ + {frags_before:?} to {frags_after:?}, so this test never ran the move" + ); + + let batch = ds + .scan() + .project(&["id", "category"]) + .unwrap() + .try_into_batch() + .await + .unwrap(); + assert_eq!( + batch.num_rows(), + total_rows, + "row count changed after an out-of-order covered move" + ); + let ids_col = batch["id"].as_string::(); + let cats_col = batch["category"].as_string::(); + let mut seen = HashSet::new(); + for i in 0..batch.num_rows() { + let id = ids_col.value(i); + assert!( + seen.insert(id.to_string()), + "duplicate id {id} after covered move" + ); + assert_eq!( + cats_col.value(i), + format!("updated-{}", &id[3..]), + "stale or missing covered value for {id}" + ); + } + assert_eq!(seen.len(), total_rows); + } + + /// The covered row-move must publish the moved rows as `affected_rows`, exactly as the + /// full-schema `RewriteRows` paths do. Without them `check_update_txn` short-circuits on + /// `affected_rows.is_none()` and turns any concurrent delete/update touching the same + /// fragments into a retryable conflict, so covered datasets would silently lose the + /// merge_insert/delete concurrency that non-covered ones keep. + #[tokio::test] + async fn test_covered_update_move_publishes_affected_rows() { + let dim = 2i32; + let total_rows = 64usize; + + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Utf8, false), + Field::new("category", DataType::Utf8, false), + Field::new( + "vec", + DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Float32, true)), dim), + false, + ), + ])); + let ids: Vec = (0..total_rows).map(|j| format!("id-{j:04}")).collect(); + let cats: Vec = (0..total_rows).map(|j| format!("cat-{j:04}")).collect(); + let values: Vec = (0..total_rows * dim as usize).map(|i| i as f32).collect(); + let vectors = + FixedSizeListArray::try_new_from_values(Float32Array::from(values), dim).unwrap(); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(StringArray::from(ids)), + Arc::new(StringArray::from(cats)), + Arc::new(vectors), + ], + ) + .unwrap(); + let reader = Box::new(RecordBatchIterator::new([Ok(batch)], schema.clone())); + let mut ds = Dataset::write(reader, "memory://covered_move_affected_rows", None) + .await + .unwrap(); + // Required to reach `rewrite_covered_update_as_move`: without a scalar index on every + // join key, `can_use_create_plan` routes the merge to v2, which publishes affected rows + // of its own and would satisfy this test's assertion without the move running at all. + ds.create_index( + &["id"], + IndexType::BTree, + None, + &ScalarIndexParams::default(), + false, + ) + .await + .unwrap(); + let mut params = VectorIndexParams::ivf_flat(2, MetricType::L2); + params.covering_columns(vec!["category".to_string()]); + ds.create_index(&["vec"], IndexType::Vector, None, ¶ms, false) + .await + .unwrap(); + let ds = Arc::new(ds); + + // Single fragment written in id order, so row N sits at address N. + let targets = [2usize, 5, 7]; + let sub_schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Utf8, false), + Field::new("category", DataType::Utf8, false), + ])); + let update_batch = RecordBatch::try_new( + sub_schema.clone(), + vec![ + Arc::new(StringArray::from( + targets + .iter() + .map(|n| format!("id-{n:04}")) + .collect::>(), + )), + Arc::new(StringArray::from( + targets + .iter() + .map(|n| format!("updated-{n:04}")) + .collect::>(), + )), + ], + ) + .unwrap(); + let reader = Box::new(RecordBatchIterator::new([Ok(update_batch)], sub_schema)); + let UncommittedMergeInsert { affected_rows, .. } = + MergeInsertBuilder::try_new(ds, vec!["id".to_string()]) + .unwrap() + .when_matched(WhenMatched::UpdateAll) + .when_not_matched(WhenNotMatched::DoNothing) + .try_build() + .unwrap() + .execute_uncommitted(reader) + .await + .unwrap(); + + let expected = + RowAddrTreeMap::from(RoaringTreemap::from_iter(targets.iter().map(|n| *n as u64))); + assert_eq!( + affected_rows, + Some(expected), + "covered row-move must publish the moved addresses so a concurrent delete can rebase" + ); + } + + /// Builds a 3-fragment dataset with an IVF_PQ index covering `category`, then does a + /// partial merge_insert that rewrites only `category` for fragment 1. That update is a + /// row-MOVE: fragment 1's rows are tombstoned and rewritten into a new fragment, while + /// the old segment's storage still physically holds their pre-update covering values. + /// Returns `(post-move dataset, total_rows, rows_per_frag, dim)`. + /// The covered move must report EVERY field id -- leaves included -- in + /// `fields_for_preserving_frag_bitmap`. The consumer compares that list against + /// each index's leaf-expanded dependency set, so a top-level-only list would let + /// an index keyed on a nested field (its dependency set holds the LEAF id) + /// escape the exclude-everything contract and get its bitmap wrongly extended + /// onto the moved fragment. Unobservable through query results today (the move + /// runs only on address-style datasets, where the consumer skips every index), + /// so this pins the committed transaction payload itself. + #[tokio::test] + async fn test_covered_move_reports_leaf_field_ids_for_preserving() { + use crate::io::commit::read_transaction_file; + + let dim = 4i32; + let rows_per_frag = 256usize; // >= 256 rows so PQ has enough to train + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Utf8, false), + Field::new("category", DataType::Utf8, false), + Field::new( + "meta", + DataType::Struct(vec![Field::new("code", DataType::Int32, false)].into()), + false, + ), + Field::new( + "vec", + DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Float32, true)), dim), + false, + ), + ])); + let make_batch = |frag_idx: usize| { + let start = frag_idx * rows_per_frag; + let ids: Vec = (start..start + rows_per_frag) + .map(|j| format!("id-{j:04}")) + .collect(); + let cats: Vec = (start..start + rows_per_frag) + .map(|j| format!("cat-{j:04}")) + .collect(); + let codes = + Int32Array::from_iter_values((start..start + rows_per_frag).map(|j| j as i32)); + let meta = StructArray::from(vec![( + Arc::new(Field::new("code", DataType::Int32, false)), + Arc::new(codes) as arrow_array::ArrayRef, + )]); + let values: Vec = (0..rows_per_frag * dim as usize) + .map(|i| (start * dim as usize + i) as f32) + .collect(); + let vectors = + FixedSizeListArray::try_new_from_values(Float32Array::from(values), dim).unwrap(); + RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(StringArray::from(ids)), + Arc::new(StringArray::from(cats)), + Arc::new(meta), + Arc::new(vectors), + ], + ) + .unwrap() + }; + let reader = Box::new(RecordBatchIterator::new( + [Ok(make_batch(0))], + schema.clone(), + )); + let mut ds = Dataset::write(reader, "memory://covered_move_preserving", None) + .await + .unwrap(); + let reader = Box::new(RecordBatchIterator::new( + [Ok(make_batch(1))], + schema.clone(), + )); + ds.append(reader, None).await.unwrap(); + // Scalar index on the join key so the merge joins via the indexed scan and + // reaches `rewrite_covered_update_as_move` (the full-scan join commits its + // move through the insert exec instead, which is not the producer under + // test here). + ds.create_index( + &["id"], + IndexType::BTree, + None, + &ScalarIndexParams::default(), + false, + ) + .await + .unwrap(); + let mut params = VectorIndexParams::ivf_pq(1, 8, 2, MetricType::L2, 2); + params.covering_columns(vec!["category".to_string()]); + ds.create_index(&["vec"], IndexType::Vector, None, ¶ms, false) + .await + .unwrap(); + let ds = Arc::new(ds); + + // Partial covered-column update over fragment 1 -> takes the row-move path. + let frag1_start = rows_per_frag; + let ids: Vec = (frag1_start..frag1_start + rows_per_frag) + .map(|j| format!("id-{j:04}")) + .collect(); + let new_cats: Vec = (frag1_start..frag1_start + rows_per_frag) + .map(|j| format!("updated-{j:04}")) + .collect(); + let sub_schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Utf8, false), + Field::new("category", DataType::Utf8, false), + ])); + let update_batch = RecordBatch::try_new( + sub_schema.clone(), + vec![ + Arc::new(StringArray::from(ids)), + Arc::new(StringArray::from(new_cats)), + ], + ) + .unwrap(); + let reader = Box::new(RecordBatchIterator::new([Ok(update_batch)], sub_schema)); + let (ds, _) = MergeInsertBuilder::try_new(ds, vec!["id".to_string()]) + .unwrap() + .when_matched(WhenMatched::UpdateAll) + .when_not_matched(WhenNotMatched::DoNothing) + .try_build() + .unwrap() + .execute_reader(reader) + .await + .unwrap(); + + let tx_path = ds + .manifest() + .transaction_file + .clone() + .expect("the covered move writes a transaction file"); + let tx = read_transaction_file(ds.object_store.as_ref(), &ds.base, &tx_path) + .await + .unwrap(); + let Operation::Update { + fields_for_preserving_frag_bitmap, + update_mode, + .. + } = &tx.operation + else { + panic!("expected Operation::Update, got: {:?}", tx.operation); + }; + assert_eq!(*update_mode, Some(RewriteRows), "the move path committed"); + let reported: HashSet = fields_for_preserving_frag_bitmap.iter().copied().collect(); + let expected: HashSet = ds + .schema() + .fields_pre_order() + .map(|f| f.id as u32) + .collect(); + let leaf_id = ds.schema().field("meta").unwrap().children[0].id as u32; + assert!( + reported.contains(&leaf_id), + "the preserving list must include nested LEAF ids (missing {leaf_id}): {reported:?}" + ); + assert_eq!( + reported, expected, + "the covered move must report every field id, leaves included" + ); + } + + async fn build_moved_covering_dataset(uri: &str) -> (Arc, usize, usize, i32) { + let dim = 4i32; + let rows_per_frag = 256usize; + let num_frags = 3usize; + + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Utf8, false), + Field::new("category", DataType::Utf8, false), + Field::new( + "vec", + DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Float32, true)), dim), + false, + ), + ])); + + let make_batch = |frag_idx: usize| { + let start = frag_idx * rows_per_frag; + let ids: Vec = (start..start + rows_per_frag) + .map(|j| format!("id-{j:04}")) + .collect(); + let cats: Vec = (start..start + rows_per_frag) + .map(|j| format!("cat-{j:04}")) + .collect(); + let values: Vec = (0..rows_per_frag * dim as usize) + .map(|i| (start * dim as usize + i) as f32) + .collect(); + let vectors = + FixedSizeListArray::try_new_from_values(Float32Array::from(values), dim).unwrap(); + RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(StringArray::from(ids)), + Arc::new(StringArray::from(cats)), + Arc::new(vectors), + ], + ) + .unwrap() + }; + + let reader = Box::new(RecordBatchIterator::new( + [Ok(make_batch(0))], + schema.clone(), + )); + let mut ds = Dataset::write(reader, uri, None).await.unwrap(); + for frag_idx in 1..num_frags { + let reader = Box::new(RecordBatchIterator::new( + [Ok(make_batch(frag_idx))], + schema.clone(), + )); + ds.append(reader, None).await.unwrap(); + } + + // Scalar index on the join key so the merge joins via the indexed scan. The write + // strategy is unaffected: a covered-column update is always a row-move. + ds.create_index( + &["id"], + IndexType::BTree, + None, + &ScalarIndexParams::default(), + false, + ) + .await + .unwrap(); + let mut params = VectorIndexParams::ivf_pq(1, 8, 2, MetricType::L2, 2); + params.covering_columns(vec!["category".to_string()]); + ds.create_index(&["vec"], IndexType::Vector, None, ¶ms, false) + .await + .unwrap(); + let ds = Arc::new(ds); + + // Partial merge_insert rewriting ONLY `category` for fragment 1 -> prunes + // fragment 1 from the vector index's coverage bitmap. + let frag1_start = rows_per_frag; + let ids: Vec = (frag1_start..frag1_start + rows_per_frag) + .map(|j| format!("id-{j:04}")) + .collect(); + let new_cats: Vec = (frag1_start..frag1_start + rows_per_frag) + .map(|j| format!("updated-{j:04}")) + .collect(); + let sub_schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Utf8, false), + Field::new("category", DataType::Utf8, false), + ])); + let update_batch = RecordBatch::try_new( + sub_schema.clone(), + vec![ + Arc::new(StringArray::from(ids)), + Arc::new(StringArray::from(new_cats)), + ], + ) + .unwrap(); + let reader = Box::new(RecordBatchIterator::new([Ok(update_batch)], sub_schema)); + let (ds, _) = MergeInsertBuilder::try_new(ds, vec!["id".to_string()]) + .unwrap() + .when_matched(WhenMatched::UpdateAll) + .when_not_matched(WhenNotMatched::DoNothing) + .try_build() + .unwrap() + .execute_reader(reader) + .await + .unwrap(); + + // Precondition every caller relies on: the covering column really is recorded on + // the index metadata, since the move/prune logic keys on it. + let category_field_id = ds.schema().field("category").unwrap().id; + let indices = ds.load_indices().await.unwrap(); + let vec_idx = indices.iter().find(|i| i.name == "vec_idx").unwrap(); + assert_eq!(vec_idx.covering_fields, vec![category_field_id]); + + (ds, num_frags * rows_per_frag, rows_per_frag, dim) + } + + /// A covered query over every row must return each row exactly once (no stale duplicate + /// from the old storage) with the fresh `category` value. + async fn assert_covered_query_no_duplicates_fresh( + ds: &Dataset, + total_rows: usize, + rows_per_frag: usize, + dim: i32, + ) { + use arrow_array::cast::AsArray; + use arrow_array::types::UInt64Type; + use lance_core::ROW_ID; + use std::collections::HashSet; + + let frag1_start = rows_per_frag; + let query: Float32Array = (0..dim) + .map(|i| (frag1_start * dim as usize + i as usize) as f32) + .collect(); + let results = ds + .scan() + .nearest("vec", &query, total_rows) + .unwrap() + .with_row_id() + .project(&["id", "category"]) + .unwrap() + .try_into_batch() + .await + .unwrap(); + + let rowids = results[ROW_ID].as_primitive::(); + let unique: HashSet = rowids.values().iter().copied().collect(); + assert_eq!( + unique.len(), + rowids.len(), + "optimize double-counted the pruned fragment: duplicate row ids in covered results" + ); + assert_eq!( + results.num_rows(), + total_rows, + "covered query returned {} rows, expected {total_rows} (dropped/duplicated rows)", + results.num_rows() + ); + + let ids_col = results["id"].as_string::(); + let cats = results["category"].as_string::(); + for i in 0..ids_col.len() { + let id = ids_col.value(i); + let n: usize = id[3..].parse().unwrap(); + let expected = if (frag1_start..frag1_start + rows_per_frag).contains(&n) { + format!("updated-{n:04}") + } else { + format!("cat-{n:04}") + }; + assert_eq!(cats.value(i), expected, "stale covering value for row {id}"); + } + } + + /// After a covered-column update moves rows out of a fragment, an optimize must not + /// resurrect the stale copies still sitting in the old segment's storage. `merge` + /// re-scans the moved rows as unindexed data *and* folds in the old segment, so it must + /// not double-count them; the default (append) optimize leaves the old segment in place, + /// so the deletion mask must keep hiding them. + #[rstest::rstest] + #[case::merge(lance_index::optimize::OptimizeOptions::merge(1))] + #[case::append_default(lance_index::optimize::OptimizeOptions::default())] + #[tokio::test] + async fn test_optimize_after_covering_move_no_duplicate_rows( + #[case] options: lance_index::optimize::OptimizeOptions, + ) { + let (ds, total_rows, rows_per_frag, dim) = + build_moved_covering_dataset("memory://covered_opt").await; + let mut ds = Arc::try_unwrap(ds).unwrap_or_else(|arc| (*arc).clone()); + ds.optimize_indices(&options).await.unwrap(); + assert_covered_query_no_duplicates_fresh(&ds, total_rows, rows_per_frag, dim).await; + } + + /// The covered move cannot carry stable row ids onto its new fragments -- it never rechunks + /// the captured `row_id_sequence` the way the full-schema path does, so `assign_row_ids` + /// would mint fresh ids for every moved row and break the one guarantee stable row ids exist + /// to provide. The move is an optimisation over an already-correct in-place patch, so this + /// must degrade to that patch rather than fail the merge. + /// + /// The assertions are the property the old rejection protected, stated positively: the merge + /// SUCCEEDS, every row keeps the row id it had (the in-place patch moves no rows, so ids are + /// preserved by construction -- strictly safer here than the move), and confinement was the + /// only thing given up, i.e. the updated fragment left the vector index's bitmap. Run against + /// a real directory rather than `memory://` so the ids are checked after a genuine reopen. + #[tokio::test] + async fn test_merge_insert_covered_update_falls_back_on_stable_row_ids() { + use arrow_array::types::UInt64Type; + use lance_core::ROW_ID; + + /// `id -> _rowid` for the whole dataset, so "ids were preserved" is provable per row + /// rather than as a count. + async fn row_ids_by_key(ds: &Dataset) -> HashMap { + let batch = ds + .scan() + .with_row_id() + .project(&["id"]) + .unwrap() + .try_into_batch() + .await + .unwrap(); + let ids = batch["id"].as_string::(); + let row_ids = batch[ROW_ID].as_primitive::(); + (0..batch.num_rows()) + .map(|i| (ids.value(i).to_string(), row_ids.value(i))) + .collect() + } + + let test_uri = TempStrDir::default(); + let dim = 4i32; + let n = 256usize; // enough rows to train PQ + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Utf8, false), + Field::new("category", DataType::Utf8, false), + Field::new( + "vec", + DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Float32, true)), dim), + false, + ), + ])); + let ids: Vec = (0..n).map(|j| format!("id-{j:04}")).collect(); + let cats: Vec = (0..n).map(|j| format!("cat-{j:04}")).collect(); + let values: Vec = (0..n * dim as usize).map(|i| i as f32).collect(); + let vectors = + FixedSizeListArray::try_new_from_values(Float32Array::from(values), dim).unwrap(); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(StringArray::from(ids.clone())), + Arc::new(StringArray::from(cats)), + Arc::new(vectors), + ], + ) + .unwrap(); + let reader = Box::new(RecordBatchIterator::new([Ok(batch)], schema.clone())); + let mut ds = Dataset::write( + reader, + &test_uri, + Some(WriteParams { + enable_stable_row_ids: true, + ..Default::default() + }), + ) + .await + .unwrap(); + ds.create_index( + &["id"], + IndexType::BTree, + None, + &ScalarIndexParams::default(), + false, + ) + .await + .unwrap(); + let mut params = VectorIndexParams::ivf_pq(1, 8, 2, MetricType::L2, 2); + params.covering_columns(vec!["category".to_string()]); + ds.create_index(&["vec"], IndexType::Vector, None, ¶ms, false) + .await + .unwrap(); + + // Partial merge_insert updating only the covered `category`. + let sub_schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Utf8, false), + Field::new("category", DataType::Utf8, false), + ])); + let update = RecordBatch::try_new( + sub_schema.clone(), + vec![ + Arc::new(StringArray::from(ids)), + Arc::new(StringArray::from( + (0..n).map(|j| format!("new-{j:04}")).collect::>(), + )), + ], + ) + .unwrap(); + let reader = Box::new(RecordBatchIterator::new([Ok(update)], sub_schema)); + let row_ids_before = row_ids_by_key(&ds).await; + let updated_frag = ds.get_fragments().iter().map(|f| f.id()).max().unwrap(); + let (ds, _) = MergeInsertBuilder::try_new(Arc::new(ds), vec!["id".to_string()]) + .unwrap() + .when_matched(WhenMatched::UpdateAll) + .when_not_matched(WhenNotMatched::DoNothing) + .try_build() + .unwrap() + .execute_reader(reader) + .await + .expect("a covered update on a stable-row-id dataset must fall back, not fail"); + + // Reopen so the ids are read back from committed bytes, not from the in-memory handle. + let reopened = Dataset::open(test_uri.as_ref()).await.unwrap(); + let after = reopened + .scan() + .project(&["id", "category"]) + .unwrap() + .try_into_batch() + .await + .unwrap(); + assert_eq!(after.num_rows(), n, "the fallback must not lose rows"); + let cats = after["category"].as_string::(); + assert!( + (0..cats.len()).all(|i| cats.value(i).starts_with("new-")), + "the fallback must still apply the update" + ); + + // The property the old rejection existed to protect: no row changed identity. The + // in-place patch moves no rows, so every id survives, paired with the same key. + let row_ids_after = row_ids_by_key(&reopened).await; + assert_eq!( + row_ids_before, row_ids_after, + "falling back to the in-place patch must preserve every stable row id" + ); + + // And the only thing forfeited is confinement: the updated fragment left the bitmap. + let indices = ds.load_indices().await.unwrap(); + let vec_idx = indices.iter().find(|i| i.name == "vec_idx").unwrap(); + let bitmap = vec_idx.fragment_bitmap.as_ref().unwrap(); + assert!( + !bitmap.contains(updated_frag as u32), + "the in-place fallback should have evicted fragment {updated_frag} from the vector \ + index; if it is still covered the merge did not take the fallback and this test is \ + not exercising it: {bitmap:?}" + ); + } + + /// A partial-schema merge that does NOT update matched rows (DoNothing + InsertAll) + /// must not be treated as a covered-column update, even when the source carries a + /// covered column -- matched rows are untouched and unmatched rows enter unindexed + /// fragments, so nothing is patched into the index. + #[tokio::test] + async fn test_merge_insert_covered_donothing_insert_not_rejected() { + let dim = 4i32; + let n = 20usize; + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Utf8, false), + Field::new("category", DataType::Utf8, false), + Field::new( + "vec", + // nullable so InsertAll with a partial source is permitted + DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Float32, true)), dim), + true, + ), + ])); + let ids: Vec = (0..n).map(|j| format!("id-{j:04}")).collect(); + let cats: Vec = (0..n).map(|j| format!("cat-{j:04}")).collect(); + let values: Vec = (0..n * dim as usize).map(|i| i as f32).collect(); + let vectors = + FixedSizeListArray::try_new_from_values(Float32Array::from(values), dim).unwrap(); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(StringArray::from(ids)), + Arc::new(StringArray::from(cats)), + Arc::new(vectors), + ], + ) + .unwrap(); + let reader = Box::new(RecordBatchIterator::new([Ok(batch)], schema.clone())); + let mut ds = Dataset::write(reader, "memory://covered_donothing", None) + .await + .unwrap(); + ds.create_index( + &["id"], + IndexType::BTree, + None, + &ScalarIndexParams::default(), + false, + ) + .await + .unwrap(); + let mut params = VectorIndexParams::ivf_flat(2, MetricType::L2); + params.covering_columns(vec!["category".to_string()]); + ds.create_index(&["vec"], IndexType::Vector, None, ¶ms, false) + .await + .unwrap(); + + // DoNothing on matched + InsertAll unmatched, partial source (id + category) + // including new rows to insert. + let sub_schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Utf8, false), + Field::new("category", DataType::Utf8, false), + ])); + let new_ids: Vec = (n..n + 5).map(|j| format!("id-{j:04}")).collect(); + let new_cats: Vec = (n..n + 5).map(|j| format!("cat-{j:04}")).collect(); + let update = RecordBatch::try_new( + sub_schema.clone(), + vec![ + Arc::new(StringArray::from(new_ids)), + Arc::new(StringArray::from(new_cats)), + ], + ) + .unwrap(); + let reader = Box::new(RecordBatchIterator::new([Ok(update)], sub_schema)); + let result = MergeInsertBuilder::try_new(Arc::new(ds), vec!["id".to_string()]) + .unwrap() + .when_matched(WhenMatched::DoNothing) + .when_not_matched(WhenNotMatched::InsertAll) + .try_build() + .unwrap() + .execute_reader(reader) + .await; + assert!( + result.is_ok(), + "DoNothing + InsertAll with a covered column in a partial source must not be \ + rejected as a covered update; got {:?}", + result.err() + ); + } + + /// An upsert touching a covered column -- `UpdateAll` matched plus `InsertAll` unmatched -- + /// must fall back to the in-place patch. The move reads row addresses through the null mask, + /// and the unmatched group arrives with a NULL `_rowaddr`, so entering the move would + /// tombstone unrelated rows at address 0. Never entering it prevents that as completely as + /// refusing the merge did, and `update_fragments` already routes the NULL group to + /// `handle_new_fragments`. + /// + /// The data-loss assertion is the point: the row count must be exactly the original plus the + /// inserted row, and every pre-existing row must still be there. That is what a bogus + /// address-0 tombstone would break. + #[tokio::test] + async fn test_merge_insert_covered_update_with_insert_falls_back() { + let dim = 4i32; + let n = 20usize; + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Utf8, false), + Field::new("category", DataType::Utf8, false), + Field::new( + "vec", + // nullable so InsertAll with a partial source is permitted + DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Float32, true)), dim), + true, + ), + ])); + let ids: Vec = (0..n).map(|j| format!("id-{j:04}")).collect(); + let cats: Vec = (0..n).map(|j| format!("cat-{j:04}")).collect(); + let values: Vec = (0..n * dim as usize).map(|i| i as f32).collect(); + let vectors = + FixedSizeListArray::try_new_from_values(Float32Array::from(values), dim).unwrap(); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(StringArray::from(ids)), + Arc::new(StringArray::from(cats)), + Arc::new(vectors), + ], + ) + .unwrap(); + let reader = Box::new(RecordBatchIterator::new([Ok(batch)], schema.clone())); + let mut ds = Dataset::write(reader, "memory://covered_update_insert", None) + .await + .unwrap(); + ds.create_index( + &["id"], + IndexType::BTree, + None, + &ScalarIndexParams::default(), + false, + ) + .await + .unwrap(); + let mut params = VectorIndexParams::ivf_flat(2, MetricType::L2); + params.covering_columns(vec!["category".to_string()]); + ds.create_index(&["vec"], IndexType::Vector, None, ¶ms, false) + .await + .unwrap(); + + // Partial source mixing an update to an existing row with a brand-new row. + let sub_schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Utf8, false), + Field::new("category", DataType::Utf8, false), + ])); + let update = RecordBatch::try_new( + sub_schema.clone(), + vec![ + Arc::new(StringArray::from(vec![ + "id-0001".to_string(), + format!("id-{:04}", n), + ])), + Arc::new(StringArray::from(vec![ + "updated-0001".to_string(), + "brand-new".to_string(), + ])), + ], + ) + .unwrap(); + let reader = Box::new(RecordBatchIterator::new([Ok(update)], sub_schema)); + let updated_frag = ds.get_fragments().iter().map(|f| f.id()).max().unwrap(); + let (ds, stats) = MergeInsertBuilder::try_new(Arc::new(ds), vec!["id".to_string()]) + .unwrap() + .when_matched(WhenMatched::UpdateAll) + .when_not_matched(WhenNotMatched::InsertAll) + .try_build() + .unwrap() + .execute_reader(reader) + .await + .expect("a covered update carrying inserts must fall back, not fail"); + assert_eq!(stats.num_updated_rows, 1); + assert_eq!(stats.num_inserted_rows, 1); + + // No row was tombstoned at a bogus address: exactly the original rows plus the insert, + // with every original key still present. + let after = ds + .scan() + .project(&["id", "category"]) + .unwrap() + .try_into_batch() + .await + .unwrap(); + assert_eq!( + after.num_rows(), + n + 1, + "the fallback must insert the new row and tombstone nothing else" + ); + let ids_col = after["id"].as_string::(); + let cats_col = after["category"].as_string::(); + let found: HashMap<&str, &str> = (0..after.num_rows()) + .map(|i| (ids_col.value(i), cats_col.value(i))) + .collect(); + for j in 0..n { + assert!( + found.contains_key(format!("id-{j:04}").as_str()), + "the fallback dropped pre-existing row id-{j:04}" + ); + } + assert_eq!( + found["id-0001"], "updated-0001", + "the update was not applied" + ); + assert_eq!( + found[format!("id-{n:04}").as_str()], + "brand-new", + "the insert was not applied" + ); + + // Confinement is what was forfeited. + let indices = ds.load_indices().await.unwrap(); + let vec_idx = indices.iter().find(|i| i.name == "vec_idx").unwrap(); + let bitmap = vec_idx.fragment_bitmap.as_ref().unwrap(); + assert!( + !bitmap.contains(updated_frag as u32), + "the in-place fallback should have evicted fragment {updated_frag} from the vector \ + index; still covered means the fallback was not taken: {bitmap:?}" + ); + } + + /// The covered-column row-move re-reads full rows via a take, which cannot round-trip + /// legacy (v1) blob columns as binary, so a dataset carrying one must fall back to the + /// in-place patch -- which never reads the blob column at all. The assertion that matters + /// is therefore that the blob survives the merge intact, not that the merge was refused. + #[tokio::test] + async fn test_merge_insert_covered_update_falls_back_with_v1_blob() { + use arrow_array::LargeBinaryArray; + use lance_arrow::BLOB_META_KEY; + + let dim = 4i32; + let n = 20usize; + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Utf8, false), + Field::new("category", DataType::Utf8, false), + Field::new( + "vec", + DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Float32, true)), dim), + true, + ), + Field::new("blobs", DataType::LargeBinary, true).with_metadata(HashMap::from([( + BLOB_META_KEY.to_string(), + "true".to_string(), + )])), + ])); + let ids: Vec = (0..n).map(|j| format!("id-{j:04}")).collect(); + let cats: Vec = (0..n).map(|j| format!("cat-{j:04}")).collect(); + let blobs: Vec> = (0..n).map(|_| Some(b"payload".as_slice())).collect(); + let values: Vec = (0..n * dim as usize).map(|i| i as f32).collect(); + let vectors = + FixedSizeListArray::try_new_from_values(Float32Array::from(values), dim).unwrap(); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(StringArray::from(ids)), + Arc::new(StringArray::from(cats)), + Arc::new(vectors), + Arc::new(LargeBinaryArray::from(blobs)), + ], + ) + .unwrap(); + let reader = Box::new(RecordBatchIterator::new([Ok(batch)], schema.clone())); + let mut ds = Dataset::write( + reader, + "memory://covered_v1_blob", + Some(WriteParams { + data_storage_version: Some(LanceFileVersion::V2_1), + ..Default::default() + }), + ) + .await + .unwrap(); + ds.create_index( + &["id"], + IndexType::BTree, + None, + &ScalarIndexParams::default(), + false, + ) + .await + .unwrap(); + let mut params = VectorIndexParams::ivf_flat(2, MetricType::L2); + params.covering_columns(vec!["category".to_string()]); + ds.create_index(&["vec"], IndexType::Vector, None, ¶ms, false) + .await + .unwrap(); + + // Partial source updating the covered `category` on matched rows only. + let sub_schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Utf8, false), + Field::new("category", DataType::Utf8, false), + ])); + let upd_ids: Vec = (0..3).map(|j| format!("id-{j:04}")).collect(); + let upd_cats: Vec = (0..3).map(|j| format!("new-{j:04}")).collect(); + let update = RecordBatch::try_new( + sub_schema.clone(), + vec![ + Arc::new(StringArray::from(upd_ids)), + Arc::new(StringArray::from(upd_cats)), + ], + ) + .unwrap(); + let reader = Box::new(RecordBatchIterator::new([Ok(update)], sub_schema)); + let updated_frag = ds.get_fragments().iter().map(|f| f.id()).max().unwrap(); + // A v1 blob column scans as a position/size description struct, not binary -- writing + // those descriptors back as data is exactly the corruption the move would cause. Capture + // them so "the blob was left alone" is provable rather than assumed. + let blobs_before = ds + .scan() + .project(&["blobs"]) + .unwrap() + .try_into_batch() + .await + .unwrap(); + let (ds, _) = MergeInsertBuilder::try_new(Arc::new(ds), vec!["id".to_string()]) + .unwrap() + .when_matched(WhenMatched::UpdateAll) + .when_not_matched(WhenNotMatched::DoNothing) + .try_build() + .unwrap() + .execute_reader(reader) + .await + .expect("a covered update on a v1-blob dataset must fall back, not fail"); + + // The update landed, and the blob the move would have corrupted is untouched. Read the + // blobs through the binary path the v1 column is meant to be read by. + let after = ds + .scan() + .project(&["id", "category"]) + .unwrap() + .try_into_batch() + .await + .unwrap(); + assert_eq!(after.num_rows(), n); + let ids_col = after["id"].as_string::(); + let cats_col = after["category"].as_string::(); + for i in 0..after.num_rows() { + let key: usize = ids_col.value(i)[3..].parse().unwrap(); + let expected = if key < 3 { + format!("new-{key:04}") + } else { + format!("cat-{key:04}") + }; + assert_eq!(cats_col.value(i), expected, "wrong category after fallback"); + } + let blobs_after = ds + .scan() + .project(&["blobs"]) + .unwrap() + .try_into_batch() + .await + .unwrap(); + assert_eq!( + blobs_before, blobs_after, + "the in-place fallback must leave the v1 blob column untouched; the move would have \ + rewritten these descriptors into a new fragment as ordinary struct data" + ); + // And the bytes are still readable through the blob API they were written with. + let blob_handles = ds + .take_blobs_by_indices(&(0..n as u64).collect::>(), "blobs") + .await + .unwrap(); + assert_eq!(blob_handles.len(), n); + assert!( + blob_handles.iter().all(|b| b + .as_ref() + .is_some_and(|b| b.size() == b"payload".len() as u64)), + "every v1 blob must still read back at its original size" + ); + + // Confinement is what was forfeited. + let indices = ds.load_indices().await.unwrap(); + let vec_idx = indices.iter().find(|i| i.name == "vec_idx").unwrap(); + let bitmap = vec_idx.fragment_bitmap.as_ref().unwrap(); + assert!( + !bitmap.contains(updated_frag as u32), + "the in-place fallback should have evicted fragment {updated_frag} from the vector \ + index; still covered means the fallback was not taken: {bitmap:?}" + ); + } + + // Regression test: after a partial-schema merge_insert drops a fragment from the FTS + // index bitmap, a full text search should not return duplicate rows. The stale inverted + // index data still references the dropped fragment, and the scanner also flat-scans + // unindexed fragments, causing the same rows to appear from both paths. + #[tokio::test] + async fn test_partial_merge_insert_stale_fts_index_duplicates() { + let rows_per_frag = 10usize; + let num_frags = 3usize; + + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Utf8, false), + Field::new("category", DataType::Utf8, false), + Field::new("text", DataType::Utf8, false), + ])); + + let make_batch = |frag_idx: usize| { + let start = frag_idx * rows_per_frag; + let ids: Vec = (start..start + rows_per_frag) + .map(|j| format!("id-{j:04}")) + .collect(); + let cats: Vec<&str> = vec!["A"; rows_per_frag]; + // Every row contains "common" so we can search for it and expect all rows + let texts: Vec = (start..start + rows_per_frag) + .map(|j| format!("common unique{j:04}")) + .collect(); + RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(StringArray::from(ids)), + Arc::new(StringArray::from(cats)), + Arc::new(StringArray::from(texts)), + ], + ) + .unwrap() + }; + + // Write 3 fragments + let batch0 = make_batch(0); + let reader = Box::new(RecordBatchIterator::new([Ok(batch0)], schema.clone())); + let mut ds = Dataset::write(reader, "memory://fts_stale_test", None) + .await + .unwrap(); + for frag_idx in 1..num_frags { + let batch = make_batch(frag_idx); + let reader = Box::new(RecordBatchIterator::new([Ok(batch)], schema.clone())); + ds.append(reader, None).await.unwrap(); + } + + // Create inverted index on text + let params = InvertedIndexParams::default(); + ds.create_index(&["text"], IndexType::Inverted, None, ¶ms, true) + .await + .unwrap(); + + let ds = Arc::new(ds); + + // Partial merge_insert with (id, text) on fragment 1 rows. + // Text still contains "common" so FTS will find them via both paths. + // This drops fragment 1 from the inverted index bitmap. + let frag1_start = rows_per_frag; + let ids: Vec = (frag1_start..frag1_start + rows_per_frag) + .map(|j| format!("id-{j:04}")) + .collect(); + let texts: Vec = (frag1_start..frag1_start + rows_per_frag) + .map(|j| format!("common updated{j:04}")) + .collect(); + let sub_schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Utf8, false), + Field::new("text", DataType::Utf8, false), + ])); + let update_batch = RecordBatch::try_new( + sub_schema.clone(), + vec![ + Arc::new(StringArray::from(ids)), + Arc::new(StringArray::from(texts)), + ], + ) + .unwrap(); + let reader = Box::new(RecordBatchIterator::new([Ok(update_batch)], sub_schema)); + let (ds, _) = MergeInsertBuilder::try_new(ds, vec!["id".to_string()]) + .unwrap() + .when_matched(WhenMatched::UpdateAll) + .when_not_matched(WhenNotMatched::DoNothing) + .try_build() + .unwrap() + .execute_reader(reader) + .await + .unwrap(); + + // FTS search for "common" — every row should match exactly once + let query = FullTextSearchQuery::new("common".to_string()); + let results = ds + .scan() + .full_text_search(query) + .unwrap() + .try_into_batch() + .await + .unwrap(); + // Check no duplicate ids let ids = results .column_by_name("id") @@ -14921,11 +17071,11 @@ MergeInsert: on=[id], when_matched=DoNothing, when_not_matched=InsertAll, when_n #[tokio::test] async fn test_merge_insert_with_blob_v1_source_provides_blob() { use arrow_array::LargeBinaryArray; - use arrow_schema::Schema as ArrowSchema; + use arrow_schema::Schema; use lance_arrow::BLOB_META_KEY; let test_dir = TempStrDir::default(); - let schema = Arc::new(ArrowSchema::new(vec![ + let schema = Arc::new(Schema::new(vec![ Field::new("blobs", DataType::LargeBinary, true).with_metadata(HashMap::from([( BLOB_META_KEY.to_string(), "true".to_string(), @@ -15000,10 +17150,10 @@ MergeInsert: on=[id], when_matched=DoNothing, when_not_matched=InsertAll, when_n #[tokio::test] async fn test_merge_insert_with_blob_v2_source_provides_blob() { use crate::{BlobArrayBuilder, blob_field}; - use arrow_schema::Schema as ArrowSchema; + use arrow_schema::Schema; let test_dir = TempStrDir::default(); - let schema = Arc::new(ArrowSchema::new(vec![ + let schema = Arc::new(Schema::new(vec![ blob_field("blobs", true), Field::new("id", DataType::Int64, true), Field::new("other", DataType::Int64, true), diff --git a/rust/lance/src/index.rs b/rust/lance/src/index.rs index b5803624520..1f2218c1c17 100644 --- a/rust/lance/src/index.rs +++ b/rust/lance/src/index.rs @@ -78,6 +78,7 @@ use vector::utils::get_vector_type; mod api; pub(crate) mod append; +pub(crate) mod covering; mod create; pub mod frag_reuse; pub mod mem_wal; @@ -137,7 +138,7 @@ fn validate_segment_metadata(index_name: &str, segments: &[IndexMetadata]) -> Re Ok(()) } -fn collect_subtree_field_ids(field: &Field, field_ids: &mut HashSet) { +pub(crate) fn collect_subtree_field_ids(field: &Field, field_ids: &mut HashSet) { field_ids.insert(field.id); for child in &field.children { collect_subtree_field_ids(child, field_ids); @@ -163,7 +164,16 @@ struct PhysicalDataFileIdentity<'a> { base_id: Option, base_binding: PhysicalBaseBinding<'a>, path: &'a str, - fields: &'a [i32], + // Deliberately NOT the file's `fields` list. This identity is built per field, but + // `fields` is file-wide, so a tombstone of any *other* field in a shared file (an + // unrelated `RewriteColumns` turns its id negative) would change every field's + // identity and prune coverage for indexes whose data did not move. A tombstone of + // *this* field is still caught: its id stops matching and it drops out of the map + // `fragment_field_files` builds. `column_indices` covers an actual layout change. + // + // Upstream cannot hit this without a covering producer -- with no payload written, + // two indexes never share a file with different field sets -- which is why the guard + // held until this commit introduced one. column_indices: &'a [i32], file_major_version: u32, file_minor_version: u32, @@ -185,7 +195,6 @@ impl<'a> PhysicalDataFileIdentity<'a> { base_id: file.base_id, base_binding, path: &file.path, - fields: file.fields.as_ref(), column_indices: file.column_indices.as_ref(), file_major_version: file.file_major_version, file_minor_version: file.file_minor_version, @@ -265,7 +274,40 @@ async fn prune_stale_segment_coverage( .iter_mut() .filter(|segment| segment.dataset_version() == version) { - let indexed_field_ids = segment_indexed_field_ids(dataset, segment)?; + // `segment_indexed_field_ids` already walks every entry of `segment.fields()`, + // keyed and carried alike -- `covering_fields` is always the trailing slice of + // `fields` -- so no separate expansion over `covering_fields()` is needed here. + let relevant_field_ids = segment_indexed_field_ids(dataset, segment)?; + for included_id in segment.covering_fields() { + // No exemption for a field the segment is also keyed on: carried refine + // vectors declare the indexed column here, and index storage holds that copy + // under its build-time name, so a rename leaves storage the rebuild paths + // cannot resolve. The `alter_columns` preflight (`schema_evolution`) and the + // commit-boundary guard (`reject_covered_field_subtree_change`) refuse it for + // the same reason; all three must agree or one of them forbids what another + // accepts. + // + // A covered field whose subtree changed between the segment's build version + // and now -- e.g. a struct that gained a child via `add_columns` (which + // commits while this segment is still uncommitted, so the Merge-commit guard + // never saw it) -- makes the segment's stored covering payload type- + // incompatible with the current schema for *every* row it covers. No + // per-fragment file-path prune below can repair a global type change, so + // reject the commit; the segment must be rebuilt against the current schema. + if Transaction::covered_field_subtree_changed( + historical.schema(), + dataset.schema(), + *included_id, + ) { + return Err(Error::invalid_input(format!( + "CreateIndex: covered (included) field id {included_id} changed its \ + subtree between the segment's build version {version} and the current \ + dataset version {}; the segment's stored covering payload no longer \ + matches the schema. Rebuild the segment against the current schema.", + dataset.manifest.version + ))); + } + } let stale_fragments = segment .fragment_bitmap() .iter() @@ -277,9 +319,9 @@ async fn prune_stale_segment_coverage( return true; }; let historical_files = - fragment_field_files(&historical, historical_fragment, &indexed_field_ids); + fragment_field_files(&historical, historical_fragment, &relevant_field_ids); let current_files = - fragment_field_files(dataset, current_fragment, &indexed_field_ids); + fragment_field_files(dataset, current_fragment, &relevant_field_ids); let changed_files = historical_files.is_none() || historical_files != current_files; let changed_overlays = prune_newer_overlays @@ -289,7 +331,7 @@ async fn prune_stale_segment_coverage( .data_file .fields .iter() - .any(|field_id| indexed_field_ids.contains(field_id)) + .any(|field_id| relevant_field_ids.contains(field_id)) }); changed_files || changed_overlays }) @@ -393,6 +435,10 @@ pub(crate) async fn build_index_metadata_from_segments( covered_fragments |= segment.fragment_bitmap().clone(); } + // The declaration is a dependency list, so validate it against the dataset schema and + // prune stale fragment coverage before committing. Physical covering payload is not a + // commit precondition: readers independently verify every segment's storage capability + // and fall back to a base-table take for any declared field that is absent. prune_stale_segment_coverage(dataset, &mut segments, false, false).await?; let new_indices = futures::stream::iter(segments.into_iter().map(|segment| async move { @@ -1170,34 +1216,21 @@ pub(crate) async fn remap_index( // withdrawal below would otherwise swallow it as an ordinary covered index. matched.validate_covering_fields()?; - // A covered index cannot survive a remap yet. Nothing writes the carried - // values in the first place, and each index type's `remap` rewrites only the - // schema that type knows about, so the remapped segment would still declare - // payload its storage no longer holds. Withdraw it instead, exactly as an - // index whose type reports `can_remap() == false` is withdrawn below: an - // absent index costs a fallback scan, whereas a surviving false declaration - // is answered from data that is not there. Erroring here would instead block - // compaction of the whole table. - if !matched.covering_fields.is_empty() { - log::warn!( - "Index '{}' declares covering fields {:?}, which no index builder \ - writes or preserves yet. Index will be dropped during compaction \ - and must be rebuilt.", - matched.name, - matched.covering_fields, - ); - return Ok(RemapResult::Drop); - } - - // With covering withdrawn above, `fields` is entirely keyed, so more than one - // entry means a genuinely composite index. - if matched.fields.len() > 1 { + // Remap handles a single keyed field. Carried fields are irrelevant to the + // remap itself: the index is opened by its keyed field, and each storage + // format's `remap` preserves its covering columns row-aligned (the format + // commit dropped covered indexes here because nothing wrote carried values + // yet; the builders in this commit do). Only a genuinely composite index is + // rejected -- `keyed_fields` fails closed on malformed metadata, but the + // validation above already rejected that. + if matched.keyed_fields().len() > 1 { return Err(Error::index(format!( - "Remapping index '{}' is not supported: it has {} keyed fields {:?}; \ - only one keyed field is supported", + "Remapping index '{}' is not supported: it has {} keyed fields {:?} \ + (carried fields {:?}); only one keyed field is supported", matched.name, - matched.fields.len(), + matched.keyed_fields().len(), matched.fields, + matched.covering_fields, ))); } @@ -1238,6 +1271,24 @@ pub(crate) async fn remap_index( let created_index = match generic.index_type() { it if it.is_scalar() => { + // Withdraw rather than remap a covered scalar index. The vector storages' + // `remap` carries their covering columns through row-aligned, which is what + // lets a covered vector index be remapped in place; `BTreeIndex::remap` and + // friends rewrite only `(value, row_id)`. Remapping would republish a + // declaration the new files do not back -- and because carried fields live in + // `fields`, the index would then shed fragment coverage on every update to a + // column it never held. Withdrawing follows the same shape as the + // `can_remap()` check below and the open failure above: a table-level + // compaction degrades one index rather than failing outright. + if !matched.covering_fields.is_empty() { + log::warn!( + "Withdrawing covered index '{}' during remap: its storage format does \ + not carry covering columns through a remap, so the declaration would \ + outlive the payload.", + matched.name + ); + return Ok(RemapResult::Drop); + } let new_store = LanceIndexStore::from_dataset_for_new(dataset, &new_id)?; let scalar_index = dataset @@ -2311,42 +2362,6 @@ impl DatasetIndexExt for Dataset { continue; } - // Optimizing a covered index would republish its declaration on a - // segment rebuilt without the carried values: `scan_vector_fragments` - // projects the keyed field and `_rowid` only, and the scalar merges - // reconstruct value plus row id. - // - // What decides is the caller's intent, not whether this group is - // stale. An unfiltered `optimize_indices()` is a table-wide - // maintenance request, and erroring aborts the loop before the - // replacements accumulated for the other groups are committed -- so - // one index this build cannot rebuild would leave every other index - // on the table stale. Skip it with a warning instead. - // - // A caller that listed this index in `index_names` asked for it - // specifically, so refuse out loud. The loop is already filtered by - // that list, so reaching here with it set means this group was named. - if let Some(covered) = deltas - .iter() - .find(|index| !index.covering_fields.is_empty()) - { - if options.index_names.is_none() { - log::warn!( - "Skipping index '{}': it declares covering fields {:?}, \ - which no index builder writes or preserves yet.", - covered.name, - covered.covering_fields, - ); - continue; - } - return Err(Error::index(format!( - "Optimizing index '{}' is not supported: it declares \ - covering fields {:?}, which no index builder writes or \ - preserves yet", - covered.name, covered.covering_fields, - ))); - } - // Optimizing a name means replacing its segments with one that // covers their union, which this build cannot compute when it // cannot read one of them: the merged segment would overlap the @@ -3637,6 +3652,18 @@ impl DatasetIndexInternalExt for Dataset { ))); } + // `field_paths` is already the keyed prefix -- it was built from `keyed_fields()` + // above -- so a carried column is deliberately NOT resolved or type-checked here. + // That is the point: the rebuild writes fresh, non-covered metadata and never reads + // the carried payload, so demanding those columns exist in the target would reject + // a target that can hold this index perfectly well. + // + // `initialize_vector_index` then re-resolves the carried columns itself, by name, + // against the target schema. `initialize_scalar_index` does NOT -- it recreates + // with default params and never reads `covering_fields`, so a covered scalar index + // would land on the target with covering silently dropped. Unreachable today: + // `covering_columns` exists only on `VectorIndexParams`, so nothing produces a + // covered scalar index. Whoever adds a scalar producer must fix that path too. let field_names = field_paths.iter().map(String::as_str).collect::>(); if let Some(index_details) = &source_index.index_details { let index_details_wrapper = IndexDetails(index_details.clone()); @@ -6567,6 +6594,90 @@ mod tests { assert_eq!(desc.rows_indexed(), 10); } + /// `describe_indices` with a column criteria must still find a covered index. + /// + /// `index_matches_criteria` is documented to receive only the column named by + /// `for_column` -- it rejects anything else with `fields.len() != 1`. A covered + /// index lists its carried columns in `fields` too, so resolving all of `fields` + /// silently fails that check and drops the index from the description with no + /// error: `describe_indices(for_column)` reports "no index on this column" for a + /// column that is indexed. + #[tokio::test] + async fn test_describe_indices_for_column_finds_covered_index() { + use lance_datagen::{BatchCount, RowCount, array}; + + let test_dir = tempfile::tempdir().unwrap(); + let test_uri = test_dir.path().to_str().unwrap(); + let reader = lance_datagen::gen_batch() + .col("id", array::step::()) + .col( + "vector", + array::rand_vec::(8.into()), + ) + .into_reader_rows(RowCount::from(20), BatchCount::from(1)); + let mut dataset = Dataset::write(reader, test_uri, None).await.unwrap(); + + let mut params = crate::index::vector::VectorIndexParams::ivf_flat( + 1, + lance_linalg::distance::MetricType::L2, + ); + params.covering_columns(vec!["id".to_string()]); + dataset + .create_index( + &["vector"], + IndexType::Vector, + Some("covered_idx".to_string()), + ¶ms, + false, + ) + .await + .unwrap(); + let covered = dataset.load_indices_by_name("covered_idx").await.unwrap(); + assert_eq!(covered[0].fields.len(), 2, "index should be covered"); + + let descriptions = dataset + .describe_indices(Some(IndexCriteria::default().for_column("vector"))) + .await + .unwrap(); + assert_eq!( + descriptions.len(), + 1, + "a covered index must still be described for its keyed column" + ); + assert_eq!(descriptions[0].name(), "covered_idx"); + + // The carried column is not the indexed column, so it must not match. + let by_carried = dataset + .describe_indices(Some(IndexCriteria::default().for_column("id"))) + .await + .unwrap(); + assert!( + by_carried.is_empty(), + "a carried column is not indexed and must not match `for_column`" + ); + + // `field_ids()` is a capability statement -- "the fields the index is built on" -- + // so it must report the keyed column only. Reporting the carried column here tells + // every consumer (Python `describe_indices().fields`, Java `getFieldIds()`, the + // namespace REST `columns` list) that this index can serve `id`; Python's + // `_default_vector_index_for_column` then hands back `vec`'s IVF centroids for the + // int32 column `id` instead of raising. + let vector_field_id = dataset.schema().field("vector").unwrap().id as u32; + let id_field_id = dataset.schema().field("id").unwrap().id as u32; + let desc = &descriptions[0]; + assert_eq!( + desc.field_ids(), + &[vector_field_id], + "field_ids() must report the keyed column only, not the carried one" + ); + // The carried set is not lost -- it stays reachable per segment. + assert_eq!( + desc.metadata()[0].covering_fields, + vec![id_field_id as i32], + "the carried column must remain reachable via metadata()[i].covering_fields" + ); + } + #[tokio::test] async fn test_describe_indices_derives_type_from_url_without_plugin() { // When index details exist but no plugin is registered for the type @@ -7959,6 +8070,68 @@ mod tests { ); } + /// `initialize_index` must copy a covered index, not reject it. + /// + /// It derives the column names it initializes from `source_index.fields`, and a + /// covered index lists its carried columns there too. Both initializers accept + /// exactly one name (they re-resolve the carried columns themselves from + /// `covering_fields`), so passing all of `fields` makes every covered index + /// un-copyable -- and `initialize_indices` then fails the whole dataset copy. + #[tokio::test] + async fn test_initialize_index_copies_covered_index() { + use arrow_array::types::Float32Type; + use lance_datagen::{BatchCount, RowCount, array}; + + let test_dir = TempStrDir::default(); + let source_uri = format!("{}/{}", test_dir, "source"); + let target_uri = format!("{}/{}", test_dir, "target"); + + let make_reader = || { + lance_datagen::gen_batch() + .col("id", array::step::()) + .col("vector", array::rand_vec::(8.into())) + .into_reader_rows(RowCount::from(300), BatchCount::from(1)) + }; + + let mut source_dataset = Dataset::write(make_reader(), &source_uri, None) + .await + .unwrap(); + let mut params = crate::index::vector::VectorIndexParams::ivf_flat( + 4, + lance_linalg::distance::MetricType::L2, + ); + params.covering_columns(vec!["id".to_string()]); + source_dataset + .create_index( + &["vector"], + IndexType::Vector, + Some("covered_idx".to_string()), + ¶ms, + false, + ) + .await + .unwrap(); + let source_dataset = Dataset::open(&source_uri).await.unwrap(); + + let mut target_dataset = Dataset::write(make_reader(), &target_uri, None) + .await + .unwrap(); + target_dataset + .initialize_index(&source_dataset, "covered_idx") + .await + .unwrap(); + + let target_indices = target_dataset.load_indices().await.unwrap(); + assert_eq!(target_indices.len(), 1); + let id_field_id = target_dataset.schema().field("id").unwrap().id; + assert_eq!( + target_indices[0].covering_fields, + vec![id_field_id], + "the copied index must keep its covering declaration" + ); + target_indices[0].validate_covering_fields().unwrap(); + } + #[rstest] #[case::simple("value", "data.value")] #[case::quoted("value.with.dot", "data.`value.with.dot`")] @@ -8831,8 +9004,12 @@ mod tests { ); } + /// The logical covering declaration and the physical payload are independent at commit. + /// Transitional segments may declare a dependency before their writer emits the values, + /// while physical payload without a declaration remains an unused implementation detail. #[tokio::test] - async fn test_commit_existing_index_segments_rejects_duplicate_segment_ids() { + async fn test_commit_existing_index_segments_preserves_covering_declaration_independently() { + use crate::index::vector::VectorIndexParams; use lance_datagen::{BatchCount, RowCount, array}; let test_dir = tempfile::tempdir().unwrap(); @@ -8844,163 +9021,605 @@ mod tests { "vector", array::rand_vec::(8.into()), ) - .into_reader_rows(RowCount::from(10), BatchCount::from(1)); - + .into_reader_rows(RowCount::from(256), BatchCount::from(1)); let mut dataset = Dataset::write(reader, test_uri, None).await.unwrap(); + let id_field_id = dataset.schema().field("id").unwrap().id; - let field_id = dataset.schema().field("vector").unwrap().id; - let base = write_vector_segment_metadata( - &dataset, - "vector_idx", - field_id, - Uuid::new_v4(), - [0_u32], - b"base", - ) - .await; - - let err = dataset + // Direction 1: storage built WITHOUT covering, declaration fabricated. + let params = VectorIndexParams::ivf_pq(2, 8, 4, lance_linalg::distance::MetricType::L2, 2); + let mut plain_segment = dataset + .create_index_builder(&["vector"], IndexType::Vector, ¶ms) + .name("vector_idx".to_string()) + .execute_uncommitted() + .await + .unwrap(); + assert!(plain_segment.covering_fields.is_empty()); + // Carried columns live at the end of `fields`, so fabricating a declaration means + // extending both lists, not just `covering_fields`. + plain_segment.fields.push(id_field_id); + plain_segment.covering_fields = vec![id_field_id]; + dataset .commit_existing_index_segments( "vector_idx", "vector", - vec![ - segment_from_metadata(&base), - segment_from_metadata(&IndexMetadata { - fragment_bitmap: Some(std::iter::once(1_u32).collect()), - ..base - }), - ], + vec![segment_from_metadata(&plain_segment)], ) .await - .unwrap_err(); - assert!(err.to_string().contains("duplicate segment uuid")); - } - - #[tokio::test] - async fn test_commit_existing_index_segments_rejects_empty_segments() { - use lance_datagen::{BatchCount, RowCount, array}; - - let test_dir = tempfile::tempdir().unwrap(); - let test_uri = test_dir.path().to_str().unwrap(); - - let reader = lance_datagen::gen_batch() - .col("id", array::step::()) - .col( + .unwrap(); + let committed = dataset.load_indices_by_name("vector_idx").await.unwrap(); + assert_eq!(committed[0].covering_fields, vec![id_field_id]); + + // Direction 2: storage built WITH covering, declaration dropped. + let mut covered_params = + VectorIndexParams::ivf_pq(2, 8, 4, lance_linalg::distance::MetricType::L2, 2); + covered_params.covering_columns(vec!["id".to_string()]); + let mut covered_segment = dataset + .create_index_builder(&["vector"], IndexType::Vector, &covered_params) + .name("vector_idx".to_string()) + .replace(true) + .execute_uncommitted() + .await + .unwrap(); + assert_eq!(covered_segment.covering_fields, vec![id_field_id]); + // Dropping the declaration drops the carried suffix from `fields` too, leaving a + // segment that claims to key on `vector` alone while its storage still holds `id`. + covered_segment.fields.pop(); + covered_segment.covering_fields = Vec::new(); + dataset + .commit_existing_index_segments( + "vector_idx", "vector", - array::rand_vec::(8.into()), + vec![segment_from_metadata(&covered_segment)], ) - .into_reader_rows(RowCount::from(10), BatchCount::from(1)); - - let mut dataset = Dataset::write(reader, test_uri, None).await.unwrap(); - - let err = dataset - .commit_existing_index_segments("vector_idx", "vector", Vec::::new()) .await - .unwrap_err(); - assert!(err.to_string().contains("at least one index segment")); + .unwrap(); + let committed = dataset.load_indices_by_name("vector_idx").await.unwrap(); + assert!(committed[0].covering_fields.is_empty()); + + // Sanity: the honest covered segment commits fine, and the declaration survives the + // round trip. That matters beyond bookkeeping -- the segment files still hold the + // payload, so losing the declaration would silently drop covered queries back to a + // base-table take. + covered_segment.fields.push(id_field_id); + covered_segment.covering_fields = vec![id_field_id]; + dataset + .commit_existing_index_segments( + "vector_idx", + "vector", + vec![segment_from_metadata(&covered_segment)], + ) + .await + .unwrap(); + let committed = dataset.load_indices_by_name("vector_idx").await.unwrap(); + assert_eq!(committed.len(), 1); + assert_eq!(committed[0].covering_fields, vec![id_field_id]); } - /// A segment may carry extra columns; `build_index_metadata_from_segments` - /// must still commit it, keyed on the field the index is being built for. - /// Calls the guarded function directly (not through - /// `commit_existing_index_segments`) so this test's outcome depends only - /// on this guard, not on the independent duplicate a few lines into - /// `commit_existing_index_segments`. + /// A second commit whose incoming segment covers a different column set than an + /// existing, disjoint segment we would *retain* must be rejected: all segments of one + /// logical index must agree on their declaration so they expose one dependency set and + /// one declaration-ordered query schema. Their physical capabilities may still differ. #[tokio::test] - async fn test_build_index_metadata_from_segments_accepts_carried_fields() { + async fn test_commit_existing_index_segments_rejects_inconsistent_retained_covering() { use lance_datagen::{BatchCount, RowCount, array}; let test_dir = tempfile::tempdir().unwrap(); + let test_uri = test_dir.path().to_str().unwrap(); + let reader = lance_datagen::gen_batch() .col("id", array::step::()) .col( "vector", array::rand_vec::(8.into()), ) - .into_reader_rows(RowCount::from(10), BatchCount::from(1)); - let dataset = Dataset::write(reader, test_dir.path().to_str().unwrap(), None) - .await - .unwrap(); - - let vector_field_id = dataset.schema().field("vector").unwrap().id; - let id_field_id = dataset.schema().field("id").unwrap().id; - let mut metadata = write_vector_segment_metadata( - &dataset, - "vector_idx", - vector_field_id, - Uuid::new_v4(), - [0_u32], - b"segment", - ) - .await; - // Carry `id` alongside the keyed `vector` field. - metadata.fields = vec![vector_field_id, id_field_id]; - metadata.covering_fields = vec![id_field_id]; + .into_reader_rows(RowCount::from(20), BatchCount::from(1)); - let new_indices = build_index_metadata_from_segments( - &dataset, - "vector_idx", - vector_field_id, - vec![segment_from_metadata(&metadata)], + let mut dataset = Dataset::write( + reader, + test_uri, + Some(WriteParams { + max_rows_per_file: 10, + max_rows_per_group: 10, + ..Default::default() + }), ) .await .unwrap(); - assert_eq!(new_indices.len(), 1); - assert_eq!(new_indices[0].fields, vec![vector_field_id, id_field_id]); - assert_eq!(new_indices[0].covering_fields, vec![id_field_id]); - } - - /// A segment's carried columns can go stale independently of its keyed column, - /// so the staleness check has to walk every entry of `fields`, not just the - /// keyed subtree. Here the carried column did not exist when the segment was - /// built, so the segment cannot be carrying its values and its coverage of that - /// fragment must be pruned. Walking only the keyed subtree sees no change and - /// leaves the fragment covered, which is the bug. - #[tokio::test] - async fn test_prune_stale_coverage_notices_a_changed_carried_column() { - use crate::dataset::NewColumnTransform; - use lance_datagen::{BatchCount, RowCount, array}; + let id_field_id = dataset.schema().field("id").unwrap().id; - let test_dir = tempfile::tempdir().unwrap(); - let reader = lance_datagen::gen_batch() - .col("id", array::step::()) - .col( + // First commit: a real segment covering `id`, over fragment 0 only. + let mut covered_params = crate::index::vector::VectorIndexParams::ivf_flat( + 1, + lance_linalg::distance::MetricType::L2, + ); + covered_params.covering_columns(vec!["id".to_string()]); + let seg0 = dataset + .create_index_builder(&["vector"], IndexType::Vector, &covered_params) + .name("vector_idx".to_string()) + .fragments(vec![0]) + .execute_uncommitted() + .await + .unwrap(); + assert_eq!(seg0.covering_fields, vec![id_field_id]); + dataset + .commit_existing_index_segments( + "vector_idx", "vector", - array::rand_vec::(8.into()), + vec![segment_from_metadata(&seg0)], ) - .into_reader_rows(RowCount::from(10), BatchCount::from(1)); - let mut dataset = Dataset::write(reader, test_dir.path().to_str().unwrap(), None) .await .unwrap(); - let vector_field_id = dataset.schema().field("vector").unwrap().id; - // Built against the schema as it stands now, before `payload` exists. - let metadata = write_vector_segment_metadata( - &dataset, - "vector_idx", - vector_field_id, - Uuid::new_v4(), - [0_u32], - b"segment", + // Second commit: a disjoint segment (fragment 1) that covers *nothing*. It would + // be retained alongside seg0, whose covering is `[id]` -- an inconsistency. + let plain_params = crate::index::vector::VectorIndexParams::ivf_flat( + 1, + lance_linalg::distance::MetricType::L2, + ); + let seg1 = dataset + .create_index_builder(&["vector"], IndexType::Vector, &plain_params) + .name("vector_idx".to_string()) + // The name already exists (seg0 was committed); replace(true) only skips + // the duplicate-name check for this uncommitted build. + .replace(true) + .fragments(vec![1]) + .execute_uncommitted() + .await + .unwrap(); + assert!(seg1.covering_fields.is_empty()); + let result = dataset + .commit_existing_index_segments( + "vector_idx", + "vector", + vec![segment_from_metadata(&seg1)], + ) + .await; + let err = + result.expect_err("mismatched covering across retained segments must be rejected"); + assert!( + err.to_string() + .contains("a logical index cannot mix declarations") + && err.to_string().contains("covering_fields"), + "expected a covering-consistency rejection, got: {err}" + ); + } + + /// When a segment covers a column whose data file was rewritten in place after the + /// segment was built, `prune_stale_segment_coverage` must drop the affected fragment + /// from the segment's coverage -- even though the *indexed* column's file is untouched. + /// Otherwise a covered query reads the segment's stale copy of that column. + #[tokio::test] + async fn test_commit_existing_index_segments_prunes_stale_covered_fragment() { + use crate::dataset::{MergeInsertBuilder, WhenMatched, WhenNotMatched}; + + let test_dir = tempfile::tempdir().unwrap(); + let test_uri = test_dir.path().to_str().unwrap(); + + let dim = 8i32; + let n = 20i32; + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new("tag", DataType::Int32, true), + Field::new( + "vector", + DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Float32, true)), dim), + false, + ), + ])); + let vectors = FixedSizeListArray::try_new_from_values( + Float32Array::from((0..n * dim).map(|v| v as f32).collect::>()), + dim, ) - .await; + .unwrap(); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(Int32Array::from((0..n).collect::>())), + Arc::new(Int32Array::from((0..n).collect::>())), + Arc::new(vectors), + ], + ) + .unwrap(); + // Two fragments (ids 0..10 -> frag 0, 10..20 -> frag 1). + let mut dataset = Dataset::write( + RecordBatchIterator::new([Ok(batch)], schema.clone()), + test_uri, + Some(WriteParams { + max_rows_per_file: 10, + max_rows_per_group: 10, + ..Default::default() + }), + ) + .await + .unwrap(); - // `payload` lands in a new data file on the same fragment, so the fragment's - // file layout changes for `payload` while `vector`'s file is untouched. + // A scalar index on the merge key routes the partial-schema update through the + // in-place column-rewrite path (RewriteColumns) instead of a row-move, so the + // covered column's file changes while the fragment id is preserved. dataset - .add_columns( - NewColumnTransform::SqlExpressions(vec![("payload".into(), "id * 2".into())]), - None, + .create_index( + &["id"], + IndexType::BTree, None, + &ScalarIndexParams::default(), + false, ) .await .unwrap(); - let payload_field_id = dataset.schema().field("payload").unwrap().id; - let mut covered = metadata.clone(); - covered.fields = vec![vector_field_id, payload_field_id]; - covered.covering_fields = vec![payload_field_id]; + let tag_field_id = dataset.schema().field("tag").unwrap().id; + + // Build a real covering segment BEFORE the update, so it is stamped at the + // pre-update version and covers `tag`, whose files still hold the old values. + let mut params = crate::index::vector::VectorIndexParams::ivf_flat( + 1, + lance_linalg::distance::MetricType::L2, + ); + params.covering_columns(vec!["tag".to_string()]); + let seg = dataset + .create_index_builder(&["vector"], IndexType::Vector, ¶ms) + .name("vector_idx".to_string()) + .execute_uncommitted() + .await + .unwrap(); + assert_eq!(seg.covering_fields, vec![tag_field_id]); + + // In-place rewrite of the covered `tag` column for fragment 0 only (no covering + // index committed yet, so this patches the column file in place and keeps the + // fragment id). The indexed `vector` column's file is untouched. + let sub_schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new("tag", DataType::Int32, true), + ])); + let update = RecordBatch::try_new( + sub_schema.clone(), + vec![ + Arc::new(Int32Array::from((0..5).collect::>())), + Arc::new(Int32Array::from((100..105).collect::>())), + ], + ) + .unwrap(); + let (updated, _) = MergeInsertBuilder::try_new(Arc::new(dataset), vec!["id".to_string()]) + .unwrap() + .when_matched(WhenMatched::UpdateAll) + .when_not_matched(WhenNotMatched::DoNothing) + .try_build() + .unwrap() + .execute_reader(Box::new(RecordBatchIterator::new([Ok(update)], sub_schema))) + .await + .unwrap(); + let mut dataset = Arc::try_unwrap(updated).unwrap_or_else(|arc| (*arc).clone()); + + dataset + .commit_existing_index_segments( + "vector_idx", + "vector", + vec![segment_from_metadata(&seg)], + ) + .await + .unwrap(); + + let committed = dataset.load_indices_by_name("vector_idx").await.unwrap(); + assert_eq!(committed.len(), 1); + let bitmap = committed[0].fragment_bitmap.as_ref().unwrap(); + assert!( + !bitmap.contains(0), + "fragment 0's covered column was rewritten; its stale coverage must be pruned" + ); + assert!( + bitmap.contains(1), + "fragment 1's covered column was untouched; it stays covered" + ); + } + + /// A covered *struct* that grows a child (via `add_columns`) after a segment was built + /// makes the segment's stored covering payload type-incompatible for every covered row -- + /// a global change no per-fragment prune can repair. Since the covering index is not + /// committed while the segment is staged, the Merge-commit guard never sees the add; + /// `prune_stale_segment_coverage` must reject the segment commit instead. + #[tokio::test] + async fn test_commit_existing_index_segments_rejects_grown_covered_struct() { + use arrow_array::StructArray; + use arrow_schema::Fields; + use lance_file::version::LanceFileVersion; + + use crate::dataset::NewColumnTransform; + + let test_dir = tempfile::tempdir().unwrap(); + let test_uri = test_dir.path().to_str().unwrap(); + + let dim = 8i32; + let n = 10i32; + let meta_fields = Fields::from(vec![Field::new("a", DataType::Int32, false)]); + let schema = Arc::new(Schema::new(vec![ + Field::new("meta", DataType::Struct(meta_fields.clone()), true), + Field::new( + "vector", + DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Float32, true)), dim), + false, + ), + ])); + let a = Arc::new(Int32Array::from((0..n).collect::>())); + let meta = Arc::new(StructArray::new(meta_fields, vec![a], None)); + let vectors = FixedSizeListArray::try_new_from_values( + Float32Array::from((0..n * dim).map(|v| v as f32).collect::>()), + dim, + ) + .unwrap(); + let batch = RecordBatch::try_new(schema.clone(), vec![meta, Arc::new(vectors)]).unwrap(); + let mut dataset = Dataset::write( + RecordBatchIterator::new([Ok(batch)], schema.clone()), + test_uri, + Some(WriteParams { + data_storage_version: Some(LanceFileVersion::V2_2), + ..Default::default() + }), + ) + .await + .unwrap(); + + let meta_field_id = dataset.schema().field("meta").unwrap().id; + + // Build a real covering segment stamped at the current (pre-add) version, + // covering the struct column `meta`. + let mut params = crate::index::vector::VectorIndexParams::ivf_flat( + 1, + lance_linalg::distance::MetricType::L2, + ); + params.covering_columns(vec!["meta".to_string()]); + let seg = dataset + .create_index_builder(&["vector"], IndexType::Vector, ¶ms) + .name("vector_idx".to_string()) + .execute_uncommitted() + .await + .unwrap(); + assert_eq!(seg.covering_fields, vec![meta_field_id]); + + // Grow the covered struct AFTER the segment was built. No covering index is + // committed yet, so the Merge-commit guard does not fire and the add succeeds. + let add = + NewColumnTransform::AllNulls(Arc::new(arrow_schema::Schema::new(vec![Field::new( + "meta", + DataType::Struct(Fields::from(vec![Field::new("b", DataType::Int32, true)])), + true, + )]))); + dataset.add_columns(add, None, None).await.unwrap(); + + // Committing the now-stale segment must be rejected: its stored covering payload + // still holds the pre-child struct type. + let err = dataset + .commit_existing_index_segments( + "vector_idx", + "vector", + vec![segment_from_metadata(&seg)], + ) + .await + .expect_err("committing a segment covering a since-grown struct must be rejected"); + assert!( + err.to_string().contains("changed its subtree"), + "expected a covered-subtree-change rejection, got: {err}" + ); + } + + #[tokio::test] + async fn test_commit_existing_index_segments_rejects_duplicate_segment_ids() { + use lance_datagen::{BatchCount, RowCount, array}; + + let test_dir = tempfile::tempdir().unwrap(); + let test_uri = test_dir.path().to_str().unwrap(); + + let reader = lance_datagen::gen_batch() + .col("id", array::step::()) + .col( + "vector", + array::rand_vec::(8.into()), + ) + .into_reader_rows(RowCount::from(10), BatchCount::from(1)); + + let mut dataset = Dataset::write(reader, test_uri, None).await.unwrap(); + + let field_id = dataset.schema().field("vector").unwrap().id; + let base = write_vector_segment_metadata( + &dataset, + "vector_idx", + field_id, + Uuid::new_v4(), + [0_u32], + b"base", + ) + .await; + + let err = dataset + .commit_existing_index_segments( + "vector_idx", + "vector", + vec![ + segment_from_metadata(&base), + segment_from_metadata(&IndexMetadata { + fragment_bitmap: Some(std::iter::once(1_u32).collect()), + ..base + }), + ], + ) + .await + .unwrap_err(); + assert!(err.to_string().contains("duplicate segment uuid")); + } + + #[tokio::test] + async fn test_commit_existing_index_segments_rejects_empty_segments() { + use lance_datagen::{BatchCount, RowCount, array}; + + let test_dir = tempfile::tempdir().unwrap(); + let test_uri = test_dir.path().to_str().unwrap(); + + let reader = lance_datagen::gen_batch() + .col("id", array::step::()) + .col( + "vector", + array::rand_vec::(8.into()), + ) + .into_reader_rows(RowCount::from(10), BatchCount::from(1)); + + let mut dataset = Dataset::write(reader, test_uri, None).await.unwrap(); + + let err = dataset + .commit_existing_index_segments("vector_idx", "vector", Vec::::new()) + .await + .unwrap_err(); + assert!(err.to_string().contains("at least one index segment")); + } + + /// A segment may carry extra columns; `build_index_metadata_from_segments` + /// must still commit it, keyed on the field the index is being built for. + /// Calls the guarded function directly (not through + /// `commit_existing_index_segments`) so this test's outcome depends only + /// on this guard, not on the independent duplicate a few lines into + /// `commit_existing_index_segments`. + #[tokio::test] + async fn test_build_index_metadata_from_segments_accepts_carried_fields() { + use crate::index::vector::VectorIndexParams; + use lance_datagen::{BatchCount, RowCount, array}; + + let test_dir = tempfile::tempdir().unwrap(); + let reader = lance_datagen::gen_batch() + .col("id", array::step::()) + .col( + "vector", + array::rand_vec::(8.into()), + ) + .into_reader_rows(RowCount::from(256), BatchCount::from(1)); + let mut dataset = Dataset::write(reader, test_dir.path().to_str().unwrap(), None) + .await + .unwrap(); + + let vector_field_id = dataset.schema().field("vector").unwrap().id; + let id_field_id = dataset.schema().field("id").unwrap().id; + // A genuinely covered segment, not a hand-edited declaration: the commit boundary + // now cross-checks `covering_fields` against the payload in the segment's auxiliary + // storage, so `id` has to really be there. + let mut params = + VectorIndexParams::ivf_pq(2, 8, 4, lance_linalg::distance::MetricType::L2, 2); + params.covering_columns(vec!["id".to_string()]); + let metadata = dataset + .create_index_builder(&["vector"], IndexType::Vector, ¶ms) + .name("vector_idx".to_string()) + .execute_uncommitted() + .await + .unwrap(); + // Carried `id` sits at the end of `fields`, behind the keyed `vector` field. + assert_eq!(metadata.fields, vec![vector_field_id, id_field_id]); + assert_eq!(metadata.covering_fields, vec![id_field_id]); + + let new_indices = build_index_metadata_from_segments( + &dataset, + "vector_idx", + vector_field_id, + vec![segment_from_metadata(&metadata)], + ) + .await + .unwrap(); + + assert_eq!(new_indices.len(), 1); + assert_eq!(new_indices[0].fields, vec![vector_field_id, id_field_id]); + assert_eq!(new_indices[0].covering_fields, vec![id_field_id]); + } + + /// A segment's carried columns can go stale independently of its keyed column, + /// so the staleness check has to walk every entry of `fields`, not just the + /// keyed subtree. Here an in-place update rewrites the carried column's data + /// file after the segment was built, so the segment's stored copy is stale and + /// its coverage of that fragment must be pruned. Walking only the keyed subtree + /// sees no change (the vector file is untouched) and leaves the fragment + /// covered, which is the bug. Both segments are genuinely built (the commit + /// boundary cross-checks a covering declaration against the auxiliary storage, + /// so hand-doctored metadata cannot reach the prune). + #[tokio::test] + async fn test_prune_stale_coverage_notices_a_changed_carried_column() { + use crate::index::vector::VectorIndexParams; + use lance_datagen::{BatchCount, RowCount, array}; + + let test_dir = tempfile::tempdir().unwrap(); + let reader = lance_datagen::gen_batch() + .col("id", array::step::()) + .col("payload", array::step::()) + .col( + "vector", + array::rand_vec::(8.into()), + ) + .into_reader_rows(RowCount::from(256), BatchCount::from(1)); + let mut dataset = Dataset::write(reader, test_dir.path().to_str().unwrap(), None) + .await + .unwrap(); + + let vector_field_id = dataset.schema().field("vector").unwrap().id; + + // A real covered segment carrying `payload`, and a plain control segment, + // both built against the current file layout. + let mut covered_params = + VectorIndexParams::ivf_pq(2, 8, 4, lance_linalg::distance::MetricType::L2, 2); + covered_params.covering_columns(vec!["payload".to_string()]); + let covered = dataset + .create_index_builder(&["vector"], IndexType::Vector, &covered_params) + .name("vector_idx".to_string()) + .execute_uncommitted() + .await + .unwrap(); + let plain_params = + VectorIndexParams::ivf_pq(2, 8, 4, lance_linalg::distance::MetricType::L2, 2); + let plain = dataset + .create_index_builder(&["vector"], IndexType::Vector, &plain_params) + .name("vector_idx".to_string()) + .execute_uncommitted() + .await + .unwrap(); + + // Rewrite `payload`'s backing file IN PLACE (same fragment id, new data + // file for `payload` only); `vector`'s file is untouched. A row-moving + // update would replace the fragment and prune BOTH segments' coverage, + // which the control below would catch. + let update_schema = Arc::new(arrow_schema::Schema::new(vec![ + arrow_schema::Field::new("id", arrow_schema::DataType::Int32, false), + arrow_schema::Field::new("payload", arrow_schema::DataType::Int32, false), + ])); + let update_batch = RecordBatch::try_new( + update_schema.clone(), + vec![ + Arc::new(arrow_array::Int32Array::from(vec![2, 4, 6])), + Arc::new(arrow_array::Int32Array::from(vec![99, 99, 99])), + ], + ) + .unwrap(); + let right: Box = + Box::new(arrow_array::RecordBatchIterator::new( + vec![Ok(update_batch)].into_iter(), + update_schema, + )); + let mut frag0 = dataset.get_fragment(0).unwrap(); + let (updated_fragment, fields_modified) = + frag0.update_columns(right, "id", "id").await.unwrap(); + let op = Operation::Update { + removed_fragment_ids: vec![], + updated_fragments: vec![updated_fragment], + new_fragments: vec![], + fields_modified, + compacted_sstables: Vec::new(), + fields_for_preserving_frag_bitmap: vec![], + update_mode: Some(crate::dataset::transaction::UpdateMode::RewriteColumns), + inserted_rows_filter: None, + updated_fragment_offsets: None, + }; + let read_version = dataset.manifest.version; + let mut dataset = dataset; + dataset + .apply_commit( + Transaction::new(read_version, op, None), + &Default::default(), + &Default::default(), + ) + .await + .unwrap(); let new_indices = build_index_metadata_from_segments( &dataset, @@ -9020,12 +9639,12 @@ mod tests { // Control: with nothing carried, the same segment over the same fragment // keeps its coverage -- so the assertion above is about the carried column, - // not about `add_columns` invalidating everything. + // not about the update invalidating everything. let plain_indices = build_index_metadata_from_segments( &dataset, "vector_idx", vector_field_id, - vec![segment_from_metadata(&metadata)], + vec![segment_from_metadata(&plain)], ) .await .unwrap(); @@ -9035,7 +9654,7 @@ mod tests { .as_ref() .unwrap() .is_empty(), - "a non-covering segment must keep its coverage across the same add_columns" + "a non-covering segment must keep its coverage across the same update" ); } @@ -9153,6 +9772,109 @@ mod tests { ); } + /// `covering_fields` is a format-level declaration shared by every index type. A + /// multi-field `fields` with a carried suffix must therefore remain legal even when + /// that index implementation has no physical covering capability yet. + /// A covered SCALAR index must be withdrawn by remap, not remapped. + /// + /// `remap_index` stopped withdrawing covered indexes on the grounds that "each storage + /// format's `remap` preserves its covering columns row-aligned". That holds for the four + /// vector storages; `BTreeIndex::remap` rewrites only `(value, row_id)`. Remapping one + /// republishes a declaration nothing backs, and because carried fields live in `fields` + /// the index then loses fragment coverage on every update to a column it never held. + #[tokio::test] + async fn test_remap_withdraws_a_covered_scalar_index() { + use lance_datagen::{BatchCount, RowCount, array}; + + let test_dir = tempfile::tempdir().unwrap(); + let reader = lance_datagen::gen_batch() + .col("id", array::step::()) + .col("other", array::step::()) + .into_reader_rows(RowCount::from(64), BatchCount::from(1)); + let mut dataset = Dataset::write(reader, test_dir.path().to_str().unwrap(), None) + .await + .unwrap(); + + let id_field_id = dataset.schema().field("id").unwrap().id; + let other_field_id = dataset.schema().field("other").unwrap().id; + + let btree_params = ScalarIndexParams::for_builtin(BuiltinIndexType::BTree); + dataset + .create_index( + &["id"], + IndexType::BTree, + Some("id_idx".into()), + &btree_params, + true, + ) + .await + .unwrap(); + + // Actually COMMIT the covered declaration -- mutating a loaded copy proves nothing, + // because `remap_index` re-loads the indices from the dataset itself. + crate::utils::test::covering::declare_covering(&mut dataset, "id", "other").await; + let indices = dataset.load_indices().await.unwrap(); + let committed = indices.iter().find(|i| i.name == "id_idx").unwrap(); + assert_eq!( + committed.covering_fields, + vec![other_field_id], + "precondition: the covered declaration must be committed, or this test is vacuous" + ); + assert_eq!(committed.fields, vec![id_field_id, other_field_id]); + let index_id = committed.uuid; + + let result = remap_index(&dataset, &index_id, &RowAddrRemap::empty()).await; + assert!( + matches!(result, Ok(RemapResult::Drop)), + "a covered scalar index must be withdrawn rather than remapped without its \ + payload; got {result:?}" + ); + } + + #[tokio::test] + async fn test_build_index_metadata_from_segments_accepts_carried_fields_on_scalar() { + use lance_datagen::{BatchCount, RowCount, array}; + + let test_dir = tempfile::tempdir().unwrap(); + let reader = lance_datagen::gen_batch() + .col("id", array::step::()) + .col("other", array::step::()) + .into_reader_rows(RowCount::from(64), BatchCount::from(1)); + let mut dataset = Dataset::write(reader, test_dir.path().to_str().unwrap(), None) + .await + .unwrap(); + + let id_field_id = dataset.schema().field("id").unwrap().id; + let other_field_id = dataset.schema().field("other").unwrap().id; + + let btree_params = ScalarIndexParams::for_builtin(BuiltinIndexType::BTree); + let mut metadata = dataset + .create_index_builder(&["id"], IndexType::BTree, &btree_params) + .name("id_idx".to_string()) + .execute_uncommitted() + .await + .unwrap(); + // No producer builds covered scalar indexes (`covering_columns` lives only on + // `VectorIndexParams`), so the declaration is hand-constructed on a real BTree + // segment. That is the shape a distributed caller can hand to the public + // `commit_existing_index_segments`, which is what this guard has to tolerate. + metadata.fields = vec![id_field_id, other_field_id]; + metadata.covering_fields = vec![other_field_id]; + + let new_indices = build_index_metadata_from_segments( + &dataset, + "id_idx", + id_field_id, + vec![segment_from_metadata(&metadata)], + ) + .await + .expect("a carried-field declaration must not require physical payload at commit"); + + assert_eq!(new_indices.len(), 1); + assert_eq!(new_indices[0].fields, vec![id_field_id, other_field_id]); + assert_eq!(new_indices[0].covering_fields, vec![other_field_id]); + } + /// A segment keyed on the wrong field must still be rejected. Calls /// `build_index_metadata_from_segments` directly so the assertion is /// evidence for *this* guard specifically -- going through @@ -9243,6 +9965,75 @@ mod tests { assert!(error.to_string().contains("are not among its fields")); } + /// Two segments passed to a single `build_index_metadata_from_segments` + /// call must agree on `covering_fields`: the read path (`knn.rs`, + /// `scanner.rs`) derives its declared exec output schema from a single + /// segment and assumes every sibling of the same logical index agrees. + /// Each segment individually satisfies the per-segment suffix rule + /// (`validate_covering_fields`), so only a cross-segment comparison + /// catches the disagreement. Calls `build_index_metadata_from_segments` + /// directly, not through `commit_existing_index_segments`, so the + /// assertion is evidence for this guard specifically. + #[tokio::test] + async fn test_build_index_metadata_from_segments_rejects_covering_fields_disagreement() { + use lance_datagen::{BatchCount, RowCount, array}; + + let test_dir = tempfile::tempdir().unwrap(); + let reader = lance_datagen::gen_batch() + .col("id", array::step::()) + .col( + "vector", + array::rand_vec::(8.into()), + ) + .into_reader_rows(RowCount::from(10), BatchCount::from(1)); + let dataset = Dataset::write(reader, test_dir.path().to_str().unwrap(), None) + .await + .unwrap(); + + let vector_field_id = dataset.schema().field("vector").unwrap().id; + let id_field_id = dataset.schema().field("id").unwrap().id; + + let mut covered = write_vector_segment_metadata( + &dataset, + "vector_idx", + vector_field_id, + Uuid::new_v4(), + [0_u32], + b"seg0", + ) + .await; + covered.fields = vec![vector_field_id, id_field_id]; + covered.covering_fields = vec![id_field_id]; + + let uncovered = write_vector_segment_metadata( + &dataset, + "vector_idx", + vector_field_id, + Uuid::new_v4(), + [1_u32], + b"seg1", + ) + .await; + + let error = build_index_metadata_from_segments( + &dataset, + "vector_idx", + vector_field_id, + vec![ + segment_from_metadata(&covered), + segment_from_metadata(&uncovered), + ], + ) + .await + .unwrap_err(); + assert!( + error + .to_string() + .contains("every segment of one index must declare the same columns"), + "unexpected error: {error}" + ); + } + #[tokio::test] async fn test_commit_existing_index_segments_rejects_wrong_field_provenance() { use lance_datagen::{BatchCount, RowCount, array}; diff --git a/rust/lance/src/index/append.rs b/rust/lance/src/index/append.rs index 8bfb32aec75..c7e6e420752 100644 --- a/rust/lance/src/index/append.rs +++ b/rust/lance/src/index/append.rs @@ -663,6 +663,7 @@ async fn build_fresh_vector_segment( logical_index: &LogicalVectorIndex, field_path: &str, fragment_bitmap: &RoaringBitmap, + covering_columns: &[String], progress: Arc, ) -> Result { let (reference_metadata, reference_index) = logical_index.iter().last().ok_or_else(|| { @@ -671,7 +672,12 @@ async fn build_fresh_vector_segment( logical_index.name() )) })?; - let params = fresh_vector_segment_params(reference_metadata, reference_index.as_ref())?; + let mut params = fresh_vector_segment_params(reference_metadata, reference_index.as_ref())?; + // A retrain rebuilds the storage from scratch, so re-declare the covering + // ("included") columns -- otherwise the fresh files omit the payload while the + // committed metadata still advertises `covering_fields`, and covered queries + // then expect columns the storage cannot emit. + crate::index::vector::redeclare_covering(&mut params, covering_columns.to_vec(), field_path); let mut build_dataset = dataset.clone(); CreateIndexBuilder::new( &mut build_dataset, @@ -692,12 +698,26 @@ async fn scan_vector_fragments( field_path: &str, column_nullable: bool, fragments: &[Fragment], + covering_columns: &[String], ) -> Result { let mut scanner = dataset.scan(); + // Project the vector column plus any covering ("included") columns so rows from the + // newly-indexed fragments carry the same columns as the existing partitions. + // The indexed column is already the first entry, and can also appear in + // `covering_columns`: carried refine vectors live in storage under that column's own + // name, so a rebuild reads it back as covering. The projection builder rejects a + // repeated name outright. + let mut projection: Vec<&str> = vec![field_path]; + projection.extend( + covering_columns + .iter() + .map(String::as_str) + .filter(|name| *name != field_path), + ); scanner .with_fragments(fragments.to_vec()) .with_row_id() - .project(&[field_path])?; + .project(&projection)?; if column_nullable { let column_expr = lance_datafusion::logical_expr::field_path_to_expr(field_path)?; scanner.filter_expr(column_expr.is_not_null()); @@ -788,6 +808,29 @@ pub async fn merge_indices_with_unindexed_frags<'a>( .as_ref() .map(|resolved| resolved.canonical_path.clone()) .unwrap_or(raw_field_path); + // Covering ("included") columns recorded on the index, resolved to names. + // Threaded into the merge builder and projected into the new-data scan so + // they survive optimize/merge and partition split/join. Resolution goes + // through `effective_covering` -- the single authority for covering + // resolution -- so it is top-level only and fails loudly on an unresolvable + // id, exactly like the read path: silently rebuilding without the column + // would leave the committed metadata advertising a column the storage + // cannot emit, and a recursive lookup would resolve a (corrupt or foreign) + // nested id to a leaf name the read path errors on. + let covering_columns: Vec = crate::index::covering::effective_covering( + &old_indices[0].covering_fields, + None, + dataset.schema(), + ) + .map_err(|e| { + Error::index(format!( + "Append index: covering declaration of index '{}' cannot be resolved: {e}", + old_indices[0].name + )) + })? + .into_iter() + .map(|field| field.name().clone()) + .collect(); let first_is_vector_index = metadata_is_vector_index(dataset.as_ref(), old_indices[0]).await?; for idx in old_indices.iter().skip(1) { let is_vector_index = metadata_is_vector_index(dataset.as_ref(), idx).await?; @@ -878,6 +921,7 @@ pub async fn merge_indices_with_unindexed_frags<'a>( &logical_index, &field_path, &fragment_bitmap, + &covering_columns, options.progress.clone(), ) .await?; @@ -920,6 +964,7 @@ pub async fn merge_indices_with_unindexed_frags<'a>( &field_path, column.nullable, unindexed, + &covering_columns, ) .await?; let mut append_options = options.clone(); @@ -931,6 +976,7 @@ pub async fn merge_indices_with_unindexed_frags<'a>( &field_path, &reference_ivf_view, &append_options, + &covering_columns, ) .boxed() .await?; @@ -1012,6 +1058,7 @@ pub async fn merge_indices_with_unindexed_frags<'a>( &field_path, &selected_ivf_view, options, + &covering_columns, )) .await?; if indices_merged == 0 { @@ -1067,6 +1114,7 @@ pub async fn merge_indices_with_unindexed_frags<'a>( &field_path, column.nullable, unindexed, + &covering_columns, ) .await?, ) @@ -1078,6 +1126,7 @@ pub async fn merge_indices_with_unindexed_frags<'a>( &field_path, &merge_ivf_view, options, + &covering_columns, ) .boxed() .await?; @@ -2954,6 +3003,69 @@ mod tests { ); } + /// An index whose `covering_fields` records a field id missing from the current + /// schema is metadata/schema-inconsistent. Optimize must fail loudly rather than + /// silently rebuild without the covered column: the merged delta would be committed + /// still advertising the column, and every covered query would then fail at read + /// time. + #[tokio::test] + async fn test_merge_indices_errors_on_unresolvable_covered_field() { + const DIM: usize = 16; + const TOTAL: usize = 256; + + let test_dir = TempStrDir::default(); + let test_uri = test_dir.as_str(); + + let vectors = generate_random_array(TOTAL * DIM); + let schema = Arc::new(Schema::new(vec![ + Field::new( + "vector", + DataType::FixedSizeList( + Arc::new(Field::new("item", DataType::Float32, true)), + DIM as i32, + ), + true, + ), + Field::new("id", DataType::UInt32, false), + ])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(FixedSizeListArray::try_new_from_values(vectors, DIM as i32).unwrap()), + Arc::new(UInt32Array::from_iter_values(0..TOTAL as u32)), + ], + ) + .unwrap(); + let batches = RecordBatchIterator::new(vec![batch].into_iter().map(Ok), schema.clone()); + let mut dataset = Dataset::write(batches, test_uri, None).await.unwrap(); + let index_params = VectorIndexParams::ivf_pq(2, 8, 4, MetricType::L2, 2); + dataset + .create_index(&["vector"], IndexType::Vector, None, &index_params, true) + .await + .unwrap(); + + let indices = dataset.load_indices_by_name("vector_idx").await.unwrap(); + let mut stale_index = indices[0].clone(); + stale_index.covering_fields = vec![9999]; + let old_indices = vec![&stale_index]; + let result = merge_indices_with_unindexed_frags( + Arc::new(dataset), + &old_indices, + &[], + &OptimizeOptions::merge(1), + ) + .await; + let err = result.unwrap_err(); + assert!( + matches!(err, Error::Index { .. }), + "unexpected error: {err}" + ); + assert!( + err.to_string().contains("9999"), + "error should name the unresolvable covered field id: {err}" + ); + } + #[tokio::test] async fn test_optimize_btree_multi_segment_optimize_default() { async fn query_id_count(dataset: &Dataset, id: &str) -> usize { @@ -3095,6 +3207,86 @@ mod tests { ); } + /// Covering resolution is top-level only (creation rejects dotted paths, so + /// `covering_fields` can only legally hold top-level ids). A NESTED field id -- + /// possible only in corrupt or foreign metadata -- must therefore fail loudly + /// here too, not resolve through the recursive `Schema::field_by_id` to a leaf + /// name and rebuild storage under a name the read path errors on: commit-time + /// and query-time must not disagree about the same index. + #[tokio::test] + async fn test_merge_indices_rejects_nested_covered_field_id() { + const DIM: usize = 16; + const TOTAL: usize = 256; + + let test_dir = TempStrDir::default(); + let test_uri = test_dir.as_str(); + + let vectors = generate_random_array(TOTAL * DIM); + let schema = Arc::new(Schema::new(vec![ + Field::new( + "vector", + DataType::FixedSizeList( + Arc::new(Field::new("item", DataType::Float32, true)), + DIM as i32, + ), + true, + ), + Field::new( + "meta", + DataType::Struct(vec![Field::new("code", DataType::UInt32, false)].into()), + false, + ), + ])); + let codes = Arc::new(UInt32Array::from_iter_values(0..TOTAL as u32)); + let meta = Arc::new(arrow_array::StructArray::from(vec![( + Arc::new(Field::new("code", DataType::UInt32, false)), + codes as arrow_array::ArrayRef, + )])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(FixedSizeListArray::try_new_from_values(vectors, DIM as i32).unwrap()), + meta, + ], + ) + .unwrap(); + let batches = RecordBatchIterator::new(vec![batch].into_iter().map(Ok), schema.clone()); + let mut dataset = Dataset::write(batches, test_uri, None).await.unwrap(); + let index_params = VectorIndexParams::ivf_pq(2, 8, 4, MetricType::L2, 2); + dataset + .create_index(&["vector"], IndexType::Vector, None, &index_params, true) + .await + .unwrap(); + + // The LEAF id of meta.code -- resolvable by the recursive lookup, but not a + // top-level field. + let leaf_id = dataset + .schema() + .field("meta") + .unwrap() + .children + .first() + .unwrap() + .id; + + let indices = dataset.load_indices_by_name("vector_idx").await.unwrap(); + let mut nested_covered = indices[0].clone(); + nested_covered.covering_fields = vec![leaf_id]; + let old_indices = vec![&nested_covered]; + let err = merge_indices_with_unindexed_frags( + Arc::new(dataset), + &old_indices, + &[], + &OptimizeOptions::merge(1), + ) + .await + .unwrap_err(); + assert!( + err.to_string().contains("top-level"), + "a nested covered field id must be rejected as non-top-level: {err}" + ); + } + #[tokio::test] async fn test_optimize_fmindex_default_rebuilds_old_and_new_rows() { let test_dir = TempStrDir::default(); diff --git a/rust/lance/src/index/covering.rs b/rust/lance/src/index/covering.rs new file mode 100644 index 00000000000..5f4bb92605b --- /dev/null +++ b/rust/lance/src/index/covering.rs @@ -0,0 +1,283 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +//! Read-path resolution of an index's covering ("included") columns. +//! +//! # Declaration versus capability +//! +//! Covering has two deliberately independent sources of truth: +//! +//! * the **manifest declaration** — `IndexMetadata::covering_fields`, a dependency list of +//! dataset field ids in declaration order; +//! * the **physical capability** — columns the selected segment's storage schema proves it +//! can serve, including their source dataset field ids. +//! +//! The declaration is not evidence that the payload exists. Transitional writers and +//! lifecycle operations may preserve the dependency while omitting or withdrawing the +//! physical values. Readers may skip a base-table take only for the declaration/capability +//! intersection proven across every selected segment. +//! +//! This module resolves only the manifest side. Storage-aware planners must narrow its +//! result through [`VectorIndex::physical_covering_fields`] before adding columns to an +//! executable output schema. +//! +//! Nothing here is vector-specific: `covering_fields` lives on `IndexMetadata`, so the +//! FTS execs resolve covering exactly the same way and must not re-derive it. +//! +//! [`VectorIndex::physical_covering_fields`]: lance_index::vector::VectorIndex::physical_covering_fields + +use arrow_schema::Field as ArrowField; +use lance_core::datatypes::Schema; +use lance_core::{Error, Result}; + +/// The logically requested covering columns for one query: everything the manifest +/// declares, narrowed to that query's covering projection. +/// +/// `projection` is [`Query::covering_projection`], whose three states this function is +/// the primary consumer of: +/// +/// | `projection` | result | +/// |---------------|----------------------------------------------| +/// | `None` | every declared covering column | +/// | `Some(&[])` | empty — the caller must do no covering work | +/// | `Some(cols)` | the declared columns named in `cols` | +/// +/// An empty result is a *contract*, not an incidental outcome: it means the caller skips +/// the covering read altogether rather than projecting an already-loaded batch down to +/// zero columns. Note that `None` and `Some(&[])` therefore differ enormously in cost +/// while being indistinguishable in results. +/// +/// Returns resolved Arrow **fields**, not ids. This is not yet an executable capability: +/// callers that plan a storage read must intersect these fields with the physical segment +/// schemas before declaring them in an output plan. +/// +/// # Order +/// +/// The result follows `covering_fields`, never the order of `projection`. That is the +/// logical output order. Physical storage must prove that it can emit the selected fields +/// in a compatible order; otherwise the reader falls back to the base table. +/// +/// # Errors +/// +/// A declared id absent from the dataset schema's top-level fields is an error, and is +/// reported even when `projection` would have narrowed that column away. The manifest is +/// authoritative for dependencies: such a mismatch means index metadata and schema +/// disagree, and letting a narrowing decision suppress it would turn a corrupt index into +/// one that appears to work for some queries. For the same reason callers must run any cross-segment +/// agreement checks on the *declared* sets before calling this — narrowing first can +/// reduce genuinely disagreeing segments to a subset on which they happen to agree. +/// +/// [`Query::covering_projection`]: lance_index::vector::Query::covering_projection +pub fn effective_covering( + covering_fields: &[i32], + projection: Option<&[String]>, + schema: &Schema, +) -> Result> { + Ok( + effective_covering_with_ids(covering_fields, projection, schema)? + .into_iter() + .map(|(_, field)| field) + .collect(), + ) +} + +/// [`effective_covering`] with each resolved Arrow field still paired to the manifest +/// field id that selected it. Storage-capability checks use this form so logical identity +/// is never reconstructed from a column name. +pub fn effective_covering_with_ids( + covering_fields: &[i32], + projection: Option<&[String]>, + schema: &Schema, +) -> Result> { + if covering_fields.is_empty() { + return Ok(Vec::new()); + } + let mut fields = Vec::with_capacity(covering_fields.len()); + for id in covering_fields { + // Top-level lookup only, deliberately not `Schema::field_by_id`: that recurses + // into struct children, and a nested field matched here would be emitted under + // its leaf name -- a name a covered projection cannot address, and one that can + // collide with an unrelated top-level column. Covering columns are top-level by + // construction; `validate_covering_columns` rejects dotted paths at create time. + let field = schema + .fields + .iter() + .find(|field| field.id == *id) + .ok_or_else(|| { + Error::index(format!( + "index declares covering field id {id}, which is not present as a \ + top-level field in the current dataset schema; index metadata and \ + schema are inconsistent" + )) + })?; + // Convert just this field. `ArrowSchema::from(schema)` is exactly this conversion + // mapped over every field, but it walks the whole table (including nested + // structs) on every covered plan. + if projection.is_none_or(|wanted| wanted.iter().any(|name| name == &field.name)) { + fields.push((*id, ArrowField::from(field))); + } + } + Ok(fields) +} + +#[cfg(test)] +mod tests { + use super::*; + use arrow_schema::{DataType, Field as ArrowField, Fields, Schema as ArrowSchema}; + + /// `price` and `payload` are the covering candidates; `nested.price` exists to prove + /// that a nested field sharing a top-level name is never what gets resolved. + fn schema() -> Schema { + let arrow = ArrowSchema::new(vec![ + ArrowField::new("vec", DataType::Int32, false), + ArrowField::new("price", DataType::Int32, true), + ArrowField::new("payload", DataType::Utf8, true), + ArrowField::new( + "nested", + DataType::Struct(Fields::from(vec![ArrowField::new( + "price", + DataType::Float64, + true, + )])), + true, + ), + ]); + Schema::try_from(&arrow).unwrap() + } + + fn names(fields: &[ArrowField]) -> Vec<&str> { + fields.iter().map(|f| f.name().as_str()).collect() + } + + fn field_id(schema: &Schema, name: &str) -> i32 { + schema.field(name).unwrap().id + } + + #[test] + fn no_projection_keeps_every_declared_column() { + let s = schema(); + let ids = vec![field_id(&s, "price"), field_id(&s, "payload")]; + let got = effective_covering(&ids, None, &s).unwrap(); + assert_eq!( + names(&got), + vec!["price", "payload"], + "`None` means no narrowing was computed, so every declared column is emitted" + ); + } + + #[test] + fn projection_narrows_to_what_the_query_needs() { + let s = schema(); + let ids = vec![field_id(&s, "price"), field_id(&s, "payload")]; + let projection = vec!["payload".to_string()]; + let got = effective_covering(&ids, Some(&projection), &s).unwrap(); + assert_eq!(names(&got), vec!["payload"]); + } + + /// The state the whole feature exists for. It is distinct from `None`, which yields + /// the full declared set -- conflating them silently restores full materialization. + #[test] + fn empty_projection_is_not_the_same_as_no_projection() { + let s = schema(); + let ids = vec![field_id(&s, "price"), field_id(&s, "payload")]; + let narrowed = effective_covering(&ids, Some(&[]), &s).unwrap(); + assert!( + narrowed.is_empty(), + "`Some(&[])` means this query needs no covering column at all" + ); + let unnarrowed = effective_covering(&ids, None, &s).unwrap(); + assert_eq!( + names(&unnarrowed), + vec!["price", "payload"], + "`None` must NOT behave like `Some(&[])`; if these two agree the seam is dead" + ); + } + + /// Emission order is the index's declaration order, never the caller's request order + /// and never schema order. Both alternatives pair values with the wrong names. + /// + /// The declared order below is deliberately the reverse of the schema's ascending + /// field-id order (vec=0, price=1, payload=2, ...), and the requested order is + /// deliberately different again. An implementation that walked `schema.fields()` and + /// filtered by membership, or that echoed the request order, produces `["price", + /// "payload"]` and fails. Do not "tidy" either list back into agreement. + #[test] + fn preserves_declaration_order_not_schema_or_request_order() { + let s = schema(); + let declared = vec![field_id(&s, "payload"), field_id(&s, "price")]; + let requested = vec!["price".to_string(), "payload".to_string()]; + let got = effective_covering(&declared, Some(&requested), &s).unwrap(); + assert_eq!( + names(&got), + vec!["payload", "price"], + "declaration order is storage order; any other order mismatches the batch \ + the storage produces" + ); + + // Permuting only the declaration must permute only the result: that is what + // makes this test sensitive to order rather than to membership. + let permuted = vec![field_id(&s, "price"), field_id(&s, "payload")]; + let got = effective_covering(&permuted, Some(&requested), &s).unwrap(); + assert_eq!(names(&got), vec!["price", "payload"]); + } + + #[test] + fn declared_columns_absent_from_the_projection_are_dropped_entirely() { + let s = schema(); + let ids = vec![field_id(&s, "price")]; + let projection = vec!["vec".to_string()]; + assert!( + effective_covering(&ids, Some(&projection), &s) + .unwrap() + .is_empty() + ); + } + + #[test] + fn a_plain_index_declares_nothing_in_every_projection_state() { + let s = schema(); + for projection in [None, Some(&[][..]), Some(&["price".to_string()][..])] { + assert!(effective_covering(&[], projection, &s).unwrap().is_empty()); + } + } + + /// A nested field can carry the same leaf name as a top-level column. Resolving it + /// would emit a `Float64` column named `price` where the covered projection expects + /// the top-level `Int32` one. + #[test] + fn resolves_top_level_fields_only() { + let s = schema(); + let nested_price_id = s + .field("nested") + .unwrap() + .children + .iter() + .find(|f| f.name == "price") + .unwrap() + .id; + let err = effective_covering(&[nested_price_id], None, &s) + .expect_err("a nested field id is not a valid covering declaration"); + assert!(matches!(err, Error::Index { .. })); + + let top_level = effective_covering(&[field_id(&s, "price")], None, &s).unwrap(); + assert_eq!(top_level[0].data_type(), &DataType::Int32); + } + + /// Metadata that disagrees with the schema must fail loudly even when the narrowing + /// would have discarded the offending column: a corrupt index must not appear to work + /// for the subset of queries that happen not to ask for the broken column. + #[test] + fn unresolvable_declared_id_errors_even_when_narrowed_away() { + let s = schema(); + let ids = vec![9999, field_id(&s, "price")]; + for projection in [None, Some(&[][..]), Some(&["price".to_string()][..])] { + let err = effective_covering(&ids, projection, &s) + .expect_err("a declared covering id absent from the schema must be rejected"); + assert!(matches!(err, Error::Index { .. })); + assert!( + err.to_string().contains("9999"), + "error names the offending id, got: {err}" + ); + } + } +} diff --git a/rust/lance/src/index/create.rs b/rust/lance/src/index/create.rs index 8cd006a59a1..f50f694da83 100644 --- a/rust/lance/src/index/create.rs +++ b/rust/lance/src/index/create.rs @@ -597,11 +597,52 @@ impl<'a> CreateIndexBuilder<'a> { } }; + // Resolve any covering ("included") columns to field ids for the generic + // IndexMetadata. The build path validates these names first, so a failure here is + // defence-in-depth -- but this is the site that decides what lands in the manifest, + // so it refuses rather than dropping: silently skipping a name would commit an + // index whose `fields` under-declare the payload its storage actually carries, + // leaving that column out of every invalidation rule keyed on `fields`. + let covering_fields = match self.params.as_any().downcast_ref::() { + Some(vp) => { + let mut ids = vp + .covering_columns + .iter() + .map(|name| { + self.dataset + .schema() + .field(name) + .map(|f| f.id) + .ok_or_else(|| { + Error::index(format!( + "covering column '{name}' is not present in the dataset schema" + )) + }) + }) + .collect::>>()?; + // Carried refine vectors are the indexed column, stored a second time in + // full precision. Declaring them makes the search emit them and `refine` + // re-rank from the index; leaving the declaration off would store the copy + // and never read it. Appended last on purpose: the covering ids come first + // and the carried vectors close the list (see the `fields` contract below). + if vp.store_vectors_for_refine { + ids.push(field.id); + } + ids + } + None => Vec::new(), + }; + // `covering_fields` must be the trailing entries of `fields` (see + // `IndexMetadata::validate_covering_fields`), so the keyed field comes first + // and the covering fields are appended, not tracked separately. + let mut fields = vec![field.id]; + fields.extend_from_slice(&covering_fields); + Ok(IndexMetadata { uuid: output_index_uuid, name: index_name, - fields: vec![field.id], - covering_fields: vec![], + fields, + covering_fields, dataset_version: self.dataset.manifest.version, fragment_bitmap: if train { match &self.fragments { diff --git a/rust/lance/src/index/vector.rs b/rust/lance/src/index/vector.rs index d9e749d44ad..537339ef06e 100644 --- a/rust/lance/src/index/vector.rs +++ b/rust/lance/src/index/vector.rs @@ -72,6 +72,33 @@ use crate::{Error, Result, dataset::Dataset, index::pb::vector_index_stage::Stag pub const LANCE_VECTOR_INDEX: &str = "__lance_vector_index"; +/// Names a user column can never take inside vector index storage: the union of every +/// storage's exported internal set (row id, partition id, each quantizer's code/factor +/// columns -- extending any `*_INTERNAL_COLUMNS` const extends this set automatically) +/// plus the pipeline-transient names that never reach storage (the distance column and +/// the IVF partition transform's `__centroid_dist`). Used both to reject reserved names +/// in `covering_columns` at create time and to classify which columns of an index +/// storage schema are covering ("included") payload. +pub(crate) static RESERVED_STORAGE_COLUMNS: std::sync::LazyLock< + std::collections::HashSet<&'static str>, +> = std::sync::LazyLock::new(|| { + lance_index::vector::pq::storage::PQ_INTERNAL_COLUMNS + .iter() + .chain(lance_index::vector::sq::storage::SQ_INTERNAL_COLUMNS) + .chain(lance_index::vector::flat::storage::FLAT_INTERNAL_COLUMNS) + .chain(lance_index::vector::bq::storage::RABIT_INTERNAL_COLUMNS) + // Carried by the storage but never read per partition -- reserved against user + // column names and never classified as covering payload, exactly like the + // per-partition internals above. + .chain(lance_index::vector::storage::DEFERRED_INTERNAL_COLUMNS) + .chain(&[ + lance_index::vector::DIST_COL, + lance_index::vector::CENTROID_DIST_COLUMN, + ]) + .copied() + .collect() +}); + /// A materialized snapshot of one logical vector index and all of its segments. #[derive(Debug)] pub struct LogicalVectorIndex { @@ -276,6 +303,30 @@ pub struct VectorIndexParams { /// Keys use reverse-DNS namespacing (e.g., "lance.ivf.max_iters"). /// Populated by the build path and merged into VectorIndexDetails at creation time. pub runtime_hints: HashMap, + + /// Columns to co-locate ("include") in the index storage alongside the row id + /// and the quantization code. Empty by default. + pub covering_columns: Vec, + + /// Carry full-precision vectors in index storage so `refine` re-ranks from the index + /// instead of taking them from the base table. Off by default: it stores a second copy + /// of the vector column, which is the largest column in most datasets. + /// + /// Deliberately not expressible through [`Self::covering_columns`]. The build pipeline + /// rewrites the indexed column in place -- residual for L2 and cosine, normalized first + /// for cosine -- so a covering column under that name would serve values that are not + /// the user's vectors, and cosine's normalization discards magnitude irrecoverably. The + /// copy taken here predates those transforms. It is declared as a covering column under + /// the indexed column's own field id, which is what lets the search emit it and `refine` + /// find it by name. + /// + /// Because the copy lives in index storage under that column's name, **the indexed column + /// can no longer be renamed or have its nullability changed while this index exists** -- + /// `alter_columns` refuses it, since nothing renames the stored copy and both + /// `optimize_indices` and `compact_files` would break permanently. Drop the index, alter + /// the column, and rebuild. Dropping the column outright stays allowed: that removes the + /// index with it. A cast is refused for an indexed column either way. + pub store_vectors_for_refine: bool, } impl VectorIndexParams { @@ -289,6 +340,19 @@ impl VectorIndexParams { self } + /// Carry full-precision vectors for `refine`; see + /// [`VectorIndexParams::store_vectors_for_refine`]. + pub fn store_vectors_for_refine(&mut self, enabled: bool) -> &mut Self { + self.store_vectors_for_refine = enabled; + self + } + + /// Set the columns to co-locate ("include") in the index storage. + pub fn covering_columns(&mut self, columns: Vec) -> &mut Self { + self.covering_columns = columns; + self + } + pub fn ivf_flat(num_partitions: usize, metric_type: MetricType) -> Self { let ivf_params = IvfBuildParams::new(num_partitions); let stages = vec![StageParams::Ivf(ivf_params)]; @@ -298,6 +362,8 @@ impl VectorIndexParams { version: IndexFileVersion::V3, skip_transpose: false, runtime_hints: HashMap::new(), + covering_columns: Vec::new(), + store_vectors_for_refine: false, } } @@ -309,6 +375,8 @@ impl VectorIndexParams { version: IndexFileVersion::V3, skip_transpose: false, runtime_hints: HashMap::new(), + covering_columns: Vec::new(), + store_vectors_for_refine: false, } } @@ -344,6 +412,8 @@ impl VectorIndexParams { version: IndexFileVersion::V3, skip_transpose: false, runtime_hints: HashMap::new(), + covering_columns: Vec::new(), + store_vectors_for_refine: false, } } @@ -371,6 +441,8 @@ impl VectorIndexParams { version: IndexFileVersion::V3, skip_transpose: false, runtime_hints: HashMap::new(), + covering_columns: Vec::new(), + store_vectors_for_refine: false, } } @@ -387,6 +459,8 @@ impl VectorIndexParams { version: IndexFileVersion::V3, skip_transpose: false, runtime_hints: HashMap::new(), + covering_columns: Vec::new(), + store_vectors_for_refine: false, } } @@ -402,6 +476,8 @@ impl VectorIndexParams { version: IndexFileVersion::V3, skip_transpose: false, runtime_hints: HashMap::new(), + covering_columns: Vec::new(), + store_vectors_for_refine: false, } } @@ -417,6 +493,8 @@ impl VectorIndexParams { version: IndexFileVersion::V3, skip_transpose: false, runtime_hints: HashMap::new(), + covering_columns: Vec::new(), + store_vectors_for_refine: false, } } @@ -432,6 +510,8 @@ impl VectorIndexParams { version: IndexFileVersion::V3, skip_transpose: false, runtime_hints: HashMap::new(), + covering_columns: Vec::new(), + store_vectors_for_refine: false, } } @@ -454,6 +534,8 @@ impl VectorIndexParams { version: IndexFileVersion::V3, skip_transpose: false, runtime_hints: HashMap::new(), + covering_columns: Vec::new(), + store_vectors_for_refine: false, } } @@ -476,6 +558,8 @@ impl VectorIndexParams { version: IndexFileVersion::V3, skip_transpose: false, runtime_hints: HashMap::new(), + covering_columns: Vec::new(), + store_vectors_for_refine: false, } } @@ -599,6 +683,255 @@ async fn prepare_vector_segment_build( Ok((element_type, index_type, ivf_params, shuffler)) } +/// Whether this configuration carries values the search must read per partition rather +/// than through one global top-k heap -- see [`warn_if_covering_is_pessimised`]. +/// +/// Carried refine vectors count: they are an ordinary covering column, and the widest one +/// a table has, so an HNSW index pays the per-partition read on every probe for them too. +fn covering_is_pessimised(index_type: IndexType, params: &VectorIndexParams) -> bool { + let carries_anything = !params.covering_columns.is_empty() || params.store_vectors_for_refine; + carries_anything + && matches!( + index_type, + IndexType::IvfHnswFlat | IndexType::IvfHnswPq | IndexType::IvfHnswSq + ) +} + +/// Warn when covering is declared on an index type whose search is a pessimisation for it. +/// +/// Covering pays for itself only where the search settles into ONE global top-k heap: there +/// the covering read is bounded by the survivors (see `gather_survivor_covering`). Every HNSW +/// index, and any query with `query_parallelism > 1`, emits per partition instead, so the +/// covering read scales with partitions probed rather than `k` and is issued serially -- +/// measured at 2.8x a plain index's latency warm at nprobe 32 (1.18x cold). That is slower +/// than the base-table take covering exists to elide, so warn at build time rather than let +/// the regression be found in production. Batching the per-partition reads globally would +/// remove this; until then the combination is legal but not recommended. +/// +/// Called after `validate_covering_columns` at every build entry point: warning first would +/// emit this for configurations that are about to be rejected outright, and warning on only +/// one entry point would let a distributed build silently take the shape this exists to flag. +fn warn_if_covering_is_pessimised(index_type: IndexType, column: &str, params: &VectorIndexParams) { + if !covering_is_pessimised(index_type, params) { + return; + } + log::warn!( + "Index on column '{}' carries values in the index ({:?}, refine vectors: {}) on a \ + {:?} index. HNSW \ + search emits per partition rather than through one global top-k heap, so \ + covering reads scale with partitions probed rather than k and are issued \ + serially -- about 2.8x a plain index's latency warm at nprobe 32. Prefer \ + IVF_PQ or IVF_FLAT for covering, or expect covered queries on this index to \ + be slower than uncovered ones.", + column, + params.covering_columns, + params.store_vectors_for_refine, + index_type + ); +} + +/// Reject `store_vectors_for_refine` where the build cannot honour it. +/// +/// Runs before the `covering_columns.is_empty()` short-circuit below, because the option +/// is expressed as a flag rather than a covering name and would otherwise skip every check +/// the covering set gets. +/// +/// Two configurations cannot pay off. Multivector: `refine` runs *after* +/// `MultivectorScoringExec` has re-grouped sub-vectors back to rows, so it scores the row-shaped `List` column, while index storage +/// holds one `FixedSizeList` per sub-vector -- the shape the search itself needs. The carried +/// copy therefore cannot substitute for the refine take, and would be a second copy of the +/// widest column in the table that nothing can read. And a flat quantizer already stores the +/// vectors verbatim, so the copy would be a duplicate of a column the index has. An error +/// beats silently doubling the index. +fn validate_store_vectors_for_refine( + params: &VectorIndexParams, + dataset: &Dataset, + column: &str, +) -> Result<()> { + if !params.store_vectors_for_refine { + return Ok(()); + } + + // Only the V3 covering-aware storages keep an extra column row-aligned with the code. + // A legacy build ignores the option entirely while `CreateIndexBuilder` still records + // the vector in `covering_fields`, so the index would advertise a payload its storage + // never wrote and every query would quietly fall back to the base table. + if !matches!(params.version, IndexFileVersion::V3) { + return Err(Error::invalid_input(format!( + "store_vectors_for_refine requires index file version V3, but got {:?}", + params.version + ))); + } + + // A flat quantizer is an identity copy of the vector column (`FlatQuantizer::quantize`) + // and `IvfTransformer::new_flat` adds no `ResidualTransform`, so under L2 the carried + // copy is byte-identical to the `FLAT_COLUMN` the index already stores -- and there is + // nothing for `refine` to improve on either, since flat distances are exact. Checked + // here, before the checks that need the dataset, because it reads only `params`. + let index_type = params.index_type(); + if matches!(index_type, IndexType::IvfFlat | IndexType::IvfHnswFlat) { + return Err(Error::invalid_input(format!( + "store_vectors_for_refine is not supported on a {index_type:?} index: a flat \ + quantizer already stores the full-precision vectors, so the carried copy would \ + duplicate a column the index has and refine has nothing to re-rank" + ))); + } + + // The copy is written and read back as an ordinary top-level covering column, and + // covering resolution is top-level only by construction (`effective_covering`), so a + // nested indexed column has nowhere to put it. Plain indexing of such a column keeps + // working; only this option is refused. + if column.contains('.') { + return Err(Error::invalid_input(format!( + "store_vectors_for_refine is not supported on the nested indexed column \ + '{column}': the carried vectors are stored as a top-level covering column" + ))); + } + + // The copy travels under a reserved internal name until the transform chain has + // consumed the indexed column. A column already using that name collides with the copy + // mid-flight, and the classifier that decides what storage carries as covering payload + // never counts that name either. + if column == lance_index::vector::storage::REFINE_VECTOR_COLUMN { + return Err(Error::invalid_input(format!( + "store_vectors_for_refine is not supported on the indexed column '{column}': \ + that name is reserved for the internal column the carried copy travels under" + ))); + } + + // Precomputed shuffle buffers arrive already quantized and without the raw vector + // column, so by the time the build reads them there is nothing left to copy aside -- + // the same reason covering columns are refused with them. + if params.stages.iter().any( + |stage| matches!(stage, StageParams::Ivf(ivf) if ivf.precomputed_shuffle_buffers.is_some()), + ) { + return Err(Error::invalid_input( + "store_vectors_for_refine is not supported with precomputed shuffle buffers: \ + the buffers are already quantized and do not carry the raw vectors" + .to_string(), + )); + } + + let (vector_type, _) = get_vector_type(dataset.schema(), column)?; + if matches!(vector_type, DataType::List(_)) { + return Err(Error::invalid_input(format!( + "store_vectors_for_refine is not supported on the multivector column '{column}': \ + refine scores rows after sub-vectors are re-grouped, but the index stores one \ + vector per sub-vector, so the carried copy could never serve it" + ))); + } + Ok(()) +} + +/// Split a committed covering declaration into the covering columns proper and the +/// carried-refine-vectors flag. +/// +/// An index records carried refine vectors by naming the indexed column in its +/// `covering_fields`, and nothing else may name it there. Every rebuild reads that +/// declaration back -- through [`VectorIndexParams`] on the retrain and cross-dataset +/// paths, straight into the builder on the merge/split/join path -- so both funnel through +/// here instead of each re-deriving the rule. +pub(crate) fn split_carried_refine( + covering: Vec, + indexed_column: &str, +) -> (Vec, bool) { + let (refine, columns): (Vec, Vec) = covering + .into_iter() + .partition(|name| name.as_str() == indexed_column); + (columns, !refine.is_empty()) +} + +/// Re-declare on `params` the covering payload a committed index carries. +/// +/// The covering setter is the user-facing one, and `validate_covering_columns` rejects the +/// indexed column there, so a declaration read back from metadata cannot be handed to it +/// verbatim -- see [`split_carried_refine`]. +pub(crate) fn redeclare_covering( + params: &mut VectorIndexParams, + covering: Vec, + indexed_column: &str, +) { + let (columns, carries_refine) = split_carried_refine(covering, indexed_column); + params.covering_columns(columns); + params.store_vectors_for_refine(carries_refine); +} + +/// Validate `covering_columns` (covering columns) at the API boundary, before any build. +/// Rejects non-V3 format, nested/dotted names, the indexed vector column itself, names that +/// collide with a build pipeline's internal storage/transform columns, duplicates, missing +/// columns, and blob columns. Shared by the single-node and distributed vector build paths so +/// the reserved-name list lives in one place. With no covering columns configured this does +/// nothing beyond [`validate_store_vectors_for_refine`], which is checked either way. +fn validate_covering_columns( + params: &VectorIndexParams, + dataset: &Dataset, + column: &str, +) -> Result<()> { + validate_store_vectors_for_refine(params, dataset, column)?; + + if params.covering_columns.is_empty() { + return Ok(()); + } + + // Covering ("included") columns require the V3 index file format: only the V3 + // covering-aware storages preserve the extra columns row-aligned and report them via + // `covering_field_indices`. Every vector index type supports covering at V3, so reject + // only legacy/non-V3 builds up front -- otherwise the build silently ignores the option + // while `covering_fields` still lands in the index metadata, and the search exec then + // declares a covered schema the storage cannot emit. + if !matches!(params.version, IndexFileVersion::V3) { + return Err(Error::invalid_input(format!( + "covering_columns (covering columns) requires index file version V3, but got {:?}", + params.version + ))); + } + + // A covering column is stored inline in the index, row-aligned with the code, and + // projected by name from the source batch. Reject columns that cannot satisfy that + // contract up front (otherwise the build fails deep in projection, or -- worse -- + // stores unusable data): nested/dotted paths (the covering gather projects only + // top-level names) and blob columns (which store out-of-line descriptors, not inline + // data). + let mut seen_include = std::collections::HashSet::with_capacity(params.covering_columns.len()); + for name in ¶ms.covering_columns { + if name.contains('.') { + return Err(Error::invalid_input(format!( + "covering_columns: nested/dotted covering column '{name}' is not supported; only \ + top-level columns can be covered" + ))); + } + if name == column { + return Err(Error::invalid_input(format!( + "covering_columns: covering column '{name}' is the indexed vector column itself; it \ + is stored as the quantization code, not a covered payload" + ))); + } + if RESERVED_STORAGE_COLUMNS.contains(&name.as_str()) { + return Err(Error::invalid_input(format!( + "covering_columns: covering column '{name}' collides with a reserved index storage \ + column name" + ))); + } + if !seen_include.insert(name.as_str()) { + return Err(Error::invalid_input(format!( + "covering_columns: duplicate covering column '{name}'" + ))); + } + let field = dataset.schema().field(name).ok_or_else(|| { + Error::invalid_input(format!( + "covering_columns: covering column '{name}' does not exist in the dataset schema" + )) + })?; + if field.is_blob() { + return Err(Error::invalid_input(format!( + "covering_columns: covering column '{name}' is a blob column; blob columns cannot be \ + covered by a vector index" + ))); + } + } + Ok(()) +} + /// Build a Distributed Vector Index for specific fragments #[allow(clippy::too_many_arguments)] #[instrument(level = "debug", skip(dataset))] @@ -612,6 +945,15 @@ pub(crate) async fn build_distributed_vector_index( fragment_ids: &[u32], progress: Arc, ) -> Result<(Uuid, Vec)> { + // Each shard reuses `IvfIndexBuilder`, which threads `covering_columns` through its + // projection and per-partition storage exactly like the single-node build (the shard is + // just a precomputed-IVF, fragment-filtered build). Validate the covering columns up + // front, same as the single-node path, then thread them into each shard builder below so + // the segment storage carries the payload the committed `covering_fields` advertises. + // The cross-shard segment merger carries covering too (`merge_partial_vector_auxiliary_files` + // appends the covering fields to each writer's schema and rejects shards whose covering sets + // disagree), so both per-segment commit and merged builds are covered. + validate_covering_columns(params, dataset, column)?; let (element_type, index_type, ivf_params, shuffler) = prepare_vector_segment_build( dataset, column, @@ -622,6 +964,7 @@ pub(crate) async fn build_distributed_vector_index( Some(fragment_ids), ) .await?; + warn_if_covering_is_pessimised(index_type, column, params); let stages = ¶ms.stages; let ivf_centroids = ivf_params @@ -689,6 +1032,10 @@ pub(crate) async fn build_distributed_vector_index( frag_reuse_index, )? .with_ivf(ivf_model) + .with_covering_payload( + params.covering_columns.clone(), + params.store_vectors_for_refine, + ) .with_fragment_filter(fragment_filter) .with_progress(progress.clone()) .build() @@ -710,6 +1057,10 @@ pub(crate) async fn build_distributed_vector_index( frag_reuse_index, )? .with_ivf(ivf_model) + .with_covering_payload( + params.covering_columns.clone(), + params.store_vectors_for_refine, + ) .with_fragment_filter(fragment_filter) .with_progress(progress.clone()) .build() @@ -759,6 +1110,10 @@ pub(crate) async fn build_distributed_vector_index( // For distributed shards, keep PQ codes in row-major layout. // A single transpose is performed in the distributed merge stage. .with_transpose(false) + .with_covering_payload( + params.covering_columns.clone(), + params.store_vectors_for_refine, + ) .with_fragment_filter(fragment_filter) .with_progress(progress.clone()) .build() @@ -786,6 +1141,10 @@ pub(crate) async fn build_distributed_vector_index( (), frag_reuse_index, )? + .with_covering_payload( + params.covering_columns.clone(), + params.store_vectors_for_refine, + ) .with_fragment_filter(fragment_filter) .with_progress(progress.clone()) .build() @@ -814,6 +1173,10 @@ pub(crate) async fn build_distributed_vector_index( hnsw_params.clone(), frag_reuse_index, )? + .with_covering_payload( + params.covering_columns.clone(), + params.store_vectors_for_refine, + ) .with_fragment_filter(fragment_filter) .with_progress(progress.clone()) .build() @@ -832,6 +1195,10 @@ pub(crate) async fn build_distributed_vector_index( hnsw_params.clone(), frag_reuse_index, )? + .with_covering_payload( + params.covering_columns.clone(), + params.store_vectors_for_refine, + ) .with_fragment_filter(fragment_filter) .with_progress(progress.clone()) .build() @@ -874,6 +1241,10 @@ pub(crate) async fn build_distributed_vector_index( // For distributed shards, keep PQ codes in row-major layout. // A single transpose is performed in the distributed merge stage. .with_transpose(false) + .with_covering_payload( + params.covering_columns.clone(), + params.store_vectors_for_refine, + ) .with_fragment_filter(fragment_filter) .with_progress(progress.clone()) .build() @@ -905,6 +1276,10 @@ pub(crate) async fn build_distributed_vector_index( hnsw_params.clone(), frag_reuse_index, )? + .with_covering_payload( + params.covering_columns.clone(), + params.store_vectors_for_refine, + ) .with_fragment_filter(fragment_filter) .with_progress(progress.clone()) .build() @@ -937,6 +1312,10 @@ pub(crate) async fn build_distributed_vector_index( // For distributed shards, keep RQ codes in row-major layout. // A single packing pass is performed in the distributed merge stage. .with_transpose(false) + .with_covering_payload( + params.covering_columns.clone(), + params.store_vectors_for_refine, + ) .with_fragment_filter(fragment_filter) .with_progress(progress.clone()) .build() @@ -1025,6 +1404,9 @@ async fn build_vector_index_impl( .await?; let stages = ¶ms.stages; + validate_covering_columns(params, dataset, column)?; + warn_if_covering_is_pessimised(index_type, column, params); + match index_type { IndexType::IvfFlat => match element_type { DataType::Float16 | DataType::Float32 | DataType::Float64 => { @@ -1039,6 +1421,10 @@ async fn build_vector_index_impl( (), frag_reuse_index, )? + .with_covering_payload( + params.covering_columns.clone(), + params.store_vectors_for_refine, + ) .with_optional_fragment_filter(fragment_ids) .with_progress(progress.clone()) .build() @@ -1057,6 +1443,10 @@ async fn build_vector_index_impl( (), frag_reuse_index, )? + .with_covering_payload( + params.covering_columns.clone(), + params.store_vectors_for_refine, + ) .with_optional_fragment_filter(fragment_ids) .with_progress(progress.clone()) .build() @@ -1112,6 +1502,10 @@ async fn build_vector_index_impl( )?; let summary = builder + .with_covering_payload( + params.covering_columns.clone(), + params.store_vectors_for_refine, + ) .with_transpose(!params.skip_transpose) .with_optional_fragment_filter(fragment_ids) .with_progress(progress.clone()) @@ -1140,6 +1534,10 @@ async fn build_vector_index_impl( (), frag_reuse_index, )? + .with_covering_payload( + params.covering_columns.clone(), + params.store_vectors_for_refine, + ) .with_optional_fragment_filter(fragment_ids) .with_progress(progress.clone()) .build() @@ -1167,6 +1565,10 @@ async fn build_vector_index_impl( )?; let summary = builder + .with_covering_payload( + params.covering_columns.clone(), + params.store_vectors_for_refine, + ) .with_transpose(!params.skip_transpose) .with_optional_fragment_filter(fragment_ids) .with_progress(progress.clone()) @@ -1194,6 +1596,10 @@ async fn build_vector_index_impl( hnsw_params.clone(), frag_reuse_index, )? + .with_covering_payload( + params.covering_columns.clone(), + params.store_vectors_for_refine, + ) .with_optional_fragment_filter(fragment_ids) .with_progress(progress.clone()) .build() @@ -1212,6 +1618,10 @@ async fn build_vector_index_impl( hnsw_params.clone(), frag_reuse_index, )? + .with_covering_payload( + params.covering_columns.clone(), + params.store_vectors_for_refine, + ) .with_optional_fragment_filter(fragment_ids) .with_progress(progress.clone()) .build() @@ -1244,6 +1654,10 @@ async fn build_vector_index_impl( hnsw_params.clone(), frag_reuse_index, )? + .with_covering_payload( + params.covering_columns.clone(), + params.store_vectors_for_refine, + ) .with_optional_fragment_filter(fragment_ids) .with_progress(progress.clone()) .build() @@ -1274,6 +1688,10 @@ async fn build_vector_index_impl( hnsw_params.clone(), frag_reuse_index, )? + .with_covering_payload( + params.covering_columns.clone(), + params.store_vectors_for_refine, + ) .with_optional_fragment_filter(fragment_ids) .with_progress(progress.clone()) .build() @@ -1355,6 +1773,11 @@ pub(crate) async fn build_vector_index_incremental( // Determine the index type and build incrementally let (sub_index_type, quantization_type) = existing_index.sub_index_type(); + // Every vector index type supports covering ("included") columns, so the + // incremental/copy path needs no per-type gate here -- each branch below threads + // `covering_columns` into its builder, and the source index was already built with a + // covering-aware storage. + match (sub_index_type, quantization_type) { // IVF_FLAT (SubIndexType::Flat, QuantizationType::Flat) => { @@ -1370,6 +1793,10 @@ pub(crate) async fn build_vector_index_incremental( )? .with_ivf(ivf_model) .with_quantizer(quantizer.try_into()?) + .with_covering_payload( + params.covering_columns.clone(), + params.store_vectors_for_refine, + ) .with_progress(progress.clone()) .build() .await?; @@ -1388,6 +1815,10 @@ pub(crate) async fn build_vector_index_incremental( )? .with_ivf(ivf_model) .with_quantizer(quantizer.try_into()?) + .with_covering_payload( + params.covering_columns.clone(), + params.store_vectors_for_refine, + ) .with_progress(progress.clone()) .build() .await?; @@ -1409,6 +1840,10 @@ pub(crate) async fn build_vector_index_incremental( .with_ivf(ivf_model) .with_quantizer(quantizer.try_into()?) .with_transpose(!params.skip_transpose) + .with_covering_payload( + params.covering_columns.clone(), + params.store_vectors_for_refine, + ) .with_progress(progress.clone()) .build() .await?; @@ -1428,6 +1863,10 @@ pub(crate) async fn build_vector_index_incremental( )? .with_ivf(ivf_model) .with_quantizer(quantizer.try_into()?) + .with_covering_payload( + params.covering_columns.clone(), + params.store_vectors_for_refine, + ) .with_progress(progress.clone()) .build() .await?; @@ -1448,6 +1887,10 @@ pub(crate) async fn build_vector_index_incremental( let summary = builder .with_ivf(ivf_model) .with_quantizer(quantizer.try_into()?) + .with_covering_payload( + params.covering_columns.clone(), + params.store_vectors_for_refine, + ) .with_transpose(!params.skip_transpose) .with_progress(progress.clone()) .build() @@ -1477,6 +1920,10 @@ pub(crate) async fn build_vector_index_incremental( )? .with_ivf(ivf_model) .with_quantizer(quantizer.try_into()?) + .with_covering_payload( + params.covering_columns.clone(), + params.store_vectors_for_refine, + ) .with_progress(progress.clone()) .build() .await?; @@ -1495,6 +1942,10 @@ pub(crate) async fn build_vector_index_incremental( )? .with_ivf(ivf_model) .with_quantizer(quantizer.try_into()?) + .with_covering_payload( + params.covering_columns.clone(), + params.store_vectors_for_refine, + ) .with_progress(progress.clone()) .build() .await?; @@ -1513,6 +1964,10 @@ pub(crate) async fn build_vector_index_incremental( )? .with_ivf(ivf_model) .with_quantizer(quantizer.try_into()?) + .with_covering_payload( + params.covering_columns.clone(), + params.store_vectors_for_refine, + ) .with_progress(progress.clone()) .build() .await?; @@ -1531,6 +1986,10 @@ pub(crate) async fn build_vector_index_incremental( )? .with_ivf(ivf_model) .with_quantizer(quantizer.try_into()?) + .with_covering_payload( + params.covering_columns.clone(), + params.store_vectors_for_refine, + ) .with_progress(progress.clone()) .build() .await?; @@ -1864,7 +2323,7 @@ pub async fn initialize_vector_index( let (sub_index_type, quantization_type) = source_vector_index.sub_index_type(); let ivf_params = derive_ivf_params(ivf_model); - let params = match (sub_index_type, quantization_type) { + let mut params = match (sub_index_type, quantization_type) { (SubIndexType::Flat, QuantizationType::Flat) | (SubIndexType::Flat, QuantizationType::FlatBin) => { VectorIndexParams::with_ivf_flat_params(metric_type, ivf_params) @@ -1919,6 +2378,41 @@ pub async fn initialize_vector_index( } }; + // Carry the source index's covering ("included") columns so the rebuilt + // target storage keeps them. covering_fields are field ids in the source; + // resolve to names (the same columns must exist in the target dataset). + // Through `effective_covering` -- the single authority for covering resolution -- and + // not `Schema::field_by_id`: that recurses into struct children while `Field::name` is + // unqualified, so a (corrupt or foreign) nested id would resolve to a bare leaf name, + // and an unrelated top-level column of that name in the target would then be covered + // silently in its place. `append.rs` routes its own rebuild the same way. + let covering_columns: Vec = crate::index::covering::effective_covering( + &source_index.covering_fields, + None, + source_dataset.schema(), + ) + .map_err(|e| { + Error::index(format!( + "covering declaration of index '{}' cannot be resolved against the source \ + dataset schema: {e}", + source_index.name + )) + })? + .into_iter() + .map(|field| field.name().clone()) + .collect(); + redeclare_covering(&mut params, covering_columns, column_name); + // Validate against the TARGET dataset, not the source: the names resolve there, but + // the column they resolve to may be a different kind (a blob, a reserved storage name, + // a non-V3 build). Without this, the two ordinary build entry points would reject a + // covering set that this cross-dataset path silently accepts. + validate_covering_columns(¶ms, target_dataset, column_name)?; + warn_if_covering_is_pessimised( + vector_index_type(source_vector_index.as_ref()), + column_name, + ¶ms, + ); + let new_uuid = Uuid::new_v4(); let frag_reuse_index = target_dataset .open_frag_reuse_index(&NoOpMetricsCollector) @@ -1944,11 +2438,47 @@ pub async fn initialize_vector_index( let fragment_bitmap = Some(target_dataset.fragment_bitmap.as_ref().clone()); + // Field ids are per-dataset: re-resolve the covering column names against + // the TARGET schema (like `fields` below), never copy the source's ids. + let covering_fields = params + .covering_columns + .iter() + .map(|name| { + target_dataset + .schema() + .field(name) + .map(|f| f.id) + .ok_or_else(|| { + Error::index(format!( + "covering column '{}' not found in target dataset schema", + name + )) + }) + }) + .collect::>>()?; + // Carried refine vectors are declared the same way the create path declares them + // (see `CreateIndexBuilder`): as the indexed column's own field id, appended last. + // `redeclare_covering` split it out of `covering_columns` on the way in, so it has to + // be put back here or the copied index stores a payload no query can reach. + let covering_fields = match params.store_vectors_for_refine { + true => { + let mut ids = covering_fields; + ids.push(field.id); + ids + } + false => covering_fields, + }; + // `covering_fields` must be the trailing entries of `fields` (see + // `IndexMetadata::validate_covering_fields`), so the keyed field comes first + // and the covering fields are appended, not tracked separately. + let mut fields = vec![field.id]; + fields.extend_from_slice(&covering_fields); + let new_idx = IndexMetadata { uuid: new_uuid, name: source_index.name.clone(), - fields: vec![field.id], - covering_fields: vec![], + fields, + covering_fields, dataset_version: target_dataset.manifest.version, fragment_bitmap, index_details: source_index.index_details.clone(), @@ -2184,12 +2714,40 @@ pub(crate) fn fresh_vector_segment_params( #[cfg(test)] mod tests { + /// Carried refine vectors are a covering column -- the widest one a table has -- so an + /// HNSW index pays the same per-partition read for them that the covering warning + /// exists to flag. Gating the warning on `covering_columns` alone would let the option + /// take that regression silently. + #[test] + fn covering_pessimisation_counts_carried_refine_vectors() { + let mut params = VectorIndexParams::ivf_pq(4, 8, 4, MetricType::L2, 50); + params.store_vectors_for_refine(true); + assert!(params.covering_columns.is_empty()); + + assert!( + covering_is_pessimised(IndexType::IvfHnswPq, ¶ms), + "carrying refine vectors on HNSW must be flagged like any other covering column" + ); + assert!( + !covering_is_pessimised(IndexType::IvfPq, ¶ms), + "IVF_PQ settles into one global heap, so it is not pessimised" + ); + + let mut plain = VectorIndexParams::ivf_pq(4, 8, 4, MetricType::L2, 50); + plain.store_vectors_for_refine(false); + assert!( + !covering_is_pessimised(IndexType::IvfHnswPq, &plain), + "an index carrying nothing is not pessimised" + ); + } + use super::*; use crate::dataset::Dataset; use crate::index::DatasetIndexExt; use arrow_array::Array; use arrow_array::RecordBatch; - use arrow_array::types::{Float32Type, Int32Type}; + use arrow_array::cast::AsArray; + use arrow_array::types::{Float32Type, Int32Type, UInt64Type}; use arrow_schema::{DataType as ArrowDataType, Field, Schema as ArrowSchema}; use lance_core::utils::tempfile::TempStrDir; use lance_datagen::{BatchCount, RowCount, array}; @@ -2525,6 +3083,605 @@ mod tests { assert_eq!(results.num_rows(), 10, "Should return 10 nearest neighbors"); } + /// `covering_columns` requires the V3 index file format. A legacy/non-V3 build must be + /// rejected at create time: the build would silently ignore the option while + /// `create_index` still records `covering_fields` in the manifest, so the exec would + /// declare a covered schema the storage cannot emit and every query on the index would + /// fail. This restriction is about the file format, not the index type -- it applies + /// even though covering now works for every IVF vector index type (see + /// `test_covering_columns_supported_for_all_ivf_index_types` below). + #[tokio::test] + async fn test_covering_columns_rejected_for_legacy_index_file_version() { + let test_dir = TempStrDir::default(); + let reader = lance_datagen::gen_batch() + .col("id", array::step::()) + .col("vector", array::rand_vec::(32.into())) + .into_reader_rows(RowCount::from(256), BatchCount::from(1)); + let mut dataset = Dataset::write(reader, test_dir.as_str(), None) + .await + .unwrap(); + + let mut params = VectorIndexParams::with_ivf_pq_params( + MetricType::L2, + IvfBuildParams::new(4), + PQBuildParams::new(8, 8), + ) + .version(IndexFileVersion::Legacy) + .clone(); + params.covering_columns(vec!["id".to_string()]); + let err = dataset + .create_index(&["vector"], IndexType::Vector, None, ¶ms, true) + .await + .expect_err("covering_columns on a legacy/non-V3 IVF_PQ build must be rejected"); + assert!( + err.to_string().contains("covering_columns"), + "error should mention covering_columns, got: {err}" + ); + // Fail fast: no index metadata must have been committed. + assert!( + dataset.load_indices().await.unwrap().is_empty(), + "no index should have been created" + ); + } + + /// Covering ("included") columns now work on every IVF vector index type, not just + /// IVF_PQ: each storage format declares which of its own columns are internal (see + /// `VectorStore::covering_field_indices`), so the creation-time restriction that used to + /// reject non-PQ types here (`"covering_columns ... only supported for IVF_PQ"`) is gone. + /// + /// This is the per-type coverage of the *creation* path: `create_index` succeeds and the + /// committed `IndexMetadata.covering_fields` records exactly the requested column. The + /// read-path proof (search reads the covering column back from the index, no base-table + /// take) is covered per-type by `ivf::v2::tests::test_covered_projection_skips_take`. + #[rstest::rstest] + #[case::ivf_pq(VectorIndexParams::ivf_pq(4, 8, 4, MetricType::L2, 2))] + #[case::ivf_sq(VectorIndexParams::with_ivf_sq_params( + MetricType::L2, + IvfBuildParams::new(4), + SQBuildParams::default(), + ))] + #[case::ivf_hnsw_pq(VectorIndexParams::with_ivf_hnsw_pq_params( + MetricType::L2, + IvfBuildParams::new(4), + HnswBuildParams::default(), + PQBuildParams::new(8, 8), + ))] + #[case::ivf_hnsw_sq(VectorIndexParams::with_ivf_hnsw_sq_params( + MetricType::L2, + IvfBuildParams::new(4), + HnswBuildParams::default(), + SQBuildParams::default(), + ))] + #[case::ivf_rq(VectorIndexParams::with_ivf_rq_params( + MetricType::L2, + IvfBuildParams::new(4), + RQBuildParams::with_rotation_type(1, RQRotationType::Fast), + ))] + #[case::ivf_flat(VectorIndexParams::ivf_flat(4, MetricType::L2))] + #[case::ivf_hnsw_flat(VectorIndexParams::ivf_hnsw( + MetricType::L2, + IvfBuildParams::new(4), + HnswBuildParams::default(), + ))] + #[tokio::test] + async fn test_covering_columns_supported_for_all_ivf_index_types( + #[case] mut params: VectorIndexParams, + ) { + let test_dir = TempStrDir::default(); + let reader = lance_datagen::gen_batch() + .col("id", array::step::()) + .col("vector", array::rand_vec::(32.into())) + .into_reader_rows(RowCount::from(256), BatchCount::from(1)); + let mut dataset = Dataset::write(reader, test_dir.as_str(), None) + .await + .unwrap(); + + params.covering_columns(vec!["id".to_string()]); + dataset + .create_index(&["vector"], IndexType::Vector, None, ¶ms, true) + .await + .unwrap(); + + let indices = dataset.load_indices().await.unwrap(); + assert_eq!( + indices.len(), + 1, + "exactly one index should have been created" + ); + let id_field_id = dataset.schema().field("id").unwrap().id; + assert_eq!( + indices[0].covering_fields, + vec![id_field_id], + "covering_fields must record exactly the requested covering column" + ); + + // The metadata above is derived purely from the request (`create.rs`) and + // never consults the built storage, so it stays correct even if + // `with_covering_payload` were neutralized into a no-op. Force a read that + // only the storage can satisfy: a covered projection must be answered + // without a base-table take, and the returned values must be the row's + // true covering value (not garbage/misaligned) -- proof the storage + // actually carries the column, not just the metadata. + let q = arrow_array::Float32Array::from(vec![0.0f32; 32]); + let mut scan = dataset.scan(); + scan.nearest("vector", &q, 10).unwrap(); + scan.nprobes(4); + scan.with_row_id(); + scan.project(&["id"]).unwrap(); + let plan = scan.explain_plan(true).await.unwrap(); + assert!( + !plan.contains("LanceRead"), + "covered projection ['id'] should be answered from the index storage, \ + not a base-table take; plan:\n{plan}" + ); + let batch = scan.try_into_batch().await.unwrap(); + let ids = batch + .column_by_name("id") + .expect("covered 'id' column must be emitted by the storage") + .as_primitive::(); + let row_ids = batch + .column_by_name(lance_core::ROW_ID) + .expect("row id column") + .as_primitive::(); + assert_eq!(ids.len(), 10, "should return k=10 rows"); + // Single-fragment, step-id dataset => id == row offset == _rowid, so a + // correctly covered id equals the row id for every returned row. + for i in 0..ids.len() { + assert_eq!( + ids.value(i) as u64, + row_ids.value(i), + "covered id must be the row's true value, not a no-op/garbage storage column" + ); + } + } + + /// `initialize_vector_index` re-resolves the source index's covering columns + /// against the TARGET schema, so a name that was valid to cover in the source can + /// resolve to something that is not coverable in the target. It must run the same + /// validation the two ordinary build entry points do, or it silently materializes + /// (here) blob descriptor structs as the covered payload. + /// A covering id naming a *nested* field must not resolve to its bare leaf name. + /// + /// `Schema::field_by_id` recurses into struct children and `Field::name` is + /// unqualified, so a nested id resolves to e.g. `price` -- which can collide with an + /// unrelated top-level column in the target and silently build an index covering the + /// wrong column. `effective_covering` is top-level only for exactly this reason, and + /// the optimize path routes through it (`append.rs`); this path must too. + /// + /// Reachable because `validate_covering_columns` rejects dotted paths only at *create* + /// time -- a manifest written by another build is not covered by it. + #[tokio::test] + async fn test_initialize_vector_index_rejects_a_nested_covering_id() { + use arrow_schema::Fields; + + let test_dir = TempStrDir::default(); + let source_uri = format!("{}/source", test_dir.as_str()); + let target_uri = format!("{}/target", test_dir.as_str()); + + // Source has a struct whose child is named `price`, plus a top-level `price`. + let rows = 300usize; + let nested_price: arrow_array::ArrayRef = + Arc::new(arrow_array::Int64Array::from(vec![1i64; rows])); + let nested: arrow_array::ArrayRef = Arc::new(arrow_array::StructArray::from(vec![( + Arc::new(Field::new("price", ArrowDataType::Int64, true)), + nested_price, + )])); + let top_price: arrow_array::ArrayRef = + Arc::new(arrow_array::Int64Array::from(vec![999i64; rows])); + let vectors: arrow_array::ArrayRef = Arc::new( + arrow_array::FixedSizeListArray::try_new_from_values( + arrow_array::Float32Array::from(vec![0.5f32; rows * 32]), + 32, + ) + .unwrap(), + ); + let schema = Arc::new(ArrowSchema::new(vec![ + Field::new( + "meta", + ArrowDataType::Struct(Fields::from(vec![Field::new( + "price", + ArrowDataType::Int64, + true, + )])), + true, + ), + Field::new("price", ArrowDataType::Int64, true), + Field::new("vector", vectors.data_type().clone(), false), + ])); + let batch = RecordBatch::try_new(schema.clone(), vec![nested, top_price, vectors]).unwrap(); + let reader = arrow_array::RecordBatchIterator::new(vec![Ok(batch)], schema.clone()); + let mut source_dataset = Dataset::write(reader, &source_uri, None).await.unwrap(); + + let mut params = VectorIndexParams::ivf_pq(4, 8, 2, MetricType::L2, 20); + params.covering_columns(vec!["price".to_string()]); + source_dataset + .create_index( + &["vector"], + IndexType::Vector, + Some("vidx".to_string()), + ¶ms, + true, + ) + .await + .unwrap(); + + // Forge what another writer could commit: the declaration rebound to the NESTED + // `meta.price` id. Create-time validation never sees this. + let source_dataset = Dataset::open(&source_uri).await.unwrap(); + let nested_id = source_dataset + .schema() + .field("meta.price") + .expect("nested field must exist") + .id; + let mut source_index = source_dataset + .load_indices() + .await + .unwrap() + .iter() + .find(|i| i.name == "vidx") + .unwrap() + .clone(); + let top_id = source_dataset.schema().field("price").unwrap().id; + assert_ne!(nested_id, top_id, "the two `price` fields must differ"); + for id in source_index.fields.iter_mut() { + if *id == top_id { + *id = nested_id; + } + } + for id in source_index.covering_fields.iter_mut() { + if *id == top_id { + *id = nested_id; + } + } + + let target_reader = arrow_array::RecordBatchIterator::new( + Vec::>::new(), + schema.clone(), + ); + let mut target_dataset = Dataset::write(target_reader, &target_uri, None) + .await + .unwrap(); + + let err = initialize_vector_index( + &mut target_dataset, + &source_dataset, + &source_index, + &["vector"], + ) + .await + .expect_err("a nested covering id must not resolve to a bare leaf name"); + let msg = err.to_string(); + assert!( + msg.contains(&nested_id.to_string()), + "the error must name the unresolvable id rather than silently covering the \ + unrelated top-level `price`; was: {msg}" + ); + } + + #[tokio::test] + async fn test_initialize_vector_index_validates_covering_against_target() { + let test_dir = TempStrDir::default(); + let source_uri = format!("{}/source", test_dir.as_str()); + let target_uri = format!("{}/target", test_dir.as_str()); + + // Source: `payload` is an ordinary binary column, perfectly coverable. + let source_reader = lance_datagen::gen_batch() + .col("payload", array::rand_type(&ArrowDataType::LargeBinary)) + .col("vector", array::rand_vec::(32.into())) + .into_reader_rows(RowCount::from(300), BatchCount::from(1)); + let mut source_dataset = Dataset::write(source_reader, &source_uri, None) + .await + .unwrap(); + + let mut params = VectorIndexParams::ivf_pq(4, 8, 2, MetricType::L2, 20); + params.covering_columns(vec!["payload".to_string()]); + source_dataset + .create_index( + &["vector"], + IndexType::Vector, + Some("vidx".to_string()), + ¶ms, + true, + ) + .await + .unwrap(); + let source_dataset = Dataset::open(&source_uri).await.unwrap(); + let source_index = source_dataset + .load_indices() + .await + .unwrap() + .iter() + .find(|i| i.name == "vidx") + .unwrap() + .clone(); + + // Target: same column name, same arrow type -- but declared a blob, so it + // stores out-of-line descriptors rather than inline data. + let rows = 300usize; + let payload: arrow_array::ArrayRef = + Arc::new(arrow_array::LargeBinaryArray::from_iter_values( + (0..rows).map(|i| format!("v{i}").into_bytes()), + )); + let vectors: arrow_array::ArrayRef = Arc::new( + arrow_array::FixedSizeListArray::try_new_from_values( + arrow_array::Float32Array::from(vec![0.5f32; rows * 32]), + 32, + ) + .unwrap(), + ); + let blob_field = Field::new("payload", ArrowDataType::LargeBinary, true).with_metadata( + [(lance_arrow::BLOB_META_KEY.to_string(), "true".to_string())] + .into_iter() + .collect(), + ); + let target_schema = Arc::new(ArrowSchema::new(vec![ + blob_field, + Field::new("vector", vectors.data_type().clone(), false), + ])); + let batch = RecordBatch::try_new(target_schema.clone(), vec![payload, vectors]).unwrap(); + let target_reader = + arrow_array::RecordBatchIterator::new(vec![Ok(batch)], target_schema.clone()); + let mut target_dataset = Dataset::write(target_reader, &target_uri, None) + .await + .unwrap(); + assert!( + target_dataset.schema().field("payload").unwrap().is_blob(), + "test setup must give the target a blob 'payload', or this test proves nothing" + ); + + let err = initialize_vector_index( + &mut target_dataset, + &source_dataset, + &source_index, + &["vector"], + ) + .await + .expect_err("covering a blob column in the target must be rejected"); + let msg = err.to_string(); + assert!( + msg.contains("blob") && msg.contains("payload"), + "error should name the offending blob covering column; was: {msg}" + ); + } + + /// Copying/importing a covered index into another dataset (initialize) must + /// rebuild the target storage WITH the covering columns, not just copy the + /// `covering_fields` metadata -- otherwise the target advertises covering + /// columns its storage can't emit and covered queries fail. + /// + /// Covering works on every IVF vector index type (see + /// `test_covering_columns_supported_for_all_ivf_index_types`), and + /// `initialize_vector_index`/`build_vector_index_incremental` thread + /// `covering_columns` through every `(SubIndexType, QuantizationType)` branch with no + /// per-type gate, so this parametrizes over all of them the same way. + #[rstest::rstest] + #[case::pq(VectorIndexParams::ivf_pq(10, 8, 16, MetricType::L2, 50))] + #[case::sq(VectorIndexParams::with_ivf_sq_params( + MetricType::L2, + IvfBuildParams::new(10), + SQBuildParams::default(), + ))] + #[case::rq(VectorIndexParams::with_ivf_rq_params( + MetricType::L2, + IvfBuildParams::new(10), + RQBuildParams::with_rotation_type(1, RQRotationType::Fast), + ))] + #[case::flat(VectorIndexParams::ivf_flat(10, MetricType::L2))] + #[case::hnsw_flat(VectorIndexParams::ivf_hnsw( + MetricType::L2, + IvfBuildParams::new(10), + HnswBuildParams::default(), + ))] + #[case::hnsw_pq(VectorIndexParams::with_ivf_hnsw_pq_params( + MetricType::L2, + IvfBuildParams::new(10), + HnswBuildParams::default(), + PQBuildParams::new(8, 8), + ))] + #[case::hnsw_sq(VectorIndexParams::with_ivf_hnsw_sq_params( + MetricType::L2, + IvfBuildParams::new(10), + HnswBuildParams::default(), + SQBuildParams::default(), + ))] + #[tokio::test] + async fn test_initialize_vector_index_preserves_covering( + #[case] mut params: VectorIndexParams, + ) { + let test_dir = TempStrDir::default(); + let source_uri = format!("{}/source", test_dir.as_str()); + let target_uri = format!("{}/target", test_dir.as_str()); + + let source_reader = lance_datagen::gen_batch() + .col("id", array::step::()) + .col("vector", array::rand_vec::(32.into())) + .into_reader_rows(RowCount::from(300), BatchCount::from(1)); + let mut source_dataset = Dataset::write(source_reader, &source_uri, None) + .await + .unwrap(); + + // Covered index on source (covers `id`). + params.covering_columns(vec!["id".to_string()]); + source_dataset + .create_index( + &["vector"], + IndexType::Vector, + Some("vidx".to_string()), + ¶ms, + true, + ) + .await + .unwrap(); + + let source_dataset = Dataset::open(&source_uri).await.unwrap(); + let source_index = source_dataset + .load_indices() + .await + .unwrap() + .iter() + .find(|i| i.name == "vidx") + .unwrap() + .clone(); + + // The target has an extra leading column, so its field ids are shifted + // relative to the source's: `covering_fields` must be re-resolved + // against the TARGET schema (by name), not copied from the source. + let target_reader = lance_datagen::gen_batch() + .col("pad", array::step::()) + .col("id", array::step::()) + .col("vector", array::rand_vec::(32.into())) + .into_reader_rows(RowCount::from(300), BatchCount::from(1)); + let mut target_dataset = Dataset::write(target_reader, &target_uri, None) + .await + .unwrap(); + + initialize_vector_index( + &mut target_dataset, + &source_dataset, + &source_index, + &["vector"], + ) + .await + .unwrap(); + + // Metadata carried over, resolved to the target's field ids. + let target_index = target_dataset.load_indices().await.unwrap()[0].clone(); + let id_field = target_dataset.schema().field("id").unwrap().id; + assert_ne!( + id_field, source_index.covering_fields[0], + "test setup must give 'id' different field ids in source and target" + ); + assert_eq!(target_index.covering_fields, vec![id_field]); + + // End-to-end: a covered projection on the copied index skips the take and + // returns `id` -- which only works if the rebuilt storage carries it. + let query = arrow_array::Float32Array::from(vec![0.5f32; 32]); + let mut scan = target_dataset.scan(); + scan.nearest("vector", &query, 10).unwrap(); + scan.with_row_id(); + scan.project(&["id"]).unwrap(); + let plan = scan.explain_plan(true).await.unwrap(); + assert!( + !plan.contains("LanceRead"), + "covered projection should skip the take on the copied index; plan:\n{plan}" + ); + let batch = scan.try_into_batch().await.unwrap(); + let ids = batch + .column_by_name("id") + .expect("copied index should emit the covering column 'id'") + .as_primitive::(); + let row_ids = batch + .column_by_name(lance_core::ROW_ID) + .expect("row id column") + .as_primitive::(); + // Single-fragment, step-id target dataset => id == row offset == _rowid, so + // a correctly-rebuilt covering column matches the row's true value (not a + // stale copy of the source's storage, and not garbage from a schema + // mismatch during the rebuild). + for i in 0..ids.len() { + assert_eq!( + ids.value(i) as u64, + row_ids.value(i), + "copied index's covered id must be the row's true value, row {i}" + ); + } + } + + /// Carried refine vectors round-trip through `covering_fields` like any other covering + /// payload, except the field they name is the *indexed* column -- which the covering + /// setter rejects outright. The cross-dataset rebuild reads that declaration back, so + /// it has to split it into the flag before validation sees it. + #[tokio::test] + async fn test_initialize_vector_index_preserves_carried_refine_vectors() { + let test_dir = TempStrDir::default(); + let source_uri = format!("{}/source", test_dir.as_str()); + let target_uri = format!("{}/target", test_dir.as_str()); + + let source_reader = lance_datagen::gen_batch() + .col("id", array::step::()) + .col("vector", array::rand_vec::(32.into())) + .into_reader_rows(RowCount::from(300), BatchCount::from(1)); + let mut source_dataset = Dataset::write(source_reader, &source_uri, None) + .await + .unwrap(); + + let mut params = VectorIndexParams::ivf_pq(4, 8, 16, MetricType::L2, 50); + // `id` is covered so the user's projection is served from the index too: that + // leaves refine's own vector fetch as the only thing that could still touch the + // base table, which is what the plan assertion below turns on. + params.covering_columns(vec!["id".to_string()]); + params.store_vectors_for_refine(true); + source_dataset + .create_index( + &["vector"], + IndexType::Vector, + Some("vidx".to_string()), + ¶ms, + true, + ) + .await + .unwrap(); + + let source_dataset = Dataset::open(&source_uri).await.unwrap(); + let source_index = source_dataset + .load_indices() + .await + .unwrap() + .iter() + .find(|i| i.name == "vidx") + .unwrap() + .clone(); + + // A leading column the source lacks, so the target's field ids are shifted and the + // declaration has to be re-resolved by name rather than copied. + let target_reader = lance_datagen::gen_batch() + .col("pad", array::step::()) + .col("id", array::step::()) + .col("vector", array::rand_vec::(32.into())) + .into_reader_rows(RowCount::from(300), BatchCount::from(1)); + let mut target_dataset = Dataset::write(target_reader, &target_uri, None) + .await + .unwrap(); + + initialize_vector_index( + &mut target_dataset, + &source_dataset, + &source_index, + &["vector"], + ) + .await + .unwrap(); + + let target_index = target_dataset.load_indices().await.unwrap()[0].clone(); + let id_field = target_dataset.schema().field("id").unwrap().id; + let vector_field = target_dataset.schema().field("vector").unwrap().id; + assert_ne!( + vector_field, + source_dataset.schema().field("vector").unwrap().id, + "test setup must give 'vector' different field ids in source and target" + ); + assert_eq!( + target_index.covering_fields, + vec![id_field, vector_field], + "the copied index must still declare both the covering column and the \ + carried vectors, re-resolved against the target" + ); + + // End-to-end: refine on the copied index must not fall back to the base table, + // which it can only avoid if the rebuilt storage actually carries the vectors. + let query = arrow_array::Float32Array::from(vec![0.5f32; 32]); + let mut scan = target_dataset.scan(); + scan.nearest("vector", &query, 10).unwrap(); + scan.nprobes(4); + scan.refine(2); + scan.project(&["id"]).unwrap(); + let plan = scan.explain_plan(true).await.unwrap(); + assert!( + !plan.contains("LanceRead"), + "refine on the copied index must be served from the index; plan:\n{plan}" + ); + } + #[tokio::test] async fn test_initialize_vector_index_ivf_flat() { let test_dir = TempStrDir::default(); diff --git a/rust/lance/src/index/vector/builder.rs b/rust/lance/src/index/vector/builder.rs index 3a65d4776d1..fa45a873b47 100644 --- a/rust/lance/src/index/vector/builder.rs +++ b/rust/lance/src/index/vector/builder.rs @@ -43,7 +43,9 @@ use lance_index::vector::quantizer::{ }; use lance_index::vector::quantizer::{QuantizerMetadata, QuantizerStorage}; use lance_index::vector::shared::{SupportedIvfIndexType, write_unified_ivf_and_index_metadata}; -use lance_index::vector::storage::STORAGE_METADATA_KEY; +use lance_index::vector::storage::{ + COVERING_FIELD_IDS_KEY, PartitionColumns, REFINE_VECTOR_COLUMN, STORAGE_METADATA_KEY, +}; use lance_index::vector::transform::Flatten; use lance_index::vector::v3::shuffler::{ DEFAULT_PARTITION_WINDOW_BYTES, EmptyReader, IvfShufflerReader, create_ivf_shuffler, @@ -78,7 +80,7 @@ use lance_table::format::IndexFile; use log::info; use object_store::path::Path; use prost::Message; -use roaring::RoaringBitmap; +use roaring::{RoaringBitmap, RoaringTreemap}; use tokio::sync::{OnceCell, OwnedSemaphorePermit, Semaphore}; use tracing::{Level, instrument, span}; @@ -86,6 +88,7 @@ use crate::Dataset; use crate::dataset::ProjectionRequest; use crate::dataset::index::dataset_format_version; use crate::index::append::build_old_data_filter; +use crate::index::vector::RESERVED_STORAGE_COLUMNS; use crate::index::vector::bounded_partition_stream::{ BoundedPartitionStream, Budgeted, OrderedPartitionResults, WeightedJob, }; @@ -94,7 +97,7 @@ use crate::index::vector::utils::infer_vector_dim; use super::v2::IVFIndex; use super::{ ivf::load_precomputed_partitions_if_available, - utils::{self, get_vector_type}, + utils::{self, gather_covering_columns_by_row_id, get_vector_type}, }; // the number of partitions to evaluate for reassigning @@ -122,6 +125,56 @@ impl Default for FreshPartitionBuildLimits { } } +/// Append the covering columns from `covering` (`[_rowid, ]`) +/// onto `batch` (which must contain a `_rowid` column), gathered by row id. +/// Used to re-attach covering columns to rows that were reordered/regrouped by +/// the partition-join reassignment. +fn append_covering_by_row_id( + mut batch: RecordBatch, + covering: &RecordBatch, +) -> Result { + let target_ids = batch[ROW_ID].as_primitive::(); + let included = gather_covering_columns_by_row_id(covering, target_ids)?; + for (field, col) in included { + batch = batch.try_with_column(field.as_ref().clone(), col)?; + } + Ok(batch) +} + +/// Drop covering ("included") columns that the index metadata does not declare. +/// +/// A pre-covering writer can clear `covering_fields` while the auxiliary storage keeps the +/// payload: every commit re-serializes the whole index list, and that writer's +/// `pb::IndexMetadata` has no field 11, so prost drops the declaration. +/// `FLAG_COVERED_INDEX_METADATA` fences pre-covering builds off the dataset, but the +/// fence is best-effort -- released clients consult writer flags only on the +/// append/overwrite path, so delete/update/raw commits can still slip through it. Existing +/// partition storage then still carries those columns while freshly scanned rows do not +/// (the scan projects only what the metadata declares), and `StorageBuilder::build` rejects +/// the width mismatch, leaving a readable index that can never be merge-optimized again. +/// The manifest declaration is the source of truth here, exactly as it is on the read path, +/// so narrow storage down to it rather than widening the fresh side -- widening would +/// re-materialize a payload the metadata does not advertise. +/// +/// A no-op in the healthy case: storage and `declared` agree, so nothing is projected away. +fn drop_undeclared_covering(batch: RecordBatch, declared: &[String]) -> Result { + let schema = batch.schema(); + let keep: Vec = schema + .fields() + .iter() + .enumerate() + .filter(|(_, field)| { + RESERVED_STORAGE_COLUMNS.contains(field.name().as_str()) + || declared.iter().any(|name| name == field.name()) + }) + .map(|(idx, _)| idx) + .collect(); + if keep.len() == batch.num_columns() { + return Ok(batch); + } + Ok(batch.project(&keep)?) +} + /// Build a new centroid array that incorporates the results of partition splits. /// /// For each `(part_idx, centroid1, centroid2)` in `splits`: @@ -247,6 +300,13 @@ pub struct IvfIndexBuilder { ivf_params: Option, quantizer_params: Option, sub_index_params: Option, + /// Columns to co-locate ("include") in the index storage alongside the row + /// id and quantization code, so covered queries can avoid a take. + covering_columns: Vec, + /// Carry full-precision vectors in storage for `refine` (see + /// [`REFINE_VECTOR_COLUMN`]). Copied off the batch *before* the IVF/quantizer + /// transforms, which rewrite the indexed column in place. + store_refine_vectors: bool, _temp_dir: TempStdDir, // store this for keeping the temp dir alive and clean up after build temp_dir: Path, @@ -376,6 +436,8 @@ impl IvfIndexBuilder optimize_options: None, merged_num: 0, transpose_codes: true, + covering_columns: Vec::new(), + store_refine_vectors: false, format_version, progress: Arc::new(NoopIndexBuildProgress), }) @@ -418,6 +480,44 @@ impl IvfIndexBuilder .downcast_ref::>() .ok_or(Error::invalid_input("existing index is not IVF index"))?; + // Carry the covering ("included") columns so the remapped storage keeps + // them -- otherwise merge_partitions would write a schema without them + // while the metadata still advertises them. + // + // This set comes from *storage*, so in principle it could name a column the + // schema no longer has. Refuse rather than silently dropping it: `remap_index` + // copies `covering_fields` onto the replacement metadata verbatim + // (`dataset/optimize/remapping.rs`), so dropping the column here would commit an + // index whose declaration its storage cannot satisfy, and every later query would + // fail in `effective_covering` ("declares covering field id N ...") rather than + // falling back. A loud failure here beats a silent one at read time. + // + // Unreachable on the commit path today: covering fields live in `fields`, and + // `retain_relevant_indices` drops any index with a field missing from the schema, + // so dropping a covered column deletes the whole index instead of stranding it. + // Defence-in-depth against a manifest this build did not write. + let covering_columns = ivf_index.covering_column_names()?; + if let Some(missing) = covering_columns + .iter() + .find(|name| dataset.schema().field(name).is_none()) + { + return Err(Error::index(format!( + "cannot remap index: its storage carries covering column '{missing}', which \ + is not present in the current dataset schema; index metadata and schema \ + are inconsistent (covering columns: {covering_columns:?})" + ))); + } + + // Storage cannot tell the two apart -- the carried vectors sit under the indexed + // column's own name, exactly like an ordinary covering column -- so split them back + // into the pair every other construction path holds: covering columns proper, plus + // the flag. Leaving the indexed column in `covering_columns` would be the one state + // `split_carried_refine` exists to prevent, and it reads identically from the two + // consumers a remap reaches (`covering_arrow_fields` and `covering_field_ids` both + // append the carried vectors after the covering columns). + let (covering_columns, store_refine_vectors) = + crate::index::vector::split_carried_refine(covering_columns, &column); + let temp_dir = TempStdDir::default(); let temp_dir_path = Path::from_filesystem_path(&temp_dir)?; let format_version = dataset_format_version(&dataset); @@ -431,6 +531,7 @@ impl IvfIndexBuilder ivf_params: None, quantizer_params: None, sub_index_params: None, + store_refine_vectors, _temp_dir: temp_dir, temp_dir: temp_dir_path, ivf: Some(ivf_index.ivf_model().clone()), @@ -443,6 +544,7 @@ impl IvfIndexBuilder optimize_options: None, merged_num: 0, transpose_codes: true, + covering_columns, format_version, progress: Arc::new(NoopIndexBuildProgress), }) @@ -513,8 +615,12 @@ impl IvfIndexBuilder .as_any() .downcast_ref::>() .ok_or(Error::invalid_input("existing index is not IVF index"))?; + // Every column, not just the internal ones: the remapped storage is + // written straight into the new index file, so a covering column + // that is not read here is a covering column the remapped index + // does not have. let part = ivf_index - .load_partition(part_id, false, &NoOpMetricsCollector) + .load_partition_entry_with_covering(part_id) .await?; let storage = part.storage.remap(&mapping)?; @@ -579,6 +685,36 @@ impl IvfIndexBuilder self } + /// Set the whole covering payload: the columns to co-locate ("include") in index + /// storage so covered queries skip the base-table take, and whether full-precision + /// vectors ride along for `refine`. + /// + /// **One setter on purpose.** These are two halves of a single declaration, and when + /// they were settable separately the flag depended on the order the two were called in + /// -- an assignment in one and an or-assignment in the other, so whichever ran last won. + /// Setting both here makes that unrepresentable. + /// + /// The two sources of a declaration disagree about how the flag is spelled, so both are + /// accepted: [`VectorIndexParams`] passes it as `store_refine_vectors` alongside columns + /// that never name the indexed column (`validate_covering_columns` rejects that), while a + /// rebuild reading committed metadata gets one list in which naming the indexed column + /// *is* how carried vectors are recorded -- retrain, merge, split, join and append all + /// take that path and pass `false`. Carried vectors are not an ordinary covering column: + /// the transform chain consumes the indexed column, so they are copied aside rather than + /// projected from the table, and leaving the name in this list would ask the scan to + /// project the same column twice. + pub fn with_covering_payload( + &mut self, + columns: Vec, + store_refine_vectors: bool, + ) -> &mut Self { + let (columns, carries_refine) = + crate::index::vector::split_carried_refine(columns, &self.column); + self.covering_columns = columns; + self.store_refine_vectors = store_refine_vectors || carries_refine; + self + } + /// Set progress callback for index building pub fn with_progress(&mut self, progress: Arc) -> &mut Self { self.progress = progress; @@ -740,12 +876,24 @@ impl IvfIndexBuilder return Err(Error::invalid_input("dataset not set before shuffling")); }; + let mut from_precomputed_buffers = false; let stream = match self .ivf_params .as_ref() .and_then(|p| p.precomputed_shuffle_buffers.as_ref()) { Some((uri, _)) => { + // Precomputed shuffle buffers were materialized without the covering + // columns, so a covered build reading from them would silently omit those + // columns from partition storage. Reject the combination up front. + if !self.covering_columns.is_empty() { + return Err(Error::invalid_input( + "covering_columns (covering columns) are not supported with precomputed \ + shuffle buffers: the precomputed buffers do not carry the covering columns" + .to_string(), + )); + } + from_precomputed_buffers = true; let uri = to_local_path(uri); // the uri points to data directory, // so need to trim the "data" suffix for reading the dataset @@ -756,10 +904,19 @@ impl IvfIndexBuilder } _ => { log::info!("shuffle column {} over dataset", self.column); + // Read the vector column plus any included ("covering") columns + // so they flow through the shuffle into the partition storage. + // `covering_columns` never names the indexed column -- every constructor + // routes its declaration through `split_carried_refine`, which moves that + // name into `store_refine_vectors` -- so this cannot repeat a name (which + // the projection builder would reject outright). + let mut projection: Vec<&str> = Vec::with_capacity(1 + self.covering_columns.len()); + projection.push(self.column.as_str()); + projection.extend(self.covering_columns.iter().map(String::as_str)); let mut builder = dataset.scan(); builder .batch_readahead(get_num_compute_intensive_cpus()) - .project(&[self.column.as_str()])? + .project(&projection)? .with_row_id(); // Apply fragment filter for distributed indexing @@ -782,10 +939,16 @@ impl IvfIndexBuilder } }; - if let Some((row_id_idx, _)) = stream.schema().column_with_name("row_id") { - // When using precomputed shuffle buffers we can't use the column name _rowid - // since it is reserved. So we tolerate `row_id` as well here (and rename it - // to _rowid to match the non-precomputed path) + // Scoped to the precomputed-buffer path on purpose. Those buffers cannot name a + // column `_rowid` (reserved), so they materialize it as `row_id` and it is renamed + // back here to match the non-precomputed path. The ordinary scan above already + // produces a real `_rowid` via `with_row_id()`, and it now also projects the + // covering columns -- so renaming there would rewrite a *user* column named + // `row_id` on top of the real one and fail index creation with + // `Duplicate field name "_rowid" in schema`. + if from_precomputed_buffers + && let Some((row_id_idx, _)) = stream.schema().column_with_name("row_id") + { self.shuffle_data(Some(Self::rename_row_id(stream, row_id_idx))) .await?; } else { @@ -845,6 +1008,8 @@ impl IvfIndexBuilder }; let code_column = quantizer.column(); + let refine_vector_column = self.column.clone(); + let refine_field = self.refine_vector_field()?; let transformer = Arc::new( lance_index::vector::ivf::new_ivf_transformer_with_quantizer( @@ -869,6 +1034,8 @@ impl IvfIndexBuilder data.map(move |batch| { let partition_map = partition_map.clone(); let ivf_transformer = transformer.clone(); + let refine_vector_column = refine_vector_column.clone(); + let refine_field = refine_field.clone(); tokio::spawn(async move { let mut batch = batch?; if !partition_map.is_empty() { @@ -906,7 +1073,14 @@ impl IvfIndexBuilder // this batch is already transformed (in case of GPU training) Ok(batch) } - None => ivf_transformer.transform(&batch), + None => { + let batch = add_refine_vectors( + batch, + &refine_vector_column, + refine_field.as_ref(), + )?; + ivf_transformer.transform(&batch) + } } }) }) @@ -1072,6 +1246,19 @@ impl IvfIndexBuilder FreshPartitionBuildLimits::default(), ); } + // Covering builds may re-scan a pruned fragment that still lives in old + // storage; drop the stale duplicate row ids during the partition read. + let dedup_existing = !self.covering_columns.is_empty(); + // The covering set this build declares. Existing storage may be WIDER (see + // `drop_undeclared_covering`); narrow it so every batch handed to + // `StorageBuilder::build` agrees on width. + let declared_covering = Arc::new(self.covering_columns.clone()); + // Set only while this build carries refine vectors: it names the column whose + // stored copy has to go back under the scratch name before existing and fresh + // rows meet (see [`stash_refine_vectors`]). + let carried_vector_column: Option> = self + .store_refine_vectors + .then(|| Arc::from(self.column.as_str())); let partition_adjustment = Arc::new(partition_adjustment); let build_iter = assign_batches @@ -1086,6 +1273,9 @@ impl IvfIndexBuilder let sub_index_params = sub_index_params.clone(); let column = column.clone(); let frag_reuse_index = frag_reuse_index.clone(); + let dedup_existing = dedup_existing; + let declared_covering = declared_covering.clone(); + let carried_vector_column = carried_vector_column.clone(); let partition_adjustment = partition_adjustment.clone(); async move { let (is_affected, split_reader) = match partition_adjustment.as_ref() { @@ -1111,11 +1301,14 @@ impl IvfIndexBuilder // all data (existing + new), re-assigned with updated // centroids. For other partitions, read from existing // indices + original shuffle reader as normal. + let carried = carried_vector_column.as_deref(); let (mut batches, mut loss) = if is_affected { Self::take_partition_batches( partition, &[], Some(split_reader.as_ref().unwrap().as_ref()), + dedup_existing, + carried, ) .await? } else { @@ -1123,6 +1316,8 @@ impl IvfIndexBuilder partition, indices.as_ref(), Some(reader.as_ref()), + dedup_existing, + carried, ) .await? }; @@ -1130,13 +1325,26 @@ impl IvfIndexBuilder // For unaffected partitions during a split, vectors from // affected partitions may have been reassigned here. if !is_affected && let Some(sr) = split_reader.as_ref() { - let (extra, extra_loss) = - Self::take_partition_batches(partition, &[], Some(sr.as_ref())) - .await?; + let (extra, extra_loss) = Self::take_partition_batches( + partition, + &[], + Some(sr.as_ref()), + dedup_existing, + carried, + ) + .await?; batches.extend(extra); loss += extra_loss; } + // Existing storage can carry covering columns this build does not + // declare; narrow before the fresh rows join them below. + + batches = batches + .into_iter() + .map(|batch| drop_undeclared_covering(batch, &declared_covering)) + .collect::>>()?; + spawn_cpu(move || { // Apply assign_batch for join operations (splits no // longer use assign_batches) @@ -1424,6 +1632,8 @@ impl IvfIndexBuilder part_id: usize, existing_indices: &[ExistingIndex], reader: Option<&dyn ShuffleReader>, + dedup_by_row_id: bool, + carried_vector_column: Option<&str>, ) -> Result<(Vec, f64)> { let mut batches = Vec::new(); for source in existing_indices.iter() { @@ -1444,7 +1654,12 @@ impl IvfIndexBuilder // decoded partition per in-flight task while the rest wait on it. let old_data_filter = source.old_data_filter().await?; - let part_storage = existing_index.load_partition_storage(part_id, None).await?; + // Every column, not just the internal ones: these batches are re-written into + // the index being built, so a covering column that is not read here is a + // covering column the new index does not have. + let part_storage = existing_index + .load_partition_storage(part_id, PartitionColumns::All, None) + .await?; let mut part_batches = part_storage.to_batches()?.collect::>(); // for PQ, the PQ codes are transposed, so we need to transpose them back match Q::quantization_type() { @@ -1497,9 +1712,24 @@ impl IvfIndexBuilder } } + // Stored rows carry the refine vectors under the indexed column's own name, + // while the freshly shuffled rows appended below still have them under the + // scratch name. Both go into one `StorageBuilder::build`, so normalise the + // stored side now -- see [`stash_refine_vectors`] for why that direction. + if let Some(column) = carried_vector_column { + for batch in part_batches.iter_mut() { + *batch = stash_refine_vectors(batch.clone(), column)?; + } + } + batches.extend(part_batches); } + // Everything pushed so far came from existing-segment storage; anything + // pushed below comes from the fresh shuffle reader. The dedup at the end + // only removes stale existing rows whose row id reappears in the fresh data. + let existing_end = batches.len(); + let mut loss = 0.0; // Skip if this partition doesn't exist in the reader // This can happen after a split creates a new partition @@ -1524,9 +1754,230 @@ impl IvfIndexBuilder } } + // When rebuilding a covering index, a fragment pruned from a segment's + // coverage (e.g. an in-place rewrite of a covering column) is re-scanned as + // fresh unindexed data *and* is still present, stale, in an old segment's + // storage -- the same row id then appears in both. Drop the stale existing + // copy in favour of the fresh one. + // + // The dedup is scoped to the existing-storage <-> fresh-reader boundary: it + // removes only an *existing-storage* row whose row id reappears in the fresh + // reader data, and never dedups within a single source. That preserves + // legitimate repeated row ids -- a multivector column stores one entry per + // sub-vector, all sharing the source row's id -- which a global per-row-id + // dedup would otherwise silently collapse. + if dedup_by_row_id && existing_end > 0 { + // A roaring treemap rather than a `HashSet`: this holds every fresh row + // id in the partition, HNSW partitions target `1 << 20` rows, and + // `build_partitions` runs several partitions concurrently. + let mut fresh = RoaringTreemap::new(); + for batch in &batches[existing_end..] { + for rid in batch[ROW_ID].as_primitive::().iter().flatten() { + fresh.insert(rid); + } + } + if !fresh.is_empty() { + for batch in batches[..existing_end].iter_mut() { + if batch.num_rows() == 0 { + continue; + } + let row_ids = batch[ROW_ID].as_primitive::(); + let keep: Vec = row_ids + .iter() + .map(|row_id| match row_id { + Some(rid) => !fresh.contains(rid), + None => true, + }) + .collect(); + *batch = arrow::compute::filter_record_batch(batch, &BooleanArray::from(keep))?; + } + } + } + Ok((batches, loss)) } + /// Resolve the covering ("included") columns to their arrow fields from the dataset + /// schema, in declaration order. Empty when no covering columns are configured. Both + /// the code-storage schema and the empty-flat fallback schema append these so covered + /// storage is written consistently regardless of whether any partition held data. + fn covering_arrow_fields(&self) -> Result> { + if self.covering_columns.is_empty() && !self.store_refine_vectors { + return Ok(Vec::new()); + } + let Some(ds) = self.dataset.as_ref() else { + // `covering_columns` is non-empty here, so silently returning no + // fields would let the storage writer (which resolves covering + // columns by name) write storage without them while the committed + // metadata still advertises `covering_fields` -- the exact + // metadata/storage desync every other covering guard prevents. + return Err(Error::invalid_input( + "dataset not set before resolving covering columns".to_string(), + )); + }; + let ds_schema = arrow_schema::Schema::from(ds.schema()); + let mut fields = self + .covering_columns + .iter() + .map(|name| { + ds_schema.field_with_name(name).cloned().map_err(|e| { + Error::invalid_input(format!( + "include column '{name}' not found in dataset schema: {e}" + )) + }) + }) + .collect::>>()?; + // Carried vectors are an ordinary covering column under the dataset column's own + // name: the copy is taken before the transform chain, so its values *are* that + // column's values, and the shapes match (multivector, where they would not, is + // refused at create time). Being covering rather than storage-internal is what lets + // the search emit it, the exec declare it, and `refine` find it by name. + if let Some(field) = self.refine_vector_field()? { + fields.push(field); + } + Ok(fields) + } + + /// Refuse to re-stamp a covering payload onto a different logical field. + /// + /// A rebuild (compaction remap, or `optimize_indices` merging existing sources) copies + /// the covering payload out of existing storage **by name**, while the stamp written + /// here is resolved from the **current** dataset schema. If a covered column was + /// rebound between the two -- dropped and re-added under the same name, so the name + /// resolves to a fresh id -- those disagree, and writing the current id over payload + /// built from the old one silently converts a detectable mismatch into an undetectable + /// one. + /// + /// That matters because the read path's protection *is* this stamp: planning compares + /// the stamped source ids against the index's declared `covering_fields` + /// (`common_physical_covering`) and falls back to a base-table take when they differ. + /// Re-stamping makes the comparison succeed, so the stale payload starts being served + /// and the base-table read is elided. Refusing keeps the index in the state the read + /// path already handles correctly. + /// + /// Sources with no covering payload, or whose stamp is absent/malformed (which already + /// makes their payload unservable), are not the case this guards and are skipped. + /// + /// Nor does it fire when this rebuild writes no stamp at all (`about_to_stamp` empty). + /// That is the degraded state -- a declaration cleared by a pre-covering writer while + /// the storage still physically carries the payload -- and it converges to an ordinary + /// index: with no stamp, `physical_covering_fields` yields nothing, the payload is + /// unservable, and reads use the base table. There is no mismatch to launder, and + /// refusing here would make a degraded index permanently un-optimizable + /// (see `test_degraded_covered_index_can_still_be_optimized`). + fn validate_covering_stamp_is_inherited(&self, about_to_stamp: &[i32]) -> Result<()> { + if about_to_stamp.is_empty() { + return Ok(()); + } + for source in &self.existing_indices { + let stamped: Vec = source + .index + .physical_covering_fields()? + .into_iter() + .map(|(source_id, _)| source_id) + .collect(); + if stamped.is_empty() || stamped == about_to_stamp { + continue; + } + return Err(Error::index(format!( + "cannot rebuild this covered index: its existing storage carries covering \ + values stamped with source field ids {stamped:?}, but the covering columns \ + {:?} now resolve to {about_to_stamp:?} in the dataset schema. A covered \ + column was replaced by a different field of the same name, so the stored \ + payload belongs to the old field. Drop and recreate the index rather than \ + rebuilding it.", + self.covering_columns + ))); + } + Ok(()) + } + + /// The declared field for [`REFINE_VECTOR_COLUMN`], or `None` when refine vectors are + /// not being carried. + /// + /// The type is the *flattened* vector type: a multivector column is `List` + /// in the dataset but is flattened to one `FixedSizeList` per sub-vector before the copy + /// is taken, so the stored column matches what `Flatten` produces, not the source column. + fn refine_vector_field(&self) -> Result> { + if !self.store_refine_vectors { + return Ok(None); + } + let Some(ds) = self.dataset.as_ref() else { + return Err(Error::invalid_input( + "dataset not set before resolving the refine vector column".to_string(), + )); + }; + // Cloned from the dataset schema, not constructed: the read path's capability check + // compares name, data type, nullability AND metadata against the dataset's field + // (`covering_fields_match`), so anything synthesised here that differs in any of + // those makes the column fail to match and silently withdraws it. + // + // Safe to take the dataset field verbatim because multivector -- the one case where + // the stored shape differs from the column's (per sub-vector, not per row) -- is + // refused at create time. + // Convert just this field. `arrow_schema::Schema::from(ds.schema())` is exactly this + // conversion mapped over every field, but it walks the whole table (including nested + // structs) -- and `covering_arrow_fields`, this method's main caller, has already + // paid for that walk four lines before calling it. Top-level lookup only, which is + // what the indexed column is (a dotted name is refused at create time). + let field = ds + .schema() + .fields + .iter() + .find(|field| field.name == self.column) + .ok_or_else(|| { + Error::invalid_input(format!( + "indexed column '{}' not found in dataset schema", + self.column + )) + })?; + Ok(Some(arrow_schema::Field::from(field))) + } + + /// The *source dataset* field ids of the covering columns, in declaration order. + /// Stamped into the storage file (see [`COVERING_FIELD_IDS_KEY`]) so a distributed + /// merge can tell shards that cover the same logical fields from shards whose + /// covering columns merely share a name and type. + /// + /// Fails exactly where [`Self::covering_arrow_fields`] fails, and for the same reason: + /// the two must agree in arity or the storage file gets N covering columns and a stamp + /// of a different length, which makes `physical_covering_fields_from_schema` withdraw + /// covering for that index permanently and silently. + fn covering_field_ids(&self) -> Result> { + if self.covering_columns.is_empty() && !self.store_refine_vectors { + return Ok(Vec::new()); + } + let Some(ds) = self.dataset.as_ref() else { + return Err(Error::invalid_input( + "dataset not set before resolving covering field ids".to_string(), + )); + }; + let mut ids = self + .covering_columns + .iter() + .map(|name| { + ds.schema() + .field(name) + .map(|field| field.id) + .ok_or_else(|| { + Error::invalid_input(format!( + "include column '{name}' not found in dataset schema" + )) + }) + }) + .collect::>>()?; + if self.store_refine_vectors { + let field = ds.schema().field(&self.column).ok_or_else(|| { + Error::invalid_input(format!( + "indexed column '{}' not found in dataset schema", + self.column + )) + })?; + ids.push(field.id); + } + Ok(ids) + } + #[instrument(name = "merge_partitions", level = "debug", skip_all)] async fn merge_partitions( &mut self, @@ -1544,6 +1995,12 @@ impl IvfIndexBuilder let quantization_type = Q::quantization_type(); let is_pq = quantization_type == QuantizationType::Product; let is_rq = quantization_type == QuantizationType::Rabit; + // Flat's "code" column is an identity copy of the raw vector (see + // `FlatQuantizer::quantize`), so its true arrow value type (Float16/32/64) + // is only known from the data itself: `FlatQuantizer::field()` hardcodes + // Float32, which is wrong whenever the vector column isn't Float32. Every + // other quantizer type's `field()` is driven by its own parameters (e.g. + // SQ's `num_bits`), not the raw vector's type, and is always correct. let is_flat = quantization_type == QuantizationType::Flat; // prepare the final writers @@ -1551,12 +2008,28 @@ impl IvfIndexBuilder let index_path = self.index_dir.clone().join(INDEX_FILE_NAME); let writer_options = FileWriterOptions::default(); + // Covering columns are appended to whichever storage schema this build writes + // (code storage below, or the empty-flat fallback further down). The schema is + // always declared as `[_rowid, code, extra…, covering…]` -- row id first, never + // inferred from a batch's own column order, which can put covering columns + // before `_rowid` and make every downstream reader that trusts column position + // (e.g. the default `QuantizerStorage::remap`) misread row ids. + let covering_fields = self.covering_arrow_fields()?; + // Every storage batch is reordered into this declared schema before it is written: + // `FileWriter` picks columns by name when encoding but checks nullability by + // position, so a batch ordered `[covering…, _rowid, code]` would be checked against + // the wrong fields and a null-bearing covering column rejected as if it were the + // (non-nullable) code column. + let mut storage_arrow_schema: Option = None; let mut storage_writer = if is_flat { None } else { let mut fields = vec![ROW_ID_FIELD.clone(), quantizer.field()]; fields.extend(quantizer.extra_fields()); - let storage_schema: Schema = (&arrow_schema::Schema::new(fields)).try_into()?; + fields.extend(covering_fields.iter().cloned()); + let arrow_schema = Arc::new(arrow_schema::Schema::new(fields)); + let storage_schema: Schema = arrow_schema.as_ref().try_into()?; + storage_arrow_schema = Some(arrow_schema); Some(file_versions::create_writer( self.format_version, self.store.create(&storage_path).await?, @@ -1644,7 +2117,27 @@ impl IvfIndexBuilder } if storage_writer.is_none() { - let storage_schema: Schema = batch.schema_ref().as_ref().try_into()?; + // Only flat reaches here (every other type declared its schema + // eagerly above). Declare `[_rowid, flat, covering…]` explicitly -- + // row id always first -- rather than cloning the batch's own + // schema, whose column order can put covering columns before + // `_rowid`. The flat column's arrow type is looked up by name + // from this batch (not assumed), since it is an identity copy + // of the raw vector column and can be Float16/32/64. + let flat_field = batch + .schema_ref() + .field_with_name(lance_index::vector::flat::storage::FLAT_COLUMN) + .map_err(|e| { + Error::invalid_input(format!( + "flat storage batch missing its code column: {e}" + )) + })? + .clone(); + let mut fields = vec![ROW_ID_FIELD.as_ref().clone(), flat_field]; + fields.extend(covering_fields.iter().cloned()); + let arrow_schema = Arc::new(arrow_schema::Schema::new(fields)); + let storage_schema: Schema = arrow_schema.as_ref().try_into()?; + storage_arrow_schema = Some(arrow_schema); storage_writer = Some(file_versions::create_writer( self.format_version, self.store.create(&storage_path).await?, @@ -1652,6 +2145,25 @@ impl IvfIndexBuilder writer_options.clone(), )?); } + // Match the declared column order (see `storage_arrow_schema`). + // + // Load-bearing, not tidiness: the writer schema is declared as + // `[_rowid, code, extra..., covering...]` while a covered batch arrives + // as `[covering..., _rowid, code]`, and the file writer encodes columns + // **by name** but checks nullability **by position**. Writing the batch + // unprojected therefore matches a non-nullable declared slot against a + // nullable covering column and kills the build on the first null covering + // value. The pre-covering code derived the writer schema from the batch + // and so could not disagree; declaring it instead -- which is what stops a + // positional row-id read corrupting row ids -- is what made this necessary. + // The carried vectors travelled under an internal name so they would + // not collide with the indexed column the transform chain consumes. + // At rest they are that column, so restore its name before the + // declared-order projection selects columns by name. + batch = rename_refine_vectors(batch, &self.column)?; + if let Some(schema) = storage_arrow_schema.as_ref() { + batch = batch.project_by_schema(schema)?; + } storage_writer .as_mut() .expect("storage writer must be initialized before write") @@ -1696,12 +2208,15 @@ impl IvfIndexBuilder } if storage_writer.is_none() { + // Every partition was empty, so flat never saw a batch to learn its value + // type from above. The IVF centroids preserve that type (they are computed + // over the same raw vectors `FlatQuantizer::quantize` copies verbatim). let Some(centroids) = ivf.centroids.as_ref() else { return Err(Error::invalid_input( "flat storage writer could not infer schema from empty partitions without IVF centroids", )); }; - let flat_schema = arrow_schema::Schema::new(vec![ + let mut flat_fields = vec![ ROW_ID_FIELD.as_ref().clone(), arrow_schema::Field::new( lance_index::vector::flat::storage::FLAT_COLUMN, @@ -1715,7 +2230,11 @@ impl IvfIndexBuilder ), true, ), - ]); + ]; + // An all-empty covered flat index must still declare its covering columns, + // or the written storage schema disagrees with `covering_fields`. + flat_fields.extend(covering_fields.iter().cloned()); + let flat_schema = arrow_schema::Schema::new(flat_fields); let storage_schema: Schema = (&flat_schema).try_into()?; storage_writer = Some(file_versions::create_writer( self.format_version, @@ -1755,6 +2274,18 @@ impl IvfIndexBuilder STORAGE_METADATA_KEY, serde_json::to_string(&storage_partition_metadata)?, ); + let covering_field_ids = self.covering_field_ids()?; + self.validate_covering_stamp_is_inherited(&covering_field_ids)?; + if !covering_field_ids.is_empty() { + storage_writer.add_schema_metadata( + COVERING_FIELD_IDS_KEY, + covering_field_ids + .iter() + .map(|id| id.to_string()) + .collect::>() + .join(","), + ); + } let index_type_str = index_type_string(S::name().try_into()?, Q::quantization_type()); if let Some(idx_type) = SupportedIvfIndexType::from_index_type_str(&index_type_str) { @@ -1805,14 +2336,17 @@ impl IvfIndexBuilder // take raw vectors from the dataset // - // returns batches of schema | row_id | vector | + // returns batches of schema | row_id | vector | | async fn take_vectors( dataset: &Dataset, column: &str, store: &ObjectStore, row_ids: &[u64], + covering_columns: &[String], ) -> Result> { - let projection = Arc::new(dataset.schema().project(&[column])?); + let mut proj_columns: Vec<&str> = vec![column]; + proj_columns.extend(covering_columns.iter().map(String::as_str)); + let projection = Arc::new(dataset.schema().project(&proj_columns)?); // arrow uses i32 for index, so we chunk the row ids to avoid large batch causing overflow let mut batches = Vec::new(); let row_ids = dataset.filter_deleted_ids(row_ids).await?; @@ -1836,11 +2370,13 @@ impl IvfIndexBuilder Ok(batches) } - // helper to load row ids and vectors for a partition + // helper to load row ids and vectors for a partition, plus a covering batch + // `[_rowid, ]` (None when the index has no covering columns) + // so the columns can be re-attached to reassigned rows after quantization. async fn load_partition_raw_vectors( &self, part_idx: usize, - ) -> Result> { + ) -> Result)>> { let Some(dataset) = self.dataset.as_ref() else { return Err(Error::invalid_input( "dataset not set before split partition", @@ -1854,11 +2390,56 @@ impl IvfIndexBuilder // dedup is needed if it's multivector row_ids.dedup(); - let batches = Self::take_vectors(dataset, &self.column, &self.store, &row_ids).await?; + let batches = Self::take_vectors( + dataset, + &self.column, + &self.store, + &row_ids, + &self.covering_columns, + ) + .await?; if batches.is_empty() { return Ok(None); } let batch = arrow::compute::concat_batches(&batches[0].schema(), batches.iter())?; + // Extract the covering columns BEFORE flattening. `Flatten` replicates every + // non-partition column across a multivector row's expansion, so taking them + // afterwards would yield one copy per sub-vector; the pre-flatten batch has one + // row per row id, which is what the later gather-by-row-id expects. + // Carried refine vectors ride along as one more covering column, under the + // in-flight scratch name. The reassignment below rebuilds its rows from the base + // table and never runs the transform chain, so without this the reassigned rows + // rejoin the untouched partitions' stored copies simply missing the column, and + // the storage build has nothing to align them against. + let carried_vector_idx = self + .store_refine_vectors + .then(|| batch.schema().index_of(&self.column)) + .transpose()?; + let covering = if self.covering_columns.is_empty() && carried_vector_idx.is_none() { + None + } else { + let mut indices = vec![batch.schema().index_of(ROW_ID)?]; + for name in &self.covering_columns { + indices.push(batch.schema().index_of(name)?); + } + indices.extend(carried_vector_idx); + let projected = batch.project(&indices)?; + Some(match carried_vector_idx { + Some(_) => stash_refine_vectors(projected, &self.column)?, + None => projected, + }) + }; + // Narrow to `[_rowid, vector]` first: `Flatten` replicates every remaining column + // across the multivector expansion, so leaving the covering columns in would + // materialize a second copy of the payload that nothing reads -- the gather below + // uses the pre-flatten `covering` batch. Skipped when the vector column is not a + // plain top-level column (precomputed buffers), where `Flatten` is a no-op anyway. + let batch = match (&covering, batch.schema().index_of(&self.column)) { + (Some(_), Ok(vector_idx)) => { + batch.project(&[batch.schema().index_of(ROW_ID)?, vector_idx])? + } + _ => batch, + }; // for multivector, we need to flatten the vectors let batch = Flatten::new(&self.column).transform(&batch)?; // need to retrieve the row ids from the batch because some rows may have been deleted @@ -1872,7 +2453,7 @@ impl IvfIndexBuilder )))? .as_fixed_size_list() .clone(); - Ok(Some((row_ids, vectors))) + Ok(Some((row_ids, vectors, covering))) } // check whether need to split or join partition @@ -2064,11 +2645,16 @@ impl IvfIndexBuilder all_row_ids.sort(); all_row_ids.dedup(); - // Stream raw vectors in chunks - let projection = Arc::new(dataset.schema().project(&[self.column.as_str()])?); + // Stream raw vectors plus any covering ("included") columns so a + // partition split preserves them in the re-quantized storage. As in the shuffle + // projection above, `covering_columns` cannot name the indexed column. + let mut split_projection: Vec<&str> = vec![self.column.as_str()]; + split_projection.extend(self.covering_columns.iter().map(String::as_str)); + let projection = Arc::new(dataset.schema().project(&split_projection)?); let row_ids = dataset.filter_deleted_ids(&all_row_ids).await?; let block_size = self.store.block_size(); let column = self.column.clone(); + let refine_field = self.refine_vector_field()?; let dataset_clone = dataset.clone(); let projection_clone = projection.clone(); @@ -2082,6 +2668,7 @@ impl IvfIndexBuilder let dataset = dataset_clone.clone(); let projection = projection_clone.clone(); let column = column.clone(); + let refine_field = refine_field.clone(); async move { let batch = dataset .take_rows(&chunk, ProjectionRequest::Schema(projection)) @@ -2089,7 +2676,12 @@ impl IvfIndexBuilder let batch = batch .try_with_column(ROW_ID_FIELD.clone(), Arc::new(UInt64Array::from(chunk)))?; // For multivector, flatten - Flatten::new(&column).transform(&batch) + let batch = Flatten::new(&column).transform(&batch)?; + // Carry the vectors on this path too. A split re-streams the affected rows + // through the transform chain, which consumes the indexed column, so without + // this the rebuilt partitions come back missing a column the declared storage + // schema still names. + add_refine_vectors(batch, &column, refine_field.as_ref()) } }) .boxed(); @@ -2174,7 +2766,8 @@ impl IvfIndexBuilder row_ids = (0..sample_size).map(|i| row_ids[i * stride]).collect(); } - let batches = Self::take_vectors(dataset, &self.column, &self.store, &row_ids).await?; + // Centroid sampling only needs the vectors, not covering columns. + let batches = Self::take_vectors(dataset, &self.column, &self.store, &row_ids, &[]).await?; if batches.is_empty() { return Ok(None); } @@ -2251,8 +2844,9 @@ impl IvfIndexBuilder let new_centroids = FixedSizeListArray::try_new_from_values(new_centroids, centroids.value_length())?; - // take the raw vectors from dataset - let Some((row_ids, vectors)) = self.load_partition_raw_vectors(part_idx).await? else { + // take the raw vectors from dataset (plus any covering columns) + let Some((row_ids, vectors, covering)) = self.load_partition_raw_vectors(part_idx).await? + else { return Ok(AssignResult { assign_batches: vec![None; ivf.num_partitions() - 1], new_centroids, @@ -2266,6 +2860,7 @@ impl IvfIndexBuilder ivf, &row_ids, &vectors, + covering.as_ref(), new_centroids, ) .await @@ -2276,6 +2871,7 @@ impl IvfIndexBuilder ivf, &row_ids, &vectors, + covering.as_ref(), new_centroids, ) .await @@ -2286,6 +2882,7 @@ impl IvfIndexBuilder ivf, &row_ids, &vectors, + covering.as_ref(), new_centroids, ) .await @@ -2296,6 +2893,7 @@ impl IvfIndexBuilder ivf, &row_ids, &vectors, + covering.as_ref(), new_centroids, ) .await @@ -2313,6 +2911,10 @@ impl IvfIndexBuilder ivf: &IvfModel, row_ids: &UInt64Array, vectors: &FixedSizeListArray, + // `[_rowid, ]` for the reassigned rows, plus the carried refine + // vectors when the index has them, or None when it has neither. Re-attached to + // each assign batch by row id. + covering: Option<&RecordBatch>, new_centroids: FixedSizeListArray, ) -> Result where @@ -2356,6 +2958,21 @@ impl IvfIndexBuilder } let assign_batches = self.build_assign_batch::(&new_centroids, &assign_ops)?; + // Re-attach the covering columns to each reassigned batch, gathered by + // row id (in-memory; the rows were reordered across target partitions). + let assign_batches = match covering { + None => assign_batches, + Some(covering) => assign_batches + .into_iter() + .map(|entry| match entry { + Some((batch, deleted)) => { + Ok(Some((append_covering_by_row_id(batch, covering)?, deleted))) + } + None => Ok(None), + }) + .collect::>>()?, + }; + Ok(AssignResult { assign_batches, new_centroids, @@ -2671,6 +3288,96 @@ pub(crate) fn index_type_string(sub_index: SubIndexType, quantizer: Quantization } } +/// Copy the vector column aside as [`REFINE_VECTOR_COLUMN`] before the IVF/quantizer +/// transforms run, so `refine` can re-rank against full-precision vectors read from index +/// storage rather than taken from the base table. +/// +/// **Order matters twice.** It must run *before* the transform chain, which rewrites the +/// indexed column in place -- residual for L2 and cosine, normalized first for cosine -- +/// so a copy taken afterwards would not hold the user's vectors. And it must run *after* +/// `Flatten`, or a multivector row's copy would be the whole `List`, replicated once per +/// sub-vector; `Flatten` is idempotent on `FixedSizeList`, so the chain's own copy of it +/// is then a no-op. +/// +/// Every transform in the chain leaves an unrecognised column alone: the two that filter +/// rows (`KeepFiniteVectors`, `PartitionFilter`) use `batch.take`, so the copy stays +/// row-aligned, and the rest touch only the columns they name. +/// `refine_field` is the declared field for the carried copy (see +/// [`IvfIndexBuilder::refine_vector_field`]), and `None` when the copy is not being carried +/// at all -- the two are the same decision, so they travel as one value rather than as a +/// field plus a flag that could disagree. +fn add_refine_vectors( + batch: RecordBatch, + vector_column: &str, + refine_field: Option<&arrow_schema::Field>, +) -> Result { + let Some(refine_field) = refine_field else { + return Ok(batch); + }; + let batch = Flatten::new(vector_column).transform(&batch)?; + let Some(vectors) = batch.column_by_name(vector_column) else { + // Every source that reaches here carries the indexed column: the one build input + // that does not -- precomputed shuffle buffers -- is refused at the API boundary by + // `validate_store_vectors_for_refine`. Returning the batch unchanged would let an + // upstream defect surface much later as `project_by_schema` complaining that the + // indexed column is missing from a partition, which never mentions refine at all. + return Err(Error::index(format!( + "store_vectors_for_refine: the indexed column '{vector_column}' is missing from \ + a batch reaching the transform chain, so the full-precision copy cannot be taken" + ))); + }; + // Nullability comes from the declaration, never from `Array::is_nullable()` -- that is + // `logical_null_count() != 0`, so it would answer "this batch has a null" and make + // consecutive batches of one column declare different schemas. + let field = Field::new( + REFINE_VECTOR_COLUMN, + vectors.data_type().clone(), + refine_field.is_nullable(), + ); + Ok(batch.try_with_column(field, vectors.clone())?) +} + +/// Rename `from` to `to`, or return the batch untouched when it has no such column. +fn rename_column(batch: RecordBatch, from: &str, to: &str) -> Result { + let Ok(idx) = batch.schema().index_of(from) else { + return Ok(batch); + }; + let mut fields: Vec> = batch.schema().fields().to_vec(); + // Rename in place rather than rebuild: hardcoding a nullability here would disagree + // with what `add_refine_vectors` declared on the way in and with the storage schema + // resolved from the dataset field, and the carried copy round-trips through both + // directions of this rename. + fields[idx] = Arc::new(fields[idx].as_ref().clone().with_name(to)); + let schema = Arc::new(arrow_schema::Schema::new(fields)); + Ok(RecordBatch::try_new(schema, batch.columns().to_vec())?) +} + +/// Restore the indexed column's name on the carried full-precision vectors. +/// +/// They are copied aside as [`REFINE_VECTOR_COLUMN`] before the transform chain, because +/// the chain consumes a column of the indexed column's own name. Once the chain has run, +/// the collision is gone and the copy is exactly that column's values, so it takes the +/// name back and becomes an ordinary covering column. +fn rename_refine_vectors(batch: RecordBatch, column: &str) -> Result { + rename_column(batch, REFINE_VECTOR_COLUMN, column) +} + +/// Put carried vectors read back out of storage under the in-flight scratch name. +/// +/// The inverse of [`rename_refine_vectors`], and the reason both directions exist. Storage +/// holds the copy under the indexed column's own name, while freshly shuffled rows still +/// have it under [`REFINE_VECTOR_COLUMN`]; an optimize that merges the two hands both to +/// one `StorageBuilder::build`, which concatenates by column name and rejects sources that +/// disagree. +/// +/// The scratch name is the one they have to agree on, not the stored one: `StorageBuilder` +/// quantizes from a column of the indexed column's own name and drops it whenever a batch +/// reaches it without codes, so normalising the other way would feed the carried copy to +/// the quantizer and consume it instead of preserving it. +fn stash_refine_vectors(batch: RecordBatch, column: &str) -> Result { + rename_column(batch, column, REFINE_VECTOR_COLUMN) +} + #[cfg(test)] mod tests { use std::sync::atomic::{AtomicUsize, Ordering}; @@ -2682,6 +3389,79 @@ mod tests { ShufflePartition, ShufflePartitionWindow, ShufflePartitionWindowPlan, }; + /// The carried copy's declared nullability must come from the *schema*, not from + /// whether the batch in hand happens to hold a null. + /// + /// `Array::is_nullable()` is `logical_null_count() != 0` -- "contains nulls", not "is + /// nullable". Deriving the field from it makes a nullable vector column declare a + /// non-nullable copy for every batch that happens to have no nulls, and a nullable one + /// for the next batch that does. The shuffle stream fixes the file schema from the + /// first batch it peeks and builds every partition writer from that, so the later + /// batches disagree with the schema their own file was declared with. + #[test] + fn test_add_refine_vectors_declares_nullability_from_the_schema() { + let vectors = arrow_array::FixedSizeListArray::try_new_from_values( + Float32Array::from(vec![1.0f32, 2.0, 3.0, 4.0]), + 2, + ) + .unwrap(); + // Declared nullable, but this batch carries no nulls -- the case that makes + // reading nullability off the data disagree with the schema. + let declared = Field::new("vector", vectors.data_type().clone(), true); + let schema = Arc::new(arrow_schema::Schema::new(vec![declared.clone()])); + let batch = RecordBatch::try_new(schema, vec![Arc::new(vectors)]).unwrap(); + assert_eq!( + batch.column(0).logical_null_count(), + 0, + "precondition: the batch must hold no nulls, or the bug cannot show" + ); + + let carried = add_refine_vectors(batch, "vector", Some(&declared)).unwrap(); + + let field = carried + .schema() + .field_with_name(REFINE_VECTOR_COLUMN) + .expect("the carried copy must be present") + .clone(); + assert!( + field.is_nullable(), + "the carried copy must inherit the column's declared nullability, not the \ + null count of the batch in hand" + ); + } + + /// A batch that reaches the transform chain without the indexed column is an upstream + /// defect, not a configuration: the one input that legitimately arrives without it -- + /// precomputed shuffle buffers -- is refused at the API boundary by + /// `validate_store_vectors_for_refine`. Skipping silently would surface it much later + /// as `project_by_schema` reporting a missing column, with nothing tying it to refine. + #[test] + fn test_add_refine_vectors_errors_when_the_indexed_column_is_absent() { + let schema = Arc::new(arrow_schema::Schema::new(vec![Field::new( + "payload", + DataType::Int32, + true, + )])); + let batch = RecordBatch::try_new( + schema, + vec![Arc::new(arrow_array::Int32Array::from(vec![1, 2, 3]))], + ) + .unwrap(); + let declared = Field::new( + "vector", + DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Float32, true)), 2), + true, + ); + + let err = add_refine_vectors(batch, "vector", Some(&declared)) + .expect_err("a missing indexed column must be reported, not skipped"); + let msg = err.to_string(); + assert!( + msg.contains("store_vectors_for_refine") && msg.contains("vector"), + "the error must name the option and the column; got: {msg}" + ); + } + struct SingleBatchReader { batch: RecordBatch, partition_id: usize, @@ -3117,6 +3897,8 @@ mod tests { 0, &[], Some(&reader), + false, + None, ) .await .unwrap(); diff --git a/rust/lance/src/index/vector/fixture_test.rs b/rust/lance/src/index/vector/fixture_test.rs index 2af020e1df2..88c58cc85af 100644 --- a/rust/lance/src/index/vector/fixture_test.rs +++ b/rust/lance/src/index/vector/fixture_test.rs @@ -266,6 +266,7 @@ mod test { query_parallelism: lance_index::vector::DEFAULT_QUERY_PARALLELISM, dist_q_c: 0.0, approx_mode: Default::default(), + covering_projection: None, }; let idx = make_idx.clone()(expected_query_at_subindex, metric).await; let (partition_ids, _) = idx.find_partitions(&q).unwrap(); diff --git a/rust/lance/src/index/vector/hamming.rs b/rust/lance/src/index/vector/hamming.rs index 3240ab3021b..6a33b235a6f 100644 --- a/rust/lance/src/index/vector/hamming.rs +++ b/rust/lance/src/index/vector/hamming.rs @@ -26,7 +26,7 @@ use lance_index::vector::VectorIndex; use lance_index::vector::flat::index::{FlatBinQuantizer, FlatIndex}; use lance_index::vector::flat::storage::FLAT_COLUMN; use lance_index::vector::ivf::storage::IvfModel; -use lance_index::vector::storage::VectorStore; +use lance_index::vector::storage::{PartitionColumns, VectorStore}; use lance_linalg::distance::{ BinaryHashValues, ClusteringResult, cluster_pairwise_result, extract_binary_hashes_from_fixed_list, pairwise_hamming_distance_binary_parallel, @@ -242,9 +242,10 @@ async fn hamming_clustering_for_ivf_partition_impl( let mut hash_chunks = Vec::new(); let mut num_hashes = 0; for segment in &segments { + // Only the row ids and the binary hashes are read below, both internal columns. let storage = segment .ivf_flat_bin() - .load_partition_storage(partition_id, None) + .load_partition_storage(partition_id, PartitionColumns::Internal, None) .await?; all_row_ids.extend(storage.row_ids().copied()); for batch in storage.to_batches()? { diff --git a/rust/lance/src/index/vector/ivf.rs b/rust/lance/src/index/vector/ivf.rs index 2ac0760caf6..c534067cf4a 100644 --- a/rust/lance/src/index/vector/ivf.rs +++ b/rust/lance/src/index/vector/ivf.rs @@ -649,6 +649,9 @@ pub(crate) async fn optimize_vector_indices( vector_column: &str, logical_index: &LogicalIvfView<'_>, options: &OptimizeOptions, + // Covering ("included") column names to preserve through the merge. Empty + // for ordinary indexes. Only the v2 path supports covering columns. + covering_columns: &[String], ) -> Result<(Uuid, usize, Vec)> { let existing_indices = logical_index.indices().cloned().collect::>(); // Sanity check the indices @@ -663,8 +666,15 @@ pub(crate) async fn optimize_vector_indices( // fallback to v2 IVFIndex if it's not v1 IVFIndex if !existing_indices[0].as_any().is::() { let sources = existing_index_sources(&dataset, logical_index); - return optimize_vector_indices_v2(&dataset, unindexed, vector_column, &sources, options) - .await; + return optimize_vector_indices_v2( + &dataset, + unindexed, + vector_column, + &sources, + options, + covering_columns, + ) + .await; } let new_uuid = Uuid::new_v4(); @@ -735,6 +745,9 @@ pub(crate) async fn optimize_vector_indices_v2( vector_column: &str, existing_indices: &[ExistingIndex], options: &OptimizeOptions, + // Covering ("included") column names to preserve through the merge/split/join. + // Empty for ordinary indexes. + covering_columns: &[String], ) -> Result<(Uuid, usize, Vec)> { // Sanity check the indices if existing_indices.is_empty() { @@ -761,6 +774,8 @@ pub(crate) async fn optimize_vector_indices_v2( let shuffler = create_ivf_shuffler(temp_dir_path, num_partitions, format_version, None); let (_, element_type) = get_vector_type(dataset.schema(), vector_column)?; + // Every vector index type supports covering ("included") columns, so each branch below + // threads `covering_columns` into its builder regardless of quantizer. let summary = match index_type { // IVF_FLAT (SubIndexType::Flat, QuantizationType::Flat) => { @@ -777,6 +792,7 @@ pub(crate) async fn optimize_vector_indices_v2( )? .with_ivf(ivf_model.clone()) .with_quantizer(quantizer.try_into()?) + .with_covering_payload(covering_columns.to_vec(), false) .with_existing_index_sources(existing_indices.clone()) .with_progress(options.progress.clone()) .shuffle_data_input(unindexed) @@ -795,6 +811,7 @@ pub(crate) async fn optimize_vector_indices_v2( )? .with_ivf(ivf_model.clone()) .with_quantizer(quantizer.try_into()?) + .with_covering_payload(covering_columns.to_vec(), false) .with_existing_index_sources(existing_indices.clone()) .with_progress(options.progress.clone()) .shuffle_data_input(unindexed) @@ -818,6 +835,7 @@ pub(crate) async fn optimize_vector_indices_v2( .with_quantizer(quantizer.try_into()?) .with_existing_index_sources(existing_indices.clone()) .with_progress(options.progress.clone()) + .with_covering_payload(covering_columns.to_vec(), false) .shuffle_data_input(unindexed) .build() .await? @@ -838,6 +856,7 @@ pub(crate) async fn optimize_vector_indices_v2( .with_quantizer(quantizer.try_into()?) .with_existing_index_sources(existing_indices.clone()) .with_progress(options.progress.clone()) + .with_covering_payload(covering_columns.to_vec(), false) .shuffle_data_input(unindexed) .build() .await? @@ -858,6 +877,7 @@ pub(crate) async fn optimize_vector_indices_v2( .with_quantizer(quantizer.try_into()?) .with_existing_index_sources(existing_indices.clone()) .with_progress(options.progress.clone()) + .with_covering_payload(covering_columns.to_vec(), false) .shuffle_data_input(unindexed) .build() .await? @@ -877,6 +897,7 @@ pub(crate) async fn optimize_vector_indices_v2( .with_quantizer(quantizer.try_into()?) .with_existing_index_sources(existing_indices.clone()) .with_progress(options.progress.clone()) + .with_covering_payload(covering_columns.to_vec(), false) .shuffle_data_input(unindexed) .build() .await? @@ -896,6 +917,7 @@ pub(crate) async fn optimize_vector_indices_v2( )? .with_ivf(ivf_model.clone()) .with_quantizer(quantizer.try_into()?) + .with_covering_payload(covering_columns.to_vec(), false) .with_existing_index_sources(existing_indices.clone()) .with_progress(options.progress.clone()) .shuffle_data_input(unindexed) @@ -914,6 +936,7 @@ pub(crate) async fn optimize_vector_indices_v2( )? .with_ivf(ivf_model.clone()) .with_quantizer(quantizer.try_into()?) + .with_covering_payload(covering_columns.to_vec(), false) .with_existing_index_sources(existing_indices.clone()) .with_progress(options.progress.clone()) .shuffle_data_input(unindexed) @@ -937,6 +960,7 @@ pub(crate) async fn optimize_vector_indices_v2( .with_quantizer(quantizer.try_into()?) .with_existing_index_sources(existing_indices.clone()) .with_progress(options.progress.clone()) + .with_covering_payload(covering_columns.to_vec(), false) .shuffle_data_input(unindexed) .build() .await? @@ -957,6 +981,7 @@ pub(crate) async fn optimize_vector_indices_v2( .with_quantizer(quantizer.try_into()?) .with_existing_index_sources(existing_indices.clone()) .with_progress(options.progress.clone()) + .with_covering_payload(covering_columns.to_vec(), false) .shuffle_data_input(unindexed) .build() .await? @@ -4870,6 +4895,7 @@ mod tests { use lance_index::metrics::NoOpMetricsCollector; use lance_index::scalar::OldIndexDataFilter; use lance_index::vector::sq::builder::SQBuildParams; + use lance_index::vector::storage::PartitionColumns; use lance_linalg::distance::l2_distance_batch; use lance_testing::datagen::{ generate_random_array, generate_random_array_with_range, generate_random_array_with_seed, @@ -4967,6 +4993,7 @@ mod tests { "vector", &sources, &OptimizeOptions::new(), + &[], ) .await .unwrap(); @@ -4982,6 +5009,7 @@ mod tests { "vector", &sources, &OptimizeOptions::merge(1), + &[], ) .await .unwrap(); @@ -5323,6 +5351,7 @@ mod tests { query_parallelism: lance_index::vector::DEFAULT_QUERY_PARALLELISM, dist_q_c: 0.0, approx_mode: Default::default(), + covering_projection: None, }; let (partitions, _) = index.find_partitions(&query).unwrap(); let nearest_partition_id = partitions.value(0) as usize; @@ -6851,7 +6880,10 @@ mod tests { ); // PQ code is on residual space - let pq_store = ivf_idx.load_partition_storage(0, None).await.unwrap(); + let pq_store = ivf_idx + .load_partition_storage(0, PartitionColumns::Internal, None) + .await + .unwrap(); pq_store .codebook() .values() diff --git a/rust/lance/src/index/vector/ivf/v2.rs b/rust/lance/src/index/vector/ivf/v2.rs index e369ad53888..fb0c2a0cc31 100644 --- a/rust/lance/src/index/vector/ivf/v2.rs +++ b/rust/lance/src/index/vector/ivf/v2.rs @@ -8,19 +8,24 @@ use std::marker::PhantomData; use std::{ any::Any, borrow::Cow, - collections::{BinaryHeap, HashMap}, + collections::{BinaryHeap, HashMap, HashSet}, sync::{ Arc, LazyLock, Mutex, OnceLock, atomic::{AtomicBool, Ordering}, }, }; -use crate::index::vector::{IndexFileVersion, builder::index_type_string}; +use crate::index::vector::{ + IndexFileVersion, + builder::index_type_string, + utils::{gather_covering_columns_by_row_id, row_id_take_indices}, +}; use crate::index::{PreFilter, vector::VectorIndex}; use arrow::compute::concat_batches; use arrow_arith::numeric::sub; -use arrow_array::{ArrayRef, Float32Array, RecordBatch, UInt32Array, UInt64Array}; -use arrow_schema::DataType; +use arrow_array::{ArrayRef, Float32Array, RecordBatch, UInt32Array, UInt64Array, cast::AsArray}; +use arrow_schema::{DataType, Field}; +use arrow_select::take::take_record_batch; use async_trait::async_trait; use datafusion::error::{DataFusionError, Result as DataFusionResult}; use datafusion::execution::SendableRecordBatchStream; @@ -68,8 +73,11 @@ use lance_index::{ INDEX_AUXILIARY_FILE_NAME, INDEX_FILE_NAME, Index, IndexType, pb, vector::{ DISTANCE_TYPE_KEY, PartitionSearchControl, PreparedPartitionSearchHandle, Query, - VECTOR_RESULT_SCHEMA, ivf::storage::IVF_METADATA_KEY, quantizer::Quantization, - storage::IvfQuantizationStorage, v3::subindex::IvfSubIndex, + VECTOR_RESULT_SCHEMA, + ivf::storage::IVF_METADATA_KEY, + quantizer::Quantization, + storage::{IvfQuantizationStorage, PartitionColumns}, + v3::subindex::IvfSubIndex, }, }; use lance_index::{INDEX_METADATA_SCHEMA_KEY, IndexMetadata}; @@ -165,6 +173,10 @@ struct PreparedPartitionSearch { query: Query, pre_filter: Arc, partition_id: usize, + /// Rows this partition occupies in the storage file. Compared against the loaded + /// partition's row count to decide whether a storage position is also a file + /// position (see [`CoveringGather::WholeRange`]). + partition_rows: usize, partition_centroid: Option, rq_search_cache: Option>, raw_query_context: Option>, @@ -172,6 +184,45 @@ struct PreparedPartitionSearch { _marker: PhantomData<(S, Q)>, } +/// The covering ("included") columns one query needs from one index, resolved once per +/// search. +/// +/// Absent (`None` at the call sites that hold this) for an ordinary index **and** for a +/// query whose covering projection names none of the declared columns: both mean the same +/// thing to the search path -- emit `[_distance, _rowid]` and do no covering read at all. +/// Collapsing the second case into "project an already-read batch down to nothing" is +/// exactly the regression [`Query::covering_projection`] documents. +struct QueryCovering { + /// `[_rowid, ]`, in index declaration order. Both the declared stream + /// schema and the emitted batches are built from this. + schema: arrow_schema::SchemaRef, + /// The same included columns by name and in the same order, for the storage read. + columns: Vec, +} + +/// What the covering gather must read for one partition's search survivors. +enum CoveringGather { + /// The query needs no covering column from this index. + NotNeeded, + /// Read exactly these positions within the partition's row range. Strictly ascending + /// and deduplicated, which is what the reader's take path requires. + Positions(Vec), + /// Positions are not derivable: a deferred fragment-reuse remap drops rows from the + /// loaded partition, so a storage position is no longer the file position it was read + /// from. Read the partition's whole range instead and match by row id. + WholeRange, +} + +/// Where one heap survivor's covering values live, recorded while its partition is still +/// loaded so the gather after the heap settles is a bounded read rather than a re-search. +#[derive(Debug, Clone, Copy)] +struct CoveringLocation { + partition_id: usize, + /// Position within the partition's row range, or `None` when positions are not + /// derivable for that partition (see [`CoveringGather::WholeRange`]). + position: Option, +} + #[derive(Debug)] pub(crate) struct RabitSearchCache { rotated_centroids: Vec, @@ -832,6 +883,7 @@ impl IVFIndex { query: query.clone(), pre_filter, partition_id, + partition_rows: self.storage.partition_size(partition_id), partition_centroid: self.ivf.centroid(partition_id), rq_search_cache: self.rq_search_cache.clone(), raw_query_context, @@ -856,6 +908,7 @@ impl IVFIndex { query: query.clone(), pre_filter, partition_id, + partition_rows: self.storage.partition_size(partition_id), partition_centroid: self.ivf.centroid(partition_id), rq_search_cache: self.rq_search_cache.clone(), raw_query_context, @@ -864,17 +917,23 @@ impl IVFIndex { }) } + /// The CPU half of a prepared partition search: score the partition and locate its + /// survivors' covering rows. Reading those rows is I/O and happens in the caller's + /// async context (see [`IVFIndex::append_covering`]) -- this runs on the CPU pool, + /// where awaiting is not an option. fn run_prepared_partition_search( use_query_residual: bool, use_residual_scratch: bool, prepared: PreparedPartitionSearch, + want_covering: bool, metrics: &dyn MetricsCollector, scratch: &mut QueryScratch, - ) -> Result { + ) -> Result<(RecordBatch, CoveringGather)> { let PreparedPartitionSearch { query, pre_filter, partition_id, + partition_rows, partition_centroid, rq_search_cache, raw_query_context, @@ -911,7 +970,65 @@ impl IVFIndex { residual, scratch, )?; - Ok(batch) + let gather = match want_covering { + true => Self::survivor_positions(&batch, &part_entry.storage, partition_rows)?, + false => CoveringGather::NotNeeded, + }; + Ok((batch, gather)) + } + + /// Locate the rows of `batch` (a `[_distance, _rowid]` partition result) within the + /// partition's row range in the storage file. + /// + /// The scan walks the partition's row ids once against the (`<= k`)-sized survivor set + /// with an early exit, rather than building a partition-sized side map per probe. + /// Positions come out ascending and deduplicated because each storage row is visited + /// at most once, which is what the reader's take path requires. + fn survivor_positions( + batch: &RecordBatch, + storage: &Q::Storage, + partition_rows: usize, + ) -> Result { + // A deferred fragment-reuse remap filters rows out of the partition as it is + // loaded, so from that point a storage position is no longer the file position it + // was read from and cannot address the file. + if storage.len() != partition_rows { + return Ok(CoveringGather::WholeRange); + } + let row_ids = batch + .column_by_name(ROW_ID) + .ok_or_else(|| Error::internal("search result missing row id".to_string()))? + .as_primitive::(); + let mut needed: HashSet = row_ids.values().iter().copied().collect(); + let mut positions = Vec::with_capacity(needed.len()); + Self::locate_survivors(storage, &mut needed, |_, position| positions.push(position)); + Ok(CoveringGather::Positions(positions)) + } + + /// Scan a partition's row ids once against the (`<= k`)-sized `needed` set with an + /// early exit, calling `on_hit(row_id, storage_position)` for each located row. + /// + /// The one shared core of [`Self::survivor_positions`] (per-partition search) and + /// the locator inside `accumulate_prepared_partition_search` (global-heap search). + /// The two callers differ in how they react to a storage/partition misalignment -- + /// early whole-range fallback vs. scan-anyway with `position: None` -- but the scan + /// itself must stay identical: a future change to when positions are valid (e.g. a + /// new remap variant) goes through here, so the two paths cannot silently diverge + /// on the deferred-remap case. Positions come out ascending and deduplicated + /// because each storage row is visited at most once. + fn locate_survivors( + storage: &Q::Storage, + needed: &mut HashSet, + mut on_hit: impl FnMut(u64, u32), + ) { + for (position, row_id) in storage.row_ids().enumerate() { + if needed.remove(row_id) { + on_hit(*row_id, position as u32); + if needed.is_empty() { + break; + } + } + } } #[allow(clippy::too_many_arguments)] @@ -920,6 +1037,8 @@ impl IVFIndex { use_residual_scratch: bool, prepared: PreparedPartitionSearch, heap: &mut BinaryHeap>, + covering_locations: &mut HashMap, + want_covering: bool, scratch: &mut QueryScratch, metrics: &dyn MetricsCollector, ) -> Result<()> { @@ -927,6 +1046,7 @@ impl IVFIndex { query, pre_filter, partition_id, + partition_rows, partition_centroid, rq_search_cache, raw_query_context, @@ -953,6 +1073,8 @@ impl IVFIndex { let param = (&query).into(); let refine_factor = query.refine_factor.unwrap_or(1) as usize; let k = query.k * refine_factor; + // Accumulate this partition's contribution to the top-k heap first, so the heap + // membership consulted below already reflects it. part_entry.index.accumulate_topk_with_scratch( query.key, k, @@ -963,7 +1085,51 @@ impl IVFIndex { residual, scratch, metrics, - ) + )?; + // Keep covering O(k): record only where the current heap survivors' covering rows + // live, then prune the map to the heap's membership so rows this partition evicted + // drop their entry. Nothing is read here -- the read happens once, after the heap + // has settled, for the survivors that actually made it (see + // `gather_survivor_covering`); buffering every probed partition's covering to serve + // k winners is exactly the O(nprobe * partition_size) cost this phase removes. + // Every heap survivor is located: a row is a survivor only if it was in the top-k + // when its own partition was processed (the heap never re-admits an evicted row), + // so it was located then -- or it is new from this partition. + // The lookup scans the partition's row ids against the (<= k)-sized needed set + // with an early exit, instead of building a partition-sized side map per probe. + if want_covering { + let heap_ids: HashSet = heap.iter().map(|node| node.id).collect(); + let mut needed: HashSet = heap_ids + .iter() + .filter(|id| !covering_locations.contains_key(id)) + .copied() + .collect(); + if !needed.is_empty() { + // A deferred fragment-reuse remap filters rows out of the partition as it + // is loaded, so from that point a storage position no longer addresses the + // file and this partition must be re-read as a whole range. + let aligned = part_entry.storage.len() == partition_rows; + Self::locate_survivors(&part_entry.storage, &mut needed, |row_id, position| { + covering_locations.insert( + row_id, + CoveringLocation { + partition_id, + position: aligned.then_some(position), + }, + ); + }); + } + covering_locations.retain(|id, _| heap_ids.contains(id)); + // Invariant: the map locates exactly the heap's distinct row ids -- never more + // (that would break the O(k) bound) and never fewer (a survivor would be + // missing its covering at emit time). + debug_assert_eq!( + covering_locations.len(), + heap_ids.len(), + "covering locations must track exactly the heap's survivors (O(k))" + ); + } + Ok(()) } fn query_context_for_scratch<'a>( @@ -991,17 +1157,240 @@ impl IVFIndex { } } - fn global_heap_to_batch(heap: BinaryHeap>) -> Result { + fn global_heap_to_batch( + heap: BinaryHeap>, + // `[_rowid, ]` covering the heap survivors and nothing else, + // gathered after the heap settled by `gather_survivor_covering`. `None` when the + // gather was not run (no survivor was located). + covering: Option<&RecordBatch>, + // `[_rowid, ]` schema for the covering columns this query needs, + // or None when the index has none or the query needs none. This is a per-query + // constant, so it -- not whether any covering was gathered -- decides whether to + // emit the wider covered schema. That keeps the emitted schema equal to the schema + // the exec declares even when zero partitions were searched (heap empty). + covering_schema: Option<&arrow_schema::Schema>, + ) -> Result { let (row_ids, dists): (Vec<_>, Vec<_>) = heap.into_iter().map(|r| (r.id, r.dist.0)).unzip(); + let dist_arr: ArrayRef = Arc::new(Float32Array::from(dists)); + let row_id_arr: ArrayRef = Arc::new(UInt64Array::from(row_ids)); + let Some(covering_schema) = covering_schema else { + // Ordinary index: `[_distance, _rowid]`. + return Ok(RecordBatch::try_new( + VECTOR_RESULT_SCHEMA.clone(), + vec![dist_arr, row_id_arr], + )?); + }; + // Covered index: emit `[_distance, _rowid, ]`. + let mut fields: Vec = VECTOR_RESULT_SCHEMA.fields().to_vec(); + let mut columns: Vec = vec![dist_arr, row_id_arr.clone()]; + if row_id_arr.is_empty() { + // No survivors (heap empty / zero partitions searched). Emit the covering + // columns as empty arrays so the schema still matches the declared covered + // schema. + for field in covering_schema.fields() { + if field.name() == ROW_ID { + continue; + } + fields.push(field.clone()); + columns.push(arrow_array::new_null_array(field.data_type(), 0)); + } + } else { + // Survivors but nothing gathered for them. Reported rather than filled with + // nulls: a null fill is indistinguishable from genuine nulls once a covering + // column is nullable, so the query would return the right rows with silently + // wrong values. + let covering = covering.ok_or_else(|| { + Error::index(format!( + "index declares covering columns {:?} but none were gathered for the \ + {} result rows", + covering_schema + .fields() + .iter() + .map(|f| f.name().as_str()) + .filter(|name| *name != ROW_ID) + .collect::>(), + row_id_arr.len(), + )) + })?; + // Align the gathered values to the final row ids by row id, never by + // position: the gather returns each partition's rows in file order, which is + // not the heap's order. A survivor the gather did not return is an error + // there, so this cannot silently drop a row. + let row_id_u64 = row_id_arr.as_primitive::(); + let included = gather_covering_columns_by_row_id(covering, row_id_u64)?; + for (field, array) in included { + fields.push(field); + columns.push(array); + } + } Ok(RecordBatch::try_new( - VECTOR_RESULT_SCHEMA.clone(), - vec![ - Arc::new(Float32Array::from(dists)), - Arc::new(UInt64Array::from(row_ids)), - ], + Arc::new(arrow_schema::Schema::new(fields)), + columns, + )?) + } + + /// The schema search results carry: `[_distance, _rowid]` plus the covering + /// ("included") columns of `covering_schema`, in storage order. Used to declare + /// stream schemas that stay consistent with the (possibly widened) batches. + fn covered_result_schema( + covering_schema: Option<&arrow_schema::Schema>, + ) -> arrow_schema::SchemaRef { + let Some(covering_schema) = covering_schema else { + return VECTOR_RESULT_SCHEMA.clone(); + }; + let mut fields: Vec = VECTOR_RESULT_SCHEMA.fields().to_vec(); + for field in covering_schema.fields() { + if field.name() == ROW_ID { + continue; + } + fields.push(field.clone()); + } + Arc::new(arrow_schema::Schema::new(fields)) + } + + /// The covering columns `query` needs from this index, or `None` when there are none + /// to emit. + /// + /// `None` covers two cases that are the same instruction to the search path -- do no + /// covering work at all: an ordinary index, and a covered index whose covering + /// projection this query narrows to nothing. See [`Query::covering_projection`] for + /// why the second must not degrade into "read everything and project it away". + fn query_covering(&self, query: &Query) -> Result> { + let columns = self + .storage + .covering_columns(query.covering_projection.as_deref())?; + if columns.is_empty() { + return Ok(None); + } + Ok(Some(QueryCovering { + schema: self.storage.covering_read_schema(&columns)?, + columns, + })) + } + + /// Widen a `[_distance, _rowid]` partition result with its survivors' covering values, + /// read from the storage file by position. + /// + /// This is the per-partition emit site. It runs in the caller's async context, after + /// the CPU phase has settled which rows survived, so the read is proportional to `k` + /// rather than to the partition size -- and it is deliberately not cached, being a + /// different set of rows for every query. + async fn append_covering( + &self, + partition_id: usize, + batch: RecordBatch, + gather: CoveringGather, + covering: Option<&QueryCovering>, + io_stats: Option, + ) -> Result { + let positions = match &gather { + CoveringGather::NotNeeded => return Ok(batch), + CoveringGather::Positions(positions) => Some(positions.as_slice()), + CoveringGather::WholeRange => None, + }; + let Some(covering) = covering else { + return Ok(batch); + }; + let row_ids = batch + .column_by_name(ROW_ID) + .ok_or_else(|| Error::internal("search result missing row id".to_string()))? + .as_primitive::(); + let mut fields: Vec = batch.schema().fields().to_vec(); + let mut columns: Vec = batch.columns().to_vec(); + if row_ids.is_empty() { + // Nothing survived this partition, so there is nothing to read. Still emit the + // covering columns (empty) so every batch on the stream carries the schema the + // exec declared. + for field in covering.schema.fields() { + if field.name() == ROW_ID { + continue; + } + columns.push(arrow_array::new_null_array(field.data_type(), 0)); + fields.push(field.clone()); + } + } else { + let gathered = self + .storage + .take_covering(partition_id, positions, &covering.columns, io_stats) + .await?; + // By row id, never by position: the gather returns the partition's rows in + // file order while the search result is in distance order. A survivor the + // gather did not return is an error there rather than a silent null. + for (field, array) in gather_covering_columns_by_row_id(&gathered, row_ids)? { + fields.push(field); + columns.push(array); + } + } + Ok(RecordBatch::try_new( + Arc::new(arrow_schema::Schema::new(fields)), + columns, )?) } + /// Read `[_rowid, ]` for the heap's survivors: one bounded read per + /// partition that still owns one, after the heap has settled. + /// + /// Each partition's covering is narrowed to its own survivors before the next + /// partition is read, so what the loop accumulates is `O(survivors)` and only one + /// partition's covering is ever alive at a time. That matters on the whole-range + /// fallback (a pending fragment-reuse remap puts every partition on it): keeping each + /// partition's whole covering to concatenate at the end would make peak memory + /// `partitions x partition size`, per concurrent query and outside the shared, + /// evictable partition cache -- the opposite of what bounding this read is for. + /// + /// Returns `None` when no survivor was located, which for a covered query means the + /// heap is empty -- [`Self::global_heap_to_batch`] turns "survivors but nothing + /// gathered" into an error rather than a null fill. + async fn gather_survivor_covering( + &self, + locations: &HashMap, + covering: &QueryCovering, + io_stats: Option, + ) -> Result> { + // Group by partition so each contributing partition is read once, not once per + // survivor. `None` positions mark a partition whose positions are not derivable. + let mut by_partition: HashMap>, Vec)> = HashMap::new(); + for (row_id, location) in locations { + let (positions, row_ids) = by_partition + .entry(location.partition_id) + .or_insert_with(|| (Some(Vec::new()), Vec::new())); + row_ids.push(*row_id); + match (positions.as_mut(), location.position) { + (Some(positions), Some(position)) => positions.push(position), + _ => *positions = None, + } + } + let mut batches = Vec::with_capacity(by_partition.len()); + for (partition_id, (mut positions, row_ids)) in by_partition { + if let Some(positions) = positions.as_mut() { + // The reader's take path requires strictly ascending indices; survivors + // arrive in heap order, which is neither sorted nor partition-local. + positions.sort_unstable(); + positions.dedup(); + } + let gathered = self + .storage + .take_covering( + partition_id, + positions.as_deref(), + &covering.columns, + io_stats.clone(), + ) + .await?; + // By row id, never by position: on the fallback `gathered` is the partition's + // whole range in file order, and even on the scattered read the caller's + // survivors are in heap order. A survivor the read did not return is an error + // here rather than a row quietly dropped from the result. + let survivors = UInt64Array::from(row_ids); + let take_idx = row_id_take_indices(&gathered, &survivors)?; + batches.push(take_record_batch(&gathered, &take_idx)?); + } + if batches.is_empty() { + return Ok(None); + } + Ok(Some(concat_batches(&covering.schema, batches.iter())?)) + } + fn preprocess_partition_query( use_query_residual: bool, use_residual_scratch: bool, @@ -1270,8 +1659,15 @@ impl IVFIndex { .get_or_insert_with_key_hit(cache_key, || async { info!(target: TRACE_IO_EVENTS, r#type=IO_TYPE_LOAD_VECTOR_PART, index_type="ivf", part_id=partition_id); metrics.record_part_load(); - self.load_partition_entry(partition_id, metrics.io_stats()) - .await + // Internal columns only: this entry is cached under a partition id + // alone, so it must not depend on which covering columns the loading + // query wanted. + self.load_partition_entry( + partition_id, + PartitionColumns::Internal, + metrics.io_stats(), + ) + .await }) .await; match &result { @@ -1289,15 +1685,44 @@ impl IVFIndex { info!(target: TRACE_IO_EVENTS, r#type=IO_TYPE_LOAD_VECTOR_PART, index_type="ivf", part_id=partition_id); metrics.record_part_load(); Ok(Arc::new( - self.load_partition_entry(partition_id, metrics.io_stats()) - .await?, + self.load_partition_entry( + partition_id, + PartitionColumns::Internal, + metrics.io_stats(), + ) + .await?, )) } } + /// Load a partition entry carrying every column its storage file holds, bypassing + /// the partition cache in both directions. + /// + /// Rewrite paths need this: the entry they load is written straight back out, so a + /// covering column left unread is a covering column the rewritten index no longer + /// has. The cache is bypassed rather than consulted because its key + /// ([`IVFPartitionKey`]) has no covering component -- a hit would hand back the + /// codes-only entry the search path stores, and an insert would hand a covering-laden + /// entry to the search path. + pub(crate) async fn load_partition_entry_with_covering( + &self, + partition_id: usize, + ) -> Result> { + if partition_id >= self.ivf.num_partitions() { + return Err(Error::index(format!( + "partition id {} is out of range of {} partitions", + partition_id, + self.ivf.num_partitions() + ))); + } + self.load_partition_entry(partition_id, PartitionColumns::All, None) + .await + } + async fn load_partition_entry( &self, partition_id: usize, + columns: PartitionColumns, io_stats: Option, ) -> Result> { // `concat_batches` indexes the batches by this schema's field positions @@ -1354,16 +1779,39 @@ impl IVFIndex { self.sub_index_metadata[partition_id].clone(), )?; let idx = S::load(batch)?; - let storage = self.load_partition_storage(partition_id, io_stats).await?; + let storage = self + .load_partition_storage(partition_id, columns, io_stats) + .await?; Ok(PartitionEntry::new(idx, storage)) } pub async fn load_partition_storage( &self, partition_id: usize, + columns: PartitionColumns, io_stats: Option, ) -> Result { - self.storage.load_partition(partition_id, io_stats).await + self.storage + .load_partition(partition_id, columns, io_stats) + .await + } + + /// Names of this index's covering ("included") columns, in storage order, or + /// empty if the index has none. Used by remap/rebuild paths to re-project the + /// covering columns so they survive into the rewritten storage. + pub(crate) fn covering_column_names(&self) -> Result> { + Ok(self + .storage + .covering_schema()? + .map(|schema| { + schema + .fields() + .iter() + .filter(|f| f.name() != ROW_ID) + .map(|f| f.name().to_string()) + .collect() + }) + .unwrap_or_default()) } /// preprocess the query vector given the partition id. @@ -1542,6 +1990,10 @@ impl VectorIndex for IVFInd self.ivf.num_partitions() } + fn physical_covering_fields(&self) -> Result> { + Ok(self.storage.physical_covering_fields()) + } + #[instrument(level = "debug", skip(self, pre_filter, metrics))] async fn search_in_partition( &self, @@ -1569,7 +2021,10 @@ impl VectorIndex for IVFInd let scratch_pool = self.scratch_pool.clone(); let use_query_residual = self.use_query_residual; let use_residual_scratch = self.use_residual_scratch; - let (batch, local_metrics) = spawn_cpu(move || { + let covering = self.query_covering(&query)?; + let want_covering = covering.is_some(); + let partition_rows = self.storage.partition_size(partition_id); + let (batch, gather, local_metrics) = spawn_cpu(move || { let param = (&query).into(); let refine_factor = query.refine_factor.unwrap_or(1) as usize; let k = query.k * refine_factor; @@ -1596,13 +2051,26 @@ impl VectorIndex for IVFInd scratch, ) })?; - Result::Ok((batch, local_metrics)) + // Locate the survivors' covering rows while the partition is still loaded; + // reading them is I/O and cannot happen on the CPU pool. + let gather = match want_covering { + true => Self::survivor_positions(&batch, &part_entry.storage, partition_rows)?, + false => CoveringGather::NotNeeded, + }; + Result::Ok((batch, gather, local_metrics)) }) .await?; local_metrics.dump_into(metrics); - Ok(batch) + self.append_covering( + partition_id, + batch, + gather, + covering.as_ref(), + metrics.io_stats(), + ) + .await } async fn prepare_partition_search( @@ -1619,6 +2087,15 @@ impl VectorIndex for IVFInd )) } + /// The synchronous phase of a prepared partition search. + /// + /// A covered query cannot be served here: since phase 8 the covering values are read + /// from the storage file once the survivors are known, and this entry point exists + /// precisely to run on the CPU pool where that read cannot be awaited. Callers that + /// need covering columns use [`VectorIndex::search_in_partition`] or + /// [`VectorIndex::search_partitions`], both of which own an async context; this + /// reports the mismatch rather than returning rows with their covering columns + /// silently missing. fn search_prepared_partition( &self, prepared: PreparedPartitionSearchHandle, @@ -1627,19 +2104,40 @@ impl VectorIndex for IVFInd let prepared = prepared .downcast::>() .map_err(|_| Error::internal("failed to downcast prepared partition search"))?; - self.scratch_pool.with_scratch(|scratch| { + if let Some(covering) = self.query_covering(&prepared.query)? { + return Err(Error::index(format!( + "search_prepared_partition cannot materialize the covering columns {:?} \ + this query needs: the covering values are read from index storage after \ + scoring, and this is the synchronous phase of a prepared search. Use \ + search_in_partition or search_partitions instead.", + covering.columns, + ))); + } + let (batch, _) = self.scratch_pool.with_scratch(|scratch| { Self::run_prepared_partition_search( self.use_query_residual, self.use_residual_scratch, *prepared, + false, metrics, scratch, ) - }) + })?; + Ok(batch) } + /// False for a covered index, because [`Self::search_prepared_partition`] rejects the + /// queries such an index exists to serve (the covering values are read after scoring, + /// which is I/O, and that method is the synchronous phase). + /// + /// The flag carries no query, so it answers for the index as a whole: a covered index + /// that a particular query narrows to no covering column would in fact be servable + /// there, but advertising `true` and then failing would strand a dispatcher after it + /// had already paid for the partition load. Saying `false` lets it choose + /// [`VectorIndex::search_in_partition`] or [`VectorIndex::search_partitions`] up + /// front. An unreadable covering schema answers `false` for the same reason. fn supports_prepared_partition_search(&self) -> bool { - true + matches!(self.storage.covering_schema(), Ok(None)) } fn auto_query_parallelism(&self, cpu_pool_size: usize) -> usize { @@ -1678,6 +2176,9 @@ impl VectorIndex for IVFInd let prepare_parallelism = get_num_compute_intensive_cpus().max(1); let raw_query_context = self.prepare_rq_raw_query_context(&query.key)?; + // Resolved once per search, before `query` is moved into the prepare stream. + let covering = self.query_covering(&query)?; + let want_covering = covering.is_some(); if control.is_none() && S::supports_global_topk_heap() { let heap_capacity = query.k * query.refine_factor.unwrap_or(1) as usize; @@ -1714,28 +2215,57 @@ impl VectorIndex for IVFInd let use_residual_scratch = self.use_residual_scratch; let search_metrics = metrics.clone(); let scratch_pool = self.scratch_pool.clone(); - let batch = spawn_cpu(move || -> DataFusionResult { - let mut heap = BinaryHeap::with_capacity(heap_capacity); - scratch_pool.with_scratch(|scratch| -> DataFusionResult<()> { - for prepared in prepared { - Self::accumulate_prepared_partition_search( - use_query_residual, - use_residual_scratch, - prepared, - &mut heap, - scratch, - search_metrics.as_ref(), - ) - .map_err(DataFusionError::from)?; - } - Ok(()) - })?; - Self::global_heap_to_batch(heap).map_err(DataFusionError::from) - }) + let (heap, covering_locations) = spawn_cpu( + move || -> DataFusionResult<( + BinaryHeap>, + HashMap, + )> { + let mut heap = BinaryHeap::with_capacity(heap_capacity); + // O(k) side map locating the heap survivors' covering rows, recorded + // in-flight while each partition is loaded (empty for ordinary + // indexes). Nothing is read until the heap has settled. + let mut covering_locations: HashMap = HashMap::new(); + scratch_pool.with_scratch(|scratch| -> DataFusionResult<()> { + for prepared in prepared { + Self::accumulate_prepared_partition_search( + use_query_residual, + use_residual_scratch, + prepared, + &mut heap, + &mut covering_locations, + want_covering, + scratch, + search_metrics.as_ref(), + ) + .map_err(DataFusionError::from)?; + } + Ok(()) + })?; + Ok((heap, covering_locations)) + }, + ) .await?; + // The gather is I/O, so it runs here rather than on the CPU pool: one bounded + // read per contributing partition, for the survivors only. + let gathered = match covering.as_ref() { + Some(covering) => { + self.gather_survivor_covering(&covering_locations, covering, metrics.io_stats()) + .await? + } + None => None, + }; + let batch = Self::global_heap_to_batch( + heap, + gathered.as_ref(), + covering.as_ref().map(|covering| covering.schema.as_ref()), + )?; + + // Schema may be wider than VECTOR_RESULT_SCHEMA when covering columns + // are emitted; take it from the produced batch so they stay consistent. + let result_schema = batch.schema(); return Ok(Box::pin(RecordBatchStreamAdapter::new( - VECTOR_RESULT_SCHEMA.clone(), + result_schema, stream::once(async move { Ok(batch) }), ))); } @@ -1788,6 +2318,16 @@ impl VectorIndex for IVFInd let search_metrics = metrics.clone(); let search_control = control.clone(); let scratch_pool = self.scratch_pool.clone(); + // The covering gather runs in the async half of the search loop, so it needs the + // index and the metrics sink there as well as inside the CPU closure. + let gather_index = self.clone(); + let gather_metrics = metrics.clone(); + // The per-partition batches are widened with the covering columns this query needs + // (`append_covering`), so declare the matching schema -- the global-heap branch + // above already emits its batch's own (covered) schema. Resolved before the search + // loop takes ownership of the covering set. + let result_schema = + Self::covered_result_schema(covering.as_ref().map(|covering| covering.schema.as_ref())); // Search prepared partitions in batches. Each batch is searched in a single // `spawn_cpu` dispatch (amortizing the per-dispatch overhead the single-worker // design in #6475 avoided), but the channel `recv`/`send` stay in async code so @@ -1850,8 +2390,13 @@ impl VectorIndex for IVFInd // cancellable, so abandoning the await leaves the work running.) let cancel_probe = batch_tx.clone(); let search_output = spawn_cpu(move || { - let mut outputs: Vec> = - Vec::with_capacity(prepared_batch.len()); + // Each output carries the partition it came from and where its + // survivors sit in that partition, so the covering read below can + // be a bounded take rather than a whole-partition scan. The read + // itself is I/O and stays out of this CPU closure. + let mut outputs: Vec< + DataFusionResult<(usize, RecordBatch, CoveringGather)>, + > = Vec::with_capacity(prepared_batch.len()); // `stopped` means the whole search should end (an error, an // early-stop signal, or cancellation), not just this batch. let mut stopped = false; @@ -1865,20 +2410,22 @@ impl VectorIndex for IVFInd stopped = true; break; } + let partition_id = prepared.partition_id; match Self::run_prepared_partition_search( use_query_residual, use_residual_scratch, prepared, + want_covering, search_metrics.as_ref(), scratch, ) .map_err(DataFusionError::from) { - Ok(batch) => { + Ok((batch, gather)) => { if let Some(control) = search_control.as_ref() { control.record_batch(&batch); } - outputs.push(Ok(batch)); + outputs.push(Ok((partition_id, batch, gather))); } Err(err) => { outputs.push(Err(err)); @@ -1903,6 +2450,19 @@ impl VectorIndex for IVFInd } }; for output in outputs { + let output = match output { + Ok((partition_id, batch, gather)) => gather_index + .append_covering( + partition_id, + batch, + gather, + covering.as_ref(), + gather_metrics.io_stats(), + ) + .await + .map_err(DataFusionError::from), + Err(err) => Err(err), + }; if batch_tx.send(output).await.is_err() { return; } @@ -1923,7 +2483,7 @@ impl VectorIndex for IVFInd }); Ok(Box::pin(RecordBatchStreamAdapter::new( - VECTOR_RESULT_SCHEMA.clone(), + result_schema, ReceiverStream::new(batch_rx), ))) } @@ -2117,15 +2677,19 @@ mod tests { storage::{RABIT_BLOCKED_EX_CODE_COLUMN, RabitQuantizationMetadata, RabitQueryEstimator}, transform::{EX_ADD_FACTORS_COLUMN, EX_SCALE_FACTORS_COLUMN}, }; - use lance_index::vector::storage::VectorStore; + use lance_index::vector::storage::{PartitionColumns, VectorStore}; use lance_index::vector::v3::subindex::IvfSubIndex; - use crate::dataset::{InsertBuilder, UpdateBuilder, WriteMode, WriteParams}; + use crate::dataset::{ + InsertBuilder, NewColumnTransform, UpdateBuilder, WriteMode, WriteParams, + }; use crate::index::DatasetIndexExt; use crate::index::DatasetIndexInternalExt; use crate::index::vector::ivf::v2::{ - IVFPartitionKey, IvfFlatIndex, IvfHnswSqIndex, IvfPq, IvfStateEntryBox, PartitionEntry, + CoveringLocation, IVFPartitionKey, IvfFlatIndex, IvfHnswSqIndex, IvfPq, IvfStateEntryBox, + PartitionEntry, }; + use crate::index::vector::utils::gather_covering_columns_by_row_id; use crate::utils::test::copy_test_data_to_tmp; use crate::{ Dataset, @@ -2335,6 +2899,36 @@ mod tests { assert_eq!(capacity.u32, 0); } + /// Vector values laid out as `num_clusters` well-separated clusters: row `r` lands in + /// cluster `r % num_clusters`, centred at `cluster * 1000.0` with a tiny per-row offset + /// so no two rows coincide. + /// + /// **The separation is load-bearing, not cosmetic.** `early_pruning` RAISES + /// `minimum_nprobes` to the number of centroids within `dists[0] * factor` (7.0 for + /// `k` in 2..=10, 81.0 for `k >= 11`). With evenly-spread vectors every centroid + /// survives pruning, so `minimum_nprobes` reaches `maximum_nprobes`, `late_search` + /// returns at its first guard, and the no-rows shortcut -- and therefore the covered + /// recovery path -- is never reached. A test targeting either then passes vacuously + /// against a live bug, which has happened repeatedly on this feature. With these + /// clusters a query at the origin is ~10^6 times closer to cluster 0 than to any + /// other, so pruning yields 1 and unsearched partitions remain. + /// + /// See also `generate_clustered_batch`, the pre-existing generator used by the + /// non-covering partition-split tests; it emits a different schema and is not a + /// drop-in substitute here. + fn clustered_vector_values( + rows: impl IntoIterator, + dim: i32, + num_clusters: usize, + ) -> Vec { + rows.into_iter() + .flat_map(|r| { + let center = (r % num_clusters) as f32 * 1000.0; + (0..dim as usize).map(move |d| center + (r * dim as usize + d) as f32 * 1e-3) + }) + .collect() + } + async fn generate_test_dataset( test_uri: &str, range: Range, @@ -2945,1233 +3539,7450 @@ mod tests { } } - async fn shrink_smallest_partition( - dataset: &mut Dataset, - index_name: &str, - expected_after_join: usize, - next_id: &mut u64, - ) -> (usize, usize, usize) { - const ROWS_TO_APPEND_FOR_JOIN: usize = 32; - let row_count_before = dataset.count_all_rows().await.unwrap(); - let index_ctx = load_vector_index_context(dataset, "vector", index_name).await; - let partitions = index_ctx.stats()["indices"][0]["partitions"] - .as_array() - .expect("partitions should be present"); - let (partition_idx, _size) = partitions - .iter() - .enumerate() - .filter_map(|(idx, part)| part["size"].as_u64().map(|size| (idx, size))) - .filter(|(_, size)| *size > 1) - .min_by_key(|(_, size)| *size) - .expect("should have at least one partition with joinable rows"); - - let row_ids = load_partition_row_ids(index_ctx.ivf(), partition_idx).await; - assert!( - row_ids.len() > 1, - "Partition {} should have removable rows", - partition_idx - ); + /// End-to-end test for index-included ("covering") columns : an + /// `covering_columns` request on the build params must survive the full + /// build -> shuffle -> persist -> reopen path and land in the IVF_PQ + /// partition storage, so covered queries can avoid a take. + #[tokio::test] + async fn test_ivf_pq_covering_columns_roundtrip() { + const INDEX_NAME: &str = "vector_idx"; + let test_dir = TempStrDir::default(); + let test_uri = test_dir.as_str(); + // `generate_test_dataset` produces an `id` (UInt64) column besides `vector`. + let (mut dataset, _) = generate_test_dataset::(test_uri, 0.0..1.0).await; - let rows = dataset - .take_rows(&row_ids, dataset.schema().clone()) + let mut params = VectorIndexParams::ivf_pq(4, 8, 4, DistanceType::L2, 2); + params.covering_columns(vec!["id".to_string()]); + dataset + .create_index( + &["vector"], + IndexType::Vector, + Some(INDEX_NAME.to_string()), + ¶ms, + true, + ) .await .unwrap(); - let ids = rows["id"].as_primitive::().values(); - let template_values = rows["vector"] - .as_fixed_size_list() - .value(0) - .as_primitive::() - .values() - .to_vec(); - delete_ids(dataset, &ids[1..]).await; - compact_after_deletions(dataset).await; - - append_constant_vector_with_start_id( - dataset, - ROWS_TO_APPEND_FOR_JOIN, - &template_values, - next_id, - ) - .await; - dataset - .optimize_indices(&OptimizeOptions::new()) + let ctx = load_vector_index_context(&dataset, "vector", INDEX_NAME).await; + let storage = ctx + .ivf() + .load_partition_storage(0, PartitionColumns::All, None) .await .unwrap(); + assert!( + storage.batch().column_by_name("id").is_some(), + "included column 'id' should be co-located in IVF_PQ partition storage" + ); - let post_ctx = load_vector_index_context(dataset, "vector", index_name).await; - let post_partitions = post_ctx.num_partitions(); + // The covering columns are also declared generically on IndexMetadata + // (by field id), so the read path can discover them for any index type. + let id_field_id = dataset.schema().field("id").unwrap().id; + let indices = dataset.load_indices().await.unwrap(); + let idx = indices + .iter() + .find(|i| i.name == INDEX_NAME) + .expect("index should exist"); assert_eq!( - post_partitions, - expected_after_join, - "Expected partitions to be at most {} after join, got stats: {}", - expected_after_join, - post_ctx.stats_json() + idx.covering_fields, + vec![id_field_id], + "IndexMetadata.covering_fields should record the covered column's field id" ); - let row_count_after = dataset.count_all_rows().await.unwrap(); - debug_assert!( - row_count_before + ROWS_TO_APPEND_FOR_JOIN >= row_count_after, - "row count should not increase after delete + append" + // The REAL producer must raise the reader+writer fence, not just the + // doctored-metadata commit that upstream's + // `test_covering_commit_fences_the_table_with_a_feature_flag` exercises: a + // pre-covering build that opened this dataset would select the index by + // membership of `fields` and answer a query on the carried column with + // an index keyed on another one. + use lance_table::feature_flags::FLAG_COVERED_INDEX_METADATA; + assert_ne!( + dataset.manifest.reader_feature_flags & FLAG_COVERED_INDEX_METADATA, + 0, + "a covered create must raise the reader fence" + ); + assert_ne!( + dataset.manifest.writer_feature_flags & FLAG_COVERED_INDEX_METADATA, + 0, + "a covered create must raise the writer fence" ); - let deleted_rows = row_count_before + ROWS_TO_APPEND_FOR_JOIN - row_count_after; - - (deleted_rows, ROWS_TO_APPEND_FOR_JOIN, post_partitions) } - async fn append_constant_vector_with_start_id( - dataset: &mut Dataset, - rows: usize, - template: &[f32], - next_id: &mut u64, - ) { - append_constant_vector_batch(dataset, rows, template, *next_id, None).await; - *next_id += rows as u64; - } + /// A multivector column stores one entry per sub-vector (all sharing the source + /// row id). The covering build must preserve every sub-vector -- the stale-row + /// dedup must not collapse legitimate repeated row ids, or recall silently drops. + #[tokio::test] + async fn test_ivf_pq_covered_multivector_preserves_all_subvectors() { + const INDEX_NAME: &str = "vector_idx"; + const SUBVECS_PER_ROW: usize = 3; + let test_dir = TempStrDir::default(); + let test_uri = test_dir.as_str(); + let (mut dataset, _) = + generate_multivec_test_dataset::(test_uri, 0.0..1.0).await; - async fn append_partition_templates( - dataset: &mut Dataset, - rows_per_template: usize, - templates: &[Vec], - ) { - assert!( - !templates.is_empty(), - "at least one template is required for append" + // Multivector requires cosine. Cover the `id` column. + let mut params = VectorIndexParams::ivf_pq(4, 8, 4, DistanceType::Cosine, 2); + params.covering_columns(vec!["id".to_string()]); + dataset + .create_index( + &["vector"], + IndexType::Vector, + Some(INDEX_NAME.to_string()), + ¶ms, + true, + ) + .await + .unwrap(); + + let ctx = load_vector_index_context(&dataset, "vector", INDEX_NAME).await; + let mut total_stored = 0usize; + let mut has_id = false; + for p in 0..ctx.num_partitions() { + let storage = ctx + .ivf() + .load_partition_storage(p, PartitionColumns::All, None) + .await + .unwrap(); + total_stored += storage.batch().num_rows(); + has_id |= storage.batch().column_by_name("id").is_some(); + } + assert!( + has_id, + "multivector covered storage should carry covering column 'id'" ); - for template in templates { + assert_eq!( + total_stored, + NUM_ROWS * SUBVECS_PER_ROW, + "covering multivector build must keep every sub-vector, not collapse per row id" + ); + } + + /// Read-side payoff for multivector: a covered projection is carried through + /// `MultivectorScoringExec` (which re-groups sub-vectors back to rows), so no + /// `TakeExec` against the base table is needed. + #[tokio::test] + async fn test_ivf_pq_covered_multivector_projection_skips_take() { + use arrow_array::types::UInt64Type; + + const INDEX_NAME: &str = "vector_idx"; + let test_dir = TempStrDir::default(); + let (mut dataset, vectors) = + generate_multivec_test_dataset::(test_dir.as_str(), 0.0..1.0).await; + + let mut params = VectorIndexParams::ivf_pq(4, 8, 4, DistanceType::Cosine, 2); + params.covering_columns(vec!["id".to_string()]); + dataset + .create_index( + &["vector"], + IndexType::Vector, + Some(INDEX_NAME.to_string()), + ¶ms, + true, + ) + .await + .unwrap(); + + let query = vectors.value(0); + let mut scan = dataset.scan(); + scan.nearest("vector", &query, 10).unwrap(); + scan.minimum_nprobes(4); + scan.with_row_id(); + scan.project(&["id"]).unwrap(); + + let plan = scan.explain_plan(true).await.unwrap(); + assert!( + !plan.contains("LanceRead"), + "covered multivector projection ['id'] should skip TakeExec; plan was:\n{plan}" + ); + + let batch = scan.try_into_batch().await.unwrap(); + let ids = batch + .column_by_name("id") + .expect("covered 'id' must be emitted through multivector scoring") + .as_primitive::(); + let row_ids = batch + .column_by_name(ROW_ID) + .expect("row id column") + .as_primitive::(); + assert!(batch.num_rows() > 0, "query should return rows"); + // Single-fragment, step-id dataset => id == row offset == _rowid, so a correctly + // carried covered value equals the row id for every returned row (payload + // re-attached to the right row through scoring, not misaligned or stale). + for i in 0..ids.len() { assert_eq!( - template.len(), - DIM, - "Template vector should have {} dimensions", - DIM + ids.value(i), + row_ids.value(i), + "covered 'id' must stay row-aligned through multivector scoring" ); } + } - let start_id = dataset.count_all_rows().await.unwrap() as u64; - let total_rows = rows_per_template * templates.len(); - let ids = Arc::new(UInt64Array::from_iter_values( - start_id..start_id + total_rows as u64, - )); - let mut appended_values = Vec::with_capacity(total_rows * DIM); - for template in templates { - for _ in 0..rows_per_template { - appended_values.extend_from_slice(template); - } + /// A dotted/nested include path cannot be covered (the covering gather projects only + /// top-level names); reject it at build time. + #[tokio::test] + async fn test_ivf_pq_rejects_dotted_include_column() { + let test_dir = TempStrDir::default(); + let (mut dataset, _) = + generate_test_dataset::(test_dir.as_str(), 0.0..1.0).await; + let mut params = VectorIndexParams::ivf_pq(4, 8, 4, DistanceType::L2, 2); + params.covering_columns(vec!["id.sub".to_string()]); + let err = dataset + .create_index(&["vector"], IndexType::Vector, None, ¶ms, true) + .await + .expect_err("a dotted include column must be rejected"); + assert!( + err.to_string().contains("nested/dotted"), + "expected a nested/dotted rejection, got: {err}" + ); + } + + /// Covering names that collide with index storage (the indexed vector column itself, + /// a reserved storage column, or a duplicate) must be rejected at build time. + #[tokio::test] + async fn test_ivf_rejects_reserved_and_duplicate_covering_columns() { + let test_dir = TempStrDir::default(); + let (mut dataset, _) = + generate_test_dataset::(test_dir.as_str(), 0.0..1.0).await; + let build = |cols: Vec| { + let mut params = VectorIndexParams::ivf_flat(2, DistanceType::L2); + params.covering_columns(cols); + params + }; + + // The indexed vector column itself. + let err = dataset + .create_index( + &["vector"], + IndexType::Vector, + None, + &build(vec!["vector".to_string()]), + true, + ) + .await + .expect_err("covering the indexed vector column must be rejected"); + assert!( + err.to_string().contains("indexed vector column itself"), + "got: {err}" + ); + + // Reserved storage/transform column names. The reserved check runs before the + // schema-existence check, so these are rejected even though the dataset lacks them. + // The RaBitQ code/factor columns and the IVF partition transform's `__centroid_dist` + // are internal to the build pipeline: covering one would advertise it in + // `covering_fields` while the storage's `covering_field_indices` drops it, so a + // covered query would declare a column storage never emits. + for internal in [ + "__pq_code", + "__sq_code", + "__ivf_part_id", + "__centroid_dist", + "__ex_codes", + "__blocked_ex_codes", + "__add_factors", + "__scale_factors", + "__error_factors", + "__add_factors_ex", + "__scale_factors_ex", + ] { + let err = dataset + .create_index( + &["vector"], + IndexType::Vector, + None, + &build(vec![internal.to_string()]), + true, + ) + .await + .expect_err("covering a reserved storage/transform name must be rejected"); + assert!( + err.to_string().contains("reserved index storage"), + "internal name '{internal}' must be reserved, got: {err}" + ); } - let vectors = Arc::new( - FixedSizeListArray::try_new_from_values( - Float32Array::from(appended_values), - DIM as i32, + + // Duplicate covering columns. + let err = dataset + .create_index( + &["vector"], + IndexType::Vector, + None, + &build(vec!["id".to_string(), "id".to_string()]), + true, ) - .unwrap(), + .await + .expect_err("duplicate covering columns must be rejected"); + assert!(err.to_string().contains("duplicate"), "got: {err}"); + } + + /// A blob column stores out-of-line descriptors, not inline data, so it cannot be + /// covered by a vector index; reject it at build time. + #[tokio::test] + async fn test_ivf_flat_rejects_blob_include_column() { + use arrow_array::{ + Float32Array, Int32Array, LargeBinaryArray, RecordBatch, RecordBatchIterator, + }; + use lance_arrow::BLOB_META_KEY; + use lance_file::version::LanceFileVersion; + use std::collections::HashMap; + + let dim = 4i32; + let n = 32usize; + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new( + "vector", + DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Float32, true)), dim), + false, + ), + Field::new("blobs", DataType::LargeBinary, true).with_metadata(HashMap::from([( + BLOB_META_KEY.to_string(), + "true".to_string(), + )])), + ])); + let vectors = FixedSizeListArray::try_new_from_values( + Float32Array::from((0..n as i32 * dim).map(|v| v as f32).collect::>()), + dim, + ) + .unwrap(); + let blobs: Vec> = (0..n).map(|_| Some(b"x".as_slice())).collect(); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(Int32Array::from((0..n as i32).collect::>())), + Arc::new(vectors), + Arc::new(LargeBinaryArray::from(blobs)), + ], + ) + .unwrap(); + let mut dataset = Dataset::write( + RecordBatchIterator::new([Ok(batch)], schema), + "memory://c1-blob-include", + Some(WriteParams { + data_storage_version: Some(LanceFileVersion::V2_1), + ..Default::default() + }), + ) + .await + .unwrap(); + let mut params = VectorIndexParams::ivf_flat(2, DistanceType::L2); + params.covering_columns(vec!["blobs".to_string()]); + let err = dataset + .create_index(&["vector"], IndexType::Vector, None, ¶ms, true) + .await + .expect_err("a blob include column must be rejected"); + assert!( + err.to_string().contains("blob column"), + "expected a blob rejection, got: {err}" ); + } + + /// Covering must never change QUERY RESULTS: for any query, a covered index + /// returns exactly the rows and distances a non-covered index returns (covering + /// only changes how projected columns are fetched). Exercises both the late-search + /// shortcut config (min < max nprobes: prefilter-matched rows without index + /// entries come back with INFINITY distance) and the all-partitions-searched + /// corner (min == max: such rows are NOT returned at all). + #[tokio::test] + async fn test_covered_prefilter_results_match_non_covered() { + use arrow_array::types::UInt64Type; + use arrow_array::{Int32Array, RecordBatchIterator, StringArray}; + use arrow_buffer::NullBuffer; + + const INDEX: &str = "vec_idx"; + let dim = 4i32; + let n = 4096usize; + let n_null = 3usize; // ids 0..3 have NULL vectors (no index entry) + let schema = Arc::new(Schema::new(vec![ - Field::new("id", DataType::UInt64, false), - Field::new("vector", vectors.data_type().clone(), false), + Field::new("id", DataType::Int32, false), + Field::new("category", DataType::Utf8, false), + Field::new( + "vector", + DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Float32, true)), dim), + true, + ), ])); - let batch = RecordBatch::try_new(schema.clone(), vec![ids, vectors]).unwrap(); - let batches = RecordBatchIterator::new(vec![Ok(batch)], schema); - dataset.append(batches, None).await.unwrap(); + let ids: Vec = (0..n as i32).collect(); + // The prefilter matches ids 0..5: three null-vector rows plus two real ones. + let cats: Vec<&str> = (0..n) + .map(|i| if i < 5 { "picked" } else { "rest" }) + .collect(); + // Two well-separated clusters (around 0.0 and 100.0). The real picked rows + // (ids 3, 4) live in cluster B; the query below probes cluster A first, so the + // late search / shortcut path is genuinely exercised at min < max nprobes. + let values: Vec = (0..n) + .flat_map(|r| { + let center = if r % 2 == 0 { 0.0f32 } else { 100.0 }; + (0..dim as usize) + .map(move |d| center + ((r * dim as usize + d) % 400) as f32 * 1e-3) + }) + .collect(); + let validity: Vec = (0..n).map(|i| i >= n_null).collect(); + let vector = FixedSizeListArray::new( + Arc::new(Field::new("item", DataType::Float32, true)), + dim, + Arc::new(Float32Array::from(values)), + Some(NullBuffer::from(validity)), + ); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(Int32Array::from(ids)), + Arc::new(StringArray::from(cats)), + Arc::new(vector), + ], + ) + .unwrap(); + + let mut dataset = Dataset::write( + RecordBatchIterator::new([Ok(batch)], schema), + "memory://covered_noncovered_parity", + None, + ) + .await + .unwrap(); + + // (rowid -> distance) results for the prefiltered query under the given probe + // configuration and index variant. + async fn run(dataset: &Dataset, min_nprobes: usize, max_nprobes: usize) -> Vec<(u64, f32)> { + let q = Float32Array::from(vec![0.0f32; 4]); + let mut scan = dataset.scan(); + scan.nearest("vector", &q, 10).unwrap(); + scan.minimum_nprobes(min_nprobes); + scan.maximum_nprobes(max_nprobes); + scan.filter("category = 'picked'").unwrap(); + scan.prefilter(true); + scan.with_row_id(); + scan.project(&["id"]).unwrap(); + let plan = scan.explain_plan(false).await.unwrap(); + assert!( + plan.contains("ANNSubIndex"), + "the parity comparison requires the index path; plan:\n{plan}" + ); + let batch = scan.try_into_batch().await.unwrap(); + let row_ids = batch[ROW_ID].as_primitive::(); + let dists = batch["_distance"].as_primitive::(); + let mut out: Vec<(u64, f32)> = row_ids + .values() + .iter() + .zip(dists.values().iter()) + .map(|(r, d)| (*r, *d)) + .collect(); + out.sort_by_key(|(r, _)| *r); + out + } + + for (min_nprobes, max_nprobes) in [(1usize, 2usize), (2, 2)] { + let centroids = || { + let vals: Vec = [0.0f32, 100.0] + .iter() + .flat_map(|c| std::iter::repeat_n(*c, dim as usize)) + .collect(); + Arc::new( + FixedSizeListArray::try_new_from_values(Float32Array::from(vals), dim).unwrap(), + ) + }; + let plain_params = VectorIndexParams::with_ivf_flat_params( + DistanceType::L2, + IvfBuildParams::try_with_centroids(2, centroids()).unwrap(), + ); + dataset + .create_index( + &["vector"], + IndexType::Vector, + Some(INDEX.to_string()), + &plain_params, + true, + ) + .await + .unwrap(); + let plain = run(&dataset, min_nprobes, max_nprobes).await; + + let mut covered_params = VectorIndexParams::with_ivf_flat_params( + DistanceType::L2, + IvfBuildParams::try_with_centroids(2, centroids()).unwrap(), + ); + covered_params.covering_columns(vec!["id".to_string()]); + dataset + .create_index( + &["vector"], + IndexType::Vector, + Some(INDEX.to_string()), + &covered_params, + true, + ) + .await + .unwrap(); + let covered = run(&dataset, min_nprobes, max_nprobes).await; + + assert_eq!( + covered, plain, + "covered results must be identical to non-covered \ + (min_nprobes={min_nprobes}, max_nprobes={max_nprobes})" + ); + } + } + + /// A row whose vector is null has no index entry, so the covered search never returns + /// it -- but a bounded, selective prefilter that admits it still expects it in the + /// results (parity with a non-covered scan). The covered path must recover it, with + /// its covering column fetched from the base table. + #[tokio::test] + async fn test_ivf_covered_recovers_null_vector_prefilter_rows() { + use arrow_array::types::{Int32Type, UInt64Type}; + use arrow_array::{Int32Array, RecordBatchIterator, StringArray}; + use arrow_buffer::NullBuffer; + + const INDEX: &str = "vec_idx"; + let dim = 4i32; + let n = 40usize; + let n_rare = 3usize; // category "rare" rows, all with NULL vectors (no index entry) + + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new("category", DataType::Utf8, false), + Field::new( + "vector", + DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Float32, true)), dim), + true, + ), + ])); + + let ids: Vec = (0..n as i32).collect(); + let cats: Vec<&str> = (0..n) + .map(|i| if i < n_rare { "rare" } else { "common" }) + .collect(); + // Two well-separated clusters so early pruning keeps minimum_nprobes at 1: + // the recovery path only fires where the non-covered shortcut would (an + // all-partitions-searched query returns only found rows on both). + let values: Vec = clustered_vector_values(0..n, dim, 2); + let validity: Vec = (0..n).map(|i| i >= n_rare).collect(); + let vector = FixedSizeListArray::new( + Arc::new(Field::new("item", DataType::Float32, true)), + dim, + Arc::new(Float32Array::from(values)), + Some(NullBuffer::from(validity)), + ); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(Int32Array::from(ids)), + Arc::new(StringArray::from(cats)), + Arc::new(vector), + ], + ) + .unwrap(); + + let mut dataset = Dataset::write( + RecordBatchIterator::new([Ok(batch)], schema), + "memory://covered_null_vec_recover", + None, + ) + .await + .unwrap(); + + let mut params = VectorIndexParams::ivf_flat(2, DistanceType::L2); + params.covering_columns(vec!["id".to_string()]); + dataset + .create_index( + &["vector"], + IndexType::Vector, + Some(INDEX.to_string()), + ¶ms, + true, + ) + .await + .unwrap(); + let scalar = lance_index::scalar::ScalarIndexParams::for_builtin( + lance_index::scalar::BuiltinIndexType::BTree, + ); + dataset + .create_index(&["category"], IndexType::BTree, None, &scalar, false) + .await + .unwrap(); + + // The prefilter admits only the "rare" rows -- all null-vector, so the search + // returns nothing; every result row must come from the recovery path. + let q = Float32Array::from(vec![0.0f32; dim as usize]); + let mut scan = dataset.scan(); + scan.nearest("vector", &q, 10).unwrap(); + // Recovery matches the non-covered shortcut, which needs unsearched + // partitions to exist (min < max nprobes); all-partitions-searched queries + // return only found rows on both covered and non-covered indexes. + scan.minimum_nprobes(1); + scan.filter("category = 'rare'").unwrap(); + scan.prefilter(true); + scan.with_row_id(); + scan.project(&["id"]).unwrap(); + + let batch = scan.try_into_batch().await.unwrap(); + assert_eq!( + batch.num_rows(), + n_rare, + "all null-vector prefilter rows must be recovered (covered parity with non-covered)" + ); + let ids = batch + .column_by_name("id") + .expect("covered 'id' must be emitted") + .as_primitive::(); + let row_ids = batch + .column_by_name(ROW_ID) + .expect("row id column") + .as_primitive::(); + for i in 0..ids.len() { + // Row-aligned covered value fetched from the base table for the recovered row. + assert_eq!(ids.value(i) as u64, row_ids.value(i)); + assert!( + (ids.value(i) as usize) < n_rare, + "only the rare rows are admitted" + ); + } + } + + /// The covered exec emits from two paths in one stream: the search path (rows with an + /// index entry, covering columns read back from index storage) and the recovery path + /// (null-vector rows, covering columns taken from the base table with the exec's declared + /// dataset-typed schema). A prefilter admitting BOTH kinds forces both paths into a single + /// result, so this guards that their batch schemas stay compatible -- if the index + /// round-trip ever changed a covered field's type or nullability, the concatenation here + /// would fail. Covered column is non-nullable to make a nullability drift observable. + #[tokio::test] + async fn test_ivf_covered_mixed_search_and_recovery_share_schema() { + use arrow_array::types::{Int32Type, UInt64Type}; + use arrow_array::{Int32Array, RecordBatchIterator, StringArray}; + use arrow_buffer::NullBuffer; + + const INDEX: &str = "vec_idx"; + let dim = 4i32; + let n = 40usize; + let n_null = 3usize; // rows 0..3 have NULL vectors (no index entry -> recovery path) + let admit_below = 5i32; // prefilter admits ids 0..5: null rows 0,1,2 + indexed rows 3,4 + + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new("category", DataType::Utf8, false), + Field::new( + "vector", + DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Float32, true)), dim), + true, + ), + ])); + + let ids: Vec = (0..n as i32).collect(); + let cats: Vec<&str> = (0..n).map(|_| "x").collect(); + // Two well-separated clusters so early pruning keeps minimum_nprobes at 1 + // (see test_ivf_covered_recovers_null_vector_prefilter_rows). + let values: Vec = clustered_vector_values(0..n, dim, 2); + let validity: Vec = (0..n).map(|i| i >= n_null).collect(); + let vector = FixedSizeListArray::new( + Arc::new(Field::new("item", DataType::Float32, true)), + dim, + Arc::new(Float32Array::from(values)), + Some(NullBuffer::from(validity)), + ); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(Int32Array::from(ids)), + Arc::new(StringArray::from(cats)), + Arc::new(vector), + ], + ) + .unwrap(); + + let mut dataset = Dataset::write( + RecordBatchIterator::new([Ok(batch)], schema), + "memory://covered_mixed_recover", + None, + ) + .await + .unwrap(); + + let mut params = VectorIndexParams::ivf_flat(2, DistanceType::L2); + params.covering_columns(vec!["id".to_string()]); + dataset + .create_index( + &["vector"], + IndexType::Vector, + Some(INDEX.to_string()), + ¶ms, + true, + ) + .await + .unwrap(); + // A btree on id makes `id < admit_below` a bounded, selective prefilter. + let scalar = lance_index::scalar::ScalarIndexParams::for_builtin( + lance_index::scalar::BuiltinIndexType::BTree, + ); + dataset + .create_index(&["id"], IndexType::BTree, None, &scalar, false) + .await + .unwrap(); + + let q = Float32Array::from(vec![0.0f32; dim as usize]); + let mut scan = dataset.scan(); + scan.nearest("vector", &q, 10).unwrap(); + // Recovery matches the non-covered shortcut, which needs unsearched + // partitions to exist (min < max nprobes); all-partitions-searched queries + // return only found rows on both covered and non-covered indexes. + scan.minimum_nprobes(1); + scan.filter(&format!("id < {admit_below}")).unwrap(); + scan.prefilter(true); + scan.with_row_id(); + scan.project(&["id"]).unwrap(); + + // Concatenating the search-path batches (ids 3,4) with the recovery batch (ids 0,1,2) + // succeeds only if both paths carry a compatible schema for the covered `id` column. + let batch = scan.try_into_batch().await.unwrap(); + assert_eq!( + batch.num_rows(), + admit_below as usize, + "both indexed and null-vector admitted rows must be returned" + ); + let ids = batch + .column_by_name("id") + .expect("covered 'id' must be emitted") + .as_primitive::(); + let row_ids = batch + .column_by_name(ROW_ID) + .expect("row id column") + .as_primitive::(); + let mut got: Vec = Vec::with_capacity(ids.len()); + for i in 0..ids.len() { + // id == row offset == _rowid, so a row-aligned covered value equals its row id. + assert_eq!(ids.value(i) as u64, row_ids.value(i)); + got.push(ids.value(i)); + } + got.sort_unstable(); + assert_eq!( + got, + (0..admit_below).collect::>(), + "search-path and recovery-path rows must together cover every admitted id" + ); + } + + /// On a STABLE-ROW-ID dataset the covered null-vector recovery must resolve stable row + /// ids through the row-id index before taking covering payload -- feeding them to an + /// address-space take (`frag = id >> 32`, `offset = id`) reads the wrong physical row or + /// errors once stable id != physical address. A second fragment makes the two diverge: + /// its rows have addresses `(1 << 32) | offset` but small monotonic stable ids. + #[tokio::test] + async fn test_ivf_covered_recovers_null_vector_stable_row_ids() { + use arrow_array::types::{Int32Type, UInt64Type}; + use arrow_array::{Int32Array, RecordBatchIterator}; + use arrow_buffer::NullBuffer; + + const INDEX: &str = "vec_idx"; + let dim = 4i32; + let frag0 = 10usize; // fragment 0: ids 0..10, all indexed + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new( + "vector", + DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Float32, true)), dim), + true, + ), + ])); + + let make_batch = |ids: Vec, null_rows: &[i32]| { + // Two well-separated clusters. Recovery is only reachable while + // `minimum_nprobes < maximum_nprobes`, and `early_pruning` raises the minimum to + // the number of centroids within `dists[0] * 7.0` (k = 10). Points strung along + // a line leave those centroids only ~7.5x apart -- a margin thin enough that a + // different kmeans outcome silently disables the path this test exists to cover. + let values: Vec = + clustered_vector_values(ids.iter().map(|id| *id as usize), dim, 2); + let validity: Vec = ids.iter().map(|id| !null_rows.contains(id)).collect(); + let vector = FixedSizeListArray::new( + Arc::new(Field::new("item", DataType::Float32, true)), + dim, + Arc::new(Float32Array::from(values)), + Some(NullBuffer::from(validity)), + ); + RecordBatch::try_new( + schema.clone(), + vec![Arc::new(Int32Array::from(ids)), Arc::new(vector)], + ) + .unwrap() + }; + + // Fragment 0: ids 0..10, all with vectors. + let batch0 = make_batch((0..frag0 as i32).collect(), &[]); + let mut dataset = Dataset::write( + RecordBatchIterator::new([Ok(batch0)], schema.clone()), + "memory://covered_stable_null_recover", + Some(WriteParams { + enable_stable_row_ids: true, + max_rows_per_file: frag0, + ..Default::default() + }), + ) + .await + .unwrap(); + + // Fragment 1 (append): ids 10..14; 10 and 11 have NULL vectors (no index entry). + // Their stable ids (10, 11) are far smaller than their physical addresses + // ((1 << 32) | 0/1), so an address-space take of id 10 lands at fragment 0 offset 10, + // which is out of range (fragment 0 has offsets 0..9) -- the bug. + let batch1 = make_batch((frag0 as i32..frag0 as i32 + 4).collect(), &[10, 11]); + dataset + .append(RecordBatchIterator::new([Ok(batch1)], schema.clone()), None) + .await + .unwrap(); + + let mut params = VectorIndexParams::ivf_flat(2, DistanceType::L2); + params.covering_columns(vec!["id".to_string()]); + dataset + .create_index( + &["vector"], + IndexType::Vector, + Some(INDEX.to_string()), + ¶ms, + true, + ) + .await + .unwrap(); + // A btree on id makes `id >= 10` a bounded, selective prefilter. + let scalar = lance_index::scalar::ScalarIndexParams::for_builtin( + lance_index::scalar::BuiltinIndexType::BTree, + ); + dataset + .create_index(&["id"], IndexType::BTree, None, &scalar, false) + .await + .unwrap(); + + // Admit ids 10..14: 12,13 come from the search path, 10,11 from the recovery path. + let q = Float32Array::from(vec![0.0f32; dim as usize]); + let mut scan = dataset.scan(); + scan.nearest("vector", &q, 10).unwrap(); + // Recovery matches the non-covered shortcut, which needs unsearched + // partitions to exist (min < max nprobes); all-partitions-searched queries + // return only found rows on both covered and non-covered indexes. + scan.minimum_nprobes(1); + scan.filter("id >= 10").unwrap(); + scan.prefilter(true); + scan.with_row_id(); + scan.project(&["id"]).unwrap(); + + let batch = scan.try_into_batch().await.unwrap(); + assert_eq!(batch.num_rows(), 4, "all admitted rows must be returned"); + let ids = batch + .column_by_name("id") + .expect("covered 'id' must be emitted") + .as_primitive::(); + let row_ids = batch + .column_by_name(ROW_ID) + .expect("row id column") + .as_primitive::(); + for i in 0..ids.len() { + // id == stable row id == _rowid, so a correctly-resolved covered value matches. + assert_eq!( + ids.value(i) as u64, + row_ids.value(i), + "covered id must be the row's true value, not an address-space misread" + ); + } + let mut got: Vec = ids.values().to_vec(); + got.sort_unstable(); + assert_eq!(got, vec![10, 11, 12, 13]); + } + + /// The covered recovery takes its payload with `MissingRowPolicy::Ignore`, which silently + /// returns fewer rows than requested when an id no longer resolves. A stale prefilter can + /// list such ids: a scalar index built over a fragment that a later delete removed keeps + /// emitting that fragment's rows, and `create_deletion_mask_impl` produces no mask at all + /// when every fragment in the VECTOR index's bitmap is intact. Pairing the requested ids + /// with the returned payload then fails with "all columns in a record batch must have the + /// same length" -- a query the same dataset answers fine without covering. + /// + /// Scope: this covers the *unresolvable-id* shortfall only -- the fragment is gone, so + /// `get_row_addrs` drops the ids before the read. The other way a take can come up short -- + /// an id that resolves to a live address whose row is tombstoned -- is not exercised here, + /// and is not reachable. + /// + /// Note those are two different masks, despite the shared name. The one above is + /// `create_deletion_mask_impl`'s, scoped to the vector index's bitmap, and it can indeed be + /// absent. The one that makes tombstoned rows unreachable is `DatasetPreFilter`'s + /// `deleted_ids`, which is intersected in separately and spans every fragment in the manifest, + /// so + /// it holds no matter which `PreFilterSource` produced the ids. + #[tokio::test] + async fn test_covered_recovery_tolerates_unresolvable_prefilter_ids() { + use arrow_array::{Int32Array, RecordBatchIterator}; + + let dim = 4i32; + let frag0 = 16usize; + let extra = 3usize; + + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new( + "vector", + DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Float32, true)), dim), + true, + ), + ])); + // Two separated clusters so `early_pruning` leaves an unsearched partition and the + // late-search shortcut (which arms the covered recovery) is reachable. + let make_batch = |ids: Vec| { + let values: Vec = + clustered_vector_values(ids.iter().map(|id| *id as usize), dim, 2); + RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(Int32Array::from(ids)), + Arc::new( + FixedSizeListArray::try_new_from_values(Float32Array::from(values), dim) + .unwrap(), + ), + ], + ) + .unwrap() + }; + + let mut dataset = Dataset::write( + RecordBatchIterator::new( + [Ok(make_batch((0..frag0 as i32).collect()))], + schema.clone(), + ), + "memory://covered_recovery_stale_prefilter", + Some(WriteParams { + enable_stable_row_ids: true, + max_rows_per_file: frag0, + ..Default::default() + }), + ) + .await + .unwrap(); + + // Covered index over fragment 0 only. + let mut params = VectorIndexParams::ivf_flat(2, DistanceType::L2); + params.covering_columns(vec!["id".to_string()]); + dataset + .create_index(&["vector"], IndexType::Vector, None, ¶ms, true) + .await + .unwrap(); + + // Fragment 1, then a scalar index spanning BOTH fragments. + let ids: Vec = (frag0 as i32..(frag0 + extra) as i32).collect(); + dataset + .append( + RecordBatchIterator::new([Ok(make_batch(ids))], schema.clone()), + None, + ) + .await + .unwrap(); + let scalar = lance_index::scalar::ScalarIndexParams::for_builtin( + lance_index::scalar::BuiltinIndexType::BTree, + ); + dataset + .create_index(&["id"], IndexType::BTree, None, &scalar, false) + .await + .unwrap(); + + // Drop fragment 1 entirely. The BTree still lists its ids, so the prefilter admits + // rows the row-id index can no longer resolve. + dataset.delete("id >= 16").await.unwrap(); + + let q = Float32Array::from(vec![0.0f32; dim as usize]); + let mut scan = dataset.scan(); + scan.nearest("vector", &q, 10).unwrap(); + scan.minimum_nprobes(1); + scan.filter("id >= 15").unwrap(); + scan.prefilter(true); + scan.project(&["id"]).unwrap(); + + let batch = scan + .try_into_batch() + .await + .expect("covered query must tolerate prefilter ids that no longer resolve"); + let ids = batch["id"].as_primitive::(); + let got: Vec = ids.values().to_vec(); + assert_eq!( + got, + vec![15], + "only the surviving admitted row may come back; ids from the deleted fragment \ + must be dropped, not paired with mismatched payload" + ); + } + + /// The covered end-of-stream recovery re-emits every prefilter row the search did not + /// emit. `DatasetPreFilter` produces a bounded ALLOW LIST from deletions alone on a + /// stable-row-id dataset, so that recovery fires on queries carrying NO filter at all + /// (`PreFilterSource::None`). Tracking the emitted row ids must therefore not be gated + /// on a prefilter source being present -- otherwise the "already emitted" set is empty + /// and every live row is emitted a second time with INFINITY distance. + #[tokio::test] + async fn test_covered_unfiltered_query_does_not_duplicate_rows() { + use arrow_array::types::{Int32Type, UInt64Type}; + use arrow_array::{Int32Array, RecordBatchIterator}; + + const INDEX: &str = "vec_idx"; + let dim = 4i32; + let n = 12usize; + + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new( + "vector", + DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Float32, true)), dim), + true, + ), + ])); + + // Two well-separated clusters with the query sitting on top of the first: the second + // centroid is orders of magnitude farther, so `early_pruning` leaves minimum_nprobes + // at 1. An unsearched partition is what makes the late-search no-rows shortcut (and + // hence the covered recovery) reachable -- with evenly spread vectors early pruning + // raises minimum_nprobes to cover every partition and the shortcut never fires. + let ids: Vec = (0..n as i32).collect(); + let values: Vec = clustered_vector_values(0..n, dim, 2); + let vector = FixedSizeListArray::new( + Arc::new(Field::new("item", DataType::Float32, true)), + dim, + Arc::new(Float32Array::from(values)), + None, + ); + let batch = RecordBatch::try_new( + schema.clone(), + vec![Arc::new(Int32Array::from(ids)), Arc::new(vector)], + ) + .unwrap(); + + let mut dataset = Dataset::write( + RecordBatchIterator::new([Ok(batch)], schema), + "memory://covered_unfiltered_no_dupes", + Some(WriteParams { + enable_stable_row_ids: true, + ..Default::default() + }), + ) + .await + .unwrap(); + + let mut params = VectorIndexParams::ivf_flat(2, DistanceType::L2); + params.covering_columns(vec!["id".to_string()]); + dataset + .create_index( + &["vector"], + IndexType::Vector, + Some(INDEX.to_string()), + ¶ms, + true, + ) + .await + .unwrap(); + + // A deletion makes the stable-row-id deletion mask a bounded ALLOW LIST, so + // `max_len()` is Some(live) even though the query below carries no filter. + dataset.delete("id = 1").await.unwrap(); + let live = n - 1; + + let q = Float32Array::from(vec![0.0f32; dim as usize]); + let mut scan = dataset.scan(); + // k > live so the "fewer than k prefilter matches" shortcut condition holds. + scan.nearest("vector", &q, live + 9).unwrap(); + scan.minimum_nprobes(1); + scan.with_row_id(); + scan.project(&["id"]).unwrap(); + + let batch = scan.try_into_batch().await.unwrap(); + + let row_ids = batch + .column_by_name(ROW_ID) + .expect("row id column") + .as_primitive::(); + let mut seen: Vec = row_ids.values().to_vec(); + seen.sort_unstable(); + let mut distinct = seen.clone(); + distinct.dedup(); + assert_eq!( + seen.len(), + distinct.len(), + "covered unfiltered query returned DUPLICATE row ids: {seen:?}" + ); + assert_eq!( + batch.num_rows(), + live, + "covered unfiltered query must return each live row exactly once" + ); + + let ids = batch + .column_by_name("id") + .expect("covered 'id' must be emitted") + .as_primitive::(); + for i in 0..ids.len() { + assert_eq!( + ids.value(i) as u64, + row_ids.value(i), + "covered id must stay row-aligned" + ); + } + let mut got: Vec = ids.values().to_vec(); + got.sort_unstable(); + assert_eq!( + got, + (0..n as i32).filter(|id| *id != 1).collect::>(), + "every live row must appear exactly once" + ); + } + + /// A covered index whose `covering_fields` a pre-covering writer cleared (prost drops + /// unknown proto field 11 on re-serialization; `FLAG_COVERED_INDEX_METADATA` fences + /// this off only best-effort, since released clients consult writer flags solely on + /// the append/overwrite path) -- the degraded state -- must stay MAINTAINABLE, not just + /// readable. Its auxiliary storage still physically carries the payload while freshly + /// scanned rows do not, so combining the two used to fail `StorageBuilder::build`'s width + /// check ("mismatched columns while merging vector storage batches: expected + /// [_rowid, __pq_code, id], got [_rowid, __pq_code]"), leaving an index that could never + /// be merge-optimized again -- recoverable only by dropping and rebuilding it. + /// `OptimizeOptions::append()` never hit this because it does not load existing storage. + #[tokio::test] + async fn test_degraded_covered_index_can_still_be_optimized() { + use crate::dataset::WriteDestination; + use crate::dataset::transaction::Operation; + use arrow_array::{Int32Array, RecordBatchIterator}; + + const DIMS: usize = 16; + const TOTAL: usize = 256; + const NPART: usize = 4; + + let test_dir = TempStrDir::default(); + let test_uri = test_dir.as_str(); + + let make = |lo: i32, hi: i32| { + let ids: Vec = (lo..hi).collect(); + let values: Vec = + clustered_vector_values(ids.iter().map(|r| *r as usize), DIMS as i32, NPART); + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new( + "vector", + DataType::FixedSizeList( + Arc::new(Field::new("item", DataType::Float32, true)), + DIMS as i32, + ), + true, + ), + ])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(Int32Array::from(ids)), + Arc::new( + FixedSizeListArray::try_new_from_values( + Float32Array::from(values), + DIMS as i32, + ) + .unwrap(), + ), + ], + ) + .unwrap(); + (schema, batch) + }; + + let (schema, batch) = make(0, TOTAL as i32); + let mut dataset = Dataset::write( + RecordBatchIterator::new([Ok(batch)], schema), + test_uri, + None, + ) + .await + .unwrap(); + + let mut params = VectorIndexParams::ivf_pq(NPART, 8, 4, DistanceType::L2, 2); + params.covering_columns(vec!["id".to_string()]); + dataset + .create_index(&["vector"], IndexType::Vector, None, ¶ms, false) + .await + .unwrap(); + + // Degrade: clear the declaration, leave the payload in the auxiliary file. + let mut cleared = dataset.load_indices_by_name("vector_idx").await.unwrap()[0].clone(); + assert!(!cleared.covering_fields.is_empty()); + cleared.covering_fields = Vec::new(); + let read_version = dataset.manifest.version; + let mut dataset = Dataset::commit( + WriteDestination::Dataset(Arc::new(dataset)), + Operation::CreateIndex { + new_indices: vec![cleared], + removed_indices: vec![], + }, + Some(read_version), + None, + None, + Arc::new(Default::default()), + false, + ) + .await + .unwrap(); + + // Fresh rows land in an unindexed fragment; the merge below must combine them with + // the wider existing storage. + let (schema2, batch2) = make(TOTAL as i32, TOTAL as i32 + 128); + dataset + .append(RecordBatchIterator::new([Ok(batch2)], schema2), None) + .await + .unwrap(); + + dataset + .optimize_indices(&OptimizeOptions::merge(1)) + .await + .expect("a metadata-degraded covered index must still merge-optimize"); + + // The rebuilt storage is narrowed to the declaration, so it no longer carries the + // undeclared payload -- and the index still answers queries. + let ctx = load_vector_index_context(&dataset, "vector", "vector_idx").await; + let storage = ctx + .ivf() + .load_partition_storage(0, PartitionColumns::All, None) + .await + .unwrap(); + assert!( + storage.batch().column_by_name("id").is_none(), + "merged storage must drop payload the metadata no longer declares" + ); + + let q = Float32Array::from(vec![0.0f32; DIMS]); + let mut scan = dataset.scan(); + scan.nearest("vector", &q, 5).unwrap(); + scan.project(&["id"]).unwrap(); + let batch = scan.try_into_batch().await.unwrap(); + assert_eq!(batch.num_rows(), 5); + assert!(batch.column_by_name("id").is_some()); + } + + /// `row_id` is an ordinary user column name -- `_rowid` is the reserved one -- so it is + /// coverable. The `row_id` -> `_rowid` rename exists only for precomputed shuffle buffers + /// (which cannot name a column `_rowid`), but it used to run on the ordinary scan path + /// too. Once covering started projecting user columns into that scan, covering a column + /// named `row_id` renamed it on top of the real `_rowid` from `with_row_id()` and index + /// creation died with `Duplicate field name "_rowid" in schema`. + #[tokio::test] + async fn test_covering_column_named_row_id_is_supported() { + use arrow_array::types::Int32Type; + use arrow_array::{Int32Array, RecordBatchIterator}; + + let dim = 4i32; + let n = 128usize; + + let schema = Arc::new(Schema::new(vec![ + Field::new("row_id", DataType::Int32, false), + Field::new( + "vector", + DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Float32, true)), dim), + false, + ), + ])); + let values: Vec = (0..n * dim as usize).map(|i| i as f32).collect(); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(Int32Array::from((0..n as i32).collect::>())), + Arc::new( + FixedSizeListArray::try_new_from_values(Float32Array::from(values), dim) + .unwrap(), + ), + ], + ) + .unwrap(); + + let mut dataset = Dataset::write( + RecordBatchIterator::new([Ok(batch)], schema), + "memory://covering_named_row_id", + None, + ) + .await + .unwrap(); + + let mut params = VectorIndexParams::ivf_flat(2, DistanceType::L2); + params.covering_columns(vec!["row_id".to_string()]); + dataset + .create_index(&["vector"], IndexType::Vector, None, ¶ms, false) + .await + .expect("a covering column named `row_id` must not collide with the virtual _rowid"); + + // And the covered payload is served correctly, not confused with the virtual column. + let q = Float32Array::from(vec![0.0f32; dim as usize]); + let mut scan = dataset.scan(); + scan.nearest("vector", &q, 5).unwrap(); + scan.project(&["row_id"]).unwrap(); + let batch = scan.try_into_batch().await.unwrap(); + assert_eq!(batch.num_rows(), 5); + let ids = batch["row_id"].as_primitive::(); + let mut got: Vec = ids.values().to_vec(); + got.sort_unstable(); + assert_eq!( + got, + vec![0, 1, 2, 3, 4], + "covered `row_id` payload must be the user's values" + ); + } + + /// A covered ("included") struct's payload schema is fixed at index build time. Growing + /// that struct via `add_columns` (which commits as `Operation::Merge`) -- even an + /// AllNulls, metadata-only child that writes no data file -- would leave the index + /// emitting the old struct type while covered queries declare the new one, an Arrow type + /// mismatch. The commit boundary must reject it (drop the index first), mirroring the + /// `Project` drop/alter guard. + #[tokio::test] + async fn test_add_columns_child_to_covered_struct_is_rejected() { + use arrow_array::{Int32Array, RecordBatchIterator, StructArray}; + use arrow_schema::Fields; + use lance_file::version::LanceFileVersion; + + use crate::dataset::NewColumnTransform; + + let n = 64usize; + let dim = 4i32; + let meta_fields = Fields::from(vec![Field::new("a", DataType::Int32, false)]); + let schema = Arc::new(Schema::new(vec![ + Field::new("meta", DataType::Struct(meta_fields.clone()), true), + Field::new( + "vector", + DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Float32, true)), dim), + true, + ), + ])); + + let a = Arc::new(Int32Array::from((0..n as i32).collect::>())); + let meta = Arc::new(StructArray::new(meta_fields, vec![a], None)); + let values: Vec = (0..n * dim as usize).map(|i| i as f32 + 1.0).collect(); + let vector = Arc::new( + FixedSizeListArray::try_new_from_values(Float32Array::from(values), dim).unwrap(), + ); + let batch = RecordBatch::try_new(schema.clone(), vec![meta, vector]).unwrap(); + + let mut dataset = Dataset::write( + RecordBatchIterator::new([Ok(batch)], schema.clone()), + "memory://covered_struct_add_child", + Some(WriteParams { + data_storage_version: Some(LanceFileVersion::V2_2), + ..Default::default() + }), + ) + .await + .unwrap(); + + let mut params = VectorIndexParams::ivf_flat(2, DistanceType::L2); + params.covering_columns(vec!["meta".to_string()]); + dataset + .create_index( + &["vector"], + IndexType::Vector, + Some("vec_idx".to_string()), + ¶ms, + true, + ) + .await + .unwrap(); + + // Add child `meta.b` -- merges into the covered struct, changing its type. + let add = + NewColumnTransform::AllNulls(Arc::new(arrow_schema::Schema::new(vec![Field::new( + "meta", + DataType::Struct(Fields::from(vec![Field::new("b", DataType::Int32, true)])), + true, + )]))); + let err = dataset + .add_columns(add, None, None) + .await + .expect_err("growing a covered struct's subtree must be rejected"); + let msg = err.to_string(); + assert!( + msg.contains("included") && msg.contains("vec_idx"), + "expected a covered-field rejection naming the index, got: {err}" + ); + } + + /// Read-side payoff for every vector index type: a query projecting only a covered + /// column is satisfied from the index -- no `TakeExec` against the base table -- + /// with row-aligned values and sane recall. + /// + /// The fixture offsets `id` away from `_rowid` on purpose. `generate_test_dataset` + /// starts ids at 0, so `id == _rowid` for every row -- and under that fixture the + /// row-alignment assertion below is satisfied by any defect that serves the row id + /// where the covered value belongs, or the reverse. That is exactly the confusion a + /// storage batch ordered `[covering..., _rowid, code]` invites, and it is silent: a + /// `UInt64` covering column substituted for `_rowid` downcasts cleanly. The offset + /// separates the two columns for all seven quantizer families. + #[rstest] + #[case::pq(VectorIndexParams::ivf_pq(4, 8, 4, DistanceType::L2, 2))] + #[case::sq(VectorIndexParams::with_ivf_sq_params( + DistanceType::L2, + IvfBuildParams::new(4), + SQBuildParams::default() + ))] + // 16 sub-vectors, not the 4 the plain `pq` case uses: at DIM=32 that is 2 dimensions + // per codebook instead of 8. HNSW *builds* its graph from PQ distances, so a coarse + // quantizer degrades construction and search together, and on this uniform-random + // fixture the true top-10 is full of near-ties -- enough that rounding differences + // between NEON and AVX reorder it, so the recall gate read 0.4 on aarch64 while + // passing on x86. Plain `pq` scans partitions exhaustively and is unaffected, which + // is why only this case needs the extra fidelity. Same reasoning as the `rq` case. + #[case::hnsw_pq(VectorIndexParams::with_ivf_hnsw_pq_params( + DistanceType::L2, + IvfBuildParams::new(4), + HnswBuildParams::default(), + PQBuildParams::new(16, 8) + ))] + #[case::hnsw_sq(VectorIndexParams::with_ivf_hnsw_sq_params( + DistanceType::L2, + IvfBuildParams::new(4), + HnswBuildParams::default(), + SQBuildParams::default() + ))] + // RQ uses 5 bits: 1-bit RaBitQ quantization is too coarse to clear the recall + // gate on this random data without a refine (which would re-add the take). + #[case::rq(VectorIndexParams::with_ivf_rq_params( + DistanceType::L2, + IvfBuildParams::new(4), + RQBuildParams::with_rotation_type(5, RQRotationType::Fast) + ))] + #[case::flat(VectorIndexParams::ivf_flat(4, DistanceType::L2))] + #[case::hnsw_flat(VectorIndexParams::ivf_hnsw( + DistanceType::L2, + IvfBuildParams::new(4), + HnswBuildParams::default() + ))] + #[tokio::test] + async fn test_covered_projection_skips_take(#[case] mut params: VectorIndexParams) { + const INDEX_NAME: &str = "vector_idx"; + // Disjoint from the row id range (0..NUM_ROWS), so no covered `id` can coincide + // with any row id. + const ID_OFFSET: u64 = 1_000_000; + let test_dir = TempStrDir::default(); + let test_uri = test_dir.as_str(); + let (batch, schema) = + generate_batch::(NUM_ROWS, Some(ID_OFFSET), 0.0..1.0, false); + let vectors = Arc::new( + batch + .column_by_name("vector") + .unwrap() + .as_fixed_size_list() + .clone(), + ); + let batches = RecordBatchIterator::new(vec![batch].into_iter().map(Ok), schema); + let mut dataset = Dataset::write( + batches, + test_uri, + Some(WriteParams { + mode: WriteMode::Overwrite, + ..Default::default() + }), + ) + .await + .unwrap(); + + params.covering_columns(vec!["id".to_string()]); + dataset + .create_index( + &["vector"], + IndexType::Vector, + Some(INDEX_NAME.to_string()), + ¶ms, + true, + ) + .await + .unwrap(); + + let q = vectors.value(0); + let q = q.as_primitive::(); + let mut scan = dataset.scan(); + scan.nearest("vector", q, 10).unwrap(); + scan.nprobes(4); + // The HNSW search beam defaults to `k + k / 2` = 15, which on this uniform-random + // fixture leaves the recall gate below a coin flip for the coarsest quantizer: + // `hnsw_pq` returned 0.4 on ~2.5% of runs. The gate is a sanity check on take + // elision, not a measurement of beam width, so give the graph a beam wide enough + // that the assertion tests what it names. Ignored by the non-HNSW cases. + scan.ef(100); + scan.with_row_id(); + scan.project(&["id"]).unwrap(); + + let plan = scan.explain_plan(true).await.unwrap(); + assert!( + !plan.contains("LanceRead"), + "covered projection ['id'] should not require a TakeExec; plan was:\n{plan}" + ); + + let batch = scan.try_into_batch().await.unwrap(); + let ids = batch + .column_by_name("id") + .expect("covered 'id' column must be emitted") + .as_primitive::(); + let row_ids = batch + .column_by_name(ROW_ID) + .expect("row id column") + .as_primitive::(); + assert_eq!(ids.len(), 10, "should return k=10 rows"); + // Single-fragment, step-id dataset => id == ID_OFFSET + row offset == ID_OFFSET + + // _rowid, so a correctly covered id is the row id plus the offset for every + // returned row (row-aligned, not stale). The offset also means a row id sourced + // from the covering column instead of `_rowid` cannot satisfy this. + for i in 0..ids.len() { + assert_eq!( + ids.value(i), + row_ids.value(i) + ID_OFFSET, + "covered id must match the row's true id (row {i})" + ); + } + + // Take elision is worthless if the covered search returns the wrong neighbors, + // so also gate on recall against brute-force ground truth. All 4 partitions are + // probed, so 0.5 sits far below any quantizer's real recall on this data. + let returned: HashSet = row_ids.values().iter().copied().collect(); + let truth = ground_truth(&dataset, "vector", q, 10, DistanceType::L2).await; + let recall = truth.intersection(&returned).count() as f32 / truth.len() as f32; + assert!( + recall >= 0.5, + "covered recall {recall} < 0.5 (returned {returned:?}, truth {truth:?})" + ); + } + + /// The most likely real query shape: a projection mixing a covered column + /// with an uncovered one. `filtered_read.rs` re-subtracts the projection + /// against the covered stream's schema, so `id` must not be re-fetched + /// while `extra` -- a column the index does not carry -- still goes + /// through a `TakeExec`. Every other covered test projects only-covered + /// or only-uncovered columns and would not catch a regression here. + #[tokio::test] + async fn test_ivf_pq_covered_partial_projection() { + const INDEX_NAME: &str = "vector_idx"; + let test_dir = TempStrDir::default(); + let test_uri = test_dir.as_str(); + let (mut dataset, vectors) = generate_test_dataset::(test_uri, 0.0..1.0).await; + + // 16 sub-vectors, not 4: at DIM=32 that is 2 dimensions per codebook instead of + // 8. This fixture is uniform-random, so the true top-10 is full of near-ties, and + // at 4 sub-vectors the PQ error is large enough that NEON-vs-AVX rounding reorders + // them -- the recall gate below read 0.4 on aarch64 while passing on x86. The gate + // is a sanity check on take elision, not a measurement of quantizer accuracy. + let mut params = VectorIndexParams::ivf_pq(4, 8, 16, DistanceType::L2, 2); + params.covering_columns(vec!["id".to_string()]); + dataset + .create_index( + &["vector"], + IndexType::Vector, + Some(INDEX_NAME.to_string()), + ¶ms, + true, + ) + .await + .unwrap(); + + // A column the index does not cover; the take path must still fetch it. + dataset + .add_columns( + NewColumnTransform::SqlExpressions(vec![( + "extra".to_string(), + "CAST(id AS BIGINT) + 1000".to_string(), + )]), + None, + None, + ) + .await + .unwrap(); + + let q = vectors.value(0); + let q = q.as_primitive::(); + let mut scan = dataset.scan(); + scan.nearest("vector", q, 10).unwrap(); + scan.nprobes(4); + scan.with_row_id(); + scan.project(&["id", "extra"]).unwrap(); + + let plan = scan.explain_plan(true).await.unwrap(); + assert!( + plan.contains("LanceRead"), + "uncovered column 'extra' must still require a TakeExec; plan was:\n{plan}" + ); + assert!( + plan.contains("projection=[extra]"), + "the take must fetch only the uncovered column, not re-fetch the \ + already-covered 'id'; plan was:\n{plan}" + ); + + let batch = scan.try_into_batch().await.unwrap(); + assert_eq!(batch.num_rows(), 10, "k=10 should return 10 rows"); + let ids = batch["id"].as_primitive::(); + let extras = batch["extra"].as_primitive::(); + for (id, extra) in ids.values().iter().zip(extras.values().iter()) { + assert_eq!( + *extra, + *id as i64 + 1000, + "uncovered 'extra' must match the base table, not a stale/misaligned value" + ); + } + + // Take elision only matters if the covered result is correct. + let returned: HashSet = batch[ROW_ID] + .as_primitive::() + .values() + .iter() + .copied() + .collect(); + let truth = ground_truth(&dataset, "vector", q, 10, DistanceType::L2).await; + let recall = truth.intersection(&returned).count() as f32 / truth.len() as f32; + assert!( + recall >= 0.5, + "covered IVF_PQ recall {recall} < 0.5 (returned {returned:?}, truth {truth:?})" + ); + } + + /// A covering declaration is not proof that every segment already stores the payload. + /// This models a transitional distributed build: both segments carry the same logical + /// declaration, but the new segment was produced by a writer that emitted only the + /// vector-index storage. Commit must preserve that declaration, and planning must see + /// the missing physical capability and fetch `id` from the base table for all results. + #[tokio::test] + async fn test_declared_but_physically_absent_covering_field_falls_back_to_take() { + const INDEX_NAME: &str = "vector_idx"; + let test_dir = TempStrDir::default(); + let test_uri = test_dir.as_str(); + let (mut dataset, vectors) = generate_test_dataset::(test_uri, 0.0..1.0).await; + + let mut covered_params = VectorIndexParams::ivf_pq(4, 8, 4, DistanceType::L2, 2); + covered_params.covering_columns(vec!["id".to_string()]); + dataset + .create_index( + &["vector"], + IndexType::Vector, + Some(INDEX_NAME.to_string()), + &covered_params, + true, + ) + .await + .unwrap(); + + let original_indices = dataset.load_indices_by_name(INDEX_NAME).await.unwrap(); + assert_eq!(original_indices.len(), 1); + let declared_fields = original_indices[0].fields.clone(); + let declared_covering_fields = original_indices[0].covering_fields.clone(); + + let original_fragment_ids: HashSet = dataset + .get_fragments() + .iter() + .map(|f| f.id() as u32) + .collect(); + // PQ training on the new segment needs >= 256 rows (2^8 codes). + append_dataset::(&mut dataset, 300, 0.0..1.0).await; + let new_fragment_ids: Vec = dataset + .get_fragments() + .into_iter() + .map(|f| f.id() as u32) + .filter(|id| !original_fragment_ids.contains(id)) + .collect(); + assert!(!new_fragment_ids.is_empty(), "append should add fragments"); + + // Build a segment for just the new fragments, without covering. + let uncovered_params = VectorIndexParams::ivf_pq(4, 8, 4, DistanceType::L2, 2); + let mut transitional_segment = dataset + .create_index_builder(&["vector"], IndexType::Vector, &uncovered_params) + .name(INDEX_NAME.to_string()) + .fragments(new_fragment_ids) + .replace(true) + .execute_uncommitted() + .await + .unwrap(); + assert!( + transitional_segment.covering_fields.is_empty(), + "the new segment must genuinely be uncovered for this repro" + ); + + // The logical declaration is intentionally independent of the payload written by + // this segment. Do not rewrite its physical auxiliary storage. + transitional_segment.fields = declared_fields; + transitional_segment.covering_fields = declared_covering_fields; + dataset + .commit_existing_index_segments(INDEX_NAME, "vector", vec![transitional_segment]) + .await + .unwrap(); + + let q = vectors.value(0); + let q = q.as_primitive::(); + let mut scan = dataset.scan(); + scan.nearest("vector", q, 10).unwrap(); + scan.nprobes(4); + scan.with_row_id(); + scan.project(&["id"]).unwrap(); + + let plan = scan.explain_plan(true).await.unwrap(); + assert!( + plan.contains("LanceRead"), + "one segment cannot physically serve the declared 'id', so the query must use a \ + base-table take; plan was:\n{plan}" + ); + + let batch = scan.try_into_batch().await.unwrap(); + let row_ids = batch[ROW_ID].as_primitive::(); + let ids = batch["id"].as_primitive::(); + let projection = crate::dataset::ProjectionRequest::from_columns(["id"], dataset.schema()); + let truth = dataset + .take_rows(row_ids.values(), projection) + .await + .unwrap(); + let truth_ids = truth["id"].as_primitive::(); + assert_eq!( + ids, truth_ids, + "fallback values must come from the base table" + ); + } + + /// A declaration may be wider than the physical payload without disabling the part + /// storage can prove. The index carries `id`; `tag` is added to the declaration only. + /// Planning must keep `id` covered and take exactly `tag` from the base table. + #[tokio::test] + async fn test_physical_covering_subset_serves_only_proven_columns() { + const INDEX_NAME: &str = "vector_idx"; + let test_dir = TempStrDir::default(); + let test_uri = test_dir.as_str(); + let (mut dataset, vectors) = generate_test_dataset::(test_uri, 0.0..1.0).await; + dataset + .add_columns( + NewColumnTransform::SqlExpressions(vec![( + "tag".to_string(), + "CAST(id AS BIGINT) + 1000".to_string(), + )]), + None, + None, + ) + .await + .unwrap(); + + let id_field_id = dataset.schema().field("id").unwrap().id; + let tag_field_id = dataset.schema().field("tag").unwrap().id; + let mut params = VectorIndexParams::ivf_flat(4, DistanceType::L2); + params.covering_columns(vec!["id".to_string()]); + let mut segment = dataset + .create_index_builder(&["vector"], IndexType::Vector, ¶ms) + .name(INDEX_NAME.to_string()) + .execute_uncommitted() + .await + .unwrap(); + assert_eq!(segment.covering_fields, vec![id_field_id]); + + // Widen only the logical dependency. The auxiliary storage remains physically + // capable of serving `id` and has no `tag` column. + segment.fields.push(tag_field_id); + segment.covering_fields.push(tag_field_id); + dataset + .commit_existing_index_segments(INDEX_NAME, "vector", vec![segment]) + .await + .unwrap(); + + let q = vectors.value(0); + let q = q.as_primitive::(); + let mut scan = dataset.scan(); + scan.nearest("vector", q, 10).unwrap(); + scan.nprobes(4); + scan.project(&["id", "tag"]).unwrap(); + + let plan = scan.explain_plan(true).await.unwrap(); + assert!( + plan.contains("LanceRead") && plan.contains("projection=[tag]"), + "only the physically absent 'tag' should use a base-table take; plan was:\n{plan}" + ); + + let batch = scan.try_into_batch().await.unwrap(); + let ids = batch["id"].as_primitive::(); + let tags = batch["tag"].as_primitive::(); + for (id, tag) in ids.values().iter().zip(tags.values()) { + assert_eq!(*tag, *id as i64 + 1000); + } + } + + /// Same payoff but forcing the BATCH/late-search path: `k` larger than one + /// partition makes the initial `minimum_nprobes` sweep under-fill, so + /// `late_search` (the run_prepared batch path) fires. Covering must hold there too. + /// + /// Four well-separated clusters, not uniform-random data: on uniform-random + /// data at this `k` (500, i.e. `k > 10`), `early_pruning`'s generous 81x + /// factor (`knn.rs`) makes the pruning heuristic select every partition + /// regardless of the caller's `minimum_nprobes(1)`, so `adjust_probes` bumps + /// `minimum_nprobes` up to `maximum_nprobes` and `late_search` returns empty + /// immediately (`max_nprobes <= min_nprobes`) -- confirmed by instrumenting + /// `late_search` on the prior uniform-random version of this test, which + /// logged `min_nprobes=4 max_nprobes=4` despite the explicit `min=1,max=4` + /// request. With separated clusters the nearest centroid is far closer than + /// the rest, so the heuristic keeps `minimum_nprobes` at 1 and `late_search` + /// genuinely runs. + #[tokio::test] + async fn test_ivf_pq_covered_projection_batch_path() { + const INDEX_NAME: &str = "vector_idx"; + const NUM_CLUSTERS: usize = 4; + const ROWS_PER_CLUSTER: usize = NUM_ROWS / NUM_CLUSTERS; + let offsets = [0.0f32, 1000.0, 2000.0, 3000.0]; + + let mut rng = StdRng::seed_from_u64(7); + let mut ids = Vec::with_capacity(NUM_ROWS); + let mut values = Vec::with_capacity(NUM_ROWS * DIM); + for (cluster_idx, offset) in offsets.iter().enumerate() { + for row in 0..ROWS_PER_CLUSTER { + ids.push((cluster_idx * ROWS_PER_CLUSTER + row) as u64); + for dim in 0..DIM { + let base = if dim == 0 { *offset } else { 0.0 }; + let noise = (rng.random::() - 0.5) * 0.02; + values.push(base + noise); + } + } + } + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::UInt64, false), + Field::new( + "vector", + DataType::FixedSizeList( + Arc::new(Field::new("item", DataType::Float32, true)), + DIM as i32, + ), + false, + ), + ])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(UInt64Array::from(ids)), + Arc::new( + FixedSizeListArray::try_new_from_values(Float32Array::from(values), DIM as i32) + .unwrap(), + ), + ], + ) + .unwrap(); + + let test_dir = TempStrDir::default(); + let test_uri = test_dir.as_str(); + let batches = RecordBatchIterator::new(vec![Ok(batch)], schema); + let mut dataset = Dataset::write(batches, test_uri, None).await.unwrap(); + + let centroids = build_centroids_for_offsets(&offsets); + let ivf_params = IvfBuildParams::try_with_centroids(NUM_CLUSTERS, centroids).unwrap(); + let mut params = VectorIndexParams::with_ivf_pq_params( + DistanceType::L2, + ivf_params, + PQBuildParams { + num_bits: 8, + num_sub_vectors: 4, + max_iters: 2, + ..Default::default() + }, + ); + params.covering_columns(vec!["id".to_string()]); + dataset + .create_index( + &["vector"], + IndexType::Vector, + Some(INDEX_NAME.to_string()), + ¶ms, + true, + ) + .await + .unwrap(); + + // A query at cluster 0's exact center: its nearest centroid is a near-zero + // distance away, the rest are ~1000+ away, so `early_pruning` keeps + // `minimum_nprobes` at 1 rather than escalating it to `maximum_nprobes`. + let q = Float32Array::from(vec![0.0f32; DIM]); + let mut scan = dataset.scan(); + // k >> one partition's rows (NUM_ROWS=512 / 4 parts) forces late expansion. + scan.nearest("vector", &q, 500).unwrap(); + scan.minimum_nprobes(1); + scan.maximum_nprobes(NUM_CLUSTERS); + scan.with_row_id(); + scan.project(&["id"]).unwrap(); + + let plan = scan.explain_plan(true).await.unwrap(); + assert!( + !plan.contains("LanceRead"), + "covered projection ['id'] should not require a TakeExec (batch path); plan:\n{plan}" + ); + let batch = scan.try_into_batch().await.unwrap(); + assert!( + batch.num_rows() > ROWS_PER_CLUSTER, + "late search should have expanded beyond the nearest partition's {ROWS_PER_CLUSTER} rows, got {}", + batch.num_rows() + ); + // Row alignment, not just presence: ids are assigned in write order over a single + // fragment, so a correctly covered `id` equals its row id on every returned row. + let ids = batch["id"].as_primitive::(); + let row_ids = batch[ROW_ID].as_primitive::(); + for i in 0..ids.len() { + assert_eq!( + ids.value(i), + row_ids.value(i), + "covered payload must stay row-aligned through the batch path" + ); + } + } + + /// A covered query combined with a selective prefilter (scalar index + + /// `prefilter`) must still skip the TakeExec and return correct, filtered + /// rows. Exercises the covering read path under a scalar-index prefilter. + #[tokio::test] + async fn test_ivf_pq_covered_with_scalar_prefilter() { + const INDEX_NAME: &str = "vector_idx"; + let test_dir = TempStrDir::default(); + let test_uri = test_dir.as_str(); + let (mut dataset, vectors) = generate_test_dataset::(test_uri, 0.0..1.0).await; + + let mut params = VectorIndexParams::ivf_pq(4, 8, 4, DistanceType::L2, 2); + params.covering_columns(vec!["id".to_string()]); + dataset + .create_index( + &["vector"], + IndexType::Vector, + Some(INDEX_NAME.to_string()), + ¶ms, + true, + ) + .await + .unwrap(); + + // A scalar index on `id` turns the prefilter into a bounded AllowList. + let scalar_params = lance_index::scalar::ScalarIndexParams::for_builtin( + lance_index::scalar::BuiltinIndexType::BTree, + ); + dataset + .create_index( + &["id"], + IndexType::BTree, + Some("id_btree".to_string()), + &scalar_params, + true, + ) + .await + .unwrap(); + + let q = vectors.value(0); + let q = q.as_primitive::(); + let mut scan = dataset.scan(); + scan.nearest("vector", q, 10).unwrap(); + scan.filter("id < 5").unwrap(); + scan.prefilter(true); + scan.project(&["id"]).unwrap(); + + let plan = scan.explain_plan(true).await.unwrap(); + assert!( + !plan.contains("LanceRead"), + "covered projection ['id'] should not require a TakeExec even with a prefilter; plan:\n{plan}" + ); + let batch = scan.try_into_batch().await.unwrap(); + assert!(batch.column_by_name("id").is_some()); + // Every id in [0, 5) is present in the dataset, so the prefilter must + // return exactly 5 rows -- not zero, which would make the loop below + // vacuously true regardless of whether filtering actually worked. + assert_eq!(batch.num_rows(), 5); + let ids = batch + .column_by_name("id") + .unwrap() + .as_primitive::(); + for v in ids.values() { + assert!( + *v < 5, + "all returned rows must satisfy the prefilter id < 5" + ); + } + } + + /// A covered index must keep its covering columns through optimize/merge. + /// After appending rows and merging them into the index, the merged + /// auxiliary storage must still carry the included column, and a covered + /// projection must still skip the take. The covering columns are threaded + /// through the incremental optimize pipeline -- the unindexed-fragment + /// shuffle, the partition-split reshuffle, and the partition-join + /// reassignment -- as passenger columns (re-gathered by row id). + #[rstest] + #[case::pq(VectorIndexParams::ivf_pq(4, 8, 4, DistanceType::L2, 2))] + #[case::sq(VectorIndexParams::with_ivf_sq_params( + DistanceType::L2, + IvfBuildParams::new(4), + SQBuildParams::default() + ))] + #[case::hnsw_pq(VectorIndexParams::with_ivf_hnsw_pq_params( + DistanceType::L2, + IvfBuildParams::new(4), + HnswBuildParams::default(), + PQBuildParams::new(4, 8) + ))] + #[case::hnsw_sq(VectorIndexParams::with_ivf_hnsw_sq_params( + DistanceType::L2, + IvfBuildParams::new(4), + HnswBuildParams::default(), + SQBuildParams::default() + ))] + #[case::rq(VectorIndexParams::with_ivf_rq_params( + DistanceType::L2, + IvfBuildParams::new(4), + RQBuildParams::with_rotation_type(1, RQRotationType::Fast) + ))] + #[case::flat(VectorIndexParams::ivf_flat(4, DistanceType::L2))] + #[case::hnsw_flat(VectorIndexParams::ivf_hnsw( + DistanceType::L2, + IvfBuildParams::new(4), + HnswBuildParams::default() + ))] + #[tokio::test] + async fn test_covered_survives_optimize(#[case] mut params: VectorIndexParams) { + const INDEX_NAME: &str = "vector_idx"; + let test_dir = TempStrDir::default(); + let test_uri = test_dir.as_str(); + let (mut dataset, vectors) = generate_test_dataset::(test_uri, 0.0..1.0).await; + + params.covering_columns(vec!["id".to_string()]); + dataset + .create_index( + &["vector"], + IndexType::Vector, + Some(INDEX_NAME.to_string()), + ¶ms, + true, + ) + .await + .unwrap(); + + // Append rows (unindexed fragments) and merge them into the index. + append_dataset::(&mut dataset, NUM_ROWS, 0.0..1.0).await; + dataset + .optimize_indices(&OptimizeOptions::new()) + .await + .unwrap(); + + // A covered projection must still be answered from the index with no take after + // the merge. This is the real proof that the appended rows kept their covering + // column: if the merge had dropped it, `try_into_batch` would error on the + // schema mismatch (the exec declares `id` from `covering_fields`, but the + // storage could not emit it). + let q = vectors.value(0); + let q = q.as_primitive::(); + let mut scan = dataset.scan(); + scan.nearest("vector", q, 10).unwrap(); + scan.with_row_id(); + scan.project(&["id"]).unwrap(); + let plan = scan.explain_plan(true).await.unwrap(); + assert!( + !plan.contains("LanceRead"), + "covered projection ['id'] should skip the take after optimize; plan:\n{plan}" + ); + let batch = scan.try_into_batch().await.unwrap(); + assert_eq!(batch.num_rows(), 10); + let ids = batch + .column_by_name("id") + .expect("covered 'id' column must be emitted") + .as_primitive::(); + let row_ids = batch + .column_by_name(ROW_ID) + .expect("row id column") + .as_primitive::(); + + // Ground truth via an independent (non-index, base-table take) path. + // The dataset now spans two fragments (the original write plus the + // appended delta), so id no longer equals _rowid by construction the + // way it does for the single-fragment `test_covered_projection_skips_take` + // -- but the covered value must still equal the row's *true* id. This + // catches a merge that carries the column but scrambles which row + // its values ended up attached to (e.g. an off-by-one row-id gather). + let row_id_vec: Vec = row_ids.values().to_vec(); + let projection = crate::dataset::ProjectionRequest::from_columns(["id"], dataset.schema()); + let truth = dataset.take_rows(&row_id_vec, projection).await.unwrap(); + let truth_ids = truth + .column_by_name("id") + .unwrap() + .as_primitive::(); + for i in 0..ids.len() { + assert_eq!( + ids.value(i), + truth_ids.value(i), + "row {i} (row_id {}): covered id {} != true id {} -- merge scrambled the \ + covering column", + row_ids.value(i), + ids.value(i), + truth_ids.value(i) + ); + } + } + + /// Covering must survive a *retrain* optimize. Retrain rebuilds the storage from + /// scratch, so the covering column must be re-materialized -- otherwise the fresh + /// files omit it while the committed metadata still advertises `covering_fields`, + /// and a covered projection expects a column the storage cannot emit. + #[tokio::test] + async fn test_ivf_pq_covered_survives_retrain() { + const INDEX_NAME: &str = "vector_idx"; + let test_dir = TempStrDir::default(); + let test_uri = test_dir.as_str(); + let (mut dataset, vectors) = generate_test_dataset::(test_uri, 0.0..1.0).await; + + let mut params = VectorIndexParams::ivf_pq(4, 8, 4, DistanceType::L2, 2); + params.covering_columns(vec!["id".to_string()]); + dataset + .create_index( + &["vector"], + IndexType::Vector, + Some(INDEX_NAME.to_string()), + ¶ms, + true, + ) + .await + .unwrap(); + + append_dataset::(&mut dataset, NUM_ROWS, 0.0..1.0).await; + dataset + .optimize_indices(&OptimizeOptions::retrain()) + .await + .unwrap(); + + // The retrained storage must re-materialize the covering column. + let ctx = load_vector_index_context(&dataset, "vector", INDEX_NAME).await; + let storage = ctx + .ivf() + .load_partition_storage(0, PartitionColumns::All, None) + .await + .unwrap(); + assert!( + storage.batch().column_by_name("id").is_some(), + "retrained storage should re-materialize covering column 'id'" + ); + + // And a covered projection is still answered from the index (no take). + let q = vectors.value(0); + let q = q.as_primitive::(); + let mut scan = dataset.scan(); + scan.nearest("vector", q, 10).unwrap(); + scan.project(&["id"]).unwrap(); + let plan = scan.explain_plan(true).await.unwrap(); + assert!( + !plan.contains("LanceRead"), + "covered projection ['id'] should skip the take after retrain; plan:\n{plan}" + ); + let batch = scan.try_into_batch().await.unwrap(); + assert!(batch.column_by_name("id").is_some()); + } + + /// The streaming partition-search branch (used by HNSW sub-indexes and controlled + /// late searches) must declare the covered schema on the stream it returns, matching + /// the widened batches it emits -- the global-heap branch already does. A stream + /// whose declared schema disagrees with its batches is a latent hazard for any + /// consumer that trusts the declaration. + #[tokio::test] + async fn test_covered_search_partitions_stream_declares_covering_schema() { + use arrow_array::UInt32Array; + use futures::TryStreamExt; + use lance_index::prefilter::NoFilter; + use lance_index::vector::{DEFAULT_QUERY_PARALLELISM, Query}; + + const INDEX_NAME: &str = "vector_idx"; + let test_dir = TempStrDir::default(); + let test_uri = test_dir.as_str(); + let (mut dataset, vectors) = generate_test_dataset::(test_uri, 0.0..1.0).await; + + let mut params = VectorIndexParams::ivf_hnsw( + DistanceType::L2, + IvfBuildParams::new(4), + HnswBuildParams::default(), + ); + params.covering_columns(vec!["id".to_string()]); + dataset + .create_index( + &["vector"], + IndexType::Vector, + Some(INDEX_NAME.to_string()), + ¶ms, + true, + ) + .await + .unwrap(); + + let ctx = load_vector_index_context(&dataset, "vector", INDEX_NAME).await; + let query = Query { + column: "vector".to_string(), + key: vectors.value(0), + k: 5, + lower_bound: None, + upper_bound: None, + minimum_nprobes: 4, + maximum_nprobes: None, + ef: None, + refine_factor: None, + metric_type: Some(DistanceType::L2), + use_index: true, + query_parallelism: DEFAULT_QUERY_PARALLELISM, + dist_q_c: 0.0, + approx_mode: Default::default(), + covering_projection: None, + }; + let partitions = Arc::new(UInt32Array::from(vec![0u32, 1, 2, 3])); + let dists = Arc::new(Float32Array::from(vec![0.0f32; 4])); + let stream = ctx + .index + .clone() + .search_partitions( + query, + partitions, + dists, + 0, + 4, + Arc::new(NoFilter), + None, + Arc::new(NoOpMetricsCollector), + ) + .await + .unwrap(); + + let declared = stream.schema(); + assert!( + declared.column_with_name("id").is_some(), + "the covered stream must declare the covering column; declared: {declared:?}" + ); + let batches: Vec = stream.try_collect().await.unwrap(); + assert!(!batches.is_empty()); + for batch in &batches { + assert_eq!( + batch.schema(), + declared, + "emitted batches must match the declared stream schema" + ); + } + } + + /// Index storage that physically carries covering columns while the manifest + /// declares none (`covering_fields` empty) must still be queryable: the extra + /// columns are dropped from search batches (with a warning) instead of emitting + /// batches wider than the plan's declared `[_distance, _rowid]` schema. This is + /// the legacy-tolerance contract: a stable-format index file with an unexpected + /// extra column used to be projected down with a warning, never a query failure. + #[tokio::test] + async fn test_undeclared_storage_covering_is_dropped_not_fatal() { + use crate::dataset::WriteDestination; + use crate::dataset::transaction::Operation; + + const DIMS: usize = 16; + const NUM_CLUSTERS: usize = 4; + const ROWS_PER_CLUSTER: usize = 64; + const TOTAL: usize = NUM_CLUSTERS * ROWS_PER_CLUSTER; + + let test_dir = TempStrDir::default(); + let test_uri = test_dir.as_str(); + + // Well-separated clusters, not uniform-random vectors: a query near one + // cluster must land far outside the early-pruning heuristic's threshold + // (`early_pruning`/`adjust_probes` in `io/exec/knn.rs`) for every other + // cluster's centroid, so `late_search` gets a real range of partitions + // left to search instead of the heuristic folding all of them into the + // early, synchronous phase. Uniform-random vectors near the dataset's + // overall centroid are roughly equidistant from every partition centroid, + // so every partition is searched up front and `late_search` never runs -- + // silently defeating the "hard case" below (a defect recorded once + // already on this project: a covered-ANN test with too few effectively- + // reachable partitions passes without ever exercising the path it claims + // to cover). + let mut rng = StdRng::seed_from_u64(7); + let mut ids = Vec::with_capacity(TOTAL); + let mut values = Vec::with_capacity(TOTAL * DIMS); + for cluster in 0..NUM_CLUSTERS { + let center = (cluster * 1000) as f32; + for row in 0..ROWS_PER_CLUSTER { + ids.push((cluster * ROWS_PER_CLUSTER + row) as i32); + for dim in 0..DIMS { + let base = if dim == 0 { center } else { 0.0 }; + values.push(base + (rng.random::() - 0.5) * 0.02); + } + } + } + let ids_arr: ArrayRef = Arc::new(arrow_array::Int32Array::from(ids)); + let vectors: ArrayRef = Arc::new( + FixedSizeListArray::try_new_from_values(Float32Array::from(values), DIMS as i32) + .unwrap(), + ); + let schema: SchemaRef = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new("vector", vectors.data_type().clone(), false), + ])); + let batch = RecordBatch::try_new(schema.clone(), vec![ids_arr, vectors]).unwrap(); + let batches = RecordBatchIterator::new(vec![batch].into_iter().map(Ok), schema); + let mut dataset = Dataset::write(batches, test_uri, None).await.unwrap(); + + let mut params = VectorIndexParams::ivf_pq(NUM_CLUSTERS, 8, 4, DistanceType::L2, 2); + params.covering_columns(vec!["id".to_string()]); + dataset + .create_index(&["vector"], IndexType::Vector, None, ¶ms, false) + .await + .unwrap(); + + // Re-commit the index metadata with the covering declaration cleared while the + // storage keeps the payload column -- the state a pre-covering writer (which + // drops the unknown `covering_fields` proto field) or a legacy index file with + // a stray extra column would produce. + let mut cleared = dataset.load_indices_by_name("vector_idx").await.unwrap()[0].clone(); + assert!(!cleared.covering_fields.is_empty()); + cleared.covering_fields = Vec::new(); + let read_version = dataset.manifest.version; + let dataset = Dataset::commit( + WriteDestination::Dataset(Arc::new(dataset)), + Operation::CreateIndex { + new_indices: vec![cleared], + removed_indices: vec![], + }, + Some(read_version), + None, + None, + Arc::new(Default::default()), + false, + ) + .await + .unwrap(); + + // Query the LAST cluster's center: its rows' ids (`3*ROWS_PER_CLUSTER..`) + // never satisfy the `id < 3` filter used below, so the "hard case" prefilter + // match (ids 0/1/2, in the first cluster) is guaranteed to be entirely + // unfound by the early search of this, the nearest, cluster/partition -- + // exercising the not-found shortcut instead of incidentally short-circuiting + // because the early search already found every matching row. + let mut q_values = vec![0.0f32; DIMS]; + q_values[0] = ((NUM_CLUSTERS - 1) * 1000) as f32; + let q = Float32Array::from(q_values); + + // The index is no longer covering, so the query takes `id` from the base + // table; the storage's undeclared payload column must be silently dropped + // from the search batches, not fail the query. + let mut scan = dataset.scan(); + scan.nearest("vector", &q, 5).unwrap(); + scan.project(&["id"]).unwrap(); + let batch = scan.try_into_batch().await.unwrap(); + assert_eq!(batch.num_rows(), 5); + assert!(batch.column_by_name("id").is_some()); + + // The hard case: a bounded selective prefilter over unsearched partitions makes + // `late_search` emit its 2-column not-found shortcut batch into the same stream as + // the search batches. Unnarrowed, those search batches still carry the undeclared + // payload column, so the widths disagree where DataFusion's TopK concatenates them + // -- `interleave_record_batch` reads every batch through the FIRST batch's schema + // and panics with an out-of-bounds column index. + // Repeated because the unnarrowed failure is order-dependent: it only panics when + // TopK happens to store a wide batch first, which is roughly 4 runs in 5. One + // iteration would make this test pass intermittently against a real regression. + for attempt in 0..5 { + let mut scan = dataset.scan(); + scan.nearest("vector", &q, 5).unwrap(); + scan.minimum_nprobes(1); + scan.maximum_nprobes(NUM_CLUSTERS); + scan.filter("id < 3").unwrap(); + scan.prefilter(true); + scan.project(&["id"]).unwrap(); + let batch = scan.try_into_batch().await.unwrap(); + assert_eq!( + batch.num_rows(), + 3, + "all prefilter-matched rows must come back (attempt {attempt})" + ); + let ids = batch["id"].as_primitive::(); + let mut got: Vec = ids.values().to_vec(); + got.sort_unstable(); + assert_eq!( + got, + vec![0, 1, 2], + "covered payload must be taken from the base table once the declaration is \ + gone (attempt {attempt})" + ); + } + } + + /// A covered index that searches zero partitions (empty heap) + /// must still emit the covered schema `[_distance, _rowid, ]`, not + /// the bare `[_distance, _rowid]` -- otherwise the produced batch mismatches + /// the wider schema the exec declares from `covering_fields`. The decision is + /// driven by the stable per-index covering schema, not the gathered `covering` batch. + #[test] + fn test_global_heap_to_batch_covered_empty_emits_covered_schema() { + use arrow_schema::{DataType, Field, Schema}; + let covering = Schema::new(vec![ + Field::new(ROW_ID, DataType::UInt64, false), + Field::new("id", DataType::UInt64, true), + ]); + // Empty heap and nothing gathered, but the query DOES want covering columns. + let batch = + IvfPq::global_heap_to_batch(std::collections::BinaryHeap::new(), None, Some(&covering)) + .unwrap(); + assert_eq!(batch.num_rows(), 0); + assert_eq!( + batch.num_columns(), + 3, + "covered empty result must keep the covering column" + ); + assert!(batch.column_by_name("id").is_some()); + + // Ordinary (non-covered) index: bare `[_distance, _rowid]`. + let plain = + IvfPq::global_heap_to_batch(std::collections::BinaryHeap::new(), None, None).unwrap(); + assert_eq!(plain.num_columns(), 2); + } + + /// A result row id that is absent from the covering buffer is an invariant + /// break (every heap winner comes from a searched partition whose covering + /// batch was captured). It must surface as an error -- never as row 0's + /// covering values silently attached to an unrelated row. + #[test] + fn test_gather_covering_by_rowid_errors_on_missing_rowid() { + use arrow_schema::{DataType, Field, Schema}; + let schema = Arc::new(Schema::new(vec![ + Field::new(ROW_ID, DataType::UInt64, false), + Field::new("id", DataType::UInt64, true), + ])); + let source = RecordBatch::try_new( + schema, + vec![ + Arc::new(UInt64Array::from(vec![1, 2, 3])), + Arc::new(UInt64Array::from(vec![10, 20, 30])), + ], + ) + .unwrap(); + + // Present row ids gather fine. + let ok = + gather_covering_columns_by_row_id(&source, &UInt64Array::from(vec![3, 1])).unwrap(); + assert_eq!(ok.len(), 1); + let (_, values) = &ok[0]; + let values = values.as_primitive::(); + assert_eq!(values.values(), &[30, 10]); + + // Row id 99 is absent from the covering source: must be an error. + let err = gather_covering_columns_by_row_id(&source, &UInt64Array::from(vec![2, 99])) + .expect_err("missing row id must error, not return unrelated values"); + assert!(err.to_string().contains("99"), "got: {err}"); + } + + /// A covered index must keep its covering columns through compaction, which + /// rewrites fragments and REMAPS the index's row ids. Delete rows to force a + /// compaction+remap; the remapped storage must still carry the covering + /// column and covered projections must still skip the take. + #[tokio::test] + async fn test_ivf_pq_covered_survives_compaction() { + use crate::dataset::optimize::{CompactionOptions, compact_files}; + const INDEX_NAME: &str = "vector_idx"; + let test_dir = TempStrDir::default(); + let test_uri = test_dir.as_str(); + let (mut dataset, vectors) = generate_test_dataset::(test_uri, 0.0..1.0).await; + + let mut params = VectorIndexParams::ivf_pq(4, 8, 4, DistanceType::L2, 2); + params.covering_columns(vec!["id".to_string()]); + dataset + .create_index( + &["vector"], + IndexType::Vector, + Some(INDEX_NAME.to_string()), + ¶ms, + true, + ) + .await + .unwrap(); + + // Delete rows so compaction rewrites fragments and remaps the index. + dataset.delete("id < 100").await.unwrap(); + compact_files(&mut dataset, CompactionOptions::default(), None) + .await + .unwrap(); + + // The remapped partition storage must still carry the covering column. + let ctx = load_vector_index_context(&dataset, "vector", INDEX_NAME).await; + let storage = ctx + .ivf() + .load_partition_storage(0, PartitionColumns::All, None) + .await + .unwrap(); + assert!( + storage.batch().column_by_name("id").is_some(), + "remapped storage should keep covering column 'id'" + ); + + // Covered projection still skips the take, and deleted rows are gone. + let q = vectors.value(0); + let q = q.as_primitive::(); + let mut scan = dataset.scan(); + scan.nearest("vector", q, 10).unwrap(); + scan.project(&["id"]).unwrap(); + let plan = scan.explain_plan(true).await.unwrap(); + assert!( + !plan.contains("LanceRead"), + "covered projection ['id'] should skip the take after compaction; plan:\n{plan}" + ); + let batch = scan.try_into_batch().await.unwrap(); + // Plenty of non-deleted rows remain for k=10 to fill; guard against an + // empty result, which would make the loop below vacuously true. + assert_eq!(batch.num_rows(), 10); + let ids = batch + .column_by_name("id") + .expect("covered projection should return id") + .as_primitive::(); + for v in ids.values() { + assert!(*v >= 100, "deleted rows (id < 100) must not be returned"); + } + } + + /// A covered index with TWO covering columns of different arrow types must + /// survive compaction (which remaps the index's row ids). Every other + /// covering test in this suite uses a single `UInt64` covering column + /// (`id`), which hid two bugs in the default `QuantizerStorage::remap` + /// (used by flat, HNSW_FLAT and SQ, unlike PQ/RQ which override it): + /// it read the row id column by *position* (`column(0)`), which is only + /// ever correct by coincidence for a single leading `UInt64` covering + /// column. With a non-`UInt64` covering column first, the downcast + /// panics; with a `UInt64` one first (as `id` always was), it silently + /// overwrites the storage's real row ids with the covering column's + /// values. Covering column order here is deliberately `[tag, id]` so + /// `tag` (`Int32`) lands first in the storage batch, reproducing the + /// panicking variant, while `id` (`UInt64`) still covers the silent one. + #[rstest] + #[case::flat(VectorIndexParams::ivf_flat(4, DistanceType::L2))] + #[case::hnsw_flat(VectorIndexParams::ivf_hnsw( + DistanceType::L2, + IvfBuildParams::new(4), + HnswBuildParams::default() + ))] + #[case::sq(VectorIndexParams::with_ivf_sq_params( + DistanceType::L2, + IvfBuildParams::new(4), + SQBuildParams::default() + ))] + #[tokio::test] + async fn test_covered_survives_compaction_multiple_covering_columns( + #[case] mut params: VectorIndexParams, + ) { + use arrow_array::Int32Array; + use arrow_array::types::Int32Type; + + const INDEX_NAME: &str = "vector_idx"; + let test_dir = TempStrDir::default(); + let test_uri = test_dir.as_str(); + + let ids = UInt64Array::from_iter_values(0..NUM_ROWS as u64); + // Distinct from `id` (negative, offset) so a bug that swaps or + // misaligns the two covering columns cannot pass by coincidence. + let tags = Int32Array::from_iter_values((0..NUM_ROWS as i32).map(|i| -i - 1)); + let vectors = generate_random_array_with_range::(NUM_ROWS * DIM, 0.0..1.0); + let fsl = + normalize_fsl(&FixedSizeListArray::try_new_from_values(vectors, DIM as i32).unwrap()) + .unwrap(); + let query_vector = fsl.value(0); + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::UInt64, false), + Field::new("tag", DataType::Int32, false), + Field::new( + "vector", + DataType::FixedSizeList( + Arc::new(Field::new("item", DataType::Float32, true)), + DIM as i32, + ), + true, + ), + ])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![Arc::new(ids), Arc::new(tags), Arc::new(fsl)], + ) + .unwrap(); + let reader = RecordBatchIterator::new(vec![Ok(batch)], schema); + let mut dataset = Dataset::write(reader, test_uri, None).await.unwrap(); + + params.covering_columns(vec!["tag".to_string(), "id".to_string()]); + dataset + .create_index( + &["vector"], + IndexType::Vector, + Some(INDEX_NAME.to_string()), + ¶ms, + true, + ) + .await + .unwrap(); + + // Delete rows so compaction rewrites fragments and remaps the index. + dataset.delete("id < 100").await.unwrap(); + compact_after_deletions(&mut dataset).await; + + let q = query_vector.as_primitive::(); + let mut scan = dataset.scan(); + scan.nearest("vector", q, 10).unwrap(); + scan.nprobes(4); + scan.with_row_id(); + scan.project(&["tag", "id"]).unwrap(); + let plan = scan.explain_plan(true).await.unwrap(); + assert!( + !plan.contains("LanceRead"), + "covered projection ['tag', 'id'] should skip the take after compaction; plan:\n{plan}" + ); + let batch = scan.try_into_batch().await.unwrap(); + // Plenty of non-deleted rows remain for k=10 to fill; guard against an + // empty result, which would make the checks below vacuously true. + assert_eq!(batch.num_rows(), 10); + + let returned_ids = batch + .column_by_name("id") + .expect("covered projection should return id") + .as_primitive::(); + let returned_tags = batch + .column_by_name("tag") + .expect("covered projection should return tag") + .as_primitive::(); + let row_ids = batch + .column_by_name(ROW_ID) + .expect("row id column") + .as_primitive::(); + for v in returned_ids.values() { + assert!(*v >= 100, "deleted rows (id < 100) must not be returned"); + } + + // Ground truth via an independent (non-index, base-table take) path, + // for BOTH covering columns: a corrupted or misaligned remap of + // either column must be caught, not just detected by its presence. + let row_id_vec: Vec = row_ids.values().to_vec(); + let projection = + crate::dataset::ProjectionRequest::from_columns(["id", "tag"], dataset.schema()); + let truth = dataset.take_rows(&row_id_vec, projection).await.unwrap(); + let truth_ids = truth + .column_by_name("id") + .unwrap() + .as_primitive::(); + let truth_tags = truth + .column_by_name("tag") + .unwrap() + .as_primitive::(); + for i in 0..returned_ids.len() { + assert_eq!( + returned_ids.value(i), + truth_ids.value(i), + "row {i} (row_id {}): covered id != true id after compaction \ + (multi-covering-column case)", + row_ids.value(i) + ); + assert_eq!( + returned_tags.value(i), + truth_tags.value(i), + "row {i} (row_id {}): covered tag != true tag after compaction \ + (multi-covering-column case)", + row_ids.value(i) + ); + } + } + + /// A current writer stores its physical covering columns in declaration order. The + /// reader now verifies this independently and falls back if a transitional segment + /// differs, but a newly built index should retain the optimization for every declared + /// column rather than relying on that safety net. + /// + /// Declaration order here is `[payload, price]`: the reverse of both the dataset's + /// field-id order (`id`, `price`, `payload`, `vector`) and the projection order the + /// query asks for. An implementation that sorted by field id, walked the schema, or + /// echoed the request order would produce `[price, payload]` and fail. + /// + /// The two covering columns carry different Arrow types and values that are + /// deliberately disjoint from the row ids and from each other (`price` is negative, + /// `payload` is a string): a fixture where a covering value equals its row id makes a + /// positional-for-by-name substitution invisible. + #[tokio::test] + async fn test_covering_declaration_and_storage_agree_on_order() { + use crate::index::covering::effective_covering; + use arrow_array::{Int32Array, StringArray, types::Int32Type}; + use lance_index::vector::storage::VectorStore; + + const INDEX_NAME: &str = "vector_idx"; + const DIMS: usize = 16; + const NUM_CLUSTERS: usize = 4; + const ROWS_PER_CLUSTER: usize = 64; + const TOTAL: usize = NUM_CLUSTERS * ROWS_PER_CLUSTER; + + let test_dir = TempStrDir::default(); + let test_uri = test_dir.as_str(); + + // Well-separated clusters, one per partition, so the ANN path is genuinely + // exercised rather than degenerating into a single-partition scan. + let mut flat = Vec::with_capacity(TOTAL * DIMS); + for row in 0..TOTAL { + let center = (row / ROWS_PER_CLUSTER) as f32 * 50.0; + for d in 0..DIMS { + flat.push(center + (row % ROWS_PER_CLUSTER) as f32 * 0.001 + d as f32 * 0.0001); + } + } + let vectors = Arc::new( + FixedSizeListArray::try_new_from_values(Float32Array::from(flat), DIMS as i32).unwrap(), + ); + let ids = Arc::new(UInt64Array::from_iter_values(0..TOTAL as u64)); + let prices = Arc::new(Int32Array::from_iter_values( + (0..TOTAL as i32).map(|i| -i - 7), + )); + let payloads = Arc::new(StringArray::from_iter_values( + (0..TOTAL).map(|i| format!("p{}", i * 3 + 11)), + )); + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::UInt64, false), + Field::new("price", DataType::Int32, false), + Field::new("payload", DataType::Utf8, false), + Field::new("vector", vectors.data_type().clone(), false), + ])); + let batch = + RecordBatch::try_new(schema.clone(), vec![ids, prices, payloads, vectors.clone()]) + .unwrap(); + let reader = RecordBatchIterator::new(vec![Ok(batch)], schema); + let mut dataset = Dataset::write(reader, test_uri, None).await.unwrap(); + + let mut params = VectorIndexParams::ivf_flat(NUM_CLUSTERS, DistanceType::L2); + params.covering_columns(vec!["payload".to_string(), "price".to_string()]); + dataset + .create_index( + &["vector"], + IndexType::Vector, + Some(INDEX_NAME.to_string()), + ¶ms, + true, + ) + .await + .unwrap(); + + // 1. The manifest records the declaration in the order it was requested. + let price_id = dataset.schema().field("price").unwrap().id; + let payload_id = dataset.schema().field("payload").unwrap().id; + assert!( + price_id < payload_id, + "the fixture needs declaration order to differ from field-id order" + ); + let index_meta = dataset + .load_indices_by_name(INDEX_NAME) + .await + .unwrap() + .into_iter() + .next() + .expect("index metadata"); + assert_eq!(index_meta.covering_fields, vec![payload_id, price_id]); + + // 2. The read path's manifest-side resolution follows that order, whichever order + // the query requests its columns in. + let requested = ["price".to_string(), "payload".to_string()]; + let declared = effective_covering( + &index_meta.covering_fields, + Some(&requested), + dataset.schema(), + ) + .unwrap(); + let declared_names: Vec<&str> = declared.iter().map(|f| f.name().as_str()).collect(); + assert_eq!(declared_names, vec!["payload", "price"]); + + // 3. The storage-side resolution -- `covering_field_indices`, derived from each + // storage's `INTERNAL_COLUMNS` -- lists the same columns in the same order. + let ctx = load_vector_index_context(&dataset, "vector", INDEX_NAME).await; + let storage = ctx + .index + .as_any() + .downcast_ref::() + .expect("expected IvfFlatIndex") + .load_partition_storage(0, PartitionColumns::All, None) + .await + .unwrap(); + let storage_schema = storage.schema().clone(); + let storage_names: Vec<&str> = storage + .covering_field_indices() + .into_iter() + .map(|i| storage_schema.field(i).name().as_str()) + .collect(); + assert_eq!( + storage_names, declared_names, + "storage order must equal declaration order; if these diverge the covered \ + values are emitted under the wrong names" + ); + + // 4. End to end: the covered values must match the base table, matched BY NAME. + let q = vectors.value(0); + let q = q.as_primitive::(); + let mut scan = dataset.scan(); + scan.nearest("vector", q, 10).unwrap(); + scan.nprobes(NUM_CLUSTERS); + scan.with_row_id(); + scan.project(&["price", "payload"]).unwrap(); + let plan = scan.explain_plan(true).await.unwrap(); + assert!( + plan.contains("ANNSubIndex"), + "the covered query must go through the index; plan was:\n{plan}" + ); + let batch = scan.try_into_batch().await.unwrap(); + assert_eq!(batch.num_rows(), 10); + + let row_ids: Vec = batch + .column_by_name(ROW_ID) + .expect("row id column") + .as_primitive::() + .values() + .to_vec(); + let truth = dataset + .take_rows( + &row_ids, + crate::dataset::ProjectionRequest::from_columns( + ["price", "payload"], + dataset.schema(), + ), + ) + .await + .unwrap(); + let got_prices = batch + .column_by_name("price") + .unwrap() + .as_primitive::(); + let truth_prices = truth + .column_by_name("price") + .unwrap() + .as_primitive::(); + let got_payloads = batch.column_by_name("payload").unwrap().as_string::(); + let truth_payloads = truth.column_by_name("payload").unwrap().as_string::(); + for (i, row_id) in row_ids.iter().enumerate() { + assert_eq!( + got_prices.value(i), + truth_prices.value(i), + "row {i} (row_id {row_id}): covered price landed under the wrong name" + ); + assert_eq!( + got_payloads.value(i), + truth_payloads.value(i), + "row {i} (row_id {row_id}): covered payload landed under the wrong name" + ); + } + } + + /// Counts the cache hit/miss signal `IVFIndex::load_partition` reports, which is how + /// the tests below observe whether two searches shared one partition entry. + #[derive(Default)] + struct CacheCountingMetrics { + hits: AtomicUsize, + misses: AtomicUsize, + } + + impl lance_index::metrics::MetricsCollector for CacheCountingMetrics { + fn record_parts_loaded(&self, _num_parts: usize) {} + fn record_index_loads(&self, _num_indexes: usize) {} + fn record_comparisons(&self, _num_comparisons: usize) {} + fn record_index_cache_hits(&self, num_hits: usize) { + self.hits.fetch_add(num_hits, Ordering::Relaxed); + } + fn record_index_cache_misses(&self, num_misses: usize) { + self.misses.fetch_add(num_misses, Ordering::Relaxed); + } + } + + /// A covered IVF_FLAT fixture with one well-separated cluster per partition and two + /// covering columns of different Arrow types. The covering values are deliberately + /// disjoint from the row ids (`price` is negative, `payload` is a string) -- a fixture + /// where a covering value equals its row id hides a positional-for-by-name + /// substitution. Row id equals row offset, so the caller can compute ground truth + /// against `vectors` directly. + async fn covered_flat_fixture( + uri: &str, + index_name: &str, + num_clusters: usize, + rows_per_cluster: usize, + ) -> (Dataset, Arc) { + use arrow_array::{Int32Array, StringArray}; + + const DIMS: usize = 16; + let total = num_clusters * rows_per_cluster; + + let mut flat = Vec::with_capacity(total * DIMS); + for row in 0..total { + let center = (row / rows_per_cluster) as f32 * 50.0; + for d in 0..DIMS { + flat.push(center + (row % rows_per_cluster) as f32 * 0.001 + d as f32 * 0.0001); + } + } + let vectors = Arc::new( + FixedSizeListArray::try_new_from_values(Float32Array::from(flat), DIMS as i32).unwrap(), + ); + let prices = Arc::new(Int32Array::from_iter_values( + (0..total as i32).map(|i| -i - 7), + )); + let payloads = Arc::new(StringArray::from_iter_values( + (0..total).map(|i| format!("p{}", i * 3 + 11)), + )); + let schema = Arc::new(Schema::new(vec![ + Field::new("price", DataType::Int32, false), + Field::new("payload", DataType::Utf8, false), + Field::new("vector", vectors.data_type().clone(), false), + ])); + let batch = + RecordBatch::try_new(schema.clone(), vec![prices, payloads, vectors.clone()]).unwrap(); + let reader = RecordBatchIterator::new(vec![Ok(batch)], schema); + let mut dataset = Dataset::write(reader, uri, None).await.unwrap(); + + let mut params = VectorIndexParams::ivf_flat(num_clusters, DistanceType::L2); + params.covering_columns(vec!["payload".to_string(), "price".to_string()]); + dataset + .create_index( + &["vector"], + IndexType::Vector, + Some(index_name.to_string()), + ¶ms, + true, + ) + .await + .unwrap(); + (dataset, vectors) + } + + /// A partition loaded for search carries only the storage's own columns, so the entry + /// it is cached under -- `IVFPartitionKey { partition_id }`, which has no covering + /// component -- is the same entry for every query. That is what makes the key correct + /// rather than accidentally correct: the naive narrowing (read whichever covering + /// columns *this* query asked for) would let the first query to touch a partition + /// decide what every later query finds there. + /// + /// Two searches asking for disjoint covering subsets run against the same index, in + /// order. The second must reuse the entry the first populated, that shared entry must + /// carry no covering column at all -- an entry holding `payload` but not `price` is + /// exactly the unsound state ruled out here -- and both must return the partition's + /// true nearest neighbours. + /// + /// The line under test is `PartitionColumns::Internal` in `load_partition_entry`. + /// Switching it to `PartitionColumns::All` makes the covering assertions below fail. + #[tokio::test] + async fn test_partition_cache_entry_is_independent_of_the_query_covering_subset() { + use lance_index::prefilter::NoFilter; + use lance_index::vector::{DEFAULT_QUERY_PARALLELISM, Query}; + + const INDEX_NAME: &str = "vector_idx"; + const NUM_CLUSTERS: usize = 4; + const ROWS_PER_CLUSTER: usize = 64; + const K: usize = 10; + + let test_dir = TempStrDir::default(); + let (dataset, vectors) = covered_flat_fixture( + test_dir.as_str(), + INDEX_NAME, + NUM_CLUSTERS, + ROWS_PER_CLUSTER, + ) + .await; + let ctx = load_vector_index_context(&dataset, "vector", INDEX_NAME).await; + let index = ctx + .index + .as_any() + .downcast_ref::() + .expect("expected IvfFlatIndex"); + + let query_key = vectors.value(0); + let query_for = |covering: &[&str]| Query { + column: "vector".to_string(), + key: query_key.clone(), + k: K, + lower_bound: None, + upper_bound: None, + minimum_nprobes: NUM_CLUSTERS, + maximum_nprobes: None, + ef: None, + refine_factor: None, + metric_type: Some(DistanceType::L2), + use_index: true, + query_parallelism: DEFAULT_QUERY_PARALLELISM, + dist_q_c: 0.0, + approx_mode: Default::default(), + covering_projection: Some(covering.iter().map(|c| c.to_string()).collect()), + }; + + let payload_metrics = CacheCountingMetrics::default(); + let payload_result = index + .search_in_partition( + 0, + &query_for(&["payload"]), + Arc::new(NoFilter), + &payload_metrics, + ) + .await + .unwrap(); + assert_eq!( + ( + payload_metrics.misses.load(Ordering::Relaxed), + payload_metrics.hits.load(Ordering::Relaxed) + ), + (1, 0), + "the first search must populate the entry rather than find one already there" + ); + + let price_metrics = CacheCountingMetrics::default(); + let price_result = index + .search_in_partition( + 0, + &query_for(&["price"]), + Arc::new(NoFilter), + &price_metrics, + ) + .await + .unwrap(); + assert_eq!( + ( + price_metrics.hits.load(Ordering::Relaxed), + price_metrics.misses.load(Ordering::Relaxed) + ), + (1, 0), + "the second search must reuse the first one's entry -- if it loaded its own, \ + the entries are query-specific and nothing below proves the key is sound" + ); + + // The one entry the two searches shared holds no covering column at all, so there + // is no subset for one query to have fixed on another's behalf. + let entry = index + .load_partition(0, true, &NoOpMetricsCollector) + .await + .unwrap(); + let part = entry.as_ref(); + assert!( + part.storage.covering_field_indices().is_empty(), + "a cached partition must carry no covering column; it carries {:?}", + part.storage.schema() + ); + assert!( + part.storage.covering_batch().unwrap().is_none(), + "covering_batch() on a codes-only partition must report none" + ); + + // Both searches answer correctly over that shared entry. Ground truth is the + // partition's own rows ranked by distance to the query, computed from the base + // table rather than from the index. + let partition_row_ids: Vec = index + .load_partition_storage(0, PartitionColumns::All, None) + .await + .unwrap() + .row_ids() + .copied() + .collect(); + assert!( + partition_row_ids.len() > K, + "partition 0 must hold more rows than k, or the ranking below is vacuous" + ); + let query_values = query_key.as_primitive::().values().to_vec(); + let mut truth: Vec<(u64, f32)> = partition_row_ids + .iter() + .map(|row_id| { + let vector = vectors.value(*row_id as usize); + let distance = vector + .as_primitive::() + .values() + .iter() + .zip(query_values.iter()) + .map(|(a, b)| (a - b) * (a - b)) + .sum::(); + (*row_id, distance) + }) + .collect(); + truth.sort_by(|a, b| a.1.total_cmp(&b.1)); + let expected: Vec = truth.iter().take(K).map(|(row_id, _)| *row_id).collect(); + + for (label, result) in [("payload", &payload_result), ("price", &price_result)] { + let mut got: Vec = result + .column_by_name(ROW_ID) + .expect("search result carries row ids") + .as_primitive::() + .values() + .to_vec(); + got.sort_unstable(); + let mut want = expected.clone(); + want.sort_unstable(); + assert_eq!( + got, want, + "the {label} query's neighbours must be the partition's true nearest k" + ); + assert!(result.column_by_name(DIST_COL).is_some()); + } + + // The covering half of "both get right answers": each query receives exactly the + // column it projected, gathered for its survivors out of a shared entry that holds + // neither. `price` is `-row_id - 7` and `payload` is `p{row_id * 3 + 11}`, so a + // value taken from the wrong row -- or from the wrong column -- cannot coincide. + for (label, result, other) in [ + ("payload", &payload_result, "price"), + ("price", &price_result, "payload"), + ] { + assert!( + result.column_by_name(other).is_none(), + "the {label} query must not materialize {other}: a covering column the \ + query never reads is the cost this narrowing exists to avoid, and it is \ + invisible in results" + ); + let row_ids = result + .column_by_name(ROW_ID) + .expect("search result carries row ids") + .as_primitive::(); + let covering = result + .column_by_name(label) + .unwrap_or_else(|| panic!("the {label} query must materialize {label}")); + for i in 0..row_ids.len() { + let row_id = row_ids.value(i); + match label { + "price" => assert_eq!( + covering + .as_primitive::() + .value(i), + -(row_id as i32) - 7, + "row {i} (row_id {row_id}): gathered price belongs to another row" + ), + _ => assert_eq!( + covering.as_string::().value(i), + format!("p{}", row_id * 3 + 11), + "row {i} (row_id {row_id}): gathered payload belongs to another row" + ), + } + } + } + } + + /// A covered IVF_FLAT fixture built for the survivor gather: one well-separated cluster + /// per partition, a **nullable** covering column, and covering values disjoint from the + /// row ids. + /// + /// Both of those properties are load-bearing and no other covered fixture in this suite + /// has them. Every other one declares its covering columns non-nullable, which is why a + /// covered query that returned a silent NULL for every survivor stayed invisible until + /// it was probed deliberately -- non-nullable columns turn it into a loud arrow error + /// instead. And a fixture whose covering value equals its row id cannot tell a by-name + /// gather from a positional one. + /// + /// `payload` is deliberately wide (~1 KB/row). The gather exists for the regime where + /// the covering payload dwarfs the quantization code; at a narrow width the scattered + /// and sequential reads cost almost the same and nothing can be measured about them. + /// + /// Returns the dataset and its vectors; row id equals row offset, so ground truth can + /// be computed against `vectors` directly. + async fn covered_gather_fixture( + uri: &str, + index_name: &str, + num_clusters: usize, + rows_per_cluster: usize, + ) -> (Dataset, Arc) { + use arrow_array::{Int32Array, StringArray}; + + const DIMS: usize = 16; + let total = num_clusters * rows_per_cluster; + + let mut flat = Vec::with_capacity(total * DIMS); + for row in 0..total { + let center = (row / rows_per_cluster) as f32 * 50.0; + for d in 0..DIMS { + flat.push(center + (row % rows_per_cluster) as f32 * 0.001 + d as f32 * 0.0001); + } + } + let vectors = Arc::new( + FixedSizeListArray::try_new_from_values(Float32Array::from(flat), DIMS as i32).unwrap(), + ); + // Decreasing, and offset by 1000, so `tag` is never equal to its row id. + let tags = Arc::new(Int32Array::from_iter_values( + (0..total as i32).map(|i| 1000 - i), + )); + // Every third row is NULL. With k = 10 the result always contains both a NULL and + // a non-NULL note, so neither case is vacuous. + let notes = Arc::new(StringArray::from_iter((0..total).map(|i| match i % 3 { + 0 => None, + _ => Some(format!("n{}", i * 7 + 3)), + }))); + let payloads = Arc::new(StringArray::from_iter_values( + (0..total).map(|i| format!("{i:04}").repeat(256)), + )); + let schema = Arc::new(Schema::new(vec![ + Field::new("tag", DataType::Int32, false), + Field::new("note", DataType::Utf8, true), + Field::new("payload", DataType::Utf8, false), + Field::new("vector", vectors.data_type().clone(), false), + ])); + let batch = + RecordBatch::try_new(schema.clone(), vec![tags, notes, payloads, vectors.clone()]) + .unwrap(); + let reader = RecordBatchIterator::new(vec![Ok(batch)], schema); + let mut dataset = Dataset::write(reader, uri, None).await.unwrap(); + + let mut params = VectorIndexParams::ivf_flat(num_clusters, DistanceType::L2); + params.covering_columns(vec![ + "note".to_string(), + "tag".to_string(), + "payload".to_string(), + ]); + dataset + .create_index( + &["vector"], + IndexType::Vector, + Some(index_name.to_string()), + ¶ms, + true, + ) + .await + .unwrap(); + (dataset, vectors) + } + + /// The gather reads the survivors' rows and nothing else -- until they stop being a + /// small fraction of the partition, at which point `Indices` degenerates into a + /// scattered read of nearly everything and the sequential read it replaced is cheaper. + /// + /// Both branches are exercised here against the same real index, and the assertion + /// distinguishes them directly rather than by cost: the scattered read returns exactly + /// the rows asked for, the sequential fallback returns the partition's whole range + /// (the caller aligns by row id either way, so both are correct). A test that passed on + /// either branch would prove nothing about the threshold, so the two row counts are + /// asserted separately and the byte counts are asserted on top. + /// + /// The line under test is `covering_read_for`'s comparison against + /// `COVERING_SCATTERED_READ_MAX_PERCENT`. Neutralise the *comparison*, not the constant: + /// `few` and `many` are derived from the same constant, so moving it drags both sides of + /// the threshold with it and the behavioural assertions stay satisfied. Collapsing + /// `covering_read_for` to `CoveringRead::Sequential` fails this test, and so does + /// collapsing it to `CoveringRead::Scattered` for every non-empty partition. + #[tokio::test] + async fn test_covering_gather_reads_only_the_survivors_rows_until_the_threshold() { + use lance_index::vector::storage::COVERING_SCATTERED_READ_MAX_PERCENT; + use lance_io::scheduler::IoStats; + + const INDEX_NAME: &str = "vector_idx"; + const NUM_CLUSTERS: usize = 4; + const ROWS_PER_CLUSTER: usize = 64; + + let test_dir = TempStrDir::default(); + let (dataset, _) = covered_gather_fixture( + test_dir.as_str(), + INDEX_NAME, + NUM_CLUSTERS, + ROWS_PER_CLUSTER, + ) + .await; + let ctx = load_vector_index_context(&dataset, "vector", INDEX_NAME).await; + let index = ctx + .index + .as_any() + .downcast_ref::() + .expect("expected IvfFlatIndex"); + + // The largest partition, not partition 0: IVF centroid initialisation is unseeded, + // so any individual partition can come out empty and `partition_size(0) == 0` makes + // both sides of the threshold the same set. Over `NUM_CLUSTERS` partitions the + // largest necessarily holds at least `ROWS_PER_CLUSTER`, which keeps the guard below + // satisfied for every assignment rather than for the lucky ones. + let partition_id = (0..index.storage.num_partitions()) + .max_by_key(|partition| index.storage.partition_size(*partition)) + .expect("the index must hold at least one partition"); + let partition_rows = index.storage.partition_size(partition_id); + assert!( + partition_rows >= 20, + "partition {partition_id} holds {partition_rows} rows; the threshold is a \ + percentage, so a tiny partition makes both sides of it the same set" + ); + // Derive both sides from the actual partition size: cluster-to-partition + // assignment is not stable across IVF training changes, so fixed counts go + // stale. `few` sits strictly below the threshold, `many` at or above it. + let threshold_rows = (partition_rows * COVERING_SCATTERED_READ_MAX_PERCENT).div_ceil(100); + let few: Vec = (0..threshold_rows.saturating_sub(1).max(1) as u32).collect(); + let many: Vec = (0..(threshold_rows + 1).min(partition_rows) as u32).collect(); + assert!( + many.len() * 100 >= partition_rows * COVERING_SCATTERED_READ_MAX_PERCENT + && few.len() * 100 < partition_rows * COVERING_SCATTERED_READ_MAX_PERCENT, + "the fixture must put `few` below the threshold and `many` above it for a \ + {partition_rows}-row partition" + ); + let columns = vec!["payload".to_string()]; + + let read = async |positions: Option<&[u32]>| { + let io_stats = IoStats::new(); + let batch = index + .storage + .take_covering(partition_id, positions, &columns, Some(io_stats.clone())) + .await + .unwrap(); + (batch, io_stats.snapshot().bytes_read) + }; + + let (few_batch, few_bytes) = read(Some(&few)).await; + let (many_batch, many_bytes) = read(Some(&many)).await; + let (whole_batch, whole_bytes) = read(None).await; + + assert_eq!( + few_batch.num_rows(), + few.len(), + "below the threshold the gather must read only the survivors' rows" + ); + assert_eq!( + many_batch.num_rows(), + partition_rows, + "above the threshold the gather must fall back to the partition's whole range" + ); + assert_eq!( + whole_batch.num_rows(), + partition_rows, + "positions the caller could not derive read the whole range" + ); + + assert!(whole_bytes > 0, "the fixture must actually perform I/O"); + assert_eq!( + many_bytes, whole_bytes, + "the fallback is the same read as the whole-range one: {many_bytes} vs \ + {whole_bytes} bytes" + ); + assert!( + few_bytes * 4 < whole_bytes, + "the scattered read must be materially cheaper than the range it replaces: \ + {few_bytes} vs {whole_bytes} bytes for {} of {partition_rows} rows", + few.len() + ); + + // Neither branch is a stub: both return the right values for the rows asked for, + // matched by row id rather than by position. + let whole_row_ids = whole_batch + .column_by_name(ROW_ID) + .expect("gathered batch carries row ids") + .as_primitive::(); + let whole_payloads = whole_batch["payload"].as_string::(); + for batch in [&few_batch, &many_batch] { + let row_ids = batch + .column_by_name(ROW_ID) + .expect("gathered batch carries row ids") + .as_primitive::(); + let payloads = batch["payload"].as_string::(); + for (i, position) in few.iter().map(|p| *p as usize).enumerate() { + assert_eq!(row_ids.value(i), whole_row_ids.value(position)); + assert_eq!( + payloads.value(i), + whole_payloads.value(position), + "gathered payload does not belong to the row it was asked for" + ); + assert_eq!( + payloads.value(i), + format!("{:04}", row_ids.value(i)).repeat(256), + "gathered payload does not match the base table" + ); + } + } + } + + /// The multi-partition gather must hold `O(survivors)`, not one whole covering batch + /// per contributing partition -- even when every partition takes the whole-range + /// fallback, which is what a pending fragment-reuse index forces for all of them. + /// + /// This is a memory bound, so it is asserted against bytes measured in the same test + /// rather than a magic number: one partition's whole-range covering read is the + /// positive control, and the gather's result over *four* such partitions must stay + /// well under it. Row counts back it up -- a gather that accumulated whole partitions + /// returns every row of every one of them, not the four survivors. + /// + /// The line under test is the narrowing in `gather_survivor_covering` + /// (`take_record_batch` on each partition's batch before the next partition is read). + /// Pushing the untouched `gathered` batch instead fails both assertions here: + /// `4 * partition_rows` rows instead of 4, and a result larger than a whole partition. + #[tokio::test] + async fn test_covering_gather_holds_only_the_survivors_when_the_read_falls_back() { + use lance_index::vector::{DEFAULT_QUERY_PARALLELISM, Query}; + + const INDEX_NAME: &str = "vector_idx"; + const NUM_CLUSTERS: usize = 4; + const ROWS_PER_CLUSTER: usize = 256; + + let test_dir = TempStrDir::default(); + let (dataset, vectors) = covered_gather_fixture( + test_dir.as_str(), + INDEX_NAME, + NUM_CLUSTERS, + ROWS_PER_CLUSTER, + ) + .await; + let ctx = load_vector_index_context(&dataset, "vector", INDEX_NAME).await; + let index = ctx + .index + .as_any() + .downcast_ref::() + .expect("expected IvfFlatIndex"); + + let query = Query { + column: "vector".to_string(), + key: vectors.value(0), + k: 10, + lower_bound: None, + upper_bound: None, + minimum_nprobes: NUM_CLUSTERS, + maximum_nprobes: None, + ef: None, + refine_factor: None, + metric_type: Some(DistanceType::L2), + use_index: true, + query_parallelism: DEFAULT_QUERY_PARALLELISM, + dist_q_c: 0.0, + approx_mode: Default::default(), + covering_projection: Some(Arc::from(vec!["payload".to_string()])), + }; + let covering = index + .query_covering(&query) + .unwrap() + .expect("the fixture declares covering columns"); + + // One survivor per partition, each with `position: None` -- the state a deferred + // fragment-reuse remap leaves every partition in, and the only one where the + // gather reads more than the survivors' own rows. + let mut locations: HashMap = HashMap::new(); + let mut expected_payloads: HashMap = HashMap::new(); + let mut whole_partition_bytes = 0; + for partition_id in 0..NUM_CLUSTERS { + let whole = index + .storage + .take_covering(partition_id, None, &covering.columns, None) + .await + .unwrap(); + assert_eq!( + whole.num_rows(), + index.storage.partition_size(partition_id), + "the control must read the partition's whole range" + ); + whole_partition_bytes = whole_partition_bytes.max(whole.get_array_memory_size()); + let row_id = whole + .column_by_name(ROW_ID) + .expect("gathered batch carries row ids") + .as_primitive::() + .value(0); + expected_payloads.insert(row_id, whole["payload"].as_string::().value(0).into()); + locations.insert( + row_id, + CoveringLocation { + partition_id, + position: None, + }, + ); + } + assert_eq!(locations.len(), NUM_CLUSTERS, "one survivor per partition"); + + let gathered = index + .gather_survivor_covering(&locations, &covering, None) + .await + .unwrap() + .expect("survivors were located, so the gather must return their covering"); + + assert_eq!( + gathered.num_rows(), + locations.len(), + "the gather must keep the survivors' rows only; holding every contributing \ + partition's whole covering would return {} rows", + (0..NUM_CLUSTERS) + .map(|p| index.storage.partition_size(p)) + .sum::() + ); + assert!( + whole_partition_bytes > 0, + "the fixture must actually hold covering bytes" + ); + assert!( + gathered.get_array_memory_size() * 4 < whole_partition_bytes, + "the gather must hold far less than a single partition's covering: {} bytes \ + for {} survivors vs {whole_partition_bytes} bytes for one partition", + gathered.get_array_memory_size(), + locations.len(), + ); + + // Bounded, and still correct: every survivor's own value, matched by row id. + let row_ids = gathered + .column_by_name(ROW_ID) + .expect("gathered batch carries row ids") + .as_primitive::(); + let payloads = gathered["payload"].as_string::(); + for (i, row_id) in row_ids.values().iter().enumerate() { + assert_eq!( + payloads.value(i), + expected_payloads[row_id], + "gathered payload does not belong to the row it was asked for" + ); + assert_eq!( + payloads.value(i), + format!("{row_id:04}").repeat(256), + "gathered payload does not match the base table" + ); + } + assert_eq!( + row_ids.values().iter().copied().collect::>(), + locations.keys().copied().collect::>(), + "the gather must return exactly the survivors it was asked for" + ); + } + + /// A covered query must return the covering values the index actually holds -- for a + /// **nullable** covering column too, where a null fill is indistinguishable from the + /// truth. + /// + /// This is the shape that hid a silent bug: when every covering fixture declares its + /// columns non-nullable, an implementation that emits nulls for the survivors fails + /// loudly on all of them and looks like a schema problem. A nullable column turns the + /// same defect into ten rows of silently wrong values, and only a comparison against + /// the base table catches it. `tag` is `1000 - offset` and `note` is `n{offset*7+3}`, + /// so neither equals its row id and a positional gather cannot coincide with a + /// by-name one. + /// + /// Both emit sites are covered: `query_parallelism = 1` merges partitions through the + /// global top-k heap (`global_heap_to_batch`), `> 1` searches each partition + /// separately (`append_covering`). Restoring only one of them fails one case. + /// + /// The absence of a `LanceRead` is what makes this an index measurement rather than a + /// scan measurement: a covering column is semantically transparent, so a plan that + /// re-fetched it from the base table would return byte-identical results. + #[rstest] + #[case::global_heap(1)] + #[case::per_partition(4)] + #[tokio::test] + async fn test_covered_query_returns_nullable_covering_values_from_the_index( + #[case] query_parallelism: i32, + ) { + const INDEX_NAME: &str = "vector_idx"; + const NUM_CLUSTERS: usize = 4; + const ROWS_PER_CLUSTER: usize = 64; + const K: usize = 10; + + let test_dir = TempStrDir::default(); + let (dataset, vectors) = covered_gather_fixture( + test_dir.as_str(), + INDEX_NAME, + NUM_CLUSTERS, + ROWS_PER_CLUSTER, + ) + .await; + + let query_key = vectors.value(0); + let mut scan = dataset.scan(); + scan.nearest("vector", query_key.as_primitive::(), K) + .unwrap(); + scan.minimum_nprobes(NUM_CLUSTERS); + scan.query_parallelism(query_parallelism); + scan.with_row_id(); + scan.project(&["tag", "note"]).unwrap(); + + let plan = scan.explain_plan(true).await.unwrap(); + assert!( + !plan.contains("LanceRead"), + "a projection of covered columns only must not fetch from the base table, or \ + this test measures the scan rather than the index; plan:\n{plan}" + ); + assert!( + !plan.contains("payload"), + "the query reads neither `payload` nor anything derived from it, so the index \ + must not be asked to materialize it; plan:\n{plan}" + ); + + let batch = scan.try_into_batch().await.unwrap(); + assert_eq!(batch.num_rows(), K, "k = {K} must return {K} rows"); + let row_ids = batch[ROW_ID].as_primitive::(); + let tags = batch["tag"].as_primitive::(); + let notes = batch["note"].as_string::(); + + let mut nulls = 0; + let mut values = 0; + for i in 0..batch.num_rows() { + let offset = row_ids.value(i) as usize; + assert_eq!( + tags.value(i), + 1000 - offset as i32, + "row {i} (row_id {offset}): covered tag belongs to another row" + ); + match offset % 3 { + 0 => { + assert!( + notes.is_null(i), + "row {i} (row_id {offset}): covered note must be NULL, got {:?}", + notes.value(i) + ); + nulls += 1; + } + _ => { + assert!( + !notes.is_null(i), + "row {i} (row_id {offset}): covered note must not be NULL -- a null \ + fill for every survivor is exactly the silent failure this fixture \ + exists to catch" + ); + assert_eq!( + notes.value(i), + format!("n{}", offset * 7 + 3), + "row {i} (row_id {offset}): covered note belongs to another row" + ); + values += 1; + } + } + } + assert!( + nulls > 0 && values > 0, + "the result must contain both NULL and non-NULL notes, or one of the two cases \ + above is vacuous (got {nulls} nulls, {values} values)" + ); + } + + /// `search_prepared_partition` is the synchronous phase of a prepared search: it runs + /// on the CPU pool, where the covering gather -- which is I/O -- cannot be awaited. + /// A query that needs covering columns must be told so, not handed rows whose covering + /// columns are quietly absent; the caller has async entry points for exactly this. + /// + /// The same call with `covering_projection = Some(&[])` must succeed and return the + /// bare `[_distance, _rowid]`. That is the state the projection narrowing exists for + /// and the one that silently degrades if it is folded into `None`: here it is the + /// difference between an error and a result. + /// + /// `supports_prepared_partition_search` must say so up front. A dispatcher that trusts + /// the flag pays for the partition load before it ever reaches the rejection, so a + /// covered index answering `true` there is a trap. The plain index built below is the + /// control: without it, `false` would be satisfied by an index that simply never + /// supported the entry point. + #[tokio::test] + async fn test_prepared_partition_search_rejects_a_query_that_needs_covering() { + use lance_index::prefilter::NoFilter; + use lance_index::vector::{DEFAULT_QUERY_PARALLELISM, Query}; + + const INDEX_NAME: &str = "vector_idx"; + const NUM_CLUSTERS: usize = 4; + const ROWS_PER_CLUSTER: usize = 64; + + let test_dir = TempStrDir::default(); + let (dataset, vectors) = covered_gather_fixture( + test_dir.as_str(), + INDEX_NAME, + NUM_CLUSTERS, + ROWS_PER_CLUSTER, + ) + .await; + let ctx = load_vector_index_context(&dataset, "vector", INDEX_NAME).await; + let index = ctx + .index + .as_any() + .downcast_ref::() + .expect("expected IvfFlatIndex"); + + let query_for = |covering: Option>| Query { + column: "vector".to_string(), + key: vectors.value(0), + k: 10, + lower_bound: None, + upper_bound: None, + minimum_nprobes: NUM_CLUSTERS, + maximum_nprobes: None, + ef: None, + refine_factor: None, + metric_type: Some(DistanceType::L2), + use_index: true, + query_parallelism: DEFAULT_QUERY_PARALLELISM, + dist_q_c: 0.0, + approx_mode: Default::default(), + covering_projection: covering.map(Arc::from), + }; + + let prepared = index + .prepare_partition_search( + 0, + &query_for(Some(vec!["tag".to_string()])), + Arc::new(NoFilter), + &NoOpMetricsCollector, + ) + .await + .unwrap(); + let err = index + .search_prepared_partition(prepared, &NoOpMetricsCollector) + .expect_err("a covered query must not be silently served without its covering"); + assert!( + err.to_string().contains("search_prepared_partition"), + "the error must name the entry point that cannot serve it, got: {err}" + ); + + let prepared = index + .prepare_partition_search( + 0, + &query_for(Some(Vec::new())), + Arc::new(NoFilter), + &NoOpMetricsCollector, + ) + .await + .unwrap(); + let batch = index + .search_prepared_partition(prepared, &NoOpMetricsCollector) + .expect("a query needing no covering column is served here as before"); + assert_eq!( + batch + .schema() + .fields() + .iter() + .map(|f| f.name().as_str()) + .collect::>(), + vec![DIST_COL, ROW_ID], + "`Some(&[])` means no covering work at all, not covering projected away" + ); + assert!(batch.num_rows() > 0, "the partition must be non-empty"); + + // The capability flag must not advertise what the rejection above denies. It + // carries no query, so it answers for the index: a covered index says no. + assert!( + !index.supports_prepared_partition_search(), + "a covered index must not advertise a prepared search it rejects" + ); + + // Control: the same index shape without covering columns still supports it. + let plain_dir = TempStrDir::default(); + let plain_schema = Arc::new(Schema::new(vec![Field::new( + "vector", + vectors.data_type().clone(), + false, + )])); + let plain_batch = + RecordBatch::try_new(plain_schema.clone(), vec![vectors.clone()]).unwrap(); + let mut plain_dataset = Dataset::write( + RecordBatchIterator::new(vec![Ok(plain_batch)], plain_schema), + plain_dir.as_str(), + None, + ) + .await + .unwrap(); + plain_dataset + .create_index( + &["vector"], + IndexType::Vector, + Some(INDEX_NAME.to_string()), + &VectorIndexParams::ivf_flat(NUM_CLUSTERS, DistanceType::L2), + true, + ) + .await + .unwrap(); + let plain_ctx = load_vector_index_context(&plain_dataset, "vector", INDEX_NAME).await; + let plain_index = plain_ctx + .index + .as_any() + .downcast_ref::() + .expect("expected IvfFlatIndex"); + assert!( + plain_index.supports_prepared_partition_search(), + "an ordinary index still supports the prepared search; without this the \ + covered assertion above would pass on an index that never supported it" + ); + } + + /// A partition read with only its internal columns must still be a well-formed + /// storage: `try_from_batch` does not require covering columns, and `covering_batch()` + /// on the result reports none rather than erroring or fabricating an empty batch. + /// + /// The same partition read with `PartitionColumns::All` is the control. Without it, + /// "reports none" would pass just as well on an index that never had covering columns, + /// and the row-id comparison would have nothing to catch a narrowed read that silently + /// shifted rows. + #[tokio::test] + async fn test_codes_only_partition_storage_is_well_formed() { + use lance_index::vector::flat::storage::FLAT_COLUMN; + + const INDEX_NAME: &str = "vector_idx"; + const NUM_CLUSTERS: usize = 4; + const ROWS_PER_CLUSTER: usize = 64; + + let test_dir = TempStrDir::default(); + let (dataset, _) = covered_flat_fixture( + test_dir.as_str(), + INDEX_NAME, + NUM_CLUSTERS, + ROWS_PER_CLUSTER, + ) + .await; + let ctx = load_vector_index_context(&dataset, "vector", INDEX_NAME).await; + let index = ctx + .index + .as_any() + .downcast_ref::() + .expect("expected IvfFlatIndex"); + + // Whichever partition k-means actually filled, not partition 0 unconditionally: + // which centroid wins which rows depends on initialisation, and on aarch64 + // partition 0 came out empty and tripped the vacuity guard. Any non-empty + // partition exercises the same storage invariants, so this asserts on the + // storage layout rather than on the clustering. + let mut chosen = None; + for candidate in 0..index.ivf_model().num_partitions() { + let storage = index + .load_partition_storage(candidate, PartitionColumns::All, None) + .await + .unwrap(); + if !storage.is_empty() { + chosen = Some((candidate, storage)); + break; + } + } + let (partition_id, all) = + chosen.expect("some partition must be non-empty or every assertion below is vacuous"); + let internal = index + .load_partition_storage(partition_id, PartitionColumns::Internal, None) + .await + .unwrap(); + let covering_names = |storage: &lance_index::vector::flat::storage::FlatFloatStorage| { + let schema = storage.schema().clone(); + storage + .covering_field_indices() + .into_iter() + .map(|i| schema.field(i).name().to_string()) + .collect::>() + }; + assert_eq!( + covering_names(&all), + vec!["payload".to_string(), "price".to_string()], + "the control read must carry both covering columns, or the narrowed arm below \ + proves nothing" + ); + assert!(all.covering_batch().unwrap().is_some()); + + assert!( + covering_names(&internal).is_empty(), + "a codes-only read must carry no covering column" + ); + assert!( + internal.covering_batch().unwrap().is_none(), + "covering_batch() on a codes-only storage reports none rather than erroring \ + or emitting an empty batch" + ); + + // Well-formed: the storage's own columns are all there, and the rows are the same + // rows in the same order. Row ids are compared through `row_ids()`, which resolves + // `_rowid` by name -- the covered layout puts covering columns before it. + assert_eq!(internal.len(), all.len()); + assert_eq!( + internal.row_ids().copied().collect::>(), + all.row_ids().copied().collect::>(), + "narrowing the read must not disturb which rows the partition holds" + ); + assert!(internal.schema().column_with_name(ROW_ID).is_some()); + assert!(internal.schema().column_with_name(FLAT_COLUMN).is_some()); + assert_eq!( + internal.schema().fields().len(), + 2, + "flat storage's internal columns are exactly [{ROW_ID}, {FLAT_COLUMN}]; schema \ + was {:?}", + internal.schema() + ); + } + + /// The rebuild path reads a partition with `PartitionColumns::All` on purpose: it + /// re-writes those batches into the index being built, so a covering column left unread + /// is a covering column the merged index no longer has -- silently, since nothing + /// downstream of the rebuild asks for it again. + /// + /// Asserted on the merged STORAGE rather than through a query, so it holds independently + /// of what the search path currently does with covering. The line under test is + /// `PartitionColumns::All` in `IvfIndexBuilder::take_partition_batches`; switching it to + /// `Internal` drops the covering columns from the merged index and fails this test. + #[tokio::test] + async fn test_merge_optimize_preserves_covering_in_storage() { + use arrow_array::{Int32Array, StringArray, types::Int32Type}; + + const INDEX_NAME: &str = "vector_idx"; + const NUM_CLUSTERS: usize = 4; + const ROWS_PER_CLUSTER: usize = 64; + const INDEXED: usize = NUM_CLUSTERS * ROWS_PER_CLUSTER; + const APPENDED: usize = 128; + const DIMS: usize = 16; + + let test_dir = TempStrDir::default(); + let test_uri = test_dir.as_str(); + let (mut dataset, _) = + covered_flat_fixture(test_uri, INDEX_NAME, NUM_CLUSTERS, ROWS_PER_CLUSTER).await; + + // Append unindexed rows, then merge them in. The merge is what re-reads the + // existing partitions and re-writes them into the new index file. + let mut flat = Vec::with_capacity(APPENDED * DIMS); + for row in 0..APPENDED { + for d in 0..DIMS { + flat.push(row as f32 * 0.01 + d as f32 * 0.0001); + } + } + let appended_vectors = Arc::new( + FixedSizeListArray::try_new_from_values(Float32Array::from(flat), DIMS as i32).unwrap(), + ); + let schema = dataset.schema().into(); + let appended = RecordBatch::try_new( + Arc::new(schema), + vec![ + Arc::new(Int32Array::from_iter_values( + (INDEXED..INDEXED + APPENDED).map(|i| -(i as i32) - 7), + )), + Arc::new(StringArray::from_iter_values( + (INDEXED..INDEXED + APPENDED).map(|i| format!("p{}", i * 3 + 11)), + )), + appended_vectors, + ], + ) + .unwrap(); + let appended_schema = appended.schema(); + dataset = Dataset::write( + RecordBatchIterator::new(vec![Ok(appended)], appended_schema), + test_uri, + Some(WriteParams { + mode: WriteMode::Append, + ..Default::default() + }), + ) + .await + .unwrap(); + dataset + .optimize_indices(&OptimizeOptions::new()) + .await + .unwrap(); + + let ctx = load_vector_index_context(&dataset, "vector", INDEX_NAME).await; + let index = ctx + .index + .as_any() + .downcast_ref::() + .expect("expected IvfFlatIndex"); + + let mut covered: HashMap = HashMap::new(); + for partition in 0..index.ivf_model().num_partitions() { + let storage = index + .load_partition_storage(partition, PartitionColumns::All, None) + .await + .unwrap(); + for batch in storage.to_batches().unwrap() { + let row_ids = batch + .column_by_name(ROW_ID) + .expect("storage batch carries row ids") + .as_primitive::(); + let payloads = batch + .column_by_name("payload") + .expect("merged storage must keep covering column 'payload'") + .as_string::(); + let prices = batch + .column_by_name("price") + .expect("merged storage must keep covering column 'price'") + .as_primitive::(); + for i in 0..batch.num_rows() { + covered.insert( + row_ids.value(i), + (payloads.value(i).to_string(), prices.value(i)), + ); + } + } + } + assert_eq!( + covered.len(), + INDEXED + APPENDED, + "the merge must index every row, or the value check below misses the rows it dropped" + ); + + // Values, not just presence: a merge that re-attached covering by position rather + // than by row id keeps the column and corrupts it. + let row_ids: Vec = covered.keys().copied().collect(); + let truth = dataset + .take_rows( + &row_ids, + crate::dataset::ProjectionRequest::from_columns( + ["payload", "price"], + dataset.schema(), + ), + ) + .await + .unwrap(); + let truth_payloads = truth.column_by_name("payload").unwrap().as_string::(); + let truth_prices = truth + .column_by_name("price") + .unwrap() + .as_primitive::(); + for (i, row_id) in row_ids.iter().enumerate() { + let (payload, price) = &covered[row_id]; + assert_eq!( + (payload.as_str(), *price), + (truth_payloads.value(i), truth_prices.value(i)), + "row_id {row_id}: merged covering value does not match the base table" + ); + } + } + + /// The remap path rewrites a partition into a new index file exactly as the merge + /// path does, so it has the same requirement: it must read every column, not just + /// the internal ones. It reaches the storage through the partition *entry* rather + /// than through `load_partition_storage`, which is why the merge test above does not + /// cover it -- and why a codes-only entry silently strips the covering columns from + /// every compacted covered index. + /// + /// Asserted on the remapped STORAGE, not through a query, so it holds independently + /// of what the search path currently does with covering (`test_ivf_pq_covered_survives_compaction` + /// asserts the same property through a query and is ignored until the survivors-only + /// gather lands). The line under test is `load_partition_entry_with_covering` in + /// `IvfIndexBuilder::remap`; the cached `load_partition` reads internal columns only + /// and fails this test. + #[tokio::test] + async fn test_compaction_remap_preserves_covering_in_storage() { + use crate::dataset::optimize::{CompactionOptions, compact_files}; + use arrow_array::types::Int32Type; + + const INDEX_NAME: &str = "vector_idx"; + const NUM_CLUSTERS: usize = 4; + const ROWS_PER_CLUSTER: usize = 64; + + let test_dir = TempStrDir::default(); + let (mut dataset, _) = covered_flat_fixture( + test_dir.as_str(), + INDEX_NAME, + NUM_CLUSTERS, + ROWS_PER_CLUSTER, + ) + .await; + + // Delete enough rows to clear the materialize-deletions threshold, so compaction + // rewrites the fragment and remaps the index rather than no-opping. + dataset.delete("price > -60").await.unwrap(); + let metrics = compact_files(&mut dataset, CompactionOptions::default(), None) + .await + .unwrap(); + assert!( + metrics.files_removed > 0, + "compaction must actually rewrite fragments, or the remap under test never runs" + ); + + let ctx = load_vector_index_context(&dataset, "vector", INDEX_NAME).await; + let index = ctx + .index + .as_any() + .downcast_ref::() + .expect("expected IvfFlatIndex"); + + let mut covered: HashMap = HashMap::new(); + for partition in 0..index.ivf_model().num_partitions() { + let storage = index + .load_partition_storage(partition, PartitionColumns::All, None) + .await + .unwrap(); + for batch in storage.to_batches().unwrap() { + let row_ids = batch + .column_by_name(ROW_ID) + .expect("storage batch carries row ids") + .as_primitive::(); + let payloads = batch + .column_by_name("payload") + .expect("remapped storage must keep covering column 'payload'") + .as_string::(); + let prices = batch + .column_by_name("price") + .expect("remapped storage must keep covering column 'price'") + .as_primitive::(); + for i in 0..batch.num_rows() { + covered.insert( + row_ids.value(i), + (payloads.value(i).to_string(), prices.value(i)), + ); + } + } + } + assert!( + !covered.is_empty(), + "the remapped index must still hold rows, or the value check below is vacuous" + ); + + // Values, not just presence: a remap that re-attached covering by position rather + // than by row id keeps the column and corrupts it, and the row ids themselves are + // the post-compaction addresses. + let row_ids: Vec = covered.keys().copied().collect(); + let truth = dataset + .take_rows( + &row_ids, + crate::dataset::ProjectionRequest::from_columns( + ["payload", "price"], + dataset.schema(), + ), + ) + .await + .unwrap(); + let truth_payloads = truth.column_by_name("payload").unwrap().as_string::(); + let truth_prices = truth + .column_by_name("price") + .unwrap() + .as_primitive::(); + for (i, row_id) in row_ids.iter().enumerate() { + let (payload, price) = &covered[row_id]; + assert_eq!( + (payload.as_str(), *price), + (truth_payloads.value(i), truth_prices.value(i)), + "row_id {row_id}: remapped covering value does not match the base table" + ); + } + } + + /// A covered FLAT/SQ index must survive a *deferred* remap: `compact_files` + /// with `defer_index_remap: true` rewrites fragments but leaves the index's + /// row ids to be remapped later via a fragment-reuse index (FRI) instead of + /// remapping them inline. `IvfQuantizationStorage::load_partition` passes + /// that FRI on every subsequent load (search, optimize, another build), so + /// a covered FLAT/SQ storage batch -- `[_rowid, code, ]`, three + /// or more columns -- must not be rejected as though it could only ever be + /// the two-column `(value, row_id)` shape a plain scalar index has. + /// + /// PQ and RQ are excluded: PQ remaps inline via `rebuild_storage_batch` and + /// RQ via its own `remap`, so neither ever reaches the shared FRI path this + /// guards. The HNSW_FLAT/HNSW_SQ variants are also excluded here: HNSW's + /// graph independently fails to survive `defer_index_remap` even without + /// any covering columns (pre-existing, unrelated to covering -- see the + /// P0 fix report), so they would fail this test for a reason this fix + /// does not address. + /// + /// This is also the only test that reaches [`CoveringGather::WholeRange`]: the FRI + /// drops rows as the partition is loaded, so a storage position no longer addresses + /// the file and the survivor gather must re-read the whole range and match by row id. + /// Removing that alignment check fails here with `row id ... missing from covering + /// source` -- loudly, because the gather is matched by row id rather than by position. + #[rstest] + #[case::flat(VectorIndexParams::ivf_flat(4, DistanceType::L2))] + #[case::sq(VectorIndexParams::with_ivf_sq_params( + DistanceType::L2, + IvfBuildParams::new(4), + SQBuildParams::default() + ))] + #[tokio::test] + async fn test_covered_survives_deferred_remap_frag_reuse( + #[case] mut params: VectorIndexParams, + ) { + const INDEX_NAME: &str = "vector_idx"; + let test_dir = TempStrDir::default(); + let test_uri = test_dir.as_str(); + let (mut dataset, vectors) = generate_test_dataset::(test_uri, 0.0..1.0).await; + + params.covering_columns(vec!["id".to_string()]); + dataset + .create_index( + &["vector"], + IndexType::Vector, + Some(INDEX_NAME.to_string()), + ¶ms, + true, + ) + .await + .unwrap(); + + // Delete rows and defer the index remap: this fragment-reuse index now + // covers the covered vector index built above, so every subsequent load + // of its storage carries it. + dataset.delete("id < 100").await.unwrap(); + compact_files( + &mut dataset, + CompactionOptions { + defer_index_remap: true, + ..Default::default() + }, + None, + ) + .await + .unwrap(); + + // A search must load the (now FRI-tagged) covered storage without + // rejecting its three-or-more-column shape, and still return correct, + // row-aligned covering values for the surviving rows. + let q = vectors.value(0); + let q = q.as_primitive::(); + let mut scan = dataset.scan(); + scan.nearest("vector", q, 10).unwrap(); + scan.nprobes(4); + scan.with_row_id(); + scan.project(&["id"]).unwrap(); + let batch = scan.try_into_batch().await.unwrap(); + assert_eq!(batch.num_rows(), 10); + let ids = batch + .column_by_name("id") + .expect("covered 'id' column must be emitted") + .as_primitive::(); + let row_ids = batch + .column_by_name(ROW_ID) + .expect("row id column") + .as_primitive::(); + + // Ground truth via an independent (non-index, base-table take) path, + // keyed by the post-compaction row id the scan returned. Compaction + // reassigns physical row addresses, so `id == _rowid` no longer holds + // here the way it does pre-compaction -- but the covered value must + // still equal the row's true id. + let row_id_vec: Vec = row_ids.values().to_vec(); + let projection = crate::dataset::ProjectionRequest::from_columns(["id"], dataset.schema()); + let truth = dataset.take_rows(&row_id_vec, projection).await.unwrap(); + let truth_ids = truth + .column_by_name("id") + .unwrap() + .as_primitive::(); + for i in 0..ids.len() { + assert!( + ids.value(i) >= 100, + "deleted rows (id < 100) must not be returned" + ); + assert_eq!( + ids.value(i), + truth_ids.value(i), + "row {i} (row_id {}): covered id != true id after deferred-remap compaction", + row_ids.value(i) + ); + } + + // Building a SECOND covered index while a fragment-reuse index already + // exists on the dataset must also succeed: `StorageBuilder::build` + // threads the FRI into `Q::Storage::try_from_batch` for every new + // build, indexed or not. + params.covering_columns(vec!["id".to_string()]); + dataset + .create_index( + &["vector"], + IndexType::Vector, + Some("vector_idx_2".to_string()), + ¶ms, + true, + ) + .await + .unwrap(); + + // ...and the index it produces must be *correct*, not merely buildable. That + // build is where a covered storage batch reaches the fragment-reuse remap in + // its build-time column order -- `[, _rowid, code]`, row id NOT + // at index 0. The search earlier in this test loads the written file instead, + // whose `[_rowid, code, ]` order makes a positional row-id lookup + // right by accident. So without querying this second index, nothing pins + // `remap_row_ids_by_name` locating the row id by name: a positional variant + // remaps the covering column, leaves the real row ids unremapped, and no + // assertion anywhere observes it. + dataset.drop_index(INDEX_NAME).await.unwrap(); + let mut scan = dataset.scan(); + scan.nearest("vector", q, 10).unwrap(); + scan.nprobes(4); + scan.with_row_id(); + scan.project(&["id"]).unwrap(); + // Pin the ANN path: a flat-KNN fallback would satisfy every assertion below + // without the remapped index being read at all. + let plan = scan.explain_plan(true).await.unwrap(); + assert!( + plan.contains("ANNSubIndex"), + "the covered query must go through the second index; plan was:\n{plan}" + ); + assert!( + !plan.contains("LanceRead"), + "covered projection ['id'] should skip the base-table take; plan was:\n{plan}" + ); + let batch = scan.try_into_batch().await.unwrap(); + assert_eq!(batch.num_rows(), 10); + let ids = batch + .column_by_name("id") + .expect("covered 'id' column must be emitted") + .as_primitive::(); + let row_ids = batch + .column_by_name(ROW_ID) + .expect("row id column") + .as_primitive::(); + let row_id_vec: Vec = row_ids.values().to_vec(); + let projection = crate::dataset::ProjectionRequest::from_columns(["id"], dataset.schema()); + let truth = dataset.take_rows(&row_id_vec, projection).await.unwrap(); + let truth_ids = truth + .column_by_name("id") + .unwrap() + .as_primitive::(); + for i in 0..ids.len() { + assert!( + ids.value(i) >= 100, + "deleted rows (id < 100) must not be returned by the second index" + ); + assert_eq!( + ids.value(i), + truth_ids.value(i), + "row {i} (row_id {}): covered id != true id for the index built \ + while a fragment-reuse index was already present", + row_ids.value(i) + ); + } + } + + /// Fix 2: under query parallelism > 1 the parallel search branch (which uses + /// `search_in_partition`, not the heap merge) must also emit covering + /// columns. On a multi-core runner this exercises the parallel path; if the + /// session resolves parallelism to 1 it still passes via the sequential path. + #[tokio::test] + async fn test_ivf_pq_covered_projection_parallel() { + const INDEX_NAME: &str = "vector_idx"; + let test_dir = TempStrDir::default(); + let test_uri = test_dir.as_str(); + let (mut dataset, vectors) = generate_test_dataset::(test_uri, 0.0..1.0).await; + let mut params = VectorIndexParams::ivf_pq(4, 8, 4, DistanceType::L2, 2); + params.covering_columns(vec!["id".to_string()]); + dataset + .create_index( + &["vector"], + IndexType::Vector, + Some(INDEX_NAME.to_string()), + ¶ms, + true, + ) + .await + .unwrap(); + + let q = vectors.value(0); + let q = q.as_primitive::(); + let mut scan = dataset.scan(); + scan.nearest("vector", q, 10).unwrap(); + scan.minimum_nprobes(4); // search several partitions... + scan.query_parallelism(4); // ...in parallel, forcing the parallel branch + scan.project(&["id"]).unwrap(); + let batch = scan.try_into_batch().await.unwrap(); + assert!(batch.column_by_name("id").is_some()); + assert_eq!(batch.num_rows(), 10, "k=10 should return 10 rows"); + } + + /// Fix 3: a covered index with unindexed fragments (rows appended but not + /// optimized) must not panic on a non-fast-search query. The flat search + /// over the unindexed rows now projects the covering columns so the union + /// with the index result succeeds. + #[tokio::test] + async fn test_ivf_pq_covered_with_unindexed_fragments() { + const INDEX_NAME: &str = "vector_idx"; + let test_dir = TempStrDir::default(); + let test_uri = test_dir.as_str(); + let (mut dataset, vectors) = generate_test_dataset::(test_uri, 0.0..1.0).await; + let mut params = VectorIndexParams::ivf_pq(4, 8, 4, DistanceType::L2, 2); + params.covering_columns(vec!["id".to_string()]); + dataset + .create_index( + &["vector"], + IndexType::Vector, + Some(INDEX_NAME.to_string()), + ¶ms, + true, + ) + .await + .unwrap(); + + // Append rows WITHOUT optimizing -> unindexed fragments; a (non-fast) + // query then goes through the index + flat-search combine path. + append_dataset::(&mut dataset, 100, 0.0..1.0).await; + + let q = vectors.value(0); + let q = q.as_primitive::(); + let mut scan = dataset.scan(); + scan.nearest("vector", q, 10).unwrap(); + scan.project(&["id"]).unwrap(); + let batch = scan.try_into_batch().await.unwrap(); + assert!(batch.column_by_name("id").is_some()); + assert_eq!(batch.num_rows(), 10, "k=10 should return 10 rows"); + } + + async fn shrink_smallest_partition( + dataset: &mut Dataset, + index_name: &str, + expected_after_join: usize, + next_id: &mut u64, + ) -> (usize, usize, usize) { + const ROWS_TO_APPEND_FOR_JOIN: usize = 32; + let row_count_before = dataset.count_all_rows().await.unwrap(); + let index_ctx = load_vector_index_context(dataset, "vector", index_name).await; + let partitions = index_ctx.stats()["indices"][0]["partitions"] + .as_array() + .expect("partitions should be present"); + let (partition_idx, _size) = partitions + .iter() + .enumerate() + .filter_map(|(idx, part)| part["size"].as_u64().map(|size| (idx, size))) + .filter(|(_, size)| *size > 1) + .min_by_key(|(_, size)| *size) + .expect("should have at least one partition with joinable rows"); + + let row_ids = load_partition_row_ids(index_ctx.ivf(), partition_idx).await; + assert!( + row_ids.len() > 1, + "Partition {} should have removable rows", + partition_idx + ); + + let rows = dataset + .take_rows(&row_ids, dataset.schema().clone()) + .await + .unwrap(); + let ids = rows["id"].as_primitive::().values(); + let template_values = rows["vector"] + .as_fixed_size_list() + .value(0) + .as_primitive::() + .values() + .to_vec(); + + delete_ids(dataset, &ids[1..]).await; + compact_after_deletions(dataset).await; + + append_constant_vector_with_start_id( + dataset, + ROWS_TO_APPEND_FOR_JOIN, + &template_values, + next_id, + ) + .await; + dataset + .optimize_indices(&OptimizeOptions::new()) + .await + .unwrap(); + + let post_ctx = load_vector_index_context(dataset, "vector", index_name).await; + let post_partitions = post_ctx.num_partitions(); + assert_eq!( + post_partitions, + expected_after_join, + "Expected partitions to be at most {} after join, got stats: {}", + expected_after_join, + post_ctx.stats_json() + ); + + let row_count_after = dataset.count_all_rows().await.unwrap(); + debug_assert!( + row_count_before + ROWS_TO_APPEND_FOR_JOIN >= row_count_after, + "row count should not increase after delete + append" + ); + let deleted_rows = row_count_before + ROWS_TO_APPEND_FOR_JOIN - row_count_after; + + (deleted_rows, ROWS_TO_APPEND_FOR_JOIN, post_partitions) + } + + async fn append_constant_vector_with_start_id( + dataset: &mut Dataset, + rows: usize, + template: &[f32], + next_id: &mut u64, + ) { + append_constant_vector_batch(dataset, rows, template, *next_id, None).await; + *next_id += rows as u64; + } + + async fn append_partition_templates( + dataset: &mut Dataset, + rows_per_template: usize, + templates: &[Vec], + ) { + assert!( + !templates.is_empty(), + "at least one template is required for append" + ); + for template in templates { + assert_eq!( + template.len(), + DIM, + "Template vector should have {} dimensions", + DIM + ); + } + + let start_id = dataset.count_all_rows().await.unwrap() as u64; + let total_rows = rows_per_template * templates.len(); + let ids = Arc::new(UInt64Array::from_iter_values( + start_id..start_id + total_rows as u64, + )); + let mut appended_values = Vec::with_capacity(total_rows * DIM); + for template in templates { + for _ in 0..rows_per_template { + appended_values.extend_from_slice(template); + } + } + let vectors = Arc::new( + FixedSizeListArray::try_new_from_values( + Float32Array::from(appended_values), + DIM as i32, + ) + .unwrap(), + ); + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::UInt64, false), + Field::new("vector", vectors.data_type().clone(), false), + ])); + let batch = RecordBatch::try_new(schema.clone(), vec![ids, vectors]).unwrap(); + let batches = RecordBatchIterator::new(vec![Ok(batch)], schema); + dataset.append(batches, None).await.unwrap(); + } + + async fn append_constant_vector_with_params( + dataset: &mut Dataset, + rows: usize, + template: &[f32], + write_params: Option, + ) { + let start_id = dataset.count_all_rows().await.unwrap() as u64; + append_constant_vector_batch(dataset, rows, template, start_id, write_params).await; + } + + async fn append_constant_vector_batch( + dataset: &mut Dataset, + rows: usize, + template: &[f32], + start_id: u64, + write_params: Option, + ) { + assert_eq!( + template.len(), + DIM, + "Template vector should have {} dimensions", + DIM + ); + + let ids = Arc::new(UInt64Array::from_iter_values( + start_id..start_id + rows as u64, + )); + let mut appended_values = Vec::with_capacity(rows * DIM); + for _ in 0..rows { + appended_values.extend_from_slice(template); + } + let vectors = Arc::new( + FixedSizeListArray::try_new_from_values( + Float32Array::from(appended_values), + DIM as i32, + ) + .unwrap(), + ); + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::UInt64, false), + Field::new("vector", vectors.data_type().clone(), false), + ])); + let batch = RecordBatch::try_new(schema.clone(), vec![ids, vectors]).unwrap(); + let batches = RecordBatchIterator::new(vec![Ok(batch)], schema); + let params = write_params.map(|mut params| { + params.mode = WriteMode::Append; + params + }); + dataset.append(batches, params).await.unwrap(); + } + + #[allow(clippy::too_many_arguments)] + async fn append_and_verify_append_phase( + dataset: &mut Dataset, + index_name: &str, + template: &[f32], + next_id: &mut u64, + rows_to_append: usize, + expected_partitions: usize, + expected_total_rows: usize, + expected_index_count: usize, + expect_split: bool, + ) { + append_constant_vector_with_start_id(dataset, rows_to_append, template, next_id).await; + dataset + .optimize_indices(&OptimizeOptions::new()) + .await + .unwrap(); + + let stats_json = dataset.index_statistics(index_name).await.unwrap(); + let stats: serde_json::Value = serde_json::from_str(&stats_json).unwrap(); + + let indices = stats["indices"] + .as_array() + .expect("indices array should exist"); + if expect_split { + assert_eq!( + indices.len(), + expected_index_count, + "Expected {} index entries after split, got {}, stats: {}", + expected_index_count, + indices.len(), + stats + ); + } else { + assert!( + indices.len() >= expected_index_count, + "Expected at least {} index entries after append, got {}, stats: {}", + expected_index_count, + indices.len(), + stats + ); + } + assert!( + stats["num_indices"].as_u64().unwrap() as usize >= expected_index_count, + "num_indices should be at least {}, stats: {}", + expected_index_count, + stats + ); + assert_eq!( + stats["num_indexed_rows"].as_u64().unwrap() as usize, + expected_total_rows, + "Total indexed rows mismatch after append" + ); + + let base_index = indices + .iter() + .max_by_key(|entry| entry["num_partitions"].as_u64().unwrap_or(0)) + .expect("at least one index entry should exist"); + assert_eq!( + base_index["num_partitions"].as_u64().unwrap() as usize, + expected_partitions, + "Partition count mismatch after append" + ); + + if expected_index_count == 1 { + let partitions = base_index["partitions"] + .as_array() + .expect("partitions should exist"); + assert_eq!( + partitions.len(), + expected_partitions, + "Expected {} partitions, found {}", + expected_partitions, + partitions.len() + ); + let partition_sizes: Vec = partitions + .iter() + .map(|part| part["size"].as_u64().unwrap() as usize) + .collect(); + let total_partition_rows: usize = partition_sizes.iter().sum(); + assert_eq!( + total_partition_rows, expected_total_rows, + "Partition sizes should sum to total rows: {:?}", + partition_sizes + ); + } else { + assert!( + !expect_split, + "Split should result in a single merged index" + ); + } + + assert_eq!( + dataset.count_all_rows().await.unwrap(), + expected_total_rows, + "Dataset row count mismatch after append" + ); + } + + async fn load_partition_row_ids(index: &IvfPq, partition_idx: usize) -> Vec { + index + .storage + .load_partition(partition_idx, PartitionColumns::Internal, None) + .await + .unwrap() + .row_ids() + .copied() + .collect() + } + + async fn load_flat_partition_row_ids(index: &IvfFlatIndex, partition_idx: usize) -> Vec { + index + .storage + .load_partition(partition_idx, PartitionColumns::Internal, None) + .await + .unwrap() + .row_ids() + .copied() + .collect() + } + + async fn delete_ids(dataset: &mut Dataset, ids: &[u64]) { + if ids.is_empty() { + return; + } + let predicate = ids + .iter() + .map(|x| x.to_string()) + .collect::>() + .join(","); + dataset + .delete(&format!("id in ({})", predicate)) + .await + .unwrap(); + } + + async fn compact_after_deletions(dataset: &mut Dataset) { + compact_files( + dataset, + CompactionOptions { + materialize_deletions_threshold: 0.0, + ..Default::default() + }, + None, + ) + .await + .unwrap(); + } + + async fn ground_truth( + dataset: &Dataset, + column: &str, + query: &dyn Array, + k: usize, + distance_type: DistanceType, + ) -> HashSet { + let batch = dataset + .scan() + .with_row_id() + .nearest(column, query, k) + .unwrap() + .distance_metric(distance_type) + .use_index(false) + .try_into_batch() + .await + .unwrap(); + batch[ROW_ID] + .as_primitive::() + .values() + .iter() + .copied() + .collect() + } + + fn multivec_ground_truth( + vectors: &ListArray, + query: &dyn Array, + k: usize, + distance_type: DistanceType, + ) -> Vec<(f32, u64)> { + let query = if let Some(list_array) = query.as_list_opt::() { + list_array.values().clone() + } else { + query.as_fixed_size_list().values().clone() + }; + multivec_distance(&query, vectors, distance_type) + .unwrap() + .into_iter() + .enumerate() + .map(|(i, dist)| (dist, i as u64)) + .sorted_by(|a, b| a.0.total_cmp(&b.0)) + .take(k) + .collect() + } + + const TWO_FRAG_NUM_ROWS: usize = 2000; + const TWO_FRAG_DIM: usize = 128; + const TWO_FRAG_NUM_PARTITIONS: usize = 4; + const TWO_FRAG_NUM_SUBVECTORS: usize = 16; + const TWO_FRAG_NUM_BITS: usize = 8; + const TWO_FRAG_SAMPLE_RATE: usize = 7; + const TWO_FRAG_MAX_ITERS: u32 = 20; + + fn make_two_fragment_batches() -> (Arc, Vec) { + let ids = Arc::new(UInt64Array::from_iter_values(0..TWO_FRAG_NUM_ROWS as u64)); + + let values = generate_random_array_with_range(TWO_FRAG_NUM_ROWS * TWO_FRAG_DIM, 0.0..1.0); + let vectors = Arc::new( + FixedSizeListArray::try_new_from_values( + Float32Array::from(values), + TWO_FRAG_DIM as i32, + ) + .unwrap(), + ); + + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::UInt64, false), + Field::new("vector", vectors.data_type().clone(), false), + ])); + let batch = RecordBatch::try_new(schema.clone(), vec![ids, vectors]).unwrap(); + + (schema, vec![batch]) + } + + /// The covering tests' own vector geometry, kept separate from the shared `TWO_FRAG_*` + /// fixture. Global codebook training costs `dim * iters * samples * centroids` and these + /// tests pay it once per case while building and merging several indexes, which is what put + /// them over the one-second local-unit-test budget. What they assert -- that covered values + /// survive every build, merge and lifecycle step and stay row-aligned with the base table -- + /// is independent of quantization resolution, so the vector shrinks instead. `dim / + /// num_sub_vectors` stays 8, as the shared fixture has it, and the codebook stays 8-bit: that + /// is the production default and the width these tests exist to cover. + /// + /// The shared fixture keeps its own dimensions: its tests assert *exact* single-vs-split + /// top-K equality over uniform-random vectors, which a coarser codebook breaks by making + /// distances tie. + const COVERED_DIM: usize = 32; + const COVERED_NUM_SUBVECTORS: usize = 4; + const COVERED_MAX_ITERS: u32 = 4; + + /// Like `make_two_fragment_batches`, but with two covering columns: a non-null `id` and a + /// **nullable** `payload` (every 3rd value null). Covered tests use this so null covering + /// values are exercised through the per-segment build and the cross-shard merge. + fn make_covered_test_batches() -> (Arc, Vec) { + let ids = Arc::new(UInt64Array::from_iter_values(0..TWO_FRAG_NUM_ROWS as u64)); + let payload = Arc::new(UInt64Array::from_iter( + (0..TWO_FRAG_NUM_ROWS as u64).map(|v| if v % 3 == 0 { None } else { Some(v + 1000) }), + )); + + // Clustered vectors: uniform-random data is pathological for IVF_PQ recall (the curse of + // dimensionality), which would make a recall gate flaky. Instead pack points into + // well-separated clusters of <= k points each (centers 4.0 apart per dim -> ~23 in L2 at + // `COVERED_DIM`, tiny within-cluster jitter), so a query drawn from a cluster has its whole + // cluster as the unambiguous nearest neighbors and IVF_PQ recall is reliably high. + const CLUSTER_SIZE: usize = 8; + let mut flat = Vec::with_capacity(TWO_FRAG_NUM_ROWS * COVERED_DIM); + for row in 0..TWO_FRAG_NUM_ROWS { + let center = (row / CLUSTER_SIZE) as f32 * 4.0; + let within = (row % CLUSTER_SIZE) as f32; + for d in 0..COVERED_DIM { + flat.push(center + within * 0.002 + d as f32 * 0.00001); + } + } + let vectors = Arc::new( + FixedSizeListArray::try_new_from_values(Float32Array::from(flat), COVERED_DIM as i32) + .unwrap(), + ); + + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::UInt64, false), + Field::new("payload", DataType::UInt64, true), + Field::new("vector", vectors.data_type().clone(), false), + ])); + let batch = RecordBatch::try_new(schema.clone(), vec![ids, payload, vectors]).unwrap(); + + (schema, vec![batch]) + } + + /// Assert `refine` is still served out of index storage rather than a base-table take. + /// + /// The index under test must cover `payload`, so the projection is served from the index + /// too and the only thing that could still read the base table is refine's own fetch of + /// full-precision vectors. Works for every index type, unlike inspecting partition + /// storage, which needs the concrete index struct. + async fn assert_refine_served_from_index(dataset: &Dataset, query: &dyn Array, what: &str) { + let mut scan = dataset.scan(); + scan.nearest("vector", query, 10).unwrap(); + scan.nprobes(TWO_FRAG_NUM_PARTITIONS); + scan.refine(2); + scan.project(&["payload"]).unwrap(); + let plan = scan.explain_plan(true).await.unwrap(); + assert!( + !plan.contains("LanceRead"), + "{what}: refine must still be served from the index; plan:\n{plan}" + ); + let rows = scan.try_into_batch().await.unwrap().num_rows(); + assert_eq!( + rows, 10, + "{what}: the query must still return its k neighbours" + ); + } + + /// Append rows in the schema `make_covered_test_batches` writes, laid out over the same + /// clusters, so every partition ends up holding existing *and* freshly indexed rows. + async fn append_covered_rows(dataset: &mut Dataset, rows: usize) { + const CLUSTER_SIZE: usize = 8; + let start = dataset.count_all_rows().await.unwrap() as u64; + let ids = Arc::new(UInt64Array::from_iter_values(start..start + rows as u64)); + let payload = Arc::new(UInt64Array::from_iter( + (0..rows as u64).map(|v| if v % 3 == 0 { None } else { Some(v + 5000) }), + )); + let mut flat = Vec::with_capacity(rows * COVERED_DIM); + for row in 0..rows { + let center = (row / CLUSTER_SIZE) as f32 * 4.0; + let within = (row % CLUSTER_SIZE) as f32; + for d in 0..COVERED_DIM { + flat.push(center + within * 0.002 + d as f32 * 0.00001); + } + } + let vectors = Arc::new( + FixedSizeListArray::try_new_from_values(Float32Array::from(flat), COVERED_DIM as i32) + .unwrap(), + ); + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::UInt64, false), + Field::new("payload", DataType::UInt64, true), + Field::new("vector", vectors.data_type().clone(), false), + ])); + let batch = RecordBatch::try_new(schema.clone(), vec![ids, payload, vectors]).unwrap(); + dataset + .append(RecordBatchIterator::new(vec![Ok(batch)], schema), None) + .await + .unwrap(); + } + + async fn write_dataset_from_batches( + test_uri: &str, + schema: Arc, + batches: Vec, + ) -> Dataset { + write_dataset_from_batches_with_max_rows(test_uri, schema, batches, 500).await + } + + async fn write_dataset_from_batches_with_max_rows( + test_uri: &str, + schema: Arc, + batches: Vec, + max_rows_per_file: usize, + ) -> Dataset { + let batches = RecordBatchIterator::new(batches.into_iter().map(Ok), schema); + + let write_params = WriteParams { + max_rows_per_file, + mode: WriteMode::Overwrite, + ..Default::default() + }; + + Dataset::write(batches, test_uri, Some(write_params)) + .await + .unwrap() + } + + async fn prepare_global_ivf_pq( + dataset: &Dataset, + vector_column: &str, + ) -> (IvfBuildParams, PQBuildParams) { + prepare_ivf_pq( + dataset, + vector_column, + TWO_FRAG_DIM, + TWO_FRAG_NUM_PARTITIONS, + TWO_FRAG_NUM_SUBVECTORS, + TWO_FRAG_NUM_BITS, + TWO_FRAG_MAX_ITERS, + TWO_FRAG_SAMPLE_RATE, + ) + .await + } + + #[allow(clippy::too_many_arguments)] + async fn prepare_ivf_pq( + dataset: &Dataset, + vector_column: &str, + expected_dimension: usize, + num_partitions: usize, + num_sub_vectors: usize, + num_bits: usize, + max_iters: u32, + sample_rate: usize, + ) -> (IvfBuildParams, PQBuildParams) { + let batch = dataset + .scan() + .project(&[vector_column.to_string()]) + .unwrap() + .try_into_batch() + .await + .unwrap(); + let vectors = batch + .column_by_name(vector_column) + .expect("vector column should exist") + .as_fixed_size_list(); + + let dim = vectors.value_length() as usize; + assert_eq!(dim, expected_dimension, "unexpected vector dimension"); + + let values = vectors.values().as_primitive::(); + + let kmeans_params = KMeansParams::new(None, max_iters, 1, DistanceType::L2); + let kmeans = + train_kmeans::(values, kmeans_params, dim, num_partitions, sample_rate) + .unwrap(); + + let centroids_flat = kmeans.centroids.as_primitive::().clone(); + let centroids_fsl = + Arc::new(FixedSizeListArray::try_new_from_values(centroids_flat, dim as i32).unwrap()); + let mut ivf_params = + IvfBuildParams::try_with_centroids(num_partitions, centroids_fsl).unwrap(); + ivf_params.max_iters = max_iters as usize; + ivf_params.sample_rate = sample_rate; + + let mut pq_train_params = PQBuildParams::new(num_sub_vectors, num_bits); + pq_train_params.max_iters = max_iters as usize; + pq_train_params.sample_rate = sample_rate; + + let pq = pq_train_params.build(vectors, DistanceType::L2).unwrap(); + let codebook_flat = pq.codebook.values().as_primitive::().clone(); + let pq_codebook: ArrayRef = Arc::new(codebook_flat); + let mut pq_params = PQBuildParams::with_codebook(num_sub_vectors, num_bits, pq_codebook); + pq_params.max_iters = max_iters as usize; + pq_params.sample_rate = sample_rate; + + (ivf_params, pq_params) + } + + /// `prepare_global_ivf_pq` for the covering fixture: same shared training path, but over + /// `COVERED_DIM` vectors and the covering codebook. + async fn prepare_covered_ivf_pq( + dataset: &Dataset, + vector_column: &str, + ) -> (IvfBuildParams, PQBuildParams) { + prepare_ivf_pq( + dataset, + vector_column, + COVERED_DIM, + TWO_FRAG_NUM_PARTITIONS, + COVERED_NUM_SUBVECTORS, + TWO_FRAG_NUM_BITS, + COVERED_MAX_ITERS, + TWO_FRAG_SAMPLE_RATE, + ) + .await + } + + /// `prepare_global_ivf` for the covering fixture. Distributed builds need every shard to + /// share one set of centroids, so the covered non-PQ cases pre-train them here rather than + /// letting each shard train its own. + async fn prepare_covered_ivf(dataset: &Dataset, vector_column: &str) -> IvfBuildParams { + let batch = dataset + .scan() + .project(&[vector_column.to_string()]) + .unwrap() + .try_into_batch() + .await + .unwrap(); + let vectors = batch + .column_by_name(vector_column) + .expect("vector column should exist") + .as_fixed_size_list(); + + let dim = vectors.value_length() as usize; + assert_eq!(dim, COVERED_DIM, "unexpected vector dimension"); + + let values = vectors.values().as_primitive::(); + let kmeans_params = KMeansParams::new(None, COVERED_MAX_ITERS, 1, DistanceType::L2); + let kmeans = train_kmeans::( + values, + kmeans_params, + dim, + TWO_FRAG_NUM_PARTITIONS, + TWO_FRAG_SAMPLE_RATE, + ) + .unwrap(); + + let centroids_flat = kmeans.centroids.as_primitive::().clone(); + let centroids_fsl = + Arc::new(FixedSizeListArray::try_new_from_values(centroids_flat, dim as i32).unwrap()); + let mut ivf_params = + IvfBuildParams::try_with_centroids(TWO_FRAG_NUM_PARTITIONS, centroids_fsl).unwrap(); + ivf_params.max_iters = COVERED_MAX_ITERS as usize; + ivf_params.sample_rate = TWO_FRAG_SAMPLE_RATE; + ivf_params + } + + /// Params for `index_type` over the covering fixture, with IVF (and PQ, where the type + /// needs it) trained from `dataset`. Shared by the covering and carried-vector lifecycle + /// tests so each of them parametrizes over the same quantizer families. + async fn params_for_index_type(dataset: &Dataset, index_type: &str) -> VectorIndexParams { + match index_type { + "IVF_PQ" => { + let (ivf_params, pq_params) = prepare_covered_ivf_pq(dataset, "vector").await; + VectorIndexParams::with_ivf_pq_params(DistanceType::L2, ivf_params, pq_params) + } + "IVF_SQ" => VectorIndexParams::with_ivf_sq_params( + DistanceType::L2, + prepare_covered_ivf(dataset, "vector").await, + SQBuildParams::default(), + ), + "IVF_RQ" => VectorIndexParams::with_ivf_rq_params( + DistanceType::L2, + prepare_covered_ivf(dataset, "vector").await, + RQBuildParams::new(1), + ), + "IVF_FLAT" => VectorIndexParams::with_ivf_flat_params( + DistanceType::L2, + prepare_covered_ivf(dataset, "vector").await, + ), + "IVF_HNSW_PQ" => { + let (ivf_params, pq_params) = prepare_covered_ivf_pq(dataset, "vector").await; + VectorIndexParams::with_ivf_hnsw_pq_params( + DistanceType::L2, + ivf_params, + lightweight_hnsw_params(), + pq_params, + ) + } + "IVF_HNSW_FLAT" => VectorIndexParams::ivf_hnsw( + DistanceType::L2, + prepare_covered_ivf(dataset, "vector").await, + lightweight_hnsw_params(), + ), + "IVF_HNSW_SQ" => VectorIndexParams::with_ivf_hnsw_sq_params( + DistanceType::L2, + prepare_covered_ivf(dataset, "vector").await, + lightweight_hnsw_params(), + SQBuildParams::default(), + ), + other => panic!("unexpected index type {other}"), + } + } + + async fn prepare_global_ivf(dataset: &Dataset, vector_column: &str) -> IvfBuildParams { + let batch = dataset + .scan() + .project(&[vector_column.to_string()]) + .unwrap() + .try_into_batch() + .await + .unwrap(); + let vectors = batch + .column_by_name(vector_column) + .expect("vector column should exist") + .as_fixed_size_list(); + + let dim = vectors.value_length() as usize; + assert_eq!(dim, TWO_FRAG_DIM, "unexpected vector dimension"); + + let values = vectors.values().as_primitive::(); + let kmeans_params = KMeansParams::new(None, TWO_FRAG_MAX_ITERS, 1, DistanceType::L2); + let kmeans = train_kmeans::( + values, + kmeans_params, + dim, + TWO_FRAG_NUM_PARTITIONS, + TWO_FRAG_SAMPLE_RATE, + ) + .unwrap(); + + let centroids_flat = kmeans.centroids.as_primitive::().clone(); + let centroids_fsl = + Arc::new(FixedSizeListArray::try_new_from_values(centroids_flat, dim as i32).unwrap()); + let mut ivf_params = + IvfBuildParams::try_with_centroids(TWO_FRAG_NUM_PARTITIONS, centroids_fsl).unwrap(); + ivf_params.max_iters = TWO_FRAG_MAX_ITERS as usize; + ivf_params.sample_rate = TWO_FRAG_SAMPLE_RATE; + ivf_params + } + + async fn build_segments_for_fragment_groups( + dataset: &mut Dataset, + fragment_groups: Vec>, // each group is a set of fragment ids + params: &VectorIndexParams, + index_name: &str, + ) -> Vec { + let mut segments = Vec::new(); + + for fragments in fragment_groups { + let mut builder = dataset.create_index_builder(&["vector"], IndexType::Vector, params); + builder = builder.name(index_name.to_string()).fragments(fragments); + segments.push(builder.execute_uncommitted().await.unwrap()); + } + + segments + } + + async fn build_ivfpq_for_fragment_groups( + dataset: &mut Dataset, + fragment_groups: Vec>, // each group is a set of fragment ids + ivf_params: &IvfBuildParams, + pq_params: &PQBuildParams, + index_name: &str, + ) { + let params = VectorIndexParams::with_ivf_pq_params( + DistanceType::L2, + ivf_params.clone(), + pq_params.clone(), + ); + + let segments = + build_segments_for_fragment_groups(dataset, fragment_groups, ¶ms, index_name).await; + let committed_segments = + build_distributed_segments(dataset, segments, params.index_type(), index_name).await; + assert!(!committed_segments.is_empty()); + } + + fn assert_centroids_equal(reference: &serde_json::Value, candidate: &serde_json::Value) { + let centroids_a = reference["centroids"] + .as_array() + .expect("centroids should be an array"); + let centroids_b = candidate["centroids"] + .as_array() + .expect("centroids should be an array"); + assert_eq!( + centroids_a.len(), + centroids_b.len(), + "num centroids mismatch", + ); + for (row_a, row_b) in centroids_a.iter().zip(centroids_b.iter()) { + let row_a = row_a + .as_array() + .unwrap_or_else(|| panic!("invalid centroid row: {:?}", row_a)); + let row_b = row_b + .as_array() + .unwrap_or_else(|| panic!("invalid centroid row: {:?}", row_b)); + assert_eq!(row_a.len(), row_b.len(), "centroid dim mismatch"); + for (va, vb) in row_a.iter().zip(row_b.iter()) { + let fa = va.as_f64().expect("centroid must be numeric") as f32; + let fb = vb.as_f64().expect("centroid must be numeric") as f32; + assert!( + (fa - fb).abs() <= 1e-4, + "centroid mismatch: {} vs {}", + fa, + fb + ); + } + } + } + + fn sum_partition_sizes(indices: &[serde_json::Value]) -> Vec { + let mut totals = Vec::new(); + for index in indices { + let partitions = index["partitions"] + .as_array() + .expect("partitions should be an array"); + if totals.is_empty() { + totals.resize(partitions.len(), 0); + } else { + assert_eq!(totals.len(), partitions.len(), "num partitions mismatch"); + } + for (total, partition) in totals.iter_mut().zip(partitions.iter()) { + *total += partition["size"].as_u64().expect("partition size"); + } + } + totals + } + + fn assert_ivf_layout_compatible(stats_a: &serde_json::Value, stats_b: &serde_json::Value) { + let indices_a = stats_a["indices"] + .as_array() + .expect("indices should be an array"); + let indices_b = stats_b["indices"] + .as_array() + .expect("indices should be an array"); + assert!( + !indices_a.is_empty() && !indices_b.is_empty(), + "indices should not be empty", + ); + + let reference = &indices_a[0]; + for index in indices_a.iter().skip(1).chain(indices_b.iter()) { + assert_centroids_equal(reference, index); + } + + let sizes_a = sum_partition_sizes(indices_a); + let sizes_b = sum_partition_sizes(indices_b); + assert_eq!(sizes_a, sizes_b, "aggregated partition sizes mismatch"); + } + + /// Commit caller-defined segment groups as one logical index. + async fn build_distributed_segments( + dataset: &mut Dataset, + segments: Vec, + _index_type: IndexType, + index_name: &str, + ) -> Vec { + dataset + .commit_existing_index_segments(index_name, "vector", segments.clone()) + .await + .unwrap(); + segments + } + + #[tokio::test] + async fn test_ivfpq_recall_performance_on_two_frags_single_vs_split() { + const INDEX_NAME: &str = "vector_idx"; + + let test_dir = TempStrDir::default(); + let base_uri = test_dir.as_str(); + + let (schema, batches) = make_two_fragment_batches(); + + let ds_single_uri = format!("{}/single", base_uri); + let ds_split_uri = format!("{}/split", base_uri); + + let mut ds_single = + write_dataset_from_batches(&ds_single_uri, schema.clone(), batches.clone()).await; + let mut ds_split = write_dataset_from_batches(&ds_split_uri, schema, batches).await; + + let fragments_single = ds_single.get_fragments(); + assert!( + fragments_single.len() >= 2, + "expected at least 2 fragments in ds_single, got {}", + fragments_single.len() + ); + let fragments_split = ds_split.get_fragments(); + assert!( + fragments_split.len() >= 2, + "expected at least 2 fragments in ds_split, got {}", + fragments_split.len() + ); + + let (ivf_params, pq_params) = prepare_global_ivf_pq(&ds_single, "vector").await; + + let group_single = vec![ + fragments_single[0].id() as u32, + fragments_single[1].id() as u32, + ]; + build_ivfpq_for_fragment_groups( + &mut ds_single, + vec![group_single], + &ivf_params, + &pq_params, + INDEX_NAME, + ) + .await; + + let group0 = vec![fragments_split[0].id() as u32]; + let group1 = vec![fragments_split[1].id() as u32]; + build_ivfpq_for_fragment_groups( + &mut ds_split, + vec![group0, group1], + &ivf_params, + &pq_params, + INDEX_NAME, + ) + .await; + + let stats_single_json = ds_single.index_statistics(INDEX_NAME).await.unwrap(); + let stats_split_json = ds_split.index_statistics(INDEX_NAME).await.unwrap(); + let stats_single: serde_json::Value = serde_json::from_str(&stats_single_json).unwrap(); + let stats_split: serde_json::Value = serde_json::from_str(&stats_split_json).unwrap(); + assert_ivf_layout_compatible(&stats_single, &stats_split); + assert_eq!( + stats_single["num_indexed_rows"], + stats_split["num_indexed_rows"] + ); + + const K: usize = 10; + const NUM_QUERIES: usize = 10; + + async fn collect_row_ids(ds: &Dataset, queries: &[Arc]) -> Vec> { + let mut ids_per_query = Vec::with_capacity(queries.len()); + for q in queries { + let result = ds + .scan() + .with_row_id() + .project(&["_rowid"] as &[&str]) + .unwrap() + .nearest("vector", q.as_ref(), K) + .unwrap() + .minimum_nprobes(TWO_FRAG_NUM_PARTITIONS) + .try_into_batch() + .await + .unwrap(); + + let row_ids = result[ROW_ID] + .as_primitive::() + .values() + .iter() + .copied() + .collect::>(); + ids_per_query.push(row_ids); + } + ids_per_query + } + + let query_batch = ds_single + .scan() + .project(&["vector"] as &[&str]) + .unwrap() + .limit(Some(NUM_QUERIES as i64), None) + .unwrap() + .try_into_batch() + .await + .unwrap(); + let vectors = query_batch["vector"].as_fixed_size_list(); + let queries: Vec> = (0..vectors.len()) + .map(|i| vectors.value(i) as Arc) + .collect(); + + let ids_single = collect_row_ids(&ds_single, &queries).await; + let ids_split = collect_row_ids(&ds_split, &queries).await; + + assert_eq!( + ids_single, ids_split, + "single vs split index returned different Top-K row ids", + ); + } + + #[rstest] + #[case::ivf_flat(IndexType::IvfFlat)] + #[case::ivf_pq(IndexType::IvfPq)] + #[case::ivf_sq(IndexType::IvfSq)] + #[case::ivf_rq(IndexType::IvfRq)] + #[tokio::test] + async fn test_distributed_vector_build_commits_multiple_segments_and_preserves_query_results( + #[case] index_type: IndexType, + ) { + const INDEX_NAME: &str = "vector_idx"; + const K: usize = 10; + const NUM_QUERIES: usize = 10; + + let test_dir = TempStrDir::default(); + let base_uri = test_dir.as_str(); + + // Generate the data once, then write it twice to two independent dataset URIs. + let (schema, batches) = make_two_fragment_batches(); + + let ds_single_uri = format!("{}/single", base_uri); + let ds_split_uri = format!("{}/split", base_uri); + + let mut ds_single = + write_dataset_from_batches(&ds_single_uri, schema.clone(), batches.clone()).await; + let mut ds_split = write_dataset_from_batches(&ds_split_uri, schema, batches).await; + + // Ensure we have at least 2 fragments. + let fragments_single = ds_single.get_fragments(); + assert!( + fragments_single.len() >= 2, + "expected at least 2 fragments in ds_single, got {}", + fragments_single.len() + ); + let fragments_split = ds_split.get_fragments(); + assert!( + fragments_split.len() >= 2, + "expected at least 2 fragments in ds_split, got {}", + fragments_split.len() + ); + + let distributed_params = match index_type { + IndexType::IvfFlat => { + let ivf_params = prepare_global_ivf(&ds_single, "vector").await; + VectorIndexParams::with_ivf_flat_params(DistanceType::L2, ivf_params) + } + IndexType::IvfPq => { + let (ivf_params, pq_params) = prepare_global_ivf_pq(&ds_single, "vector").await; + VectorIndexParams::with_ivf_pq_params(DistanceType::L2, ivf_params, pq_params) + } + IndexType::IvfSq => { + let ivf_params = prepare_global_ivf(&ds_single, "vector").await; + VectorIndexParams::with_ivf_sq_params( + DistanceType::L2, + ivf_params, + SQBuildParams::default(), + ) + } + IndexType::IvfRq => { + let ivf_params = prepare_global_ivf(&ds_single, "vector").await; + VectorIndexParams::with_ivf_rq_params( + DistanceType::L2, + ivf_params, + RQBuildParams::with_rotation_type(1, RQRotationType::Fast), + ) + } + other => panic!("unsupported test index type: {}", other), + }; + + ds_single + .create_index( + &["vector"], + IndexType::Vector, + Some(INDEX_NAME.to_string()), + &distributed_params, + true, + ) + .await + .unwrap(); + + let fragment_groups = fragments_split + .iter() + .map(|fragment| vec![fragment.id() as u32]) + .collect::>(); + let expected_segment_count = fragment_groups.len(); + let segments = build_segments_for_fragment_groups( + &mut ds_split, + fragment_groups, + &distributed_params, + INDEX_NAME, + ) + .await; + let segments = + build_distributed_segments(&mut ds_split, segments, index_type, INDEX_NAME).await; + assert_eq!(segments.len(), expected_segment_count); + for segment in &segments { + let segment_index = ds_split + .indices_dir() + .clone() + .join(segment.uuid.to_string()) + .join(crate::index::INDEX_FILE_NAME); + assert!( + ds_split + .object_store + .as_ref() + .exists(&segment_index) + .await + .unwrap(), + "segment file should exist at {}", + segment_index + ); + } + + let committed_segments = ds_split.load_indices_by_name(INDEX_NAME).await.unwrap(); + assert_eq!(committed_segments.len(), expected_segment_count); + for committed in committed_segments { + let covered_fragments = committed + .fragment_bitmap + .as_ref() + .expect("distributed segment should have fragment coverage"); + assert_eq!(covered_fragments.len(), 1); + } + + async fn collect_row_ids(ds: &Dataset, queries: &[Arc]) -> Vec> { + let mut ids_per_query = Vec::with_capacity(queries.len()); + for q in queries { + let result = ds + .scan() + .with_row_id() + .project(&["_rowid"] as &[&str]) + .unwrap() + .nearest("vector", q.as_ref(), K) + .unwrap() + .minimum_nprobes(TWO_FRAG_NUM_PARTITIONS) + .try_into_batch() + .await + .unwrap(); + + let row_ids = result[ROW_ID] + .as_primitive::() + .values() + .iter() + .copied() + .collect::>(); + ids_per_query.push(row_ids); + } + ids_per_query + } + + // Collect a deterministic query set from ds_single. + let query_batch = ds_single + .scan() + .project(&["vector"] as &[&str]) + .unwrap() + .limit(Some(NUM_QUERIES as i64), None) + .unwrap() + .try_into_batch() + .await + .unwrap(); + let vectors = query_batch["vector"].as_fixed_size_list(); + let queries: Vec> = (0..vectors.len()) + .map(|i| vectors.value(i) as Arc) + .collect(); + + let ids_single = collect_row_ids(&ds_single, &queries).await; + let ids_split = collect_row_ids(&ds_split, &queries).await; + + if index_type == IndexType::IvfRq { + for row_ids in &ids_split { + assert_eq!( + row_ids.len(), + K, + "distributed IVF_RQ query should still return exactly {K} row ids", + ); + } + } else { + assert_eq!( + ids_single, ids_split, + "single vs segmented distributed index returned different Top-K row ids", + ); + } + } + + #[rstest] + #[case::ivf_flat(IndexType::IvfFlat)] + #[case::ivf_pq(IndexType::IvfPq)] + #[case::ivf_sq(IndexType::IvfSq)] + #[tokio::test] + async fn test_distributed_vector_grouped_build_allows_concurrent_group_execution( + #[case] index_type: IndexType, + ) { + const INDEX_NAME: &str = "grouped_idx"; + const K: usize = 10; + const NUM_QUERIES: usize = 10; + + let test_dir = TempStrDir::default(); + let base_uri = test_dir.as_str(); + + let (schema, batches) = make_two_fragment_batches(); + let ds_single_uri = format!("{}/grouped_single", base_uri); + let ds_split_uri = format!("{}/grouped_split", base_uri); + + let mut ds_single = + write_dataset_from_batches(&ds_single_uri, schema.clone(), batches.clone()).await; + let mut ds_split = write_dataset_from_batches(&ds_split_uri, schema, batches).await; + + let distributed_params = match index_type { + IndexType::IvfFlat => { + let ivf_params = prepare_global_ivf(&ds_single, "vector").await; + VectorIndexParams::with_ivf_flat_params(DistanceType::L2, ivf_params) + } + IndexType::IvfPq => { + let (ivf_params, pq_params) = prepare_global_ivf_pq(&ds_single, "vector").await; + VectorIndexParams::with_ivf_pq_params(DistanceType::L2, ivf_params, pq_params) + } + IndexType::IvfSq => { + let ivf_params = prepare_global_ivf(&ds_single, "vector").await; + VectorIndexParams::with_ivf_sq_params( + DistanceType::L2, + ivf_params, + SQBuildParams::default(), + ) + } + other => panic!("unsupported test index type: {}", other), + }; + + ds_single + .create_index( + &["vector"], + IndexType::Vector, + Some(INDEX_NAME.to_string()), + &distributed_params, + true, + ) + .await + .unwrap(); + + let fragment_groups = ds_split + .get_fragments() + .into_iter() + .map(|fragment| vec![fragment.id() as u32]) + .collect::>(); + let segments = build_segments_for_fragment_groups( + &mut ds_split, + fragment_groups, + &distributed_params, + INDEX_NAME, + ) + .await; + + assert!(segments.len() >= 4); + let grouped_inputs = segments + .chunks(2) + .map(|group| group.to_vec()) + .collect::>(); + let mut expected_fragment_coverage = grouped_inputs + .iter() + .map(|group| { + group + .iter() + .flat_map(|partial| { + partial + .fragment_bitmap + .as_ref() + .expect("partial shard should have fragment coverage") + .iter() + }) + .sorted() + .collect::>() + }) + .collect::>(); + expected_fragment_coverage.sort(); + + let grouped_segments = futures::future::try_join_all( + grouped_inputs + .into_iter() + .map(|group| ds_split.merge_existing_index_segments(group)), + ) + .await + .unwrap(); + let grouped_segments = + build_distributed_segments(&mut ds_split, grouped_segments, index_type, INDEX_NAME) + .await; + assert_eq!(grouped_segments.len(), expected_fragment_coverage.len()); + let mut actual_fragment_coverage = grouped_segments + .iter() + .map(|segment| { + segment + .fragment_bitmap + .as_ref() + .expect("segment should have fragment coverage") + .iter() + .collect::>() + }) + .collect::>(); + actual_fragment_coverage.sort(); + assert_eq!( + actual_fragment_coverage, expected_fragment_coverage, + "built segment coverage should equal the union of its source partial shards", + ); + + async fn collect_row_ids(ds: &Dataset, queries: &[Arc]) -> Vec> { + let mut ids_per_query = Vec::with_capacity(queries.len()); + for q in queries { + let result = ds + .scan() + .with_row_id() + .project(&["_rowid"] as &[&str]) + .unwrap() + .nearest("vector", q.as_ref(), K) + .unwrap() + .minimum_nprobes(TWO_FRAG_NUM_PARTITIONS) + .try_into_batch() + .await + .unwrap(); + + ids_per_query.push( + result[ROW_ID] + .as_primitive::() + .values() + .iter() + .copied() + .collect(), + ); + } + ids_per_query + } + + let query_batch = ds_single + .scan() + .project(&["vector"] as &[&str]) + .unwrap() + .limit(Some(NUM_QUERIES as i64), None) + .unwrap() + .try_into_batch() + .await + .unwrap(); + let vectors = query_batch["vector"].as_fixed_size_list(); + let queries: Vec> = (0..vectors.len()) + .map(|i| vectors.value(i) as Arc) + .collect(); + + let ids_single = collect_row_ids(&ds_single, &queries).await; + let ids_split = collect_row_ids(&ds_split, &queries).await; + if matches!(index_type, IndexType::IvfSq) { + for (single, split) in ids_single.iter().zip(ids_split.iter()) { + assert_eq!(single.len(), split.len()); + let overlap = single + .iter() + .filter(|row_id| split.contains(row_id)) + .count(); + assert!( + overlap >= K / 3, + "single vs segmented distributed SQ index returned too little top-k overlap", + ); + } + } else { + assert_eq!(ids_single, ids_split); + } + } + + #[tokio::test] + async fn test_distributed_vector_plan_rejects_overlapping_fragment_coverage() { + let test_dir = TempStrDir::default(); + let base_uri = test_dir.as_str(); + let (schema, batches) = make_two_fragment_batches(); + let dataset_uri = format!("{}/overlap_fragments", base_uri); + let mut dataset = write_dataset_from_batches(&dataset_uri, schema, batches).await; + + let fragment = dataset.get_fragments()[0].id() as u32; + let params = VectorIndexParams::with_ivf_flat_params( + DistanceType::L2, + prepare_global_ivf(&dataset, "vector").await, + ); + let mut segments = Vec::new(); + + for _ in 0..2 { + let segment = dataset + .create_index_builder(&["vector"], IndexType::Vector, ¶ms) + .name("vector_idx".to_string()) + .fragments(vec![fragment]) + .execute_uncommitted() + .await + .unwrap(); + segments.push(segment); + } + + let err = dataset + .merge_existing_index_segments(segments) + .await + .unwrap_err(); + assert!(err.to_string().contains("overlapping fragment coverage")); } - async fn append_constant_vector_with_params( - dataset: &mut Dataset, - rows: usize, - template: &[f32], - write_params: Option, - ) { - let start_id = dataset.count_all_rows().await.unwrap() as u64; - append_constant_vector_batch(dataset, rows, template, start_id, write_params).await; - } + #[tokio::test] + async fn test_distributed_vector_build_supports_hnsw_variants() { + let test_dir = TempStrDir::default(); + let base_uri = test_dir.as_str(); + let (schema, batches) = make_two_fragment_batches(); + let dataset_uri = format!("{}/distributed_hnsw_supported", base_uri); + let mut dataset = write_dataset_from_batches(&dataset_uri, schema, batches).await; - async fn append_constant_vector_batch( - dataset: &mut Dataset, - rows: usize, - template: &[f32], - start_id: u64, - write_params: Option, - ) { - assert_eq!( - template.len(), - DIM, - "Template vector should have {} dimensions", - DIM + let fragments = dataset.get_fragments(); + assert!(fragments.len() >= 2); + let params = VectorIndexParams::ivf_hnsw( + DistanceType::L2, + prepare_global_ivf(&dataset, "vector").await, + HnswBuildParams::default(), ); + let mut segments = Vec::new(); - let ids = Arc::new(UInt64Array::from_iter_values( - start_id..start_id + rows as u64, - )); - let mut appended_values = Vec::with_capacity(rows * DIM); - for _ in 0..rows { - appended_values.extend_from_slice(template); + for fragment in fragments.iter().take(2) { + let segment = dataset + .create_index_builder(&["vector"], IndexType::Vector, ¶ms) + .name("vector_idx".to_string()) + .fragments(vec![fragment.id() as u32]) + .execute_uncommitted() + .await + .unwrap(); + segments.push(segment); } - let vectors = Arc::new( - FixedSizeListArray::try_new_from_values( - Float32Array::from(appended_values), - DIM as i32, - ) - .unwrap(), - ); - let schema = Arc::new(Schema::new(vec![ - Field::new("id", DataType::UInt64, false), - Field::new("vector", vectors.data_type().clone(), false), - ])); - let batch = RecordBatch::try_new(schema.clone(), vec![ids, vectors]).unwrap(); - let batches = RecordBatchIterator::new(vec![Ok(batch)], schema); - let params = write_params.map(|mut params| { - params.mode = WriteMode::Append; - params - }); - dataset.append(batches, params).await.unwrap(); - } - #[allow(clippy::too_many_arguments)] - async fn append_and_verify_append_phase( - dataset: &mut Dataset, - index_name: &str, - template: &[f32], - next_id: &mut u64, - rows_to_append: usize, - expected_partitions: usize, - expected_total_rows: usize, - expected_index_count: usize, - expect_split: bool, - ) { - append_constant_vector_with_start_id(dataset, rows_to_append, template, next_id).await; dataset - .optimize_indices(&OptimizeOptions::new()) + .commit_existing_index_segments("vector_idx", "vector", segments) .await .unwrap(); - let stats_json = dataset.index_statistics(index_name).await.unwrap(); - let stats: serde_json::Value = serde_json::from_str(&stats_json).unwrap(); + let query_batch = dataset + .scan() + .project(&["vector"] as &[&str]) + .unwrap() + .limit(Some(4), None) + .unwrap() + .try_into_batch() + .await + .unwrap(); + let q = query_batch["vector"].as_fixed_size_list().value(0); + let result = dataset + .scan() + .project(&["_rowid"] as &[&str]) + .unwrap() + .nearest("vector", q.as_ref(), 5) + .unwrap() + .try_into_batch() + .await + .unwrap(); + assert!(result.num_rows() > 0); + } - let indices = stats["indices"] - .as_array() - .expect("indices array should exist"); - if expect_split { + /// A distributed (sharded, precomputed-IVF) build must support covering ("included") + /// columns: each shard writes the covered columns into its own segment's storage, the + /// segment metadata advertises them as the trailing entries of `fields`, and a covered + /// projection over the committed index skips the base-table take while returning + /// correct, row-aligned covered values. + #[tokio::test] + async fn test_distributed_vector_build_supports_covering_columns() { + let test_dir = TempStrDir::default(); + let base_uri = test_dir.as_str(); + let (schema, batches) = make_covered_test_batches(); + let dataset_uri = format!("{}/distributed_covered", base_uri); + let mut dataset = write_dataset_from_batches(&dataset_uri, schema, batches).await; + + let fragments = dataset.get_fragments(); + assert!(fragments.len() >= 2); + let vector_field_id = dataset.schema().field("vector").unwrap().id; + let id_field_id = dataset.schema().field("id").unwrap().id; + let payload_field_id = dataset.schema().field("payload").unwrap().id; + + let (ivf_params, pq_params) = prepare_covered_ivf_pq(&dataset, "vector").await; + let mut params = + VectorIndexParams::with_ivf_pq_params(DistanceType::L2, ivf_params, pq_params); + params.covering_columns(vec!["id".to_string(), "payload".to_string()]); + + // One covered segment per fragment, over ALL fragments, so the committed index gives + // full coverage -- no uncovered flat delta-scan can mask a covering defect. Each shard + // reads the covering columns from its own fragment subset and writes them into its + // segment's storage. + let mut segments = Vec::new(); + for fragment in fragments.iter() { + let segment = dataset + .create_index_builder(&["vector"], IndexType::Vector, ¶ms) + .name("vec_idx".to_string()) + .fragments(vec![fragment.id() as u32]) + .execute_uncommitted() + .await + .unwrap(); assert_eq!( - indices.len(), - expected_index_count, - "Expected {} index entries after split, got {}, stats: {}", - expected_index_count, - indices.len(), - stats + segment.covering_fields, + vec![id_field_id, payload_field_id], + "distributed covered segment must record both covered columns' field ids" ); - } else { - assert!( - indices.len() >= expected_index_count, - "Expected at least {} index entries after append, got {}, stats: {}", - expected_index_count, - indices.len(), - stats + // Carried columns are a trailing subset of `fields` with the keyed vector field + // first; `IndexMetadata::validate_covering_fields` rejects any other shape. + assert_eq!( + segment.fields, + vec![vector_field_id, id_field_id, payload_field_id], + "covered fields must be the trailing entries of the segment's fields" ); + segments.push(segment); } + + dataset + .commit_existing_index_segments("vec_idx", "vector", segments) + .await + .unwrap(); + + let query_batch = dataset + .scan() + .project(&["vector"] as &[&str]) + .unwrap() + .limit(Some(1), None) + .unwrap() + .try_into_batch() + .await + .unwrap(); + let q = query_batch["vector"].as_fixed_size_list().value(0); + let q = q.as_primitive::(); + + let mut scan = dataset.scan(); + scan.nearest("vector", q, 10).unwrap(); + scan.nprobes(4); + scan.with_row_id(); + scan.project(&["id", "payload"]).unwrap(); + + // The base-table take renders as `LanceRead` in the explained plan, so that -- not + // the string "Take" -- is what its absence has to be asserted on. A flat-KNN + // fallback has no take either, so pin the ANN path first or the assertion below + // could pass without the index being read at all. + let plan = scan.explain_plan(true).await.unwrap(); assert!( - stats["num_indices"].as_u64().unwrap() as usize >= expected_index_count, - "num_indices should be at least {}, stats: {}", - expected_index_count, - stats + plan.contains("ANNSubIndex"), + "the covered query must go through the committed index; plan was:\n{plan}" ); - assert_eq!( - stats["num_indexed_rows"].as_u64().unwrap() as usize, - expected_total_rows, - "Total indexed rows mismatch after append" + assert!( + !plan.contains("LanceRead"), + "covered projection ['id','payload'] should skip the base-table take; plan was:\n{plan}" ); - let base_index = indices - .iter() - .max_by_key(|entry| entry["num_partitions"].as_u64().unwrap_or(0)) - .expect("at least one index entry should exist"); - assert_eq!( - base_index["num_partitions"].as_u64().unwrap() as usize, - expected_partitions, - "Partition count mismatch after append" + let covered = scan.try_into_batch().await.unwrap(); + let row_ids = covered + .column_by_name(ROW_ID) + .expect("row id column") + .as_primitive::(); + assert!(!row_ids.is_empty()); + + // Every covered column -- including the nullable `payload` with its nulls -- must be + // row-aligned with an independent base-table take. Full array equality checks values + // AND validity, so a dropped or shifted null slot fails here; the null count is + // asserted first so that claim cannot rest on an all-valid result set. + assert!( + covered + .column_by_name("payload") + .expect("covered 'payload' column") + .null_count() + > 0, + "the returned rows must include a null covered 'payload'" ); - - if expected_index_count == 1 { - let partitions = base_index["partitions"] - .as_array() - .expect("partitions should exist"); + let row_id_vec: Vec = row_ids.values().to_vec(); + let base = dataset + .take_rows( + &row_id_vec, + crate::dataset::ProjectionRequest::from_columns( + ["id", "payload"], + dataset.schema(), + ), + ) + .await + .unwrap(); + for col in ["id", "payload"] { assert_eq!( - partitions.len(), - expected_partitions, - "Expected {} partitions, found {}", - expected_partitions, - partitions.len() + covered.column_by_name(col).expect("covered column"), + base.column_by_name(col).expect("base-table column"), + "covered '{col}' must match the base-table value (incl. nulls) for each row" ); - let partition_sizes: Vec = partitions - .iter() - .map(|part| part["size"].as_u64().unwrap() as usize) - .collect(); - let total_partition_rows: usize = partition_sizes.iter().sum(); + } + + // Take elision is worthless if the covered search returns the wrong neighbors, so also + // gate on recall against brute-force ground truth (all 4 partitions probed). + let returned: HashSet = row_id_vec.iter().copied().collect(); + let truth = ground_truth(&dataset, "vector", q, 10, DistanceType::L2).await; + let recall = truth.intersection(&returned).count() as f32 / truth.len() as f32; + assert!( + recall >= 0.5, + "covered distributed build recall {recall} < 0.5 (returned {returned:?}, truth {truth:?})" + ); + } + + /// The cross-shard segment *merger* must also carry covering columns: merging covered + /// shards into one unified auxiliary index must retain the covered payload (not drop it at + /// the merger's rebuilt output schema), so a covered projection over the merged index still + /// skips the take and returns correct values. Parametrized over quantizer family so each + /// type's internal-name list (which the merger uses to detect the covering columns) is + /// exercised on the covered-merge path, including IVF_HNSW_PQ for the HNSW-variant lists. + // Note: IVF_RQ (RaBitQ) is intentionally excluded -- distributed RQ *merge* fails + // independently of covering (shards train different `fast_rotation_signs`, so the merger's + // structural-equality check rejects them); covered RQ works on the per-segment-commit path. + /// The HNSW arms build with `lightweight_hnsw_params()`: this test asserts that the merger + /// carries covering columns through each quantizer's internal-name list, which graph quality + /// has no bearing on, and a full-size graph would push the cases past the one-second budget. + /// `hnsw_flat` and `hnsw_sq` are not redundant with `flat`/`sq`: they are the only cases that + /// reach the `IvfHnswFlat` / `IvfHnswSq` arms of `covering_fields_from_shard_schema`. Without + /// them, mapping `IvfHnswSq` to `FLAT_INTERNAL_COLUMNS` would classify `__sq_code` as a + /// covering column -- a duplicate field in the writer schema -- with no test failing. + #[rstest] + #[case::pq("IVF_PQ")] + #[case::sq("IVF_SQ")] + #[case::flat("IVF_FLAT")] + #[case::hnsw_pq("IVF_HNSW_PQ")] + #[case::hnsw_flat("IVF_HNSW_FLAT")] + #[case::hnsw_sq("IVF_HNSW_SQ")] + #[tokio::test] + async fn test_distributed_vector_merge_supports_covering_columns(#[case] index_type: &str) { + let test_dir = TempStrDir::default(); + let base_uri = test_dir.as_str(); + let (schema, batches) = make_covered_test_batches(); + let dataset_uri = format!("{}/distributed_covered_merge_{index_type}", base_uri); + let mut dataset = write_dataset_from_batches(&dataset_uri, schema, batches).await; + + let fragments = dataset.get_fragments(); + assert!(fragments.len() >= 2); + let id_field_id = dataset.schema().field("id").unwrap().id; + let payload_field_id = dataset.schema().field("payload").unwrap().id; + + let mut params = params_for_index_type(&dataset, index_type).await; + params.covering_columns(vec!["id".to_string(), "payload".to_string()]); + + // All fragments, so the merged index gives full coverage (no uncovered delta scan). + let mut segments = Vec::new(); + for fragment in fragments.iter() { + let segment = dataset + .create_index_builder(&["vector"], IndexType::Vector, ¶ms) + .name("vec_idx".to_string()) + .fragments(vec![fragment.id() as u32]) + .execute_uncommitted() + .await + .unwrap(); + segments.push(segment); + } + + // Merge the covered shards into one unified segment, then commit it. + let merged = dataset + .merge_existing_index_segments(segments) + .await + .unwrap(); + assert_eq!( + merged.covering_fields, + vec![id_field_id, payload_field_id], + "merged covered segment must retain both covered columns' field ids" + ); + dataset + .commit_existing_index_segments("vec_idx", "vector", vec![merged]) + .await + .unwrap(); + + let query_batch = dataset + .scan() + .project(&["vector"] as &[&str]) + .unwrap() + .limit(Some(1), None) + .unwrap() + .try_into_batch() + .await + .unwrap(); + let q = query_batch["vector"].as_fixed_size_list().value(0); + let q = q.as_primitive::(); + + let mut scan = dataset.scan(); + scan.nearest("vector", q, 10).unwrap(); + scan.nprobes(4); + scan.with_row_id(); + scan.project(&["id", "payload"]).unwrap(); + + // As above: pin the ANN path, then assert the take (`LanceRead`) is gone. + let plan = scan.explain_plan(true).await.unwrap(); + assert!( + plan.contains("ANNSubIndex"), + "the covered query must go through the merged index; plan was:\n{plan}" + ); + assert!( + !plan.contains("LanceRead"), + "covered projection over the merged index should skip the base-table take; \ + plan was:\n{plan}" + ); + + let covered = scan.try_into_batch().await.unwrap(); + let row_ids = covered + .column_by_name(ROW_ID) + .expect("row id column") + .as_primitive::(); + assert!(!row_ids.is_empty()); + + // Both covered columns (incl. the nullable `payload` with its nulls) must be + // row-aligned with an independent base-table take over the merged index. Full array + // equality checks validity too, and the null count is asserted first so that claim + // cannot rest on an all-valid result set. + assert!( + covered + .column_by_name("payload") + .expect("covered 'payload' column") + .null_count() + > 0, + "the returned rows must include a null covered 'payload'" + ); + let row_id_vec: Vec = row_ids.values().to_vec(); + let base = dataset + .take_rows( + &row_id_vec, + crate::dataset::ProjectionRequest::from_columns( + ["id", "payload"], + dataset.schema(), + ), + ) + .await + .unwrap(); + for col in ["id", "payload"] { assert_eq!( - total_partition_rows, expected_total_rows, - "Partition sizes should sum to total rows: {:?}", - partition_sizes - ); - } else { - assert!( - !expect_split, - "Split should result in a single merged index" + covered.column_by_name(col).expect("covered column"), + base.column_by_name(col).expect("base-table column"), + "covered '{col}' over the merged index must match the base-table value for each row" ); } - assert_eq!( - dataset.count_all_rows().await.unwrap(), - expected_total_rows, - "Dataset row count mismatch after append" + // Gate on recall too, so a merge that returns the wrong neighbors is caught. + let returned: HashSet = row_id_vec.iter().copied().collect(); + let truth = ground_truth(&dataset, "vector", q, 10, DistanceType::L2).await; + let recall = truth.intersection(&returned).count() as f32 / truth.len() as f32; + assert!( + recall >= 0.5, + "covered merged {index_type} recall {recall} < 0.5 (returned {returned:?}, truth {truth:?})" ); } - async fn load_partition_row_ids(index: &IvfPq, partition_idx: usize) -> Vec { - index - .storage - .load_partition(partition_idx, None) - .await - .unwrap() - .row_ids() - .copied() - .collect() - } + /// Distributed builds must carry refine vectors too, on both shapes: committing one + /// segment per shard, and merging the shards into one unified segment first. + /// + /// The merge shape is the one with its own machinery: the cross-shard merger classifies + /// a shard's carried columns by *excluding* the storage's internal names, so carried + /// vectors -- which sit under the indexed column's own name at rest -- are only picked + /// up because that name is not internal. `payload` is covered alongside so the + /// projection is served from the index and any base-table read the plan shows can only + /// be refine's own. + #[rstest] + #[case::per_segment_commit(false)] + #[case::cross_shard_merge(true)] + #[tokio::test] + async fn test_distributed_build_carries_refine_vectors(#[case] merge_shards: bool) { + let test_dir = TempStrDir::default(); + let (schema, batches) = make_covered_test_batches(); + let query = batches[0]["vector"].as_fixed_size_list().value(0); + let uri = format!( + "{}/distributed_refine_{}", + test_dir.as_str(), + if merge_shards { "merged" } else { "segments" } + ); + let mut dataset = write_dataset_from_batches(&uri, schema, batches).await; - async fn load_flat_partition_row_ids(index: &IvfFlatIndex, partition_idx: usize) -> Vec { - index - .storage - .load_partition(partition_idx, None) - .await - .unwrap() - .row_ids() - .copied() - .collect() - } + let fragments = dataset.get_fragments(); + assert!( + fragments.len() >= 2, + "need several shards to distribute over" + ); + let vector_field_id = dataset.schema().field("vector").unwrap().id; + let payload_field_id = dataset.schema().field("payload").unwrap().id; - async fn delete_ids(dataset: &mut Dataset, ids: &[u64]) { - if ids.is_empty() { - return; + let (ivf_params, pq_params) = prepare_covered_ivf_pq(&dataset, "vector").await; + let mut params = + VectorIndexParams::with_ivf_pq_params(DistanceType::L2, ivf_params, pq_params); + params.covering_columns(vec!["payload".to_string()]); + params.store_vectors_for_refine(true); + + // One shard per fragment, over every fragment, so the committed index gives full + // coverage and no uncovered delta scan can mask a defect. + let mut segments = Vec::new(); + for fragment in fragments.iter() { + let segment = dataset + .create_index_builder(&["vector"], IndexType::Vector, ¶ms) + .name("vec_idx".to_string()) + .fragments(vec![fragment.id() as u32]) + .execute_uncommitted() + .await + .unwrap(); + assert_eq!( + segment.covering_fields, + vec![payload_field_id, vector_field_id], + "each shard must declare the covering column and the carried vectors" + ); + segments.push(segment); } - let predicate = ids - .iter() - .map(|x| x.to_string()) - .collect::>() - .join(","); + + let to_commit = match merge_shards { + true => { + let merged = dataset + .merge_existing_index_segments(segments) + .await + .unwrap(); + assert_eq!( + merged.covering_fields, + vec![payload_field_id, vector_field_id], + "the unified segment must retain the carried vectors" + ); + vec![merged] + } + false => segments, + }; dataset - .delete(&format!("id in ({})", predicate)) + .commit_existing_index_segments("vec_idx", "vector", to_commit) .await .unwrap(); - } - async fn compact_after_deletions(dataset: &mut Dataset) { - compact_files( - dataset, - CompactionOptions { - materialize_deletions_threshold: 0.0, - ..Default::default() - }, - None, - ) - .await - .unwrap(); + assert_refine_served_from_index(&dataset, query.as_ref(), "distributed").await; } - async fn ground_truth( - dataset: &Dataset, - column: &str, - query: &dyn Array, - k: usize, - distance_type: DistanceType, - ) -> HashSet { - let batch = dataset - .scan() - .with_row_id() - .nearest(column, query, k) - .unwrap() - .distance_metric(distance_type) - .use_index(false) - .try_into_batch() + /// A merged covered segment must itself be usable as an input to a later merge -- + /// the hierarchical/incremental distributed workflow. The merger requires every + /// covered shard to carry its source field ids, so the merged output has to carry + /// them too; otherwise the merger rejects its own product and tells the user to + /// rebuild a shard that only the merger can produce. + #[tokio::test] + async fn test_distributed_covered_merge_output_can_be_merged_again() { + let test_dir = TempStrDir::default(); + let base_uri = test_dir.as_str(); + let (schema, batches) = make_covered_test_batches(); + let dataset_uri = format!("{}/distributed_covered_remerge", base_uri); + let mut dataset = write_dataset_from_batches(&dataset_uri, schema, batches).await; + + let fragments = dataset.get_fragments(); + assert!( + fragments.len() >= 3, + "need >= 3 fragments so the second-stage merge has more than one input" + ); + // IVF_FLAT: a merged *PQ* segment stores transposed codes and the merger requires + // row-major shards, so PQ cannot be re-merged for reasons unrelated to covering. + // FLAT and SQ can, which is where the missing stamp actually bit. + let mut params = VectorIndexParams::with_ivf_flat_params( + DistanceType::L2, + prepare_covered_ivf(&dataset, "vector").await, + ); + params.covering_columns(vec!["payload".to_string()]); + + let mut segments = Vec::new(); + for fragment in fragments.iter() { + segments.push( + dataset + .create_index_builder(&["vector"], IndexType::Vector, ¶ms) + .name("vec_idx".to_string()) + .fragments(vec![fragment.id() as u32]) + .execute_uncommitted() + .await + .unwrap(), + ); + } + + // Stage 1: merge the first two shards. + let rest = segments.split_off(2); + let merged = dataset + .merge_existing_index_segments(segments) + .await + .unwrap(); + // Stage 2: feed that merged segment back in alongside the remaining shards. This + // is the hierarchical workflow the missing stamp used to make impossible. + let mut stage_two = vec![merged]; + stage_two.extend(rest); + let remerged = dataset + .merge_existing_index_segments(stage_two) + .await + .expect("a merged covered segment must be mergeable again"); + assert_eq!( + remerged.covering_fields, + vec![dataset.schema().field("payload").unwrap().id], + "the re-merged segment must still declare the covered field" + ); + dataset + .commit_existing_index_segments("vec_idx", "vector", vec![remerged]) .await .unwrap(); - batch[ROW_ID] - .as_primitive::() - .values() - .iter() - .copied() - .collect() - } - - fn multivec_ground_truth( - vectors: &ListArray, - query: &dyn Array, - k: usize, - distance_type: DistanceType, - ) -> Vec<(f32, u64)> { - let query = if let Some(list_array) = query.as_list_opt::() { - list_array.values().clone() - } else { - query.as_fixed_size_list().values().clone() - }; - multivec_distance(&query, vectors, distance_type) - .unwrap() - .into_iter() - .enumerate() - .map(|(i, dist)| (dist, i as u64)) - .sorted_by(|a, b| a.0.total_cmp(&b.0)) - .take(k) - .collect() } - const TWO_FRAG_NUM_ROWS: usize = 2000; - const TWO_FRAG_DIM: usize = 128; - const TWO_FRAG_NUM_PARTITIONS: usize = 4; - const TWO_FRAG_NUM_SUBVECTORS: usize = 16; - const TWO_FRAG_NUM_BITS: usize = 8; - const TWO_FRAG_SAMPLE_RATE: usize = 7; - const TWO_FRAG_MAX_ITERS: u32 = 20; + /// Same covering column *name and type* across shards, but two different logical + /// fields: `payload` is dropped and re-added between the two shard builds, so it + /// comes back with a fresh field id. Name+type comparison accepts this and + /// `concat_batches` would then stack shard A's values for the dropped field under + /// the re-added one -- covered queries over shard A's rows serving the old column. + /// + /// Carried columns live in `fields`, so the two shards disagree on `fields` as well as + /// on `covering_fields`, and the segment-level check rejects them before the auxiliary + /// merger runs. The merger's own `COVERING_FIELD_IDS_KEY` comparison guards the same + /// case one layer down, for drivers calling `merge_partial_vector_auxiliary_files` + /// directly; what this test pins is that the dataset API never lets the corruption + /// through. + #[tokio::test] + async fn test_distributed_vector_merge_rejects_covering_of_different_field_ids() { + use crate::dataset::NewColumnTransform; - fn make_two_fragment_batches() -> (Arc, Vec) { - let ids = Arc::new(UInt64Array::from_iter_values(0..TWO_FRAG_NUM_ROWS as u64)); + let test_dir = TempStrDir::default(); + let base_uri = test_dir.as_str(); - let values = generate_random_array_with_range(TWO_FRAG_NUM_ROWS * TWO_FRAG_DIM, 0.0..1.0); + let payload = Arc::new(UInt64Array::from_iter_values(0..TWO_FRAG_NUM_ROWS as u64)); + let values = generate_random_array_with_range(TWO_FRAG_NUM_ROWS * COVERED_DIM, 0.0..1.0); let vectors = Arc::new( - FixedSizeListArray::try_new_from_values( - Float32Array::from(values), - TWO_FRAG_DIM as i32, - ) - .unwrap(), + FixedSizeListArray::try_new_from_values(Float32Array::from(values), COVERED_DIM as i32) + .unwrap(), ); - let schema = Arc::new(Schema::new(vec![ - Field::new("id", DataType::UInt64, false), + Field::new("payload", DataType::UInt64, true), Field::new("vector", vectors.data_type().clone(), false), ])); - let batch = RecordBatch::try_new(schema.clone(), vec![ids, vectors]).unwrap(); + let batch = RecordBatch::try_new(schema.clone(), vec![payload, vectors]).unwrap(); + let dataset_uri = format!("{}/distributed_covered_merge_field_id", base_uri); + let mut dataset = write_dataset_from_batches(&dataset_uri, schema, vec![batch]).await; - (schema, vec![batch]) - } + let fragments = dataset.get_fragments(); + assert!(fragments.len() >= 2); - async fn write_dataset_from_batches( - test_uri: &str, - schema: Arc, - batches: Vec, - ) -> Dataset { - write_dataset_from_batches_with_max_rows(test_uri, schema, batches, 500).await - } + let (ivf_params, pq_params) = prepare_covered_ivf_pq(&dataset, "vector").await; + let mut params = + VectorIndexParams::with_ivf_pq_params(DistanceType::L2, ivf_params, pq_params); + params.covering_columns(vec!["payload".to_string()]); - async fn write_dataset_from_batches_with_max_rows( - test_uri: &str, - schema: Arc, - batches: Vec, - max_rows_per_file: usize, - ) -> Dataset { - let batches = RecordBatchIterator::new(batches.into_iter().map(Ok), schema); + // Shard A covers `payload` as it exists now. + let segment_a = dataset + .create_index_builder(&["vector"], IndexType::Vector, ¶ms) + .name("vec_idx".to_string()) + .fragments(vec![fragments[0].id() as u32]) + .execute_uncommitted() + .await + .unwrap(); + let original_field_id = dataset.schema().field("payload").unwrap().id; - let write_params = WriteParams { - max_rows_per_file, - mode: WriteMode::Overwrite, - ..Default::default() - }; + // Drop and re-add `payload`: same name, same type, new field id. The re-add is + // metadata-only (AllNulls), so no data file changes and nothing else notices. + dataset.drop_columns(&["payload"]).await.unwrap(); + dataset + .add_columns( + NewColumnTransform::AllNulls(Arc::new(arrow_schema::Schema::new(vec![ + Field::new("payload", DataType::UInt64, true), + ]))), + None, + None, + ) + .await + .unwrap(); + let new_field_id = dataset.schema().field("payload").unwrap().id; + assert_ne!( + original_field_id, new_field_id, + "re-adding the column must produce a new field id, or this test proves nothing" + ); - Dataset::write(batches, test_uri, Some(write_params)) + let fragments = dataset.get_fragments(); + let segment_b = dataset + .create_index_builder(&["vector"], IndexType::Vector, ¶ms) + .name("vec_idx".to_string()) + .fragments(vec![fragments[1].id() as u32]) + .execute_uncommitted() .await - .unwrap() - } + .unwrap(); - async fn prepare_global_ivf_pq( - dataset: &Dataset, - vector_column: &str, - ) -> (IvfBuildParams, PQBuildParams) { - prepare_ivf_pq( - dataset, - vector_column, - TWO_FRAG_DIM, - TWO_FRAG_NUM_PARTITIONS, - TWO_FRAG_NUM_SUBVECTORS, - TWO_FRAG_NUM_BITS, - TWO_FRAG_MAX_ITERS, - TWO_FRAG_SAMPLE_RATE, - ) - .await + let err = dataset + .merge_existing_index_segments(vec![segment_a, segment_b]) + .await + .expect_err("merging shards covering different field ids must fail"); + let msg = err.to_string(); + assert!( + msg.contains("identical fields"), + "error should describe the covered-field disagreement; was: {msg}" + ); } - #[allow(clippy::too_many_arguments)] - async fn prepare_ivf_pq( - dataset: &Dataset, - vector_column: &str, - expected_dimension: usize, - num_partitions: usize, - num_sub_vectors: usize, - num_bits: usize, - max_iters: u32, - sample_rate: usize, - ) -> (IvfBuildParams, PQBuildParams) { - let batch = dataset - .scan() - .project(&[vector_column.to_string()]) - .unwrap() - .try_into_batch() + #[tokio::test] + async fn test_optimize_refuses_to_restamp_rebound_covering() { + use crate::dataset::NewColumnTransform; + + let test_dir = TempStrDir::default(); + let base_uri = test_dir.as_str(); + let (schema, batches) = make_covered_test_batches(); + let query = batches[0]["vector"].as_fixed_size_list().value(0); + let dataset_uri = format!("{}/optimize_refuses_restamp", base_uri); + let mut dataset = write_dataset_from_batches(&dataset_uri, schema, batches).await; + + let all_fragments: Vec = dataset + .get_fragments() + .iter() + .map(|fragment| fragment.id() as u32) + .collect(); + + let (ivf_params, pq_params) = prepare_covered_ivf_pq(&dataset, "vector").await; + let mut params = + VectorIndexParams::with_ivf_pq_params(DistanceType::L2, ivf_params, pq_params); + params.covering_columns(vec!["payload".to_string()]); + + let mut segment = dataset + .create_index_builder(&["vector"], IndexType::Vector, ¶ms) + .name("vec_idx".to_string()) + .fragments(all_fragments) + .execute_uncommitted() + .await + .unwrap(); + let original_field_id = dataset.schema().field("payload").unwrap().id; + + // Drop and re-add `payload`: same name, same type, new field id. The re-add is + // metadata-only (AllNulls), so no data file changes and no fragment looks stale. + dataset.drop_columns(&["payload"]).await.unwrap(); + dataset + .add_columns( + NewColumnTransform::AllNulls(Arc::new(arrow_schema::Schema::new(vec![ + Field::new("payload", DataType::UInt64, true), + ]))), + None, + None, + ) .await .unwrap(); - let vectors = batch - .column_by_name(vector_column) - .expect("vector column should exist") - .as_fixed_size_list(); + let new_field_id = dataset.schema().field("payload").unwrap().id; + assert_ne!( + original_field_id, new_field_id, + "re-adding the column must produce a new field id, or this test proves nothing" + ); - let dim = vectors.value_length() as usize; - assert_eq!(dim, expected_dimension, "unexpected vector dimension"); + // The segment keeps its true build version, so the staleness pass runs and passes; + // only the id rebind is wrong. + for id in segment.fields.iter_mut() { + if *id == original_field_id { + *id = new_field_id; + } + } + for id in segment.covering_fields.iter_mut() { + if *id == original_field_id { + *id = new_field_id; + } + } - let values = vectors.values().as_primitive::(); + dataset + .commit_existing_index_segments("vec_idx", "vector", vec![segment]) + .await + .unwrap(); - let kmeans_params = KMeansParams::new(None, max_iters, 1, DistanceType::L2); - let kmeans = - train_kmeans::(values, kmeans_params, dim, num_partitions, sample_rate) - .unwrap(); + // Before optimize the mismatch is visible and the read path does the right thing: + // the stamped source id disagrees with the declaration, so payload comes from the + // base table and is all null. + let mut scan = dataset.scan(); + scan.nearest("vector", query.as_ref(), 10).unwrap(); + scan.nprobes(TWO_FRAG_NUM_PARTITIONS); + scan.project(&["payload"]).unwrap(); + let batch = scan.try_into_batch().await.unwrap(); + let payload = batch["payload"].as_primitive::(); + assert_eq!( + payload.null_count(), + payload.len(), + "precondition: the rebound declaration must fall back to the base table" + ); - let centroids_flat = kmeans.centroids.as_primitive::().clone(); - let centroids_fsl = - Arc::new(FixedSizeListArray::try_new_from_values(centroids_flat, dim as i32).unwrap()); - let mut ivf_params = - IvfBuildParams::try_with_centroids(num_partitions, centroids_fsl).unwrap(); - ivf_params.max_iters = max_iters as usize; - ivf_params.sample_rate = sample_rate; + // A rebuild copies that payload by name and would re-derive the stamp from the + // CURRENT schema, making the declaration and the stamp agree -- which would start + // serving the old field's values and elide the base-table read. Refuse instead. + let err = dataset + .optimize_indices(&OptimizeOptions::new()) + .await + .expect_err("optimize must not re-stamp a covering payload onto a different field id"); + let msg = err.to_string(); + assert!( + msg.contains("covering values stamped with source field ids") + && msg.contains(&original_field_id.to_string()) + && msg.contains(&new_field_id.to_string()), + "error must name both the stored and the current ids; got: {msg}" + ); - let mut pq_train_params = PQBuildParams::new(num_sub_vectors, num_bits); - pq_train_params.max_iters = max_iters as usize; - pq_train_params.sample_rate = sample_rate; + // And the index is left in the state the read path already handles. + let mut scan = dataset.scan(); + scan.nearest("vector", query.as_ref(), 10).unwrap(); + scan.nprobes(TWO_FRAG_NUM_PARTITIONS); + scan.project(&["payload"]).unwrap(); + let batch = scan.try_into_batch().await.unwrap(); + let payload = batch["payload"].as_primitive::(); + assert_eq!( + payload.null_count(), + payload.len(), + "after the refusal the payload must still come from the base table" + ); + } - let pq = pq_train_params.build(vectors, DistanceType::L2).unwrap(); - let codebook_flat = pq.codebook.values().as_primitive::().clone(); - let pq_codebook: ArrayRef = Arc::new(codebook_flat); - let mut pq_params = PQBuildParams::with_codebook(num_sub_vectors, num_bits, pq_codebook); - pq_params.max_iters = max_iters as usize; - pq_params.sample_rate = sample_rate; + /// Refine vectors are as wide as the vector column itself, so reading them with every + /// partition probed would undo the bound that reading covering per survivor exists to + /// give: codes are scanned for every row on every probe, refine values are needed for + /// at most `k`. They are storage-owned but deliberately NOT part of the per-partition + /// read set -- the distinction `DEFERRED_INTERNAL_COLUMNS` exists to draw. + #[tokio::test] + async fn test_refine_vectors_are_not_read_with_every_partition() { + let test_dir = TempStrDir::default(); + let (schema, batches) = make_covered_test_batches(); + let uri = format!("{}/refine_not_per_partition", test_dir.as_str()); + let mut dataset = write_dataset_from_batches(&uri, schema, batches).await; - (ivf_params, pq_params) - } + let (ivf_params, pq_params) = prepare_covered_ivf_pq(&dataset, "vector").await; + let mut params = + VectorIndexParams::with_ivf_pq_params(DistanceType::L2, ivf_params, pq_params); + params.store_vectors_for_refine(true); + dataset + .create_index( + &["vector"], + IndexType::Vector, + Some("vec_idx".into()), + ¶ms, + true, + ) + .await + .unwrap(); - async fn prepare_global_ivf(dataset: &Dataset, vector_column: &str) -> IvfBuildParams { - let batch = dataset - .scan() - .project(&[vector_column.to_string()]) - .unwrap() - .try_into_batch() + let ctx = load_vector_index_context(&dataset, "vector", "vec_idx").await; + let storage = ctx + .ivf() + .load_partition_storage(0, PartitionColumns::Internal, None) .await .unwrap(); - let vectors = batch - .column_by_name(vector_column) - .expect("vector column should exist") - .as_fixed_size_list(); + let loaded: Vec = storage + .batch() + .schema() + .fields() + .iter() + .map(|f| f.name().clone()) + .collect(); + assert!( + !loaded.iter().any(|n| n == "vector"), + "a per-partition load must leave the carried vectors on disk; loaded {loaded:?}" + ); - let dim = vectors.value_length() as usize; - assert_eq!(dim, TWO_FRAG_DIM, "unexpected vector dimension"); + // They are covering, so they are read for survivors instead -- which is the whole + // point: codes are scanned for every row on every probe, carried vectors are needed + // for at most `k`. + let physical: Vec = ctx + .ivf() + .physical_covering_fields() + .unwrap() + .iter() + .map(|(_, f)| f.name().clone()) + .collect(); + assert_eq!( + physical, + vec!["vector".to_string()], + "carried vectors must be servable as covering, just not per partition" + ); + } - let values = vectors.values().as_primitive::(); - let kmeans_params = KMeansParams::new(None, TWO_FRAG_MAX_ITERS, 1, DistanceType::L2); - let kmeans = train_kmeans::( - values, - kmeans_params, - dim, - TWO_FRAG_NUM_PARTITIONS, - TWO_FRAG_SAMPLE_RATE, - ) - .unwrap(); + /// Covering columns and refine vectors must coexist. Covering is classified by + /// name-exclusion against the storage's own internal set and then matched against the + /// `covering_field_ids` stamp by *arity*; a refine column counted as covering makes the + /// two disagree, and `physical_covering_fields_from_schema` answers "no covering at all" + /// on a mismatch -- silently withdrawing a payload the index really does carry. + #[tokio::test] + async fn test_refine_vectors_do_not_withdraw_covering_columns() { + let test_dir = TempStrDir::default(); + let (schema, batches) = make_covered_test_batches(); + let uri = format!("{}/refine_plus_covering", test_dir.as_str()); + let mut dataset = write_dataset_from_batches(&uri, schema, batches).await; - let centroids_flat = kmeans.centroids.as_primitive::().clone(); - let centroids_fsl = - Arc::new(FixedSizeListArray::try_new_from_values(centroids_flat, dim as i32).unwrap()); - let mut ivf_params = - IvfBuildParams::try_with_centroids(TWO_FRAG_NUM_PARTITIONS, centroids_fsl).unwrap(); - ivf_params.max_iters = TWO_FRAG_MAX_ITERS as usize; - ivf_params.sample_rate = TWO_FRAG_SAMPLE_RATE; - ivf_params + let (ivf_params, pq_params) = prepare_covered_ivf_pq(&dataset, "vector").await; + let mut params = + VectorIndexParams::with_ivf_pq_params(DistanceType::L2, ivf_params, pq_params); + params.covering_columns(vec!["payload".to_string()]); + params.store_vectors_for_refine(true); + dataset + .create_index( + &["vector"], + IndexType::Vector, + Some("vec_idx".into()), + ¶ms, + true, + ) + .await + .unwrap(); + + let ctx = load_vector_index_context(&dataset, "vector", "vec_idx").await; + let physical = ctx.ivf().physical_covering_fields().unwrap(); + let names: Vec = physical.iter().map(|(_, f)| f.name().clone()).collect(); + assert_eq!( + names, + vec!["payload".to_string(), "vector".to_string()], + "both the declared covering column and the carried vectors must be servable" + ); } - async fn build_segments_for_fragment_groups( - dataset: &mut Dataset, - fragment_groups: Vec>, // each group is a set of fragment ids - params: &VectorIndexParams, - index_name: &str, - ) -> Vec { - let mut segments = Vec::new(); + /// The point of carrying refine vectors: `refine` re-ranks from index storage instead + /// of taking full-precision vectors from the base table. + /// + /// Asserted against a control rather than on its own. "No base-table read" is only + /// meaningful if the same query *does* read the base table without the option -- a bare + /// absence assertion passes for any number of unrelated reasons. + #[tokio::test] + async fn test_refine_reads_vectors_from_the_index_not_the_base_table() { + async fn plan_for(store_vectors: bool, dir: &TempStrDir, name: &str) -> (String, usize) { + let (schema, batches) = make_covered_test_batches(); + let uri = format!("{}/{}", dir.as_str(), name); + let query = batches[0]["vector"].as_fixed_size_list().value(0); + let mut dataset = write_dataset_from_batches(&uri, schema, batches).await; + + let (ivf_params, pq_params) = prepare_covered_ivf_pq(&dataset, "vector").await; + let mut params = + VectorIndexParams::with_ivf_pq_params(DistanceType::L2, ivf_params, pq_params); + // `payload` is covered so the index can serve the user's projection; the only + // thing that could still need the base table is refine's own vector fetch. + params.covering_columns(vec!["payload".to_string()]); + params.store_vectors_for_refine(store_vectors); + dataset + .create_index( + &["vector"], + IndexType::Vector, + Some("vec_idx".into()), + ¶ms, + true, + ) + .await + .unwrap(); - for fragments in fragment_groups { - let mut builder = dataset.create_index_builder(&["vector"], IndexType::Vector, params); - builder = builder.name(index_name.to_string()).fragments(fragments); - segments.push(builder.execute_uncommitted().await.unwrap()); + let mut scan = dataset.scan(); + scan.nearest("vector", query.as_ref(), 10).unwrap(); + scan.nprobes(TWO_FRAG_NUM_PARTITIONS); + scan.refine(2); + // Project only `payload`: the user never asks for the vector column, so nothing + // but refine itself needs it. This is the shape the narrowing would strip. + scan.project(&["payload"]).unwrap(); + let plan = scan.explain_plan(true).await.unwrap(); + let rows = scan.try_into_batch().await.unwrap().num_rows(); + (plan, rows) } - segments - } + let test_dir = TempStrDir::default(); + let (plain_plan, plain_rows) = plan_for(false, &test_dir, "refine_plain").await; + let (carried_plan, carried_rows) = plan_for(true, &test_dir, "refine_carried").await; - async fn build_ivfpq_for_fragment_groups( - dataset: &mut Dataset, - fragment_groups: Vec>, // each group is a set of fragment ids - ivf_params: &IvfBuildParams, - pq_params: &PQBuildParams, - index_name: &str, - ) { - let params = VectorIndexParams::with_ivf_pq_params( - DistanceType::L2, - ivf_params.clone(), - pq_params.clone(), + // Control: without the option, refine must fetch vectors from the base table. The + // covered `payload` is served from the index in both variants, so a base-table read + // here can only be refine's. + assert!( + plain_plan.contains("LanceRead"), + "control: refine without carried vectors must read the base table; plan:\n{plain_plan}" + ); + assert!( + !carried_plan.contains("LanceRead"), + "refine must be served from the index when vectors are carried; plan:\n{carried_plan}" + ); + assert_eq!( + plain_rows, carried_rows, + "carrying vectors must not change the result count" ); - - let segments = - build_segments_for_fragment_groups(dataset, fragment_groups, ¶ms, index_name).await; - let committed_segments = - build_distributed_segments(dataset, segments, params.index_type(), index_name).await; - assert!(!committed_segments.is_empty()); } - fn assert_centroids_equal(reference: &serde_json::Value, candidate: &serde_json::Value) { - let centroids_a = reference["centroids"] - .as_array() - .expect("centroids should be an array"); - let centroids_b = candidate["centroids"] - .as_array() - .expect("centroids should be an array"); - assert_eq!( - centroids_a.len(), - centroids_b.len(), - "num centroids mismatch", + /// Carrying vectors buys nothing on a multivector index, so it is refused rather than + /// silently doubling the index. + /// + /// The refine step runs *after* `MultivectorScoringExec` has re-grouped sub-vectors back + /// to rows, so it scores the row-shaped `List` column. Index storage holds one + /// `FixedSizeList` per sub-vector -- the shape the search itself needs -- which cannot + /// substitute for it. A second copy of the widest column in the table, readable by + /// nothing, is worth an error rather than a surprise. + #[tokio::test] + async fn test_store_vectors_for_refine_is_rejected_on_multivector() { + let test_dir = TempStrDir::default(); + let (mut dataset, _) = + generate_multivec_test_dataset::(test_dir.as_str(), 0.0..1.0).await; + + // Multivector requires cosine. + let mut params = VectorIndexParams::ivf_pq(4, 8, 4, DistanceType::Cosine, 2); + params.store_vectors_for_refine(true); + + let err = dataset + .create_index( + &["vector"], + IndexType::Vector, + Some("vec_idx".into()), + ¶ms, + true, + ) + .await + .expect_err("carrying vectors on a multivector index must be refused"); + let msg = err.to_string(); + assert!( + msg.contains("store_vectors_for_refine") && msg.contains("multivector"), + "the error must name the option and why it does not apply; got: {msg}" ); - for (row_a, row_b) in centroids_a.iter().zip(centroids_b.iter()) { - let row_a = row_a - .as_array() - .unwrap_or_else(|| panic!("invalid centroid row: {:?}", row_a)); - let row_b = row_b - .as_array() - .unwrap_or_else(|| panic!("invalid centroid row: {:?}", row_b)); - assert_eq!(row_a.len(), row_b.len(), "centroid dim mismatch"); - for (va, vb) in row_a.iter().zip(row_b.iter()) { - let fa = va.as_f64().expect("centroid must be numeric") as f32; - let fb = vb.as_f64().expect("centroid must be numeric") as f32; - assert!( - (fa - fb).abs() <= 1e-4, - "centroid mismatch: {} vs {}", - fa, - fb - ); - } - } } - fn sum_partition_sizes(indices: &[serde_json::Value]) -> Vec { - let mut totals = Vec::new(); - for index in indices { - let partitions = index["partitions"] - .as_array() - .expect("partitions should be an array"); - if totals.is_empty() { - totals.resize(partitions.len(), 0); - } else { - assert_eq!(totals.len(), partitions.len(), "num partitions mismatch"); - } - for (total, partition) in totals.iter_mut().zip(partitions.iter()) { - *total += partition["size"].as_u64().expect("partition size"); - } - } - totals + /// Only the V3 covering-aware storages keep an extra column row-aligned with the code. + /// A legacy build ignores the option outright, but `CreateIndexBuilder` still records + /// the vector in `covering_fields`, so the index would advertise a payload its storage + /// never wrote and every query would quietly fall back to the base table. + #[tokio::test] + async fn test_store_vectors_for_refine_requires_v3() { + let test_dir = TempStrDir::default(); + let (schema, batches) = make_covered_test_batches(); + let uri = format!("{}/refine_requires_v3", test_dir.as_str()); + let mut dataset = write_dataset_from_batches(&uri, schema, batches).await; + + let (ivf_params, pq_params) = prepare_covered_ivf_pq(&dataset, "vector").await; + let mut params = + VectorIndexParams::with_ivf_pq_params(DistanceType::L2, ivf_params, pq_params); + params.store_vectors_for_refine(true); + params.version(IndexFileVersion::Legacy); + + let err = dataset + .create_index( + &["vector"], + IndexType::Vector, + Some("vec_idx".into()), + ¶ms, + true, + ) + .await + .expect_err("a legacy build cannot carry refine vectors"); + let msg = err.to_string(); + assert!( + msg.contains("store_vectors_for_refine") && msg.contains("V3"), + "the error must name the option and the version it needs; got: {msg}" + ); } - fn assert_ivf_layout_compatible(stats_a: &serde_json::Value, stats_b: &serde_json::Value) { - let indices_a = stats_a["indices"] - .as_array() - .expect("indices should be an array"); - let indices_b = stats_b["indices"] - .as_array() - .expect("indices should be an array"); + /// A flat quantizer is an identity copy of the vector column, and the flat IVF + /// transformer adds no residual step, so under L2 the carried copy is byte-identical to + /// the `FLAT_COLUMN` the index already stores -- and `refine` re-ranking flat distances + /// is a no-op, since those distances are already exact. Carrying vectors there doubles + /// the index for nothing, so it joins the other unproductive configurations the + /// validator refuses rather than being silently accepted. + #[rstest] + #[case::flat("IVF_FLAT")] + #[case::hnsw_flat("IVF_HNSW_FLAT")] + #[tokio::test] + async fn test_store_vectors_for_refine_rejects_flat_quantizers(#[case] index_type: &str) { + let test_dir = TempStrDir::default(); + let (schema, batches) = make_covered_test_batches(); + let uri = format!("{}/refine_rejects_{index_type}", test_dir.as_str()); + let mut dataset = write_dataset_from_batches(&uri, schema, batches).await; + + let mut params = params_for_index_type(&dataset, index_type).await; + params.store_vectors_for_refine(true); + + let err = dataset + .create_index( + &["vector"], + IndexType::Vector, + Some("vec_idx".into()), + ¶ms, + true, + ) + .await + .expect_err( + "a flat quantizer already stores the vectors; carrying them again \ + must be refused", + ); + let msg = err.to_string(); assert!( - !indices_a.is_empty() && !indices_b.is_empty(), - "indices should not be empty", + msg.contains("store_vectors_for_refine") && msg.contains("flat"), + "the error must name the option and the quantizer; got: {msg}" ); + } - let reference = &indices_a[0]; - for index in indices_a.iter().skip(1).chain(indices_b.iter()) { - assert_centroids_equal(reference, index); - } + /// Carried vectors are stored as a top-level covering column and read back through the + /// same top-level-only resolution every other covering column uses, so a nested + /// indexed column has nowhere to put them. Plain indexing of such a column still works, + /// so the refusal has to be the option's, not the column's. + #[tokio::test] + async fn test_store_vectors_for_refine_rejects_nested_column() { + use arrow_array::StructArray; - let sizes_a = sum_partition_sizes(indices_a); - let sizes_b = sum_partition_sizes(indices_b); - assert_eq!(sizes_a, sizes_b, "aggregated partition sizes mismatch"); - } + const ROWS: usize = 512; + const NDIM: usize = 32; + let test_dir = TempStrDir::default(); - /// Commit caller-defined segment groups as one logical index. - async fn build_distributed_segments( - dataset: &mut Dataset, - segments: Vec, - _index_type: IndexType, - index_name: &str, - ) -> Vec { + let values = generate_random_array_with_range::(ROWS * NDIM, 0.0..1.0); + let vectors = + Arc::new(FixedSizeListArray::try_new_from_values(values, NDIM as i32).unwrap()); + let inner = Field::new("embedding", vectors.data_type().clone(), false); + let data = Arc::new(StructArray::from(vec![( + Arc::new(inner), + vectors as ArrayRef, + )])); + let schema = Arc::new(Schema::new(vec![Field::new( + "data", + data.data_type().clone(), + false, + )])); + let batch = RecordBatch::try_new(schema.clone(), vec![data]).unwrap(); + let uri = format!("{}/refine_nested_column", test_dir.as_str()); + let mut dataset = write_dataset_from_batches(&uri, schema, vec![batch]).await; + + let mut params = VectorIndexParams::ivf_pq(4, 8, 4, DistanceType::L2, 2); dataset - .commit_existing_index_segments(index_name, "vector", segments.clone()) + .create_index( + &["data.embedding"], + IndexType::Vector, + Some("plain_idx".into()), + ¶ms, + true, + ) .await - .unwrap(); - segments + .expect("precondition: a nested vector column indexes normally"); + + params.store_vectors_for_refine(true); + let err = dataset + .create_index( + &["data.embedding"], + IndexType::Vector, + Some("refine_idx".into()), + ¶ms, + true, + ) + .await + .expect_err("carrying vectors for a nested indexed column must be refused"); + let msg = err.to_string(); + assert!( + msg.contains("store_vectors_for_refine") && msg.contains("data.embedding"), + "the error must name the option and the column; got: {msg}" + ); } + /// The carried copy travels under a reserved internal name until the transform chain + /// has consumed the indexed column. An indexed column already using that name collides + /// with the copy mid-flight, and the classifier that decides what storage carries as + /// covering payload never counts it either. #[tokio::test] - async fn test_ivfpq_recall_performance_on_two_frags_single_vs_split() { - const INDEX_NAME: &str = "vector_idx"; + async fn test_store_vectors_for_refine_rejects_reserved_column_name() { + use lance_index::vector::storage::REFINE_VECTOR_COLUMN; + const ROWS: usize = 512; + const NDIM: usize = 32; let test_dir = TempStrDir::default(); - let base_uri = test_dir.as_str(); - let (schema, batches) = make_two_fragment_batches(); + let values = generate_random_array_with_range::(ROWS * NDIM, 0.0..1.0); + let vectors = + Arc::new(FixedSizeListArray::try_new_from_values(values, NDIM as i32).unwrap()); + let schema = Arc::new(Schema::new(vec![Field::new( + REFINE_VECTOR_COLUMN, + vectors.data_type().clone(), + false, + )])); + let batch = RecordBatch::try_new(schema.clone(), vec![vectors]).unwrap(); + let uri = format!("{}/refine_reserved_name", test_dir.as_str()); + let mut dataset = write_dataset_from_batches(&uri, schema, vec![batch]).await; - let ds_single_uri = format!("{}/single", base_uri); - let ds_split_uri = format!("{}/split", base_uri); + let mut params = VectorIndexParams::ivf_pq(4, 8, 4, DistanceType::L2, 2); + params.store_vectors_for_refine(true); + let err = dataset + .create_index( + &[REFINE_VECTOR_COLUMN], + IndexType::Vector, + Some("vec_idx".into()), + ¶ms, + true, + ) + .await + .expect_err("an indexed column using the scratch name must be refused"); + let msg = err.to_string(); + assert!( + msg.contains("store_vectors_for_refine") && msg.contains(REFINE_VECTOR_COLUMN), + "the error must name the option and the reserved name; got: {msg}" + ); + } - let mut ds_single = - write_dataset_from_batches(&ds_single_uri, schema.clone(), batches.clone()).await; - let mut ds_split = write_dataset_from_batches(&ds_split_uri, schema, batches).await; + /// Precomputed shuffle buffers arrive already quantized and without the raw vector + /// column, so there is nothing left to copy aside by the time the build reads them -- + /// exactly the reason covering columns are refused with them. + #[tokio::test] + async fn test_store_vectors_for_refine_rejects_precomputed_shuffle_buffers() { + use object_store::path::Path; - let fragments_single = ds_single.get_fragments(); + let test_dir = TempStrDir::default(); + let (schema, batches) = make_covered_test_batches(); + let uri = format!("{}/refine_precomputed_buffers", test_dir.as_str()); + let mut dataset = write_dataset_from_batches(&uri, schema, batches).await; + + let (mut ivf_params, pq_params) = prepare_covered_ivf_pq(&dataset, "vector").await; + ivf_params.precomputed_shuffle_buffers = + Some((Path::from("shuffle/data"), vec!["part0".to_string()])); + let mut params = + VectorIndexParams::with_ivf_pq_params(DistanceType::L2, ivf_params, pq_params); + params.store_vectors_for_refine(true); + + let err = dataset + .create_index( + &["vector"], + IndexType::Vector, + Some("vec_idx".into()), + ¶ms, + true, + ) + .await + .expect_err("carrying vectors from precomputed buffers must be refused"); + let msg = err.to_string(); assert!( - fragments_single.len() >= 2, - "expected at least 2 fragments in ds_single, got {}", - fragments_single.len() + msg.contains("store_vectors_for_refine") && msg.contains("precomputed shuffle"), + "the error must name the option and the incompatible source; got: {msg}" ); - let fragments_split = ds_split.get_fragments(); + } + + /// Optimize can split an oversized partition, which re-streams the affected rows + /// through the transform chain on a different path from the initial build. That path + /// must carry the vectors too, or the rebuilt partitions lose them while the declared + /// storage schema still names the column. + #[tokio::test] + async fn test_partition_split_preserves_carried_vectors() { + const INDEX_NAME: &str = "vector_idx"; + // IVF_PQ splits above MAX_PARTITION_SIZE_FACTOR * 8192 = 32_768 rows. + const BASE_ROWS: usize = 512; + const APPEND_ROWS: usize = 33_000; + // Two clusters, but only one is grown past the split threshold. + let offsets = [-50.0, 50.0]; + + let test_dir = TempStrDir::default(); + let (batch, schema) = generate_clustered_batch(BASE_ROWS, offsets); + let batches = RecordBatchIterator::new(vec![Ok(batch)], schema.clone()); + let mut dataset = Dataset::write( + batches, + test_dir.as_str(), + Some(WriteParams { + mode: WriteMode::Overwrite, + ..Default::default() + }), + ) + .await + .unwrap(); + + let centroids = build_centroids_for_offsets(&offsets); + let ivf_params = IvfBuildParams::try_with_centroids(2, centroids).unwrap(); + let pq_params = PQBuildParams::new(4, 8); + let mut params = + VectorIndexParams::with_ivf_pq_params(DistanceType::L2, ivf_params, pq_params); + params.store_vectors_for_refine(true); + dataset + .create_index( + &["vector"], + IndexType::Vector, + Some(INDEX_NAME.to_string()), + ¶ms, + true, + ) + .await + .unwrap(); + + // Grow only the first cluster, so exactly one partition crosses the threshold. + let mut template = vec![0.0; DIM]; + template[0] = offsets[0]; + append_partition_templates(&mut dataset, APPEND_ROWS, &[template]).await; + + dataset + .optimize_indices(&OptimizeOptions::new()) + .await + .unwrap(); + + let ctx = load_vector_index_context(&dataset, "vector", INDEX_NAME).await; assert!( - fragments_split.len() >= 2, - "expected at least 2 fragments in ds_split, got {}", - fragments_split.len() + ctx.num_partitions() > 2, + "precondition: the partition must actually have split, stats: {}", + ctx.stats_json() ); + for p in 0..ctx.num_partitions() { + let storage = ctx + .ivf() + .load_partition_storage(p, PartitionColumns::All, None) + .await + .unwrap(); + assert!( + storage.batch().column_by_name("vector").is_some(), + "partition {p} lost its carried vectors across the split" + ); + } + } - let (ivf_params, pq_params) = prepare_global_ivf_pq(&ds_single, "vector").await; + /// An explicit merge hands `StorageBuilder::build` two sources at once: batches read + /// back from existing storage, where the carried vectors sit under the indexed column's + /// own name, and freshly shuffled batches, where they are still under the in-flight + /// scratch name. The two must agree on one name, or the merge fails outright. + /// + /// The split test above cannot catch this: it grows a single cluster, so the affected + /// partition is served entirely by the split shuffle reader and the untouched + /// partitions receive no fresh rows -- the two sources never meet in one build. + #[rstest] + #[case::pq("IVF_PQ")] + #[case::sq("IVF_SQ")] + #[case::rq("IVF_RQ")] + #[case::hnsw_pq("IVF_HNSW_PQ")] + #[case::hnsw_sq("IVF_HNSW_SQ")] + #[tokio::test] + async fn test_merge_keeps_carried_vectors(#[case] index_type: &str) { + const INDEX_NAME: &str = "vec_idx"; + const APPENDED_ROWS: usize = 800; - let group_single = vec![ - fragments_single[0].id() as u32, - fragments_single[1].id() as u32, - ]; - build_ivfpq_for_fragment_groups( - &mut ds_single, - vec![group_single], - &ivf_params, - &pq_params, - INDEX_NAME, - ) - .await; + let test_dir = TempStrDir::default(); + let (schema, batches) = make_covered_test_batches(); + let query = batches[0]["vector"].as_fixed_size_list().value(0); + let uri = format!( + "{}/merge_keeps_carried_vectors_{index_type}", + test_dir.as_str() + ); + let mut dataset = write_dataset_from_batches(&uri, schema, batches).await; - let group0 = vec![fragments_split[0].id() as u32]; - let group1 = vec![fragments_split[1].id() as u32]; - build_ivfpq_for_fragment_groups( - &mut ds_split, - vec![group0, group1], - &ivf_params, - &pq_params, - INDEX_NAME, - ) - .await; + let mut params = params_for_index_type(&dataset, index_type).await; + // `payload` is covered so the projection is served from the index too, which is what + // lets the plan assertion attribute any base-table read to refine alone. + params.covering_columns(vec!["payload".to_string()]); + params.store_vectors_for_refine(true); + dataset + .create_index( + &["vector"], + IndexType::Vector, + Some(INDEX_NAME.to_string()), + ¶ms, + true, + ) + .await + .unwrap(); - let stats_single_json = ds_single.index_statistics(INDEX_NAME).await.unwrap(); - let stats_split_json = ds_split.index_statistics(INDEX_NAME).await.unwrap(); - let stats_single: serde_json::Value = serde_json::from_str(&stats_single_json).unwrap(); - let stats_split: serde_json::Value = serde_json::from_str(&stats_split_json).unwrap(); - assert_ivf_layout_compatible(&stats_single, &stats_split); + append_covered_rows(&mut dataset, APPENDED_ROWS).await; + dataset + .optimize_indices(&OptimizeOptions::merge(1)) + .await + .unwrap(); + + let ctx = load_vector_index_context(&dataset, "vector", INDEX_NAME).await; + // Both preconditions matter: one segment means the appended rows were folded into + // the segment the merge read from, rather than left as a fresh delta -- which is the + // only arrangement that puts stored and freshly shuffled batches in one storage + // build, the thing this test exists to exercise. assert_eq!( - stats_single["num_indexed_rows"], - stats_split["num_indexed_rows"] + ctx.stats()["num_indices"].as_u64().unwrap(), + 1, + "precondition: the merge must leave a single segment; stats: {}", + ctx.stats_json() + ); + assert_eq!( + ctx.stats()["num_indexed_rows"].as_u64().unwrap() as usize, + TWO_FRAG_NUM_ROWS + APPENDED_ROWS, + "precondition: the merged segment must cover every row; stats: {}", + ctx.stats_json() ); + assert_refine_served_from_index(&dataset, query.as_ref(), index_type).await; + } - const K: usize = 10; - const NUM_QUERIES: usize = 10; + /// The counterpart to the split test above. Joining an undersized partition takes its + /// rows from the *base table* and reassigns them, on a path that never runs the + /// transform chain and so never sees the carried vectors. Those reassigned rows still + /// meet the untouched partitions' stored copies in one storage build, so they have to + /// arrive carrying the vectors too. + #[rstest] + #[case::pq("IVF_PQ")] + #[case::sq("IVF_SQ")] + #[case::rq("IVF_RQ")] + #[case::hnsw_pq("IVF_HNSW_PQ")] + #[case::hnsw_sq("IVF_HNSW_SQ")] + #[tokio::test] + async fn test_partition_join_keeps_carried_vectors(#[case] index_type: &str) { + const INDEX_NAME: &str = "vec_idx"; + const APPENDED_ROWS: usize = 200; - async fn collect_row_ids(ds: &Dataset, queries: &[Arc]) -> Vec> { - let mut ids_per_query = Vec::with_capacity(queries.len()); - for q in queries { - let result = ds - .scan() - .with_row_id() - .project(&["_rowid"] as &[&str]) - .unwrap() - .nearest("vector", q.as_ref(), K) - .unwrap() - .minimum_nprobes(TWO_FRAG_NUM_PARTITIONS) - .try_into_batch() - .await - .unwrap(); + let test_dir = TempStrDir::default(); + let (schema, batches) = make_covered_test_batches(); + let query = batches[0]["vector"].as_fixed_size_list().value(0); + let uri = format!( + "{}/join_keeps_carried_vectors_{index_type}", + test_dir.as_str() + ); + let mut dataset = write_dataset_from_batches(&uri, schema, batches).await; - let row_ids = result[ROW_ID] - .as_primitive::() - .values() - .iter() - .copied() - .collect::>(); - ids_per_query.push(row_ids); - } - ids_per_query - } + let mut params = params_for_index_type(&dataset, index_type).await; + // `payload` is covered so the projection is served from the index too, which is what + // lets the plan assertion attribute any base-table read to refine alone. + params.covering_columns(vec!["payload".to_string()]); + params.store_vectors_for_refine(true); + dataset + .create_index( + &["vector"], + IndexType::Vector, + Some(INDEX_NAME.to_string()), + ¶ms, + true, + ) + .await + .unwrap(); - let query_batch = ds_single - .scan() - .project(&["vector"] as &[&str]) - .unwrap() - .limit(Some(NUM_QUERIES as i64), None) - .unwrap() - .try_into_batch() + // Every partition holds far fewer than `MIN_PARTITION_SIZE_PERCENT` of IVF_PQ's + // target partition size, so the optimize below joins one away. + append_covered_rows(&mut dataset, APPENDED_ROWS).await; + dataset + .optimize_indices(&OptimizeOptions::new()) .await .unwrap(); - let vectors = query_batch["vector"].as_fixed_size_list(); - let queries: Vec> = (0..vectors.len()) - .map(|i| vectors.value(i) as Arc) - .collect(); - - let ids_single = collect_row_ids(&ds_single, &queries).await; - let ids_split = collect_row_ids(&ds_split, &queries).await; + let ctx = load_vector_index_context(&dataset, "vector", INDEX_NAME).await; + assert!( + ctx.num_partitions() < TWO_FRAG_NUM_PARTITIONS, + "precondition: a partition must actually have been joined away, stats: {}", + ctx.stats_json() + ); assert_eq!( - ids_single, ids_split, - "single vs split index returned different Top-K row ids", + ctx.stats()["num_indexed_rows"].as_u64().unwrap() as usize, + TWO_FRAG_NUM_ROWS + APPENDED_ROWS, + "the reassigned rows must all survive the join; stats: {}", + ctx.stats_json() ); + assert_refine_served_from_index(&dataset, query.as_ref(), index_type).await; } + /// A retrain rebuilds every segment from scratch and re-declares the covering payload + /// from committed metadata, where carried refine vectors appear as the indexed column + /// itself. That declaration travels back through `VectorIndexParams`, whose covering + /// setter is the user-facing one and rejects the indexed column outright, so it has to + /// be split back out into the flag before validation ever sees it. #[rstest] - #[case::ivf_flat(IndexType::IvfFlat)] - #[case::ivf_pq(IndexType::IvfPq)] - #[case::ivf_sq(IndexType::IvfSq)] - #[case::ivf_rq(IndexType::IvfRq)] + #[case::pq("IVF_PQ")] + #[case::sq("IVF_SQ")] + #[case::rq("IVF_RQ")] + #[case::hnsw_pq("IVF_HNSW_PQ")] + #[case::hnsw_sq("IVF_HNSW_SQ")] #[tokio::test] - async fn test_distributed_vector_build_commits_multiple_segments_and_preserves_query_results( - #[case] index_type: IndexType, - ) { - const INDEX_NAME: &str = "vector_idx"; - const K: usize = 10; - const NUM_QUERIES: usize = 10; + async fn test_retrain_keeps_carried_vectors(#[case] index_type: &str) { + const INDEX_NAME: &str = "vec_idx"; + const APPENDED_ROWS: usize = 200; let test_dir = TempStrDir::default(); - let base_uri = test_dir.as_str(); - - // Generate the data once, then write it twice to two independent dataset URIs. - let (schema, batches) = make_two_fragment_batches(); - - let ds_single_uri = format!("{}/single", base_uri); - let ds_split_uri = format!("{}/split", base_uri); - - let mut ds_single = - write_dataset_from_batches(&ds_single_uri, schema.clone(), batches.clone()).await; - let mut ds_split = write_dataset_from_batches(&ds_split_uri, schema, batches).await; - - // Ensure we have at least 2 fragments. - let fragments_single = ds_single.get_fragments(); - assert!( - fragments_single.len() >= 2, - "expected at least 2 fragments in ds_single, got {}", - fragments_single.len() + let (schema, batches) = make_covered_test_batches(); + let query = batches[0]["vector"].as_fixed_size_list().value(0); + let uri = format!( + "{}/retrain_keeps_carried_vectors_{index_type}", + test_dir.as_str() ); - let fragments_split = ds_split.get_fragments(); - assert!( - fragments_split.len() >= 2, - "expected at least 2 fragments in ds_split, got {}", - fragments_split.len() - ); - - let distributed_params = match index_type { - IndexType::IvfFlat => { - let ivf_params = prepare_global_ivf(&ds_single, "vector").await; - VectorIndexParams::with_ivf_flat_params(DistanceType::L2, ivf_params) - } - IndexType::IvfPq => { - let (ivf_params, pq_params) = prepare_global_ivf_pq(&ds_single, "vector").await; - VectorIndexParams::with_ivf_pq_params(DistanceType::L2, ivf_params, pq_params) - } - IndexType::IvfSq => { - let ivf_params = prepare_global_ivf(&ds_single, "vector").await; - VectorIndexParams::with_ivf_sq_params( - DistanceType::L2, - ivf_params, - SQBuildParams::default(), - ) - } - IndexType::IvfRq => { - let ivf_params = prepare_global_ivf(&ds_single, "vector").await; - VectorIndexParams::with_ivf_rq_params( - DistanceType::L2, - ivf_params, - RQBuildParams::with_rotation_type(1, RQRotationType::Fast), - ) - } - other => panic!("unsupported test index type: {}", other), - }; + let mut dataset = write_dataset_from_batches(&uri, schema, batches).await; - ds_single + let mut params = params_for_index_type(&dataset, index_type).await; + // `payload` is covered so the projection is served from the index too, which is what + // lets the plan assertion attribute any base-table read to refine alone. + params.covering_columns(vec!["payload".to_string()]); + params.store_vectors_for_refine(true); + dataset .create_index( &["vector"], IndexType::Vector, Some(INDEX_NAME.to_string()), - &distributed_params, + ¶ms, true, ) .await .unwrap(); - let fragment_groups = fragments_split + append_covered_rows(&mut dataset, APPENDED_ROWS).await; + dataset + .optimize_indices(&OptimizeOptions::retrain()) + .await + .unwrap(); + + let ctx = load_vector_index_context(&dataset, "vector", INDEX_NAME).await; + assert_eq!( + ctx.stats()["num_indexed_rows"].as_u64().unwrap() as usize, + TWO_FRAG_NUM_ROWS + APPENDED_ROWS, + "a retrain must rebuild over the whole dataset; stats: {}", + ctx.stats_json() + ); + + // The declaration has to survive too: dropping it would leave storage carrying a + // payload no query can reach. + let indices = dataset.load_indices().await.unwrap(); + let index = indices .iter() - .map(|fragment| vec![fragment.id() as u32]) - .collect::>(); - let expected_segment_count = fragment_groups.len(); - let segments = build_segments_for_fragment_groups( - &mut ds_split, - fragment_groups, - &distributed_params, - INDEX_NAME, - ) - .await; - let segments = - build_distributed_segments(&mut ds_split, segments, index_type, INDEX_NAME).await; - assert_eq!(segments.len(), expected_segment_count); - for segment in &segments { - let segment_index = ds_split - .indices_dir() - .clone() - .join(segment.uuid.to_string()) - .join(crate::index::INDEX_FILE_NAME); - assert!( - ds_split - .object_store - .as_ref() - .exists(&segment_index) - .await - .unwrap(), - "segment file should exist at {}", - segment_index - ); - } + .find(|idx| idx.name == INDEX_NAME) + .expect("index must still exist"); + let payload_field_id = dataset.schema().field("payload").unwrap().id; + let vector_field_id = dataset.schema().field("vector").unwrap().id; + assert_eq!( + index.covering_fields, + vec![payload_field_id, vector_field_id], + "the rebuilt segment must still declare both the covering column and the \ + carried vectors" + ); - let committed_segments = ds_split.load_indices_by_name(INDEX_NAME).await.unwrap(); - assert_eq!(committed_segments.len(), expected_segment_count); - for committed in committed_segments { - let covered_fragments = committed - .fragment_bitmap - .as_ref() - .expect("distributed segment should have fragment coverage"); - assert_eq!(covered_fragments.len(), 1); - } + assert_refine_served_from_index(&dataset, query.as_ref(), index_type).await; + } - async fn collect_row_ids(ds: &Dataset, queries: &[Arc]) -> Vec> { - let mut ids_per_query = Vec::with_capacity(queries.len()); - for q in queries { - let result = ds - .scan() - .with_row_id() - .project(&["_rowid"] as &[&str]) - .unwrap() - .nearest("vector", q.as_ref(), K) - .unwrap() - .minimum_nprobes(TWO_FRAG_NUM_PARTITIONS) - .try_into_batch() - .await - .unwrap(); + /// The carried copy must be the user's RAW vectors -- the whole design rests on the copy + /// being taken *before* the IVF/quantizer transform chain, which rewrites the indexed + /// column in place (residual under L2 and cosine, normalized first under cosine). + /// + /// Every other refine test asserts plan shape and row count, which a copy taken one step + /// too late would satisfy perfectly while `refine` re-ranked against residuals. This one + /// reads the values back and compares them to the source data. Because `id` is covered + /// too, the whole projection is served from index storage -- the absent `LanceRead` is + /// what proves the vectors came from the index rather than the base table, so the + /// equality below is an assertion about what the *index* holds. + /// + /// Cosine is the case with no positive coverage at all until now, and the one the design + /// note calls irrecoverable: normalization discards magnitude, so if the copy were taken + /// after it, these vectors would come back unit-length. Recall is asserted only for L2 -- + /// this fixture's rows are near-constant across dimensions, so every row points in + /// almost the same direction and cosine ranking on it would be meaningless. + #[rstest] + #[case::l2(DistanceType::L2)] + #[case::cosine(DistanceType::Cosine)] + #[tokio::test] + async fn test_carried_vectors_are_the_raw_source_values(#[case] distance_type: DistanceType) { + const INDEX_NAME: &str = "vec_idx"; + const K: usize = 10; + const CLUSTER_SIZE: u64 = 8; - let row_ids = result[ROW_ID] - .as_primitive::() - .values() - .iter() - .copied() - .collect::>(); - ids_per_query.push(row_ids); - } - ids_per_query - } + let test_dir = TempStrDir::default(); + let (schema, batches) = make_covered_test_batches(); + let source = batches[0]["vector"].as_fixed_size_list().clone(); + let query = source.value(0); + let uri = format!("{}/carried_raw_{distance_type}", test_dir.as_str()); + let mut dataset = write_dataset_from_batches(&uri, schema, batches).await; - // Collect a deterministic query set from ds_single. - let query_batch = ds_single - .scan() - .project(&["vector"] as &[&str]) - .unwrap() - .limit(Some(NUM_QUERIES as i64), None) - .unwrap() - .try_into_batch() + let (ivf_params, pq_params) = prepare_covered_ivf_pq(&dataset, "vector").await; + let mut params = + VectorIndexParams::with_ivf_pq_params(distance_type, ivf_params, pq_params); + // `id` covered so the projection below needs no base-table take either; without it + // the plan reads the base table for `id` and proves nothing about the vectors. + params.covering_columns(vec!["id".to_string()]); + params.store_vectors_for_refine(true); + dataset + .create_index( + &["vector"], + IndexType::Vector, + Some(INDEX_NAME.to_string()), + ¶ms, + true, + ) .await .unwrap(); - let vectors = query_batch["vector"].as_fixed_size_list(); - let queries: Vec> = (0..vectors.len()) - .map(|i| vectors.value(i) as Arc) - .collect(); - let ids_single = collect_row_ids(&ds_single, &queries).await; - let ids_split = collect_row_ids(&ds_split, &queries).await; + let mut scan = dataset.scan(); + scan.nearest("vector", query.as_ref(), K).unwrap(); + scan.nprobes(TWO_FRAG_NUM_PARTITIONS); + scan.refine(2); + scan.project(&["vector", "id"]).unwrap(); + let plan = scan.explain_plan(true).await.unwrap(); + assert!( + !plan.contains("LanceRead"), + "precondition: the projection must be served from the index, or the values \ + below say nothing about what the index carries; plan:\n{plan}" + ); - if index_type == IndexType::IvfRq { - for row_ids in &ids_split { - assert_eq!( - row_ids.len(), - K, - "distributed IVF_RQ query should still return exactly {K} row ids", - ); - } - } else { + let batch = scan.try_into_batch().await.unwrap(); + assert_eq!( + batch.num_rows(), + K, + "the query must return its k neighbours" + ); + let ids = batch["id"].as_primitive::(); + let got = batch["vector"].as_fixed_size_list(); + for row in 0..batch.num_rows() { + let id = ids.value(row) as usize; + let expected = source.value(id); + let expected = expected.as_primitive::(); + let actual = got.value(row); + let actual = actual.as_primitive::(); assert_eq!( - ids_single, ids_split, - "single vs segmented distributed index returned different Top-K row ids", + actual.values(), + expected.values(), + "carried vector for id {id} is not the source vector -- the copy was taken \ + after the transform chain (residual, or normalized under cosine)" + ); + } + + if distance_type == DistanceType::L2 { + // The query is row 0, whose true neighbours are its own cluster. + let hits = (0..batch.num_rows()) + .filter(|row| ids.value(*row) < CLUSTER_SIZE) + .count(); + assert!( + hits * 2 >= CLUSTER_SIZE as usize, + "recall below 0.5: only {hits} of the query's {CLUSTER_SIZE}-row cluster came \ + back in the top {K}" ); } } + /// Renaming an indexed column stays legal wherever the index stores no copy of it + /// *by name*, and the renamed index must still survive both maintenance paths. + /// + /// This is the counterpart to + /// `test_alter_of_indexed_column_refused_with_carried_refine`: that one pins the refusal + /// for a refine-carrying index, this one pins that the refusal did not spread to indexes + /// it must not cover. A plain index stores only `_rowid`/codes/part-id, and a covering + /// index adds the *covered* columns -- in neither case is the indexed column present + /// under its own name, so a rename leaves nothing stale and the field id the index is + /// keyed on survives it. #[rstest] - #[case::ivf_flat(IndexType::IvfFlat)] - #[case::ivf_pq(IndexType::IvfPq)] - #[case::ivf_sq(IndexType::IvfSq)] + #[case::plain(false)] + #[case::covering_payload(true)] #[tokio::test] - async fn test_distributed_vector_grouped_build_allows_concurrent_group_execution( - #[case] index_type: IndexType, - ) { - const INDEX_NAME: &str = "grouped_idx"; - const K: usize = 10; - const NUM_QUERIES: usize = 10; + async fn test_rename_indexed_column_survives_maintenance(#[case] covered: bool) { + use crate::dataset::ColumnAlteration; let test_dir = TempStrDir::default(); - let base_uri = test_dir.as_str(); + let (schema, batches) = make_covered_test_batches(); + let uri = format!("{}/rename_survives_{covered}", test_dir.as_str()); + let mut dataset = write_dataset_from_batches(&uri, schema, batches).await; - let (schema, batches) = make_two_fragment_batches(); - let ds_single_uri = format!("{}/grouped_single", base_uri); - let ds_split_uri = format!("{}/grouped_split", base_uri); + let (ivf_params, pq_params) = prepare_covered_ivf_pq(&dataset, "vector").await; + let mut params = + VectorIndexParams::with_ivf_pq_params(DistanceType::L2, ivf_params, pq_params); + if covered { + params.covering_columns(vec!["payload".to_string()]); + } + dataset + .create_index( + &["vector"], + IndexType::Vector, + Some("ix".into()), + ¶ms, + true, + ) + .await + .unwrap(); - let mut ds_single = - write_dataset_from_batches(&ds_single_uri, schema.clone(), batches.clone()).await; - let mut ds_split = write_dataset_from_batches(&ds_split_uri, schema, batches).await; + dataset + .alter_columns(&[ColumnAlteration::new("vector".into()).rename("embedding".into())]) + .await + .expect("renaming an indexed column stays legal when no copy is stored by name"); - let distributed_params = match index_type { - IndexType::IvfFlat => { - let ivf_params = prepare_global_ivf(&ds_single, "vector").await; - VectorIndexParams::with_ivf_flat_params(DistanceType::L2, ivf_params) - } - IndexType::IvfPq => { - let (ivf_params, pq_params) = prepare_global_ivf_pq(&ds_single, "vector").await; - VectorIndexParams::with_ivf_pq_params(DistanceType::L2, ivf_params, pq_params) - } - IndexType::IvfSq => { - let ivf_params = prepare_global_ivf(&ds_single, "vector").await; - VectorIndexParams::with_ivf_sq_params( - DistanceType::L2, - ivf_params, - SQBuildParams::default(), - ) - } - other => panic!("unsupported test index type: {}", other), - }; + // The index must still be usable, and -- the part the refine bug broke -- both + // maintenance paths must still run against it. + dataset + .optimize_indices(&OptimizeOptions::merge(1)) + .await + .expect("optimize must still work after the rename"); + dataset.delete("id < 16").await.unwrap(); + compact_files(&mut dataset, CompactionOptions::default(), None) + .await + .expect("compaction must still work after the rename"); - ds_single + let indices = dataset.load_indices().await.unwrap(); + assert_eq!(indices.len(), 1, "the index must survive the rename"); + assert_eq!( + indices[0].keyed_field(), + Some(dataset.schema().field("embedding").unwrap().id), + "the index must stay keyed on the renamed column's (unchanged) field id" + ); + } + + /// An inline remap (`compact_files` without `defer_index_remap`) rebuilds the index + /// through `new_remapper`, which reads its covering set back out of *storage* -- where + /// the carried vectors sit under the indexed column's own name, indistinguishable from + /// an ordinary covering column. That set has to be split back into covering columns + /// proper plus the carried-refine flag, or the rebuilt index disagrees with every other + /// build path about what it is carrying. + #[tokio::test] + async fn test_remap_preserves_carried_vectors() { + const INDEX_NAME: &str = "vec_idx"; + + let test_dir = TempStrDir::default(); + let (schema, batches) = make_covered_test_batches(); + let query = batches[0]["vector"].as_fixed_size_list().value(0); + let uri = format!("{}/remap_keeps_carried_vectors", test_dir.as_str()); + let mut dataset = write_dataset_from_batches(&uri, schema, batches).await; + + let (ivf_params, pq_params) = prepare_covered_ivf_pq(&dataset, "vector").await; + let mut params = + VectorIndexParams::with_ivf_pq_params(DistanceType::L2, ivf_params, pq_params); + params.covering_columns(vec!["payload".to_string()]); + params.store_vectors_for_refine(true); + dataset .create_index( &["vector"], IndexType::Vector, Some(INDEX_NAME.to_string()), - &distributed_params, + ¶ms, true, ) .await .unwrap(); - let fragment_groups = ds_split - .get_fragments() - .into_iter() - .map(|fragment| vec![fragment.id() as u32]) - .collect::>(); - let segments = build_segments_for_fragment_groups( - &mut ds_split, - fragment_groups, - &distributed_params, - INDEX_NAME, - ) - .await; - - assert!(segments.len() >= 4); - let grouped_inputs = segments - .chunks(2) - .map(|group| group.to_vec()) - .collect::>(); - let mut expected_fragment_coverage = grouped_inputs - .iter() - .map(|group| { - group - .iter() - .flat_map(|partial| { - partial - .fragment_bitmap - .as_ref() - .expect("partial shard should have fragment coverage") - .iter() - }) - .sorted() - .collect::>() - }) - .collect::>(); - expected_fragment_coverage.sort(); + // Deleting rows gives compaction something to rewrite, which is what triggers the + // remap of the index built above. + dataset.delete("id < 16").await.unwrap(); + compact_files(&mut dataset, CompactionOptions::default(), None) + .await + .unwrap(); - let grouped_segments = futures::future::try_join_all( - grouped_inputs - .into_iter() - .map(|group| ds_split.merge_existing_index_segments(group)), - ) - .await - .unwrap(); - let grouped_segments = - build_distributed_segments(&mut ds_split, grouped_segments, index_type, INDEX_NAME) - .await; - assert_eq!(grouped_segments.len(), expected_fragment_coverage.len()); - let mut actual_fragment_coverage = grouped_segments + let indices = dataset.load_indices().await.unwrap(); + let index = indices .iter() - .map(|segment| { - segment - .fragment_bitmap - .as_ref() - .expect("segment should have fragment coverage") - .iter() - .collect::>() - }) - .collect::>(); - actual_fragment_coverage.sort(); + .find(|idx| idx.name == INDEX_NAME) + .expect("the index must survive compaction"); + let payload_id = dataset.schema().field("payload").unwrap().id; + let vector_id = dataset.schema().field("vector").unwrap().id; assert_eq!( - actual_fragment_coverage, expected_fragment_coverage, - "built segment coverage should equal the union of its source partial shards", + index.covering_fields, + vec![payload_id, vector_id], + "the remapped index must still declare the covering column and the carried \ + vectors, in that order" ); - async fn collect_row_ids(ds: &Dataset, queries: &[Arc]) -> Vec> { - let mut ids_per_query = Vec::with_capacity(queries.len()); - for q in queries { - let result = ds - .scan() - .with_row_id() - .project(&["_rowid"] as &[&str]) - .unwrap() - .nearest("vector", q.as_ref(), K) - .unwrap() - .minimum_nprobes(TWO_FRAG_NUM_PARTITIONS) - .try_into_batch() - .await - .unwrap(); - - ids_per_query.push( - result[ROW_ID] - .as_primitive::() - .values() - .iter() - .copied() - .collect(), - ); - } - ids_per_query - } - - let query_batch = ds_single - .scan() - .project(&["vector"] as &[&str]) - .unwrap() - .limit(Some(NUM_QUERIES as i64), None) - .unwrap() - .try_into_batch() - .await - .unwrap(); - let vectors = query_batch["vector"].as_fixed_size_list(); - let queries: Vec> = (0..vectors.len()) - .map(|i| vectors.value(i) as Arc) - .collect(); - - let ids_single = collect_row_ids(&ds_single, &queries).await; - let ids_split = collect_row_ids(&ds_split, &queries).await; - if matches!(index_type, IndexType::IvfSq) { - for (single, split) in ids_single.iter().zip(ids_split.iter()) { - assert_eq!(single.len(), split.len()); - let overlap = single - .iter() - .filter(|row_id| split.contains(row_id)) - .count(); - assert!( - overlap >= K / 3, - "single vs segmented distributed SQ index returned too little top-k overlap", - ); - } - } else { - assert_eq!(ids_single, ids_split); - } + assert_refine_served_from_index(&dataset, query.as_ref(), "remap").await; } + /// PQ keeps only lossy codes, so `refine` re-ranks by taking full-precision vectors + /// from the base table. Storing them in the index removes that take. They are carried + /// as a storage-internal column, NOT via `covering_columns`: the transform chain + /// rewrites the indexed column in place (residual, and normalization under cosine), so + /// a user-visible copy under the column's own name would hand back values that are not + /// the user's vectors -- and under cosine the magnitude is gone for good. #[tokio::test] - async fn test_distributed_vector_plan_rejects_overlapping_fragment_coverage() { + async fn test_refine_vectors_are_stored_in_index_storage() { let test_dir = TempStrDir::default(); - let base_uri = test_dir.as_str(); - let (schema, batches) = make_two_fragment_batches(); - let dataset_uri = format!("{}/overlap_fragments", base_uri); - let mut dataset = write_dataset_from_batches(&dataset_uri, schema, batches).await; + let (schema, batches) = make_covered_test_batches(); + let uri = format!("{}/store_refine_vectors", test_dir.as_str()); + let mut dataset = write_dataset_from_batches(&uri, schema, batches).await; - let fragment = dataset.get_fragments()[0].id() as u32; - let params = VectorIndexParams::with_ivf_flat_params( - DistanceType::L2, - prepare_global_ivf(&dataset, "vector").await, - ); - let mut segments = Vec::new(); + let (ivf_params, pq_params) = prepare_covered_ivf_pq(&dataset, "vector").await; + let mut params = + VectorIndexParams::with_ivf_pq_params(DistanceType::L2, ivf_params, pq_params); + params.store_vectors_for_refine(true); - for _ in 0..2 { - let segment = dataset - .create_index_builder(&["vector"], IndexType::Vector, ¶ms) - .name("vector_idx".to_string()) - .fragments(vec![fragment]) - .execute_uncommitted() - .await - .unwrap(); - segments.push(segment); - } + dataset + .create_index( + &["vector"], + IndexType::Vector, + Some("vec_idx".into()), + ¶ms, + true, + ) + .await + .unwrap(); - let err = dataset - .merge_existing_index_segments(segments) + let indices = dataset.load_indices().await.unwrap(); + let index = dataset + .open_vector_index("vector", &indices[0].uuid, &NoOpMetricsCollector) .await - .unwrap_err(); - assert!(err.to_string().contains("overlapping fragment coverage")); + .unwrap(); + + // Physically present under the indexed column's own name -- the copy is taken + // before the transform chain, so its values are that column's values -- and stamped + // with that column's source id, which is what makes it servable rather than merely + // present. + let vector_field_id = dataset.schema().field("vector").unwrap().id; + let physical = index.physical_covering_fields().unwrap(); + assert!( + physical + .iter() + .any(|(id, field)| *id == vector_field_id && field.name() == "vector"), + "index storage must carry the vector column stamped with its source field id; \ + got {physical:?}" + ); + + // And declared, so the search emits it and `refine` can find it by name. + assert_eq!( + indices[0].covering_fields, + vec![vector_field_id], + "carried vectors must be declared covering, or nothing will read them" + ); } + /// The single-segment counterpart to the merge test above, and the case a name-and-type + /// check cannot see. A driver rebinds a segment built for a dropped `payload` onto a + /// re-added field with the same name and type but a fresh id. + /// + /// The declaration/payload contract permits the commit. At read time, however, the + /// storage's stamped source id must prevent those old values from being served under the + /// new field: planning falls back to a base-table take, whose re-added column is all null. #[tokio::test] - async fn test_distributed_vector_build_supports_hnsw_variants() { + async fn test_rebound_covering_field_id_falls_back_to_base_table() { + use crate::dataset::NewColumnTransform; + let test_dir = TempStrDir::default(); let base_uri = test_dir.as_str(); - let (schema, batches) = make_two_fragment_batches(); - let dataset_uri = format!("{}/distributed_hnsw_supported", base_uri); + let (schema, batches) = make_covered_test_batches(); + let query = batches[0]["vector"].as_fixed_size_list().value(0); + let dataset_uri = format!("{}/commit_rebind_covering_field_id", base_uri); let mut dataset = write_dataset_from_batches(&dataset_uri, schema, batches).await; - let fragments = dataset.get_fragments(); - assert!(fragments.len() >= 2); - let params = VectorIndexParams::ivf_hnsw( - DistanceType::L2, - prepare_global_ivf(&dataset, "vector").await, - HnswBuildParams::default(), - ); - let mut segments = Vec::new(); + let all_fragments: Vec = dataset + .get_fragments() + .iter() + .map(|fragment| fragment.id() as u32) + .collect(); - for fragment in fragments.iter().take(2) { - let segment = dataset - .create_index_builder(&["vector"], IndexType::Vector, ¶ms) - .name("vector_idx".to_string()) - .fragments(vec![fragment.id() as u32]) - .execute_uncommitted() - .await - .unwrap(); - segments.push(segment); - } + let (ivf_params, pq_params) = prepare_covered_ivf_pq(&dataset, "vector").await; + let mut params = + VectorIndexParams::with_ivf_pq_params(DistanceType::L2, ivf_params, pq_params); + params.covering_columns(vec!["payload".to_string()]); - dataset - .commit_existing_index_segments("vector_idx", "vector", segments) + let mut segment = dataset + .create_index_builder(&["vector"], IndexType::Vector, ¶ms) + .name("vec_idx".to_string()) + .fragments(all_fragments) + .execute_uncommitted() .await .unwrap(); + let original_field_id = dataset.schema().field("payload").unwrap().id; - let query_batch = dataset - .scan() - .project(&["vector"] as &[&str]) - .unwrap() - .limit(Some(4), None) - .unwrap() - .try_into_batch() + // Drop and re-add `payload`: same name, same type, new field id. The re-add is + // metadata-only (AllNulls), so no data file changes and no fragment looks stale. + dataset.drop_columns(&["payload"]).await.unwrap(); + dataset + .add_columns( + NewColumnTransform::AllNulls(Arc::new(arrow_schema::Schema::new(vec![ + Field::new("payload", DataType::UInt64, true), + ]))), + None, + None, + ) .await .unwrap(); - let q = query_batch["vector"].as_fixed_size_list().value(0); - let result = dataset - .scan() - .project(&["_rowid"] as &[&str]) - .unwrap() - .nearest("vector", q.as_ref(), 5) - .unwrap() - .try_into_batch() + let new_field_id = dataset.schema().field("payload").unwrap().id; + assert_ne!( + original_field_id, new_field_id, + "re-adding the column must produce a new field id, or this test proves nothing" + ); + + // The segment keeps its true build version, so the staleness pass runs and passes; + // only the id rebind is wrong. + for id in segment.fields.iter_mut() { + if *id == original_field_id { + *id = new_field_id; + } + } + for id in segment.covering_fields.iter_mut() { + if *id == original_field_id { + *id = new_field_id; + } + } + + dataset + .commit_existing_index_segments("vec_idx", "vector", vec![segment]) .await .unwrap(); - assert!(result.num_rows() > 0); + + let mut scan = dataset.scan(); + scan.nearest("vector", query.as_ref(), 10).unwrap(); + scan.nprobes(TWO_FRAG_NUM_PARTITIONS); + scan.project(&["payload"]).unwrap(); + let plan = scan.explain_plan(true).await.unwrap(); + assert!( + plan.contains("LanceRead"), + "storage was built from field id {original_field_id}, not the declaration's \ + rebound id {new_field_id}, so payload must use a base-table take; plan:\n{plan}" + ); + + let batch = scan.try_into_batch().await.unwrap(); + assert_eq!(batch.num_rows(), 10); + let payload = batch["payload"].as_primitive::(); + assert_eq!( + payload.null_count(), + payload.len(), + "the re-added all-null payload must come from the base table, not stale index values" + ); } #[rstest] @@ -5387,7 +12198,10 @@ mod tests { // Every partition, not just the first: a projection that applied // unevenly would leave the partition cache holding mixed schemas. for partition_id in 0..hnsw.ivf.num_partitions() { - let entry = hnsw.load_partition_entry(partition_id, None).await.unwrap(); + let entry = hnsw + .load_partition_entry(partition_id, PartitionColumns::Internal, None) + .await + .unwrap(); let loaded = entry.index.to_batch().unwrap(); let loaded_schema = loaded.schema(); let read = loaded_schema diff --git a/rust/lance/src/index/vector/pq.rs b/rust/lance/src/index/vector/pq.rs index 141c8b85f27..99280ff8e25 100644 --- a/rust/lance/src/index/vector/pq.rs +++ b/rust/lance/src/index/vector/pq.rs @@ -990,6 +990,7 @@ mod tests { query_parallelism: DEFAULT_QUERY_PARALLELISM, dist_q_c: 0.0, approx_mode: Default::default(), + covering_projection: None, }; let is_empty_threads = Arc::new(Mutex::new(Vec::new())); let pre_filter = Arc::new(TestPreFilter::with_thread_capture( diff --git a/rust/lance/src/index/vector/utils.rs b/rust/lance/src/index/vector/utils.rs index 012d8d3a703..53a322572df 100644 --- a/rust/lance/src/index/vector/utils.rs +++ b/rust/lance/src/index/vector/utils.rs @@ -23,6 +23,67 @@ use tokio::sync::Mutex; use crate::dataset::{Dataset, ProjectionRequest, TakeBuilder, row_offsets_to_row_addresses}; use crate::{Error, Result}; +use lance_core::ROW_ID; + +/// Row-align the non-`_rowid` columns of `source` (a `[_rowid, ]` batch) +/// to `target_rowids`, gathered by row id. Every target row id must be present in +/// `source`; a miss returns an error, because silently gathering an unrelated +/// row's values would corrupt the covered result. Returns the taken columns paired +/// with their fields, in `source` column order, with `_rowid` excluded. +/// +/// Shared by the covered-search read path (regrouping heap winners into +/// `[_distance, _rowid, ]`) and the reassignment build path +/// (re-attaching covering columns to rows regrouped by the partition join). +pub(super) fn gather_covering_columns_by_row_id( + source: &RecordBatch, + target_rowids: &arrow_array::UInt64Array, +) -> Result> { + let take_idx = row_id_take_indices(source, target_rowids)?; + let mut out = Vec::with_capacity(source.num_columns().saturating_sub(1)); + for (i, field) in source.schema().fields().iter().enumerate() { + if field.name() == ROW_ID { + continue; + } + let taken = arrow::compute::take(source.column(i), &take_idx, None)?; + out.push((field.clone(), taken)); + } + Ok(out) +} + +/// Where each of `target_rowids` sits in `source` (a batch carrying `_rowid`), as take +/// indices. Every target row id must be present; a miss is an error for the same reason +/// [`gather_covering_columns_by_row_id`] errors on one -- taking an unrelated row's +/// values would corrupt the covered result silently. +/// +/// Split out so a caller that must keep `_rowid` on the result (the survivor gather +/// narrows each partition's covering batch before reading the next one) can take the +/// whole batch instead of its non-`_rowid` columns. +pub(super) fn row_id_take_indices( + source: &RecordBatch, + target_rowids: &arrow_array::UInt64Array, +) -> Result { + let src_rowids = source + .column_by_name(ROW_ID) + .ok_or_else(|| Error::internal("covering source missing row id".to_string()))? + .as_primitive::(); + let mut pos: std::collections::HashMap = + std::collections::HashMap::with_capacity(src_rowids.len()); + for (i, rid) in src_rowids.values().iter().enumerate() { + pos.insert(*rid, i as u32); + } + Ok(target_rowids + .values() + .iter() + .map(|rid| { + pos.get(rid).copied().ok_or_else(|| { + Error::internal(format!( + "row id {rid} missing from covering source during row-id gather" + )) + }) + }) + .collect::>>()? + .into()) +} /// Helper function to extract a column from a RecordBatch, supporting nested field paths. /// diff --git a/rust/lance/src/io/commit/conflict_resolver.rs b/rust/lance/src/io/commit/conflict_resolver.rs index 4a824058640..36b2b03663b 100644 --- a/rust/lance/src/io/commit/conflict_resolver.rs +++ b/rust/lance/src/io/commit/conflict_resolver.rs @@ -10,7 +10,7 @@ use crate::{ dataset::transaction::{DataOverlayGroup, Operation, Transaction, UpdateMode}, }; use futures::{StreamExt, TryStreamExt}; -use lance_core::{Error, Result, utils::deletion::DeletionVector}; +use lance_core::{Error, Result, datatypes::Schema, utils::deletion::DeletionVector}; use lance_index::frag_reuse::FRAG_REUSE_INDEX_NAME; use lance_index::mem_wal::{CompactedSsTable, MEM_WAL_INDEX_NAME}; use lance_select::{RowAddrTreeMap, RowSetOps}; @@ -37,6 +37,57 @@ pub struct TransactionRebase<'a> { /// Compacted SSTables from conflicting UpdateMemWalState transactions. /// Used when rebasing CreateIndex of MemWalIndex. conflicting_mem_wal_compacted_sstables: Vec, + /// The dataset schema at the read version, used to expand a covering index's + /// covered (included) struct fields to their leaf subtrees when checking whether a + /// concurrent transaction touches covered data. + schema: Arc, + /// Read-version snapshot of which data file backs each covered (included) leaf + /// field, per fragment: `fragment id -> field id -> path`. Empty unless this + /// transaction creates an index that declares covering columns. + /// + /// This is what lets `check_create_index_txn` tell a concurrent `Merge` that + /// rewrote covered data *in place* from one that merely added a column. The + /// post-merge fragments alone cannot answer that: a rewrite tombstones the + /// superseded field with a single sentinel that does not name it. + covered_field_paths: Arc>>, +} + +/// Snapshot the data file backing each covered leaf field of `new_indices`, as of +/// `dataset`'s version. Covered struct fields are expanded to their whole subtree +/// because a rewritten data file lists leaf ids while `covering_fields` records only +/// the parent id. Fragments with no covered field are omitted. +fn covered_field_paths_at( + dataset: &Dataset, + new_indices: &[IndexMetadata], + schema: &Schema, +) -> HashMap> { + let mut covered: HashSet = HashSet::new(); + for index in new_indices { + for &id in index.covering_fields.iter() { + match schema.field_by_id(id) { + Some(field) => crate::index::collect_subtree_field_ids(field, &mut covered), + None => { + covered.insert(id); + } + } + } + } + if covered.is_empty() { + return HashMap::new(); + } + dataset + .manifest + .fragments + .iter() + .filter_map(|fragment| { + let paths: HashMap = Transaction::fragment_field_paths(fragment) + .into_iter() + .filter(|(field_id, _)| covered.contains(field_id)) + .map(|(field_id, path)| (field_id, path.to_string())) + .collect(); + (!paths.is_empty()).then_some((fragment.id, paths)) + }) + .collect() } /// Whether `operation` may make a nullability-affecting schema change: a @@ -99,6 +150,55 @@ impl<'a> TransactionRebase<'a> { transaction: Transaction, affected_rows: Option<&'a RowAddrTreeMap>, ) -> Result { + // Capture the schema AS OF THE TRANSACTION'S READ VERSION, matching what + // `initial_fragments_for_rebase` does for fragments. The dataset handed to us is not + // necessarily at that version: the URI commit path deliberately loads the *latest* + // version ("we are writing to the main history, and need to check out the latest + // version") for any non-detached read version. The covering checks below use this as + // the *pre-commit* side when asking whether a concurrent operation changed a covered + // field's subtree, so capturing the latest schema would compare a post-change schema + // against itself -- silently missing a covered-field rename or retype and committing + // index storage under a schema it was never built against. + // + // Only `check_create_index_txn` and `check_data_replacement_txn` consult this, so + // only they pay for the extra checkout; every other operation gets the in-hand schema + // and never reads it. If a new check starts using `self.schema`, add its operation + // here or it will silently see the wrong version. + // + // CreateIndex is narrowed further to indices that actually declare covering columns: + // comparing a covered subtree across versions is the only thing that can observe the + // difference, so an ordinary index build keeps the pre-covering behaviour (in-hand + // schema, no extra round trip) instead of paying for a manifest load on every commit + // retry. DataReplacement cannot be narrowed the same way -- the covering it has to + // respect belongs to a *concurrent* transaction's indices, which are not known until + // the check runs. + // + // `read_version == 0` is the "no prior version" sentinel (overwrite/create): there is + // nothing to check out, and attempting it fails with DatasetNotFound. + let needs_read_version_schema = match &transaction.operation { + Operation::CreateIndex { new_indices, .. } => new_indices + .iter() + .any(|idx| !idx.covering_fields.is_empty()), + Operation::DataReplacement { .. } => true, + _ => false, + }; + let checked_out; + let at_read_version = if !needs_read_version_schema + || transaction.read_version == 0 + || dataset.manifest.version == transaction.read_version + { + dataset + } else { + checked_out = dataset.checkout_version(transaction.read_version).await?; + &checked_out + }; + let schema = Arc::new(at_read_version.schema().clone()); + let covered_field_paths = Arc::new(match &transaction.operation { + Operation::CreateIndex { new_indices, .. } if needs_read_version_schema => { + covered_field_paths_at(at_read_version, new_indices, &schema) + } + _ => HashMap::new(), + }); match &transaction.operation { // These operations add new fragments or don't modify any. Operation::Append { .. } @@ -117,6 +217,8 @@ impl<'a> TransactionRebase<'a> { modified_fragment_ids: HashSet::new(), conflicting_frag_reuse_indices: Vec::new(), conflicting_mem_wal_compacted_sstables: Vec::new(), + schema: schema.clone(), + covered_field_paths: covered_field_paths.clone(), }), Operation::Delete { updated_fragments, @@ -145,6 +247,8 @@ impl<'a> TransactionRebase<'a> { affected_rows: None, conflicting_frag_reuse_indices: Vec::new(), conflicting_mem_wal_compacted_sstables: Vec::new(), + schema: schema.clone(), + covered_field_paths: covered_field_paths.clone(), }); } @@ -158,6 +262,8 @@ impl<'a> TransactionRebase<'a> { modified_fragment_ids, conflicting_frag_reuse_indices: Vec::new(), conflicting_mem_wal_compacted_sstables: Vec::new(), + schema: schema.clone(), + covered_field_paths: covered_field_paths.clone(), }) } Operation::Rewrite { groups, .. } => { @@ -176,6 +282,8 @@ impl<'a> TransactionRebase<'a> { modified_fragment_ids, conflicting_frag_reuse_indices: Vec::new(), conflicting_mem_wal_compacted_sstables: Vec::new(), + schema: schema.clone(), + covered_field_paths: covered_field_paths.clone(), }) } Operation::DataReplacement { replacements } => { @@ -191,6 +299,8 @@ impl<'a> TransactionRebase<'a> { modified_fragment_ids, conflicting_frag_reuse_indices: Vec::new(), conflicting_mem_wal_compacted_sstables: Vec::new(), + schema: schema.clone(), + covered_field_paths: covered_field_paths.clone(), }) } Operation::DataOverlay { groups } => { @@ -206,6 +316,8 @@ impl<'a> TransactionRebase<'a> { modified_fragment_ids, conflicting_frag_reuse_indices: Vec::new(), conflicting_mem_wal_compacted_sstables: Vec::new(), + schema: schema.clone(), + covered_field_paths: covered_field_paths.clone(), }) } Operation::Merge { fragments, .. } => { @@ -220,6 +332,8 @@ impl<'a> TransactionRebase<'a> { modified_fragment_ids, conflicting_frag_reuse_indices: Vec::new(), conflicting_mem_wal_compacted_sstables: Vec::new(), + schema: schema.clone(), + covered_field_paths: covered_field_paths.clone(), }) } } @@ -695,6 +809,10 @@ impl<'a> TransactionRebase<'a> { other_transaction: &Transaction, other_version: u64, ) -> Result<()> { + // Captured before the mutable borrow below so the covered-field checks can expand + // covered struct fields to their leaf subtree (an Arc clone is cheap). + let schema = self.schema.clone(); + let covered_field_paths = self.covered_field_paths.clone(); if let Operation::CreateIndex { new_indices, removed_indices, @@ -752,25 +870,89 @@ impl<'a> TransactionRebase<'a> { fields_modified, .. } => { + // Expand covered struct fields to their leaf subtree so a concurrent + // update of a covered struct's child still invalidates the fragment's + // coverage (the schema is captured at the read version above). Transaction::prune_updated_fields_from_indices( new_indices, updated_fragments, fields_modified, + schema.as_ref(), ); Ok(()) } - // Merge, reserve, and project don't change row ids. The MemWAL - // index is the exception: its install validates schema-dependent - // state, which a concurrent schema change invalidates. - Operation::Merge { .. } => { - if new_indices.iter().any(|idx| idx.name == MEM_WAL_INDEX_NAME) { + // Reserve and project don't change row ids. Merge is usually fine too, + // but two independent conditions make it unsafe here, and neither + // subsumes the other: + // + // (a) the MemWAL index's install validates schema-dependent state, which + // a concurrent schema change invalidates. + // + // (b) it can rewrite a covered column's data file *in place*, or change a + // covered field's *subtree* (notably `add_columns` growing a covered + // struct with an AllNulls child, which writes NO data file, so the + // path comparison alone is blind to it). Either leaves the rebased + // index serving values it was not built against. + // + // Merely adding an unrelated column trips none of these and keeps the + // optimistic path -- otherwise a routine `add_columns` would discard a + // covered index build that may have been running for hours. + Operation::Merge { + fragments, + schema: merged_schema, + .. + } => { + let installs_mem_wal_index = + new_indices.iter().any(|idx| idx.name == MEM_WAL_INDEX_NAME); + let rewrote_covered_data = fragments.iter().any(|fragment| { + let Some(previous) = covered_field_paths.get(&fragment.id) else { + return false; + }; + let current = Transaction::fragment_field_paths(fragment); + previous.iter().any(|(field_id, previous_path)| { + current + .get(field_id) + .is_some_and(|current_path| *current_path != previous_path.as_str()) + }) + }); + let changed_covered_subtree = new_indices.iter().any(|idx| { + idx.covering_fields.iter().any(|id| { + Transaction::covered_field_subtree_changed(&schema, merged_schema, *id) + }) + }); + if installs_mem_wal_index || rewrote_covered_data || changed_covered_subtree { Err(self.retryable_conflict_err(other_transaction, other_version)) } else { Ok(()) } } Operation::ReserveFragments { .. } => Ok(()), - Operation::Project { .. } => Ok(()), + Operation::Project { + schema: projected_schema, + .. + } => { + // A concurrent Project that changes a covered field's subtree in *any* + // way -- drop, rename, retype, nullability, child change -- leaves the + // rebased index emitting a schema the dataset no longer declares, so + // force a retry and let the index rebuild against the projected schema. + // A rename counts: `covered_field_subtree_changed` compares names, and + // the index's covering payload carries the old one. + // Non-covering indexes, and covered subtrees the Project left alone, + // keep the optimistic path (row ids are unchanged either way). + if new_indices.iter().any(|idx| { + idx.covering_fields.iter().any(|id| { + Transaction::covered_field_subtree_changed( + &schema, + projected_schema, + *id, + ) + }) + }) { + Err(self.retryable_conflict_err(other_transaction, other_version)) + } else { + Ok(()) + } + } // Should be compatible with rewrite if it didn't move the rows // we indexed. If it did, we could retry. // TODO: this will change with stable row ids. @@ -848,16 +1030,19 @@ impl<'a> TransactionRebase<'a> { } Operation::UpdateConfig { .. } => Ok(()), Operation::DataReplacement { replacements } => { - // A data replacement only conflicts if it is updating a field the - // index depends on -- whether keyed on or merely carried, since - // `fields` lists both (see `IndexMetadata::covering_fields`). - let newly_depended_fields = new_indices - .iter() - .flat_map(|idx| idx.fields.iter()) - .collect::>(); + // A data replacement conflicts if it updates a field the index + // either *indexes* or *covers*: a covered field's values are + // materialized in the index storage, so replacing them would + // leave the rebased index serving stale values. Expand each field to + // its leaf subtree, since `covering_fields` records a covered struct's + // parent id while a replacement lists leaf ids. + let mut relevant_fields: HashSet = HashSet::new(); + for idx in new_indices.iter() { + relevant_fields.extend(Transaction::index_dependent_leaf_ids(idx, &schema)); + } for replacement in replacements { for field in replacement.1.fields.iter() { - if newly_depended_fields.contains(&field) { + if relevant_fields.contains(field) { return Err( self.retryable_conflict_err(other_transaction, other_version) ); @@ -1169,6 +1354,7 @@ impl<'a> TransactionRebase<'a> { other_transaction: &Transaction, other_version: u64, ) -> Result<()> { + let schema = self.schema.clone(); if let Operation::DataReplacement { replacements } = &self.transaction.operation { match &other_transaction.operation { Operation::Append { .. } @@ -1268,13 +1454,17 @@ impl<'a> TransactionRebase<'a> { // the index's fragment bitmap, which would lead to fewer conflicts. However // this would introduce fragment bitmaps with holes which may not be well tested // yet. For now, we don't allow this case. - let newly_depended_fields = new_indices - .iter() - .flat_map(|idx| idx.fields.iter()) - .collect::>(); + // Expand covered fields to their leaf subtree so replacing a covered + // struct's subfield (a leaf id, while `covering_fields` records the + // parent) is still recognized as touching the index -- matching the + // reverse rebase direction in check_create_index_txn. + let mut relevant_fields: HashSet = HashSet::new(); + for idx in new_indices.iter() { + relevant_fields.extend(Transaction::index_dependent_leaf_ids(idx, &schema)); + } for replacement in replacements { for field in replacement.1.fields.iter() { - if newly_depended_fields.contains(&field) { + if relevant_fields.contains(field) { return Err( self.retryable_conflict_err(other_transaction, other_version) ); @@ -2329,6 +2519,14 @@ mod tests { }; use lance_table::format::DataFile; + /// An empty schema for manually-constructed rebases in tests: covered-field subtree + /// expansion is a no-op against it (each id maps to itself), preserving the exact-id + /// behavior these tests assert. Tests that exercise covered structs build the rebase + /// via `try_new`, which captures the real dataset schema. + fn empty_lance_schema() -> Arc { + Arc::new(lance_core::datatypes::Schema::default()) + } + async fn test_dataset(num_rows: usize, num_fragments: usize) -> Dataset { let write_params = WriteParams { max_rows_per_file: num_rows / num_fragments, @@ -2690,6 +2888,279 @@ mod tests { assert!(matches!(err, Error::DatasetNotFound { .. })); } + /// A concurrent Project that drops a column an in-flight index *covers* must force + /// the index build to rebase (retryable conflict); a Project that preserves the + /// covered column stays on the optimistic path. + #[tokio::test] + async fn test_create_index_conflicts_with_project_dropping_covered_field() { + use roaring::RoaringBitmap; + + let dataset = test_dataset(10, 2).await; + let a_id = dataset.schema().field("a").unwrap().id; + let b_id = dataset.schema().field("b").unwrap().id; + + // A covering index over key "a" that also covers "b". + // `covering_fields` is always the trailing entries of `fields`. + let covering_index = IndexMetadata { + uuid: Uuid::new_v4(), + name: "cov_idx".to_string(), + fields: vec![a_id, b_id], + dataset_version: 1, + fragment_bitmap: Some(RoaringBitmap::from_iter([0u32, 1])), + index_details: Some(Arc::new(prost_types::Any { + type_url: "type.googleapis.com/lance.index.VectorIndexDetails".to_string(), + value: vec![], + })), + index_version: 0, + created_at: None, + base_id: None, + files: None, + covering_fields: vec![b_id], + }; + let self_txn = Transaction::new_from_version( + 1, + Operation::CreateIndex { + new_indices: vec![covering_index], + removed_indices: vec![], + }, + ); + + // Concurrent Project that drops "b" (the covered column) -> conflict. + let dropped_schema = dataset.schema().project(&["a"]).unwrap(); + let drop_txn = Transaction::new_from_version( + 2, + Operation::Project { + schema: dropped_schema, + preserves_nullability: true, + }, + ); + let mut rebase = TransactionRebase::try_new(&dataset, self_txn.clone(), None) + .await + .unwrap(); + assert!( + rebase.check_txn(&drop_txn, 2).is_err(), + "dropping a covered column must conflict with an in-flight covering index build" + ); + + // A Project that preserves "b" -> no conflict. + let keep_txn = Transaction::new_from_version( + 2, + Operation::Project { + schema: dataset.schema().clone(), + preserves_nullability: true, + }, + ); + let mut rebase = TransactionRebase::try_new(&dataset, self_txn.clone(), None) + .await + .unwrap(); + assert!( + rebase.check_txn(&keep_txn, 2).is_ok(), + "a project that preserves the covered column must not conflict" + ); + + // A Project that renames the covered column (id stays, name changes) -> conflict: + // the built index still emits the old name. + let mut renamed_schema = dataset.schema().clone(); + renamed_schema.field_by_id_mut(b_id).unwrap().name = "b_renamed".to_string(); + let rename_txn = Transaction::new_from_version( + 2, + Operation::Project { + schema: renamed_schema, + preserves_nullability: true, + }, + ); + let mut rebase = TransactionRebase::try_new(&dataset, self_txn, None) + .await + .unwrap(); + assert!( + rebase.check_txn(&rename_txn, 2).is_err(), + "renaming a covered column must conflict with an in-flight covering index build" + ); + } + + /// A concurrent DataReplacement of a *child* of a covered struct must conflict with + /// an in-flight covering index build: `covering_fields` records the struct's parent + /// id, but the replacement lists the leaf id, so the check must expand the subtree. + #[tokio::test] + async fn test_create_index_conflicts_with_replacement_of_covered_struct_child() { + use arrow_schema::{ + DataType, Field as ArrowField, Fields as ArrowFields, Schema as ArrowSchema, + }; + use roaring::RoaringBitmap; + + let arrow = ArrowSchema::new(vec![ + ArrowField::new("vec", DataType::Int32, false), + ArrowField::new( + "s", + DataType::Struct(ArrowFields::from(vec![ + ArrowField::new("a", DataType::Int32, true), + ArrowField::new("b", DataType::Int32, true), + ])), + true, + ), + ]); + let schema = lance_core::datatypes::Schema::try_from(&arrow).unwrap(); + let vec_id = schema.field("vec").unwrap().id; + let s_id = schema.field("s").unwrap().id; + let a_id = schema.field("s.a").unwrap().id; + + // Index keyed on `vec`, covering the whole struct `s`. + // `covering_fields` is always the trailing entries of `fields`. + let covering_index = IndexMetadata { + uuid: Uuid::new_v4(), + name: "cov_idx".to_string(), + fields: vec![vec_id, s_id], + dataset_version: 1, + fragment_bitmap: Some(RoaringBitmap::from_iter([0u32])), + index_details: Some(Arc::new(prost_types::Any { + type_url: "type.googleapis.com/lance.index.VectorIndexDetails".to_string(), + value: vec![], + })), + index_version: 0, + created_at: None, + base_id: None, + files: None, + covering_fields: vec![s_id], + }; + let self_txn = Transaction::new_from_version( + 1, + Operation::CreateIndex { + new_indices: vec![covering_index], + removed_indices: vec![], + }, + ); + + // try_new needs a real dataset, so build the rebase directly with the struct schema. + let mut rebase = TransactionRebase { + transaction: self_txn, + initial_fragments: HashMap::new(), + modified_fragment_ids: HashSet::new(), + affected_rows: None, + conflicting_frag_reuse_indices: Vec::new(), + conflicting_mem_wal_compacted_sstables: Vec::new(), + covered_field_paths: Default::default(), + schema: Arc::new(schema), + }; + + // A concurrent DataReplacement rewriting the covered struct's child `s.a`. + let replace_child = Transaction::new_from_version( + 2, + Operation::DataReplacement { + replacements: vec![DataReplacementGroup( + 0, + DataFile::new_legacy_from_fields("child.lance", vec![a_id], None), + )], + }, + ); + assert!( + rebase.check_txn(&replace_child, 2).is_err(), + "replacing a child of a covered struct must conflict (subtree expansion)" + ); + } + + /// A concurrent `add_columns` (which commits as `Operation::Merge`) that grows a + /// covered struct with an AllNulls child writes **no data file**, so the covered-file + /// path comparison cannot see it. The index would otherwise commit storage built for + /// the old struct type under the new schema, and every covered query on it would fail. + /// The Merge arm must therefore also compare the covered subtree, as the Project arm + /// does -- the two checks are complementary and neither subsumes the other. + #[tokio::test] + async fn test_create_index_conflicts_with_merge_growing_covered_struct() { + use arrow_schema::{ + DataType, Field as ArrowField, Fields as ArrowFields, Schema as ArrowSchema, + }; + use roaring::RoaringBitmap; + + let struct_of = |children: Vec| { + ArrowSchema::new(vec![ + ArrowField::new("vec", DataType::Int32, false), + ArrowField::new("s", DataType::Struct(ArrowFields::from(children)), true), + ]) + }; + let before = lance_core::datatypes::Schema::try_from(&struct_of(vec![ArrowField::new( + "a", + DataType::Int32, + true, + )])) + .unwrap(); + // The AllNulls child add keeps every existing id and gives the new leaf a fresh one. + let after = lance_core::datatypes::Schema::try_from(&struct_of(vec![ + ArrowField::new("a", DataType::Int32, true), + ArrowField::new("b", DataType::Int32, true), + ])) + .unwrap(); + let vec_id = before.field("vec").unwrap().id; + let s_id = before.field("s").unwrap().id; + assert_eq!( + s_id, + after.field("s").unwrap().id, + "the struct keeps its id" + ); + + // `covering_fields` is always the trailing entries of `fields`. + let covering_index = IndexMetadata { + uuid: Uuid::new_v4(), + name: "cov_idx".to_string(), + fields: vec![vec_id, s_id], + dataset_version: 1, + fragment_bitmap: Some(RoaringBitmap::from_iter([0u32])), + index_details: Some(Arc::new(prost_types::Any { + type_url: "type.googleapis.com/lance.index.VectorIndexDetails".to_string(), + value: vec![], + })), + index_version: 0, + created_at: None, + base_id: None, + files: None, + covering_fields: vec![s_id], + }; + let create = |index: IndexMetadata| { + Transaction::new_from_version( + 1, + Operation::CreateIndex { + new_indices: vec![index], + removed_indices: vec![], + }, + ) + }; + let rebase_over = |index: IndexMetadata| TransactionRebase { + transaction: create(index), + initial_fragments: HashMap::new(), + modified_fragment_ids: HashSet::new(), + affected_rows: None, + conflicting_frag_reuse_indices: Vec::new(), + conflicting_mem_wal_compacted_sstables: Vec::new(), + // Empty: the conflict must come from the schema comparison alone, not from + // any data-file path change. + covered_field_paths: Default::default(), + schema: Arc::new(before.clone()), + }; + let merge_with = |schema: lance_core::datatypes::Schema| { + Transaction::new_from_version( + 1, + Operation::Merge { + fragments: vec![], + schema, + preserves_nullability: true, + }, + ) + }; + + let mut rebase = rebase_over(covering_index.clone()); + assert!( + rebase.check_txn(&merge_with(after), 2).is_err(), + "a Merge that grows a covered struct's subtree must conflict" + ); + + // Control: a Merge that leaves the covered subtree alone keeps the optimistic + // path, so an ordinary add_columns still does not discard a covered build. + let mut rebase = rebase_over(covering_index); + assert!( + rebase.check_txn(&merge_with(before), 2).is_ok(), + "a Merge that does not touch the covered subtree must not conflict" + ); + } + async fn apply_deletion( delete_rows: &[u32], fragment: &mut Fragment, @@ -3483,6 +3954,8 @@ mod tests { affected_rows: None, conflicting_frag_reuse_indices: Vec::new(), conflicting_mem_wal_compacted_sstables: Vec::new(), + covered_field_paths: Default::default(), + schema: empty_lance_schema(), }; for (other, expected_conflict) in other_transactions.iter().zip(expected_conflicts) { @@ -3687,6 +4160,8 @@ mod tests { affected_rows: None, conflicting_frag_reuse_indices: Vec::new(), conflicting_mem_wal_compacted_sstables: Vec::new(), + covered_field_paths: Default::default(), + schema: empty_lance_schema(), }; let other_txn = Transaction::new(0, other.clone(), None); let result = rebase.check_txn(&other_txn, 1); @@ -3746,6 +4221,8 @@ mod tests { affected_rows: None, conflicting_frag_reuse_indices: Vec::new(), conflicting_mem_wal_compacted_sstables: Vec::new(), + covered_field_paths: Default::default(), + schema: empty_lance_schema(), }; let other_txn = Transaction::new(0, other.clone(), None); let result = rebase.check_txn(&other_txn, 1); @@ -3887,6 +4364,8 @@ mod tests { affected_rows: affected_rows.as_ref(), conflicting_frag_reuse_indices: Vec::new(), conflicting_mem_wal_compacted_sstables: Vec::new(), + covered_field_paths: Default::default(), + schema: empty_lance_schema(), }; let other_txn = Transaction::new(0, other.clone(), None); let result = rebase.check_txn(&other_txn, 1); @@ -3929,6 +4408,8 @@ mod tests { affected_rows: None, conflicting_frag_reuse_indices: Vec::new(), conflicting_mem_wal_compacted_sstables: Vec::new(), + covered_field_paths: Default::default(), + schema: Arc::new(lance_core::datatypes::Schema::default()), }; let result = append_rebase.check_txn(&Transaction::new(0, merge.clone(), None), 1); assert_eq!( @@ -3944,6 +4425,8 @@ mod tests { affected_rows: None, conflicting_frag_reuse_indices: Vec::new(), conflicting_mem_wal_compacted_sstables: Vec::new(), + covered_field_paths: Default::default(), + schema: Arc::new(lance_core::datatypes::Schema::default()), }; let result = merge_rebase.check_txn(&Transaction::new(0, append, None), 1); assert!( @@ -4020,6 +4503,8 @@ mod tests { affected_rows: None, conflicting_frag_reuse_indices: Vec::new(), conflicting_mem_wal_compacted_sstables: Vec::new(), + covered_field_paths: Default::default(), + schema: Arc::new(lance_core::datatypes::Schema::default()), }; let result = rebase.check_txn(&Transaction::new(0, theirs, None), 1); assert_eq!( @@ -4154,6 +4639,8 @@ mod tests { affected_rows: None, conflicting_frag_reuse_indices: Vec::new(), conflicting_mem_wal_compacted_sstables: Vec::new(), + covered_field_paths: Default::default(), + schema: empty_lance_schema(), }; let update = Transaction::new( 1, @@ -4230,6 +4717,8 @@ mod tests { affected_rows: None, conflicting_frag_reuse_indices: Vec::new(), conflicting_mem_wal_compacted_sstables: Vec::new(), + covered_field_paths: Default::default(), + schema: Arc::new(lance_core::datatypes::Schema::default()), }; let result = rebase.check_txn(&merge, 1); assert_eq!(result.is_err(), conflicts, "{result:?}"); @@ -4251,6 +4740,8 @@ mod tests { affected_rows: None, conflicting_frag_reuse_indices: Vec::new(), conflicting_mem_wal_compacted_sstables: Vec::new(), + covered_field_paths: Default::default(), + schema: Arc::new(lance_core::datatypes::Schema::default()), }; let result = rebase.check_txn(&install, 1); assert!( @@ -4295,6 +4786,8 @@ mod tests { affected_rows: None, conflicting_frag_reuse_indices: Vec::new(), conflicting_mem_wal_compacted_sstables: Vec::new(), + covered_field_paths: Default::default(), + schema: empty_lance_schema(), }; let same_name = Transaction::new( @@ -4350,6 +4843,8 @@ mod tests { affected_rows: None, conflicting_frag_reuse_indices: Vec::new(), conflicting_mem_wal_compacted_sstables: Vec::new(), + covered_field_paths: Default::default(), + schema: empty_lance_schema(), }; let different_name_result = rebase.check_txn(&different_name, 1); assert!( @@ -4418,6 +4913,8 @@ mod tests { affected_rows: None, conflicting_frag_reuse_indices: Vec::new(), conflicting_mem_wal_compacted_sstables: Vec::new(), + covered_field_paths: Default::default(), + schema: empty_lance_schema(), }; let result = rebase.check_txn(&rewrite, 2); if expect_conflict { @@ -4765,6 +5262,124 @@ mod tests { assert!(rebase.check_txn(&txn2, 2).is_ok()); } + /// A covering-index build must treat its *covered* fields as dependencies during + /// conflict resolution: a concurrent write to a covered field would leave the + /// rebased index serving stale materialized values. + #[tokio::test] + async fn test_create_covering_index_conflicts_with_writes_to_covered_field() { + let dataset = test_dataset(10, 2).await; // fields: 0 = "a", 1 = "b" + + // `covering_fields` is always the trailing entries of `fields`. + let make_index = |name: &str, covering_fields: Vec| { + let mut fields = vec![0]; + fields.extend(covering_fields.iter().copied()); + IndexMetadata { + uuid: Uuid::new_v4(), + fields, + name: name.to_string(), + dataset_version: 1, + fragment_bitmap: Some([0, 1].into_iter().collect()), + index_details: None, + index_version: 0, + created_at: None, + base_id: None, + files: None, + covering_fields, + } + }; + let create = |idx: IndexMetadata| { + Transaction::new_from_version( + 1, + Operation::CreateIndex { + new_indices: vec![idx], + removed_indices: vec![], + }, + ) + }; + let replace_field1 = || { + Transaction::new_from_version( + 1, + Operation::DataReplacement { + replacements: vec![DataReplacementGroup( + 0, + DataFile::new_legacy_from_fields("replace.lance", vec![1], None), + )], + }, + ) + }; + // `fragment_field_paths` reads the last live file listing a field, which is how an + // in-place rewrite presents itself (the rewrite adds a file for the field and + // tombstones the old entry, and tombstones are skipped). So appending a file that + // carries `field` models a rewrite of it, and one carrying a brand-new id models + // `add_columns`. + let merge_appending_file_for = |field: i32, path: &str| { + let mut fragments = dataset.fragments().as_ref().clone(); + for fragment in fragments.iter_mut() { + fragment + .files + .push(DataFile::new_legacy_from_fields(path, vec![field], None)); + } + Transaction::new_from_version( + 1, + Operation::Merge { + fragments, + schema: dataset.schema().clone(), + preserves_nullability: true, + }, + ) + }; + let merge_rewriting_covered = || merge_appending_file_for(1, "rewritten.lance"); + let merge_adding_column = || merge_appending_file_for(2, "added.lance"); + + // Covering index (covers field 1): a replacement of field 1 conflicts... + let mut rebase = + TransactionRebase::try_new(&dataset, create(make_index("covering", vec![1])), None) + .await + .unwrap(); + assert!( + rebase.check_txn(&replace_field1(), 2).is_err(), + "covering-index build must conflict with a data replacement of its covered field" + ); + // ...and so does a merge that rewrites the covered column's data file in place. + let mut rebase = + TransactionRebase::try_new(&dataset, create(make_index("covering", vec![1])), None) + .await + .unwrap(); + assert!( + rebase.check_txn(&merge_rewriting_covered(), 2).is_err(), + "covering-index build must conflict with a merge that rewrites its covered field" + ); + // But a merge that only ADDS a column leaves the covered field's data untouched, + // so it must NOT discard the (potentially hours-long) covered index build. + let mut rebase = + TransactionRebase::try_new(&dataset, create(make_index("covering", vec![1])), None) + .await + .unwrap(); + assert!( + rebase.check_txn(&merge_adding_column(), 2).is_ok(), + "covering-index build must not conflict with an add_columns that leaves covered data alone" + ); + + // Control: a plain index (no covered fields) keeps the optimistic path for + // both a merge and a replacement of the un-indexed field 1. + let mut rebase = + TransactionRebase::try_new(&dataset, create(make_index("plain", vec![])), None) + .await + .unwrap(); + assert!( + rebase.check_txn(&merge_rewriting_covered(), 2).is_ok(), + "non-covering index build should keep the optimistic merge path" + ); + let mut rebase = + TransactionRebase::try_new(&dataset, create(make_index("plain", vec![])), None) + .await + .unwrap(); + assert!( + rebase.check_txn(&replace_field1(), 2).is_ok(), + "non-covering index build should not conflict with a replacement of an un-indexed field" + ); + } + /// Returns the IDs of fragments that have been modified by this operation. /// /// This does not include new fragments. @@ -5099,6 +5714,8 @@ mod tests { affected_rows: None, conflicting_frag_reuse_indices: Vec::new(), conflicting_mem_wal_compacted_sstables: Vec::new(), + covered_field_paths: Default::default(), + schema: empty_lance_schema(), }; let result = rebase.check_txn(&txn2, 1); @@ -5162,6 +5779,8 @@ mod tests { affected_rows: None, conflicting_frag_reuse_indices: Vec::new(), conflicting_mem_wal_compacted_sstables: Vec::new(), + covered_field_paths: Default::default(), + schema: empty_lance_schema(), }; let result = rebase.check_txn(&committed_txn, 1); @@ -5200,6 +5819,8 @@ mod tests { affected_rows: None, conflicting_frag_reuse_indices: Vec::new(), conflicting_mem_wal_compacted_sstables: Vec::new(), + covered_field_paths: Default::default(), + schema: empty_lance_schema(), }; let result = rebase.check_txn(&committed_txn, 1); @@ -5239,6 +5860,8 @@ mod tests { affected_rows: None, conflicting_frag_reuse_indices: Vec::new(), conflicting_mem_wal_compacted_sstables: Vec::new(), + covered_field_paths: Default::default(), + schema: empty_lance_schema(), }; let result = rebase.check_txn(&committed_txn, 1); @@ -5278,6 +5901,8 @@ mod tests { affected_rows: None, conflicting_frag_reuse_indices: Vec::new(), conflicting_mem_wal_compacted_sstables: Vec::new(), + covered_field_paths: Default::default(), + schema: empty_lance_schema(), }; let result = rebase.check_txn(&committed_txn, 1); @@ -5328,6 +5953,8 @@ mod tests { affected_rows: None, conflicting_frag_reuse_indices: Vec::new(), conflicting_mem_wal_compacted_sstables: Vec::new(), + covered_field_paths: Default::default(), + schema: empty_lance_schema(), }; let result = rebase.check_txn(&committed_txn, 1); @@ -5353,6 +5980,8 @@ mod tests { affected_rows: None, conflicting_frag_reuse_indices: Vec::new(), conflicting_mem_wal_compacted_sstables: Vec::new(), + covered_field_paths: Default::default(), + schema: empty_lance_schema(), }; let result_higher = rebase_higher.check_txn(&committed_txn, 1); @@ -5399,6 +6028,8 @@ mod tests { affected_rows: None, conflicting_frag_reuse_indices: Vec::new(), conflicting_mem_wal_compacted_sstables: Vec::new(), + covered_field_paths: Default::default(), + schema: empty_lance_schema(), }; // CreateIndex of MemWalIndex should be compatible with UpdateMemWalState diff --git a/rust/lance/src/io/exec/ann_proto.rs b/rust/lance/src/io/exec/ann_proto.rs index c57ad4ca1b7..7f749d6ad0d 100644 --- a/rust/lance/src/io/exec/ann_proto.rs +++ b/rust/lance/src/io/exec/ann_proto.rs @@ -25,7 +25,7 @@ use uuid::Uuid; use crate::Dataset; use crate::pb; -use super::knn::{ANNIvfPartitionExec, ANNIvfSubIndexExec}; +use super::knn::{ANNIvfPartitionExec, ANNIvfSubIndexExec, resolve_physical_covering_query}; use super::table_identifier::{resolve_dataset, table_identifier_from_dataset}; use super::utils::PreFilterSource; @@ -118,6 +118,15 @@ pub fn query_to_proto(query: &Query) -> Result { dist_q_c: Some(query.dist_q_c), query_parallelism: Some(query.query_parallelism), approx_mode: approx_mode_to_proto(query.approx_mode) as i32, + // Three states, and the message wrapper is what keeps the middle one alive across + // the wire (see `CoveringProjection` in `ann.proto`). `map` -- never + // `unwrap_or_default()` -- because an empty list means "materialize nothing", + // while absence means "materialize everything the index declares". + covering_projection: query.covering_projection.as_ref().map(|columns| { + pb::CoveringProjection { + columns: columns.to_vec(), + } + }), }) } @@ -148,6 +157,13 @@ pub fn query_from_proto(proto: pb::VectorQueryProto) -> Result { query_parallelism: proto.query_parallelism.unwrap_or(DEFAULT_QUERY_PARALLELISM), dist_q_c: proto.dist_q_c.unwrap_or(0.0), approx_mode: approx_mode_from_proto(proto.approx_mode), + // A message field, so absence survives the round trip and stays distinct from an + // empty list. A producer predating field 15 sends nothing and the reconstructed + // query narrows nothing -- the conservative direction (see + // `Query::covering_projection`: under-narrowing costs time, never correctness). + covering_projection: proto + .covering_projection + .map(|projection| Arc::from(projection.columns)), }) } @@ -288,21 +304,31 @@ pub async fn ann_ivf_sub_index_exec_from_proto( } }; - ANNIvfSubIndexExec::try_new(input, dataset, indices, query, prefilter_source) + let query = resolve_physical_covering_query(dataset.clone(), &indices, &query).await?; + ANNIvfSubIndexExec::try_new_with_resolved_covering( + input, + dataset, + indices, + query, + prefilter_source, + ) } #[cfg(test)] mod tests { use super::*; - use arrow_array::types::{Float32Type, UInt32Type}; + use arrow_array::types::{Float32Type, Int32Type, UInt32Type}; use arrow_array::{ArrayRef, Float32Array, Float64Array}; use half::f16; use lance_datagen::{array, gen_batch}; use crate::index::DatasetIndexExt; use crate::index::vector::VectorIndexParams; + use datafusion_physical_plan::ExecutionPlan; use datafusion_physical_plan::test::TestMemoryExec; + use lance_core::ROW_ID; use lance_index::IndexType; + use lance_index::vector::DIST_COL; use lance_index::vector::ivf::IvfBuildParams; use lance_index::vector::pq::PQBuildParams; @@ -354,6 +380,7 @@ mod tests { query_parallelism: -1, dist_q_c: 0.42, approx_mode: ApproxMode::Accurate, + covering_projection: None, }; let proto = query_to_proto(&query).unwrap(); @@ -394,6 +421,7 @@ mod tests { query_parallelism: DEFAULT_QUERY_PARALLELISM, dist_q_c: 0.0, approx_mode: ApproxMode::Normal, + covering_projection: None, }; let proto = query_to_proto(&query).unwrap(); @@ -408,6 +436,148 @@ mod tests { assert_eq!(back.approx_mode, ApproxMode::Normal); } + /// `covering_projection` has three states and only two of them are expressible by a + /// bare proto3 `repeated string`, so the wire format wraps the list in a message. + /// Folding `Some(&[])` into `None` here would silently restore full materialization + /// on every distributed plan while every result stayed correct. + #[test] + fn test_query_roundtrip_covering_projection_three_states() { + let query = |projection: Option>| { + let key: ArrayRef = Arc::new(Float32Array::from(vec![1.0])); + Query { + column: "v".to_string(), + key, + k: 5, + lower_bound: None, + upper_bound: None, + minimum_nprobes: 1, + maximum_nprobes: None, + ef: None, + refine_factor: None, + metric_type: None, + use_index: true, + query_parallelism: DEFAULT_QUERY_PARALLELISM, + dist_q_c: 0.0, + approx_mode: ApproxMode::Normal, + covering_projection: projection.map(Arc::from), + } + }; + let roundtrip = |projection: Option>| { + let proto = query_to_proto(&query(projection)).unwrap(); + query_from_proto(proto).unwrap().covering_projection + }; + + assert!( + roundtrip(None).is_none(), + "'not computed' must stay 'not computed'" + ); + let narrowed = roundtrip(Some(vec![])).expect( + "'this query needs no covering column' must survive; decoding it as `None` \ + turns every remote plan back into a fully materializing one", + ); + assert!(narrowed.is_empty()); + assert_eq!( + roundtrip(Some(vec!["payload".to_string(), "price".to_string()])) + .expect("a computed projection must survive") + .to_vec(), + vec!["payload".to_string(), "price".to_string()], + ); + } + + /// The schema a remote executor declares must be the schema the planner declared. + /// The surrounding nodes -- `knn_combined`'s union and its projections -- were built + /// against the planner's, so a reconstruction that widens is not merely slower. + #[tokio::test] + async fn test_sub_index_proto_preserves_the_declared_covering_schema() { + let (dataset, _dir) = make_covered_indexed_dataset().await; + let indices = dataset.load_indices_by_name("vector_idx").await.unwrap(); + assert_eq!( + indices[0].covering_fields.len(), + 2, + "the fixture must actually be a covered index" + ); + + let input: Arc = + TestMemoryExec::try_new_exec( + &[], + super::super::knn::KNN_PARTITION_SCHEMA.clone(), + None, + ) + .unwrap(); + + for (projection, expected) in [ + ( + None, + vec![ + DIST_COL.to_string(), + ROW_ID.to_string(), + "price".to_string(), + "id".to_string(), + ], + ), + (Some(vec![]), vec![DIST_COL.to_string(), ROW_ID.to_string()]), + ( + Some(vec!["id".to_string()]), + vec![DIST_COL.to_string(), ROW_ID.to_string(), "id".to_string()], + ), + ] { + let key: ArrayRef = Arc::new(Float32Array::from(vec![0.1f32; 128])); + let query = Query { + column: "vector".to_string(), + key, + k: 10, + lower_bound: None, + upper_bound: None, + minimum_nprobes: 2, + maximum_nprobes: None, + ef: None, + refine_factor: None, + metric_type: Some(DistanceType::L2), + use_index: true, + query_parallelism: DEFAULT_QUERY_PARALLELISM, + dist_q_c: 0.0, + approx_mode: ApproxMode::Normal, + covering_projection: projection.clone().map(Arc::from), + }; + let query = resolve_physical_covering_query(dataset.clone(), &indices, &query) + .await + .unwrap(); + let exec = ANNIvfSubIndexExec::try_new_with_resolved_covering( + input.clone(), + dataset.clone(), + indices.clone(), + query, + PreFilterSource::None, + ) + .unwrap(); + let names = |schema: arrow_schema::SchemaRef| { + schema + .fields() + .iter() + .map(|field| field.name().clone()) + .collect::>() + }; + // Pin the local schema too: if this fixture stopped narrowing, the equality + // below would hold trivially for the wrong reason. + assert_eq!(names(exec.schema()), expected, "local plan, {projection:?}"); + + let proto = ann_ivf_sub_index_exec_to_proto(&exec).await.unwrap(); + let back = ann_ivf_sub_index_exec_from_proto( + proto, + Some(dataset.clone()), + input.clone(), + None, + ) + .await + .unwrap(); + assert_eq!( + names(back.schema()), + expected, + "reconstructed plan, {projection:?}" + ); + } + } + async fn make_vector_dataset() -> (Arc, tempfile::TempDir) { let dir = tempfile::tempdir().unwrap(); let batch = gen_batch() @@ -429,6 +599,41 @@ mod tests { (Arc::new(ds), dir) } + /// An IVF_PQ index covering `price` and `id`, declared in the reverse of schema + /// order so the assertions above can tell declaration order from schema order. + async fn make_covered_indexed_dataset() -> (Arc, tempfile::TempDir) { + let dir = tempfile::tempdir().unwrap(); + let batch = gen_batch() + .col("id", array::step::()) + .col("price", array::step::()) + .col( + "vector", + array::rand_vec::(lance_datagen::Dimension::from(128)), + ) + .into_batch_rows(lance_datagen::RowCount::from(256)) + .unwrap(); + let path = dir.path().join("test_ann_covered.lance"); + let mut ds = Dataset::write( + arrow_array::RecordBatchIterator::new(vec![Ok(batch.clone())], batch.schema()), + path.to_str().unwrap(), + None, + ) + .await + .unwrap(); + + let mut index_params = VectorIndexParams::with_ivf_pq_params( + DistanceType::L2, + IvfBuildParams::new(2), + PQBuildParams::default(), + ); + index_params.covering_columns(vec!["price".to_string(), "id".to_string()]); + ds.create_index(&["vector"], IndexType::Vector, None, &index_params, false) + .await + .unwrap(); + let ds = Dataset::open(path.to_str().unwrap()).await.unwrap(); + (Arc::new(ds), dir) + } + async fn make_indexed_dataset() -> (Arc, tempfile::TempDir) { let (_dataset, dir) = make_vector_dataset().await; let mut ds = Dataset::open(dir.path().join("test_ann.lance").to_str().unwrap()) @@ -472,6 +677,7 @@ mod tests { query_parallelism: DEFAULT_QUERY_PARALLELISM, dist_q_c: 0.0, approx_mode: ApproxMode::Normal, + covering_projection: None, }; let exec = ANNIvfPartitionExec::try_new( @@ -520,6 +726,7 @@ mod tests { query_parallelism: DEFAULT_QUERY_PARALLELISM, dist_q_c: 0.0, approx_mode: ApproxMode::Normal, + covering_projection: None, }; // Use a TestMemoryExec as a mock input child (provides the KNN_PARTITION_SCHEMA) @@ -584,6 +791,7 @@ mod tests { query_parallelism: DEFAULT_QUERY_PARALLELISM, dist_q_c: 0.0, approx_mode: ApproxMode::Normal, + covering_projection: None, }; let input: Arc = TestMemoryExec::try_new_exec( diff --git a/rust/lance/src/io/exec/knn.rs b/rust/lance/src/io/exec/knn.rs index 7a334198610..2a1ea633ef2 100644 --- a/rust/lance/src/io/exec/knn.rs +++ b/rust/lance/src/io/exec/knn.rs @@ -47,7 +47,7 @@ use lance_datafusion::utils::{ DELTAS_SEARCHED_METRIC, ExecutionPlanMetricsSetExt, FIND_PARTITIONS_ELAPSED_METRIC, PARTITIONS_RANKED_METRIC, PARTITIONS_SEARCHED_METRIC, }; -use lance_index::metrics::MetricsCollector; +use lance_index::metrics::{MetricsCollector, NoOpMetricsCollector}; use lance_index::prefilter::PreFilter; use lance_index::vector::DIST_Q_C_COLUMN; use lance_index::vector::{ @@ -62,8 +62,9 @@ use roaring::RoaringBitmap; use tokio::sync::Notify; use uuid::Uuid; -use crate::dataset::Dataset; +use crate::dataset::{Dataset, ProjectionRequest, TakeBuilder}; use crate::index::DatasetIndexInternalExt; +use crate::index::covering::{effective_covering, effective_covering_with_ids}; use crate::index::prefilter::{DatasetPreFilter, FilterLoader}; use crate::index::vector::utils::{get_vector_type, validate_distance_type_for}; use crate::{Error, Result}; @@ -1102,11 +1103,150 @@ pub static KNN_PARTITION_SCHEMA: LazyLock = LazyLock::new(|| { ])) }); +fn covering_fields_match(expected: &Field, physical: &Field) -> bool { + expected.name() == physical.name() + && expected.data_type() == physical.data_type() + && expected.is_nullable() == physical.is_nullable() + && expected.metadata() == physical.metadata() +} + +/// Intersect a requested manifest declaration with the capabilities proven by every +/// selected physical segment. +/// +/// The output remains in declaration order because that is the schema the ANN exec +/// exposes. If a segment stores the surviving columns in a different relative order then +/// none are served: storage emits in physical order, and pairing those values with a +/// declaration-ordered schema would be incorrect. +fn common_physical_covering( + mut requested: Vec<(i32, Field)>, + physical_by_segment: &[Vec<(i32, Field)>], +) -> Vec<(i32, Field)> { + if physical_by_segment.is_empty() { + return Vec::new(); + } + + for physical in physical_by_segment { + let before = requested.len(); + requested.retain(|(expected_id, expected_field)| { + let kept = physical.iter().any(|(physical_id, physical_field)| { + expected_id == physical_id && covering_fields_match(expected_field, physical_field) + }); + if !kept { + // Withdrawal is silent by design -- results stay correct, the plan just + // grows a base-table read -- which makes a covered index that quietly + // stopped eliding takes indistinguishable from one that never worked. + // `covering_fields_match` compares name, data type, nullability AND + // metadata, so the cause is usually a drift too small to guess at. + log::debug!( + "Covering column '{}' (field id {}) is declared but not servable by a \ + segment: no physical column matches it on name, type, nullability and \ + metadata. Its values will come from a base-table read.", + expected_field.name(), + expected_id + ); + } + kept + }); + if requested.is_empty() { + if before > 0 { + log::debug!( + "No declared covering column is servable by every selected segment; \ + this query reads all of them from the base table." + ); + } + return requested; + } + } + + let requested_ids = requested.iter().map(|(id, _)| *id).collect::>(); + for physical in physical_by_segment { + let physical_ids = physical + .iter() + .filter(|(physical_id, physical_field)| { + requested.iter().any(|(expected_id, expected_field)| { + expected_id == physical_id + && covering_fields_match(expected_field, physical_field) + }) + }) + .map(|(id, _)| *id) + .collect::>(); + if requested_ids != physical_ids { + // Segments disagree on the order or set they can serve. Storage emits in + // physical order while the schema is declaration-ordered, so pairing them + // would be wrong -- none are served rather than some. + log::debug!( + "Selected segments do not agree on the covering columns they can serve \ + ({requested_ids:?} vs {physical_ids:?}); this query reads all of them from \ + the base table." + ); + return Vec::new(); + } + } + requested +} + +/// Resolve the query's manifest-requested covering columns against every selected +/// segment's physical storage capability. +/// +/// The returned query always has an explicit covering projection. A missing physical +/// column is omitted from it, causing the normal scanner take to fetch that value from +/// the base table instead. +pub async fn resolve_physical_covering_query( + dataset: Arc, + indices: &[IndexMetadata], + query: &Query, +) -> Result { + let mut resolved_query = query.clone(); + let declared_ids = indices + .first() + .map(|index| index.covering_fields.as_slice()) + .unwrap_or_default(); + if let Some(mismatch) = indices + .iter() + .find(|index| index.covering_fields != declared_ids) + { + return Err(Error::index(format!( + "ANN index delta segments of index '{}' disagree on covering fields ({:?} vs {:?}); \ + the index metadata is inconsistent -- rebuild the index", + mismatch.name, declared_ids, mismatch.covering_fields + ))); + } + + let requested = effective_covering_with_ids( + declared_ids, + query.covering_projection.as_deref(), + dataset.schema(), + )?; + if requested.is_empty() { + resolved_query.covering_projection = Some(Arc::from(Vec::::new())); + return Ok(resolved_query); + } + + let physical_by_segment = future::try_join_all(indices.iter().map(|index| { + let dataset = dataset.clone(); + let column = query.column.clone(); + let uuid = index.uuid; + async move { + dataset + .open_vector_index(&column, &uuid, &NoOpMetricsCollector) + .await? + .physical_covering_fields() + } + })) + .await?; + let names = common_physical_covering(requested, &physical_by_segment) + .into_iter() + .map(|(_, field)| field.name().clone()) + .collect::>(); + resolved_query.covering_projection = Some(Arc::from(names)); + Ok(resolved_query) +} + /// Create a new ANN execution node. `overlay_block`, when `Some`, excludes rows whose index /// entries may be stale due to a newer data overlay (see [`ANNIvfSubIndexExec::with_overlay_block`]). /// `external_mask`, when `Some`, additionally restricts the scan to a caller-supplied /// allow/block set (see [`ANNIvfSubIndexExec::with_external_mask`]). -pub fn new_knn_exec( +pub async fn new_knn_exec( dataset: Arc, indices: &[IndexMetadata], query: &Query, @@ -1114,17 +1254,18 @@ pub fn new_knn_exec( overlay_block: Option, external_mask: Option>, ) -> Result> { + let query = resolve_physical_covering_query(dataset.clone(), indices, query).await?; let ivf_node = ANNIvfPartitionExec::try_new( dataset.clone(), indices.iter().map(|idx| idx.uuid).collect_vec(), query.clone(), )?; - let mut sub_index = ANNIvfSubIndexExec::try_new( + let mut sub_index = ANNIvfSubIndexExec::try_new_with_resolved_covering( Arc::new(ivf_node), dataset, indices.to_vec(), - query.clone(), + query, prefilter_source, )?; if let Some(overlay_block) = overlay_block { @@ -1399,6 +1540,19 @@ pub struct ANNIvfSubIndexExec { /// Optional external row-address allow/block mask, combined with the /// prefilter using logical AND. external_mask: Option>, + /// Output schema: `[_distance, _rowid]` plus any included/covering columns + /// the index emits. + output_schema: SchemaRef, + + /// Whether this plan emits covering columns, resolved from the index + /// declaration at construction. Kept as its own field rather than re-derived + /// from the output schema's width at execution: width-based inference is + /// exact only while covering is the sole thing that widens the schema, and + /// the first unrelated extra column would silently arm per-query emitted-row + /// tracking and the covered recovery take on every plain query (or mask + /// covering if a column were dropped) with results staying correct -- no test + /// would catch it. + has_covered: bool, /// Datafusion Plan Properties properties: Arc, @@ -1407,7 +1561,24 @@ pub struct ANNIvfSubIndexExec { } impl ANNIvfSubIndexExec { + /// Construct an ANN sub-index exec without a verified physical covering capability. + /// + /// This synchronous entry point cannot open segment storage, so it conservatively + /// disables covering and leaves declared values to the base-table take. Dataset + /// planning and protobuf reconstruction use the storage-aware async resolution path + /// before calling `Self::try_new_with_resolved_covering` (private, so not linked). pub fn try_new( + input: Arc, + dataset: Arc, + indices: Vec, + mut query: Query, + prefilter_source: PreFilterSource, + ) -> Result { + query.covering_projection = Some(Arc::from(Vec::::new())); + Self::try_new_with_resolved_covering(input, dataset, indices, query, prefilter_source) + } + + pub(crate) fn try_new_with_resolved_covering( input: Arc, dataset: Arc, indices: Vec, @@ -1420,8 +1591,47 @@ impl ANNIvfSubIndexExec { PART_ID_COLUMN ))); } + // Start from the logical declaration, narrowed to the explicit projection already + // proven against every selected segment's physical storage. The per-partition + // search emits only this subset, so TakeExec may skip the matching base-table fields. + let mut fields = KNN_INDEX_SCHEMA.fields().to_vec(); + let included_ids = indices + .first() + .map(|idx| idx.covering_fields.clone()) + .unwrap_or_default(); + // This exec declares one output schema but executes every delta segment, + // each emitting covering columns from its own storage -- so all deltas + // must agree on the covering set. The commit boundary enforces this; the + // check here defends against already-committed inconsistent metadata, + // failing with a clear error instead of a downstream schema mismatch. + if let Some(mismatch) = indices + .iter() + .find(|idx| idx.covering_fields != included_ids) + { + return Err(Error::index(format!( + "ANNIvfSubIndexExec: delta segments of index '{}' disagree on covering fields \ + ({:?} vs {:?}); the index metadata is inconsistent -- rebuild the index", + mismatch.name, included_ids, mismatch.covering_fields + ))); + } + // Narrow only AFTER the agreement check above: narrowing first could reduce + // segments that genuinely disagree into a subset on which they happen to agree, + // masking metadata corruption as a working query. The manifest remains + // authoritative for logical identity and schema dependencies; the projection is + // authoritative for physical availability. + let covering = effective_covering( + &included_ids, + query.covering_projection.as_deref(), + dataset.schema(), + )?; + // Resolved-at-construction covered flag (see the `has_covered` field): from + // the NARROWED set, not the declaration -- a projection that needs none of + // the declared columns (`Some(&[])`) makes this plan non-covered. + let has_covered = !covering.is_empty(); + fields.extend(covering.into_iter().map(Arc::new)); + let output_schema = Arc::new(arrow_schema::Schema::new(fields)); let properties = Arc::new(PlanProperties::new( - EquivalenceProperties::new(KNN_INDEX_SCHEMA.clone()), + EquivalenceProperties::new(output_schema.clone()), Partitioning::RoundRobinBatch(1), EmissionType::Final, Boundedness::Bounded, @@ -1434,6 +1644,8 @@ impl ANNIvfSubIndexExec { prefilter_source, overlay_block: None, external_mask: None, + output_schema, + has_covered, properties, metrics: ExecutionPlanMetricsSet::new(), }) @@ -1513,6 +1725,16 @@ struct ANNIvfEarlySearchResults { deltas_remaining: AtomicUsize, all_deltas_done: Notify, took_no_rows_shortcut: AtomicBool, + /// One entry per covered delta whose late search reached the no-rows shortcut: + /// that delta's per-segment ownership mask (`None` when the delta owned the + /// whole global mask). The exec-level covered recovery emits ONLY rows these + /// scopes select, so covered result sets stay identical to non-covered ones in + /// both directions: a delta that never armed (it skipped its late search, e.g. + /// probing fewer partitions than `minimum_nprobes`) would not have emitted its + /// not-found rows on the non-covered path either, so recovery must not emit + /// them for it -- and an empty list means no delta armed, so recovery emits + /// nothing at all. + covered_recovery_scopes: Mutex>>>, } impl ANNIvfEarlySearchResults { @@ -1524,20 +1746,26 @@ impl ANNIvfEarlySearchResults { deltas_remaining: AtomicUsize::new(deltas_remaining), all_deltas_done: Notify::new(), took_no_rows_shortcut: AtomicBool::new(false), + covered_recovery_scopes: Mutex::new(Vec::new()), } } - fn record_batch(&self, batch: &RecordBatch) { + fn record_batch(&self, batch: &RecordBatch) -> DataFusionResult<()> { + let row_ids = batch.column_by_name(ROW_ID).ok_or_else(|| { + DataFusionError::Internal(format!( + "ANNIvfEarlySearchResults::record_batch: batch is missing the {ROW_ID} column" + )) + })?; let mut initial_ids = self.initial_ids.lock().unwrap(); let ids_to_record = (self.k - initial_ids.len()).min(batch.num_rows()); initial_ids.extend( - batch - .column(1) + row_ids .as_primitive::() .values() .iter() .take(ids_to_record), ); + Ok(()) } fn record_late_batch(&self, num_rows: usize) { @@ -1763,7 +1991,7 @@ impl ANNIvfSubIndexExec { } metrics.baseline_metrics.record_output(batch.num_rows()); if record_initial { - state.record_batch(&batch); + state.record_batch(&batch)?; } Ok(batch) }) @@ -1782,6 +2010,7 @@ impl ANNIvfSubIndexExec { state: Arc, target_partitions: usize, seg_mask: Option>, + has_covered: bool, ) -> impl Stream> { let stream = futures::stream::once(async move { let max_nprobes = query @@ -1826,6 +2055,29 @@ impl ANNIvfSubIndexExec { // In this case there are fewer than k results matching the prefilter so // just return the prefilter ids and don't bother searching any further + // The shortcut batch emits only [_distance, _rowid]; a covering + // index's declared schema carries more columns, and their values + // cannot be fabricated here. Emit nothing instead: the + // end-of-stream covered recovery (`covered_not_found_batch`) + // returns every unemitted prefilter row with INFINITY distance + // and its covering columns via a take bounded by the prefilter + // size -- mirroring the non-covered shortcut's semantics without + // searching (a prefilter row with no index entry, e.g. a null + // vector, could otherwise never be found and would drive the + // search through every partition). + if has_covered { + // Arm recovery with THIS delta's ownership scope: recovery may + // emit only rows the arming deltas own, or it diverges from the + // non-covered result set whenever a sibling delta skipped its + // late search without arming (see `covered_recovery_scopes`). + state + .covered_recovery_scopes + .lock() + .unwrap() + .push(seg_mask.clone()); + return futures::stream::empty().boxed(); + } + // This next if check should be true, because we wouldn't get max_results otherwise if let Some(iter_addrs) = prefilter_mask.iter_addrs() { // Emit the prefilter rows that the partition search did not reach. @@ -2022,13 +2274,146 @@ impl ANNIvfSubIndexExec { seg_mask, ) .await?; - state.record_batch(&batch); + state.record_batch(&batch)?; Ok(batch) } }) .buffered(query_parallelism) .boxed() } + + /// Recover rows that match a bounded, selective prefilter but have no index entry + /// (their vector is null or an empty multivector), which the covered search never + /// returns. The non-covered path recovers them via its not-found shortcut with + /// `INFINITY` distance; a covered index must do the same but also supply the covering + /// columns, which it fetches from the base table for these few missing rows. Only + /// rows selected by an arming delta's ownership scope in `scopes` are recovered + /// (`None` = that delta owned the whole global mask): a delta that skipped its late + /// search without arming never emits its not-found rows on the non-covered path, so + /// recovering them here would make the covered result set diverge. Returns `None` + /// when there is nothing to recover (non-covered, unbounded/non-selective prefilter, + /// or every prefilter row was already emitted or unowned). + async fn covered_not_found_batch( + has_covered: bool, + prefilter: Arc, + dataset: Arc, + output_schema: SchemaRef, + k: usize, + emitted: Arc>, + scopes: Vec>>, + ) -> DataFusionResult> { + if !has_covered { + return Ok(None); + } + prefilter + .wait_for_ready() + .await + .map_err(|e| DataFusionError::Execution(format!("prefilter not ready: {e}")))?; + let mask = prefilter.mask(); + // Only a bounded, selective prefilter (at most k allowed rows) guarantees every + // matching row belongs in the result; an unbounded prefilter has no not-found set. + let Some(max_len) = mask.max_len() else { + return Ok(None); + }; + if max_len as usize > k { + return Ok(None); + } + let Some(addrs) = mask.iter_addrs() else { + return Ok(None); + }; + let not_found: Vec = { + let emitted = emitted.lock().unwrap(); + addrs + .map(u64::from) + .filter(|addr| !emitted.contains(*addr)) + .filter(|addr| { + // Mirror the non-covered shortcut's per-segment ownership rule: + // each armed delta may recover only the rows its segment owns. + scopes + .iter() + .any(|scope| scope.as_ref().is_none_or(|m| m.selected(*addr))) + }) + .collect() + }; + if not_found.is_empty() { + return Ok(None); + } + + // Fetch the covering columns (the fields after `[_distance, _rowid]`) from the + // base table for the missing rows; the take is bounded by `not_found.len()`. + // Carry `_rowid` through as well: the take resolves ids with + // `MissingRowPolicy::Ignore` and silently returns FEWER rows than requested when an + // id no longer resolves through the row-id index -- a stale prefilter listing rows + // of a fragment a later delete removed entirely. Rebuilding the batch from + // `not_found` would then pair N distance/rowid values with < N covering values and + // fail with "all columns in a record batch must have the same length". Carrying + // `_rowid` lets the batch be rebuilt from the rows the take actually returned, + // which is what the non-covered path does when its `TakeExec` drops the same ids. + // + // A row that still resolves but is tombstoned would shorten the batch the same way. + // That is not reachable today, so the handling below is defensive for that case, not + // load-bearing. The invariant is enforced by `DatasetPreFilter`, not by the individual + // prefilter sources: its `deleted_ids` mask (`prefilter.rs`) is intersected into the final + // mask and covers every fragment in the manifest -- fragments inside the index bitmap + // contribute their deletion vectors, fragments outside it are blocked wholesale, and + // the mask is elided only when no fragment carries a deletion file at all. So a + // tombstoned address never enters `mask` regardless of which source produced it. + // Check that invariant there, not here, when adding a new `PreFilterSource`. + // + // `include_row_id()`, NOT `_rowid` in the projected column list: naming it there + // sets `must_add_row_offset`, which both arms the take's "must not target deleted + // rows" check (turning any shortfall into a hard error) and computes a row-offset + // column per batch that `project_batch` then discards. + let covered_names: Vec = output_schema + .fields() + .iter() + .skip(2) + .map(|field| field.name().clone()) + .collect(); + let projection = ProjectionRequest::from_columns(covered_names, dataset.schema()); + // `not_found` holds row ids (the `_rowid` space): stable logical ids on a + // stable-row-id dataset, physical addresses otherwise. `try_new_from_ids` + // resolves them through the row-id index when one exists and falls back to + // identity (id == address) when it does not, so it is correct in both cases -- + // unlike `try_new_from_addresses`, which would misread a stable id as + // `frag = id >> 32, offset = id` and take the wrong physical row (or error). + let covered = TakeBuilder::try_new_from_ids(dataset, not_found, projection)? + .include_row_id() + .execute() + .await?; + + // Realign on what came back, not on what was asked for. + let returned_ids = covered + .column_by_name(ROW_ID) + .ok_or_else(|| { + DataFusionError::Internal(format!( + "covered recovery take did not return the requested {ROW_ID} column" + )) + })? + .clone(); + let n = covered.num_rows(); + if n == 0 { + return Ok(None); + } + // Assemble by name. `include_row_id` appends `_rowid` to the take's output rather + // than placing it first, so no position in `covered` is contractual; the declared + // `[_distance, _rowid, ]` order comes from `output_schema` alone. + let mut columns: Vec = Vec::with_capacity(output_schema.fields().len()); + columns.push(Arc::new(Float32Array::from_value(f32::INFINITY, n))); + columns.push(returned_ids); + for field in output_schema.fields().iter().skip(2) { + let column = covered.column_by_name(field.name()).ok_or_else(|| { + DataFusionError::Internal(format!( + "covered recovery take did not return the requested column '{}'", + field.name() + )) + })?; + columns.push(column.clone()); + } + let batch = RecordBatch::try_new(output_schema.clone(), columns) + .map_err(|e| DataFusionError::ArrowError(Box::new(e), None))?; + Ok(Some(batch)) + } } impl ExecutionPlan for ANNIvfSubIndexExec { @@ -2037,7 +2422,7 @@ impl ExecutionPlan for ANNIvfSubIndexExec { } fn schema(&self) -> arrow_schema::SchemaRef { - KNN_INDEX_SCHEMA.clone() + self.output_schema.clone() } fn children(&self) -> Vec<&Arc> { @@ -2084,6 +2469,8 @@ impl ExecutionPlan for ANNIvfSubIndexExec { prefilter_source, overlay_block: self.overlay_block.clone(), external_mask: self.external_mask.clone(), + output_schema: self.output_schema.clone(), + has_covered: self.has_covered, properties: self.properties.clone(), metrics: ExecutionPlanMetricsSet::new(), } @@ -2102,6 +2489,10 @@ impl ExecutionPlan for ANNIvfSubIndexExec { ) -> DataFusionResult { let input_stream = self.input.execute(partition, context.clone())?; let schema = self.schema(); + // When the index emits covering columns, the late-search no-rows shortcut + // (which emits only [_distance, _rowid]) is disabled in favour of the + // end-of-stream covered recovery. Resolved at construction; see the field. + let has_covered = self.has_covered; let target_partitions = context.session_config().target_partitions(); let query = self.query.clone(); let ds = self.dataset.clone(); @@ -2202,6 +2593,25 @@ impl ExecutionPlan for ANNIvfSubIndexExec { let state = Arc::new(ANNIvfEarlySearchResults::new(indices.len(), query.k)); + // Covered null-vector recovery: track which row ids the search emits, then after + // it finishes emit any prefilter-matched rows that had no index entry, with their + // covering columns taken from the base table (see `covered_not_found_batch`). + // This must NOT be gated on `prefilter_source`: `DatasetPreFilter` builds a bounded + // allow-list from the deletion mask alone, so an unfiltered query over a dataset with + // deletions (or a fragment outside the index bitmap) still arms recovery. Skipping the + // bookkeeping there left `emitted` empty, and recovery re-emitted every live row on top + // of the ones the search had already returned. The cost is already confined to covered + // plans, and recovery itself stays gated behind `has_covered`. + let track_emitted = has_covered; + let emitted_row_ids = Arc::new(std::sync::Mutex::new(roaring::RoaringTreemap::new())); + let emitted_for_inspect = emitted_row_ids.clone(); + let recon_prefilter = pre_filter.clone(); + let recon_state = state.clone(); + let recon_ds = ds.clone(); + let recon_schema = schema.clone(); + let recon_k = query.k; + let declared_schema = schema.clone(); + Ok(Box::pin(RecordBatchStreamAdapter::new( schema, per_index_stream @@ -2277,6 +2687,7 @@ impl ExecutionPlan for ANNIvfSubIndexExec { state, target_partitions, seg_mask, + has_covered, ); DataFusionResult::Ok(early_search.chain(late_search).boxed()) } @@ -2285,6 +2696,66 @@ impl ExecutionPlan for ANNIvfSubIndexExec { // Each delta stream is split into an early and late search. The late search // will not start until the early search is complete across all deltas. .try_flatten_unordered(None) + // The search emits covering columns read from index STORAGE, while this node + // declares them from the manifest (`IndexMetadata.covering_fields`). The two can + // disagree: `FLAG_COVERED_INDEX_METADATA` is a best-effort writer fence, so a + // client predating it can drop the declaration while the index files keep the + // payload. Without narrowing, this stream mixes batch widths -- a wider search + // batch next to the narrower non-covered shortcut batch -- which panics inside + // DataFusion's TopK (`interleave_record_batch` reads every batch through the first + // batch's schema). Only ever narrows: the opposite desync (declared but absent from + // storage) is prevented at build time (covering requires the V3 index format and + // rejects columns the storage cannot emit) and otherwise fails loudly, which is the + // right outcome. + // + // Caveat worth knowing: the test below is cardinality-based, so it cannot detect a + // desync in which storage holds the same *number* of covering columns as the + // declaration but different ones. + .map(move |batch| { + let batch = batch?; + if batch.num_columns() > declared_schema.fields().len() { + return batch + .project_by_schema(declared_schema.as_ref()) + .map_err(|e| DataFusionError::ArrowError(Box::new(e), None)); + } + Ok(batch) + }) + .inspect_ok(move |batch| { + // Record emitted row ids so the reconciliation below can find the + // prefilter rows the search never returned (no index entry). + if track_emitted && let Some(col) = batch.column_by_name(ROW_ID) { + emitted_for_inspect + .lock() + .unwrap() + .extend(col.as_primitive::().values().iter().copied()); + } + }) + .chain( + stream::once(async move { + // Only recover when a covered late search actually took the + // no-rows shortcut; otherwise the non-covered path would not + // have emitted these rows either, and covered must match. + // Each armed delta contributes its ownership scope, and the + // recovery emits only rows those scopes select. + let scopes = std::mem::take( + &mut *recon_state.covered_recovery_scopes.lock().unwrap(), + ); + if scopes.is_empty() { + return Ok(None); + } + Self::covered_not_found_batch( + has_covered, + recon_prefilter, + recon_ds, + recon_schema, + recon_k, + emitted_row_ids, + scopes, + ) + .await + }) + .try_filter_map(|opt| future::ready(Ok(opt))), + ) .finally(move || { // Publish the exact index-file I/O measured for this query // (cache misses only) to the iops/requests/bytes_read gauges. @@ -2360,13 +2831,24 @@ pub struct MultivectorScoringExec { // the inputs are sorted ANN search results inputs: Vec>, query: Query, + /// Output schema: `[_distance, _rowid]` plus any covering columns the ANN children + /// emit. The covered payload is per source row (replicated across its sub-vectors), + /// so it is carried through the per-row-id reduction below. + output_schema: SchemaRef, properties: Arc, } impl MultivectorScoringExec { pub fn try_new(inputs: Vec>, query: Query) -> Result { + // Inherit the children's schema so covering columns flow through scoring instead + // of being dropped (a covered multivector query would otherwise still take from + // the base table). + let output_schema = inputs + .first() + .map(|input| input.schema()) + .unwrap_or_else(|| KNN_INDEX_SCHEMA.clone()); let properties = Arc::new(PlanProperties::new( - EquivalenceProperties::new(KNN_INDEX_SCHEMA.clone()), + EquivalenceProperties::new(output_schema.clone()), Partitioning::RoundRobinBatch(1), EmissionType::Final, Boundedness::Bounded, @@ -2375,6 +2857,7 @@ impl MultivectorScoringExec { Ok(Self { inputs, query, + output_schema, properties, }) } @@ -2399,7 +2882,7 @@ impl ExecutionPlan for MultivectorScoringExec { } fn schema(&self) -> arrow_schema::SchemaRef { - KNN_INDEX_SCHEMA.clone() + self.output_schema.clone() } fn children(&self) -> Vec<&Arc> { @@ -2434,11 +2917,18 @@ impl ExecutionPlan for MultivectorScoringExec { .map(|input| input.execute(partition, context.clone())) .collect::>>()?; + // Covering columns (if any) follow `[_distance, _rowid]` in the child schema. + // They are per source row -- identical across a row's sub-vectors -- so we carry + // them through the per-row-id reduction and re-attach one copy per output row. + let output_schema = self.output_schema.clone(); + let n_covered = output_schema.fields().len().saturating_sub(2); + // collect the top k results from each stream, // and max-reduce for each query, // records the minimum distance for each query as estimation. let mut reduced_inputs = stream::select_all(inputs.into_iter().map(|stream| { - stream.map(|batch| { + let output_schema = output_schema.clone(); + stream.map(move |batch| { let batch = batch?; let row_ids = batch[ROW_ID].as_primitive::(); let dists = batch[DIST_COL].as_primitive::(); @@ -2452,9 +2942,15 @@ impl ExecutionPlan for MultivectorScoringExec { .unwrap_or_default(); let mut new_row_ids = Vec::with_capacity(row_ids.len()); let mut new_sims = Vec::with_capacity(row_ids.len()); + let mut keep_indices: Vec = Vec::with_capacity(row_ids.len()); let mut visited_row_ids = HashSet::with_capacity(row_ids.len()); - for (row_id, dist) in row_ids.values().iter().zip(dists.values().iter()) { + for (i, (row_id, dist)) in row_ids + .values() + .iter() + .zip(dists.values().iter()) + .enumerate() + { // the results are sorted by distance, so we can skip if we have seen this row id before if visited_row_ids.contains(row_id) { continue; @@ -2463,14 +2959,20 @@ impl ExecutionPlan for MultivectorScoringExec { new_row_ids.push(*row_id); // it's cosine distance, so we need to convert it to similarity new_sims.push(1.0 - *dist); + keep_indices.push(i as u32); } - let new_row_ids = UInt64Array::from(new_row_ids); - let new_dists = Float32Array::from(new_sims); - let batch = RecordBatch::try_new( - KNN_INDEX_SCHEMA.clone(), - vec![Arc::new(new_dists), Arc::new(new_row_ids)], - )?; + let mut columns: Vec = vec![ + Arc::new(Float32Array::from(new_sims)), + Arc::new(UInt64Array::from(new_row_ids)), + ]; + if n_covered > 0 { + let keep = arrow_array::UInt32Array::from(keep_indices); + for covered_col in batch.columns().iter().skip(2) { + columns.push(arrow::compute::take(covered_col, &keep, None)?); + } + } + let batch = RecordBatch::try_new(output_schema.clone(), columns)?; Ok::<_, DataFusionError>((min_sim, batch)) }) @@ -2482,11 +2984,28 @@ impl ExecutionPlan for MultivectorScoringExec { let stream = stream::once(async move { // at most, we will have k * refine_factor results for each query let mut results = HashMap::with_capacity(k * refactor); + // Covering columns of every reduced batch, kept alive (cheap Arc clones) + // for one zero-copy-per-cell `interleave` gather at emit. The map records + // where each row id's covering values live: (batch position, row position), + // first-seen (the values are per source row, identical across batches). + // Materializing per-cell `ScalarValue`s here instead would copy every + // candidate cell once into the map and again through `iter_to_array` -- + // several MB per query for wide payloads. + let mut covered_batches: Vec> = Vec::new(); + let mut payload_locs: HashMap = HashMap::new(); let mut missed_sim_sum = 0.0; while let Some((min_sim, batch)) = reduced_inputs.try_next().await? { let row_ids = batch[ROW_ID].as_primitive::(); let sims = batch[DIST_COL].as_primitive::(); + if n_covered > 0 { + let batch_pos = covered_batches.len(); + for (i, row_id) in row_ids.values().iter().enumerate() { + payload_locs.entry(*row_id).or_insert((batch_pos, i)); + } + covered_batches.push(batch.columns().iter().skip(2).cloned().collect()); + } + let query_results = row_ids .values() .iter() @@ -2512,18 +3031,47 @@ impl ExecutionPlan for MultivectorScoringExec { missed_sim_sum += min_sim; } - let (row_ids, sims): (Vec<_>, Vec<_>) = results.into_iter().unzip(); + let (row_ids, sims): (Vec, Vec) = results.into_iter().unzip(); let dists = sims - .into_iter() + .iter() // it's similarity, so we need to convert it back to distance .map(|sim| num_queries - sim) .collect::>(); - let row_ids = UInt64Array::from(row_ids); - let dists = Float32Array::from(dists); - let batch = RecordBatch::try_new( - KNN_INDEX_SCHEMA.clone(), - vec![Arc::new(dists), Arc::new(row_ids)], - )?; + let mut columns: Vec = vec![ + Arc::new(Float32Array::from(dists)), + Arc::new(UInt64Array::from(row_ids.clone())), + ]; + if n_covered > 0 && !row_ids.is_empty() { + // Every emitted row id was recorded while scanning the reduced + // batches (both maps are fed from the same iteration), so a miss + // here is a results/payloads desync. Fail loudly rather than + // fall back to nulls: a silent null would masquerade as a real + // covered value and hide the desync from every result-based test. + let locations = row_ids + .iter() + .map(|row_id| { + payload_locs.get(row_id).copied().ok_or_else(|| { + DataFusionError::Internal(format!( + "multivector covering: row id {row_id} missing from \ + covering source" + )) + }) + }) + .collect::>>()?; + for covered_pos in 0..n_covered { + let arrays: Vec<&dyn Array> = covered_batches + .iter() + .map(|cols| cols[covered_pos].as_ref()) + .collect(); + columns.push(arrow::compute::interleave(&arrays, &locations)?); + } + } else { + for covered_pos in 0..n_covered { + let field = output_schema.field(2 + covered_pos); + columns.push(arrow::array::new_empty_array(field.data_type())); + } + } + let batch = RecordBatch::try_new(output_schema.clone(), columns)?; Ok::<_, DataFusionError>(batch) }); Ok(Box::pin(RecordBatchStreamAdapter::new( @@ -2597,9 +3145,334 @@ mod tests { query_parallelism: DEFAULT_QUERY_PARALLELISM, dist_q_c: 0.0, approx_mode: Default::default(), + covering_projection: None, } } + fn covering_field(id: i32, name: &str, data_type: DataType) -> (i32, Field) { + (id, Field::new(name, data_type, true)) + } + + #[test] + fn test_common_physical_covering_requires_every_segment_to_prove_the_field() { + let price = covering_field(17, "price", DataType::Int64); + let payload = covering_field(23, "payload", DataType::Utf8); + + let common = common_physical_covering( + vec![price.clone(), payload.clone()], + &[vec![price.clone(), payload.clone()], vec![price.clone()]], + ); + + assert_eq!(common, vec![price.clone()]); + + let common = common_physical_covering( + vec![price.clone(), payload.clone()], + &[vec![payload, price.clone()], vec![price.clone()]], + ); + assert_eq!( + common, + vec![price], + "order is checked after finding the plan-wide subset" + ); + } + + #[test] + fn test_common_physical_covering_rejects_wrong_identity_shape_and_order() { + let price = covering_field(17, "price", DataType::Int64); + let payload = covering_field(23, "payload", DataType::Utf8); + + assert!( + common_physical_covering( + vec![price.clone()], + &[vec![covering_field(99, "price", DataType::Int64)]], + ) + .is_empty(), + "a matching name and type cannot substitute for the source field id" + ); + assert!( + common_physical_covering( + vec![price.clone()], + &[vec![covering_field(17, "price", DataType::UInt64)]], + ) + .is_empty(), + "a matching source field id cannot substitute for the physical Arrow type" + ); + assert!( + common_physical_covering( + vec![price.clone(), payload.clone()], + &[vec![payload, price]], + ) + .is_empty(), + "a physical order that disagrees with the declared output schema is not servable" + ); + } + + /// The manifest remains authoritative for logical dependencies: if an index declares + /// a covering field id the current schema cannot resolve, construction must fail loudly + /// rather than hiding schema/metadata corruption behind physical fallback. + #[tokio::test] + async fn test_ann_sub_index_rejects_unresolvable_covering_field() { + use lance_datafusion::exec::OneShotExec; + + // Minimal input plan carrying the required PART_ID_COLUMN schema. + let input = Arc::new(OneShotExec::from_batch(RecordBatch::new_empty( + KNN_PARTITION_SCHEMA.clone(), + ))); + + // A dataset whose schema has no field with id 9999. + let schema = Arc::new(ArrowSchema::new(vec![ArrowField::new( + "id", + DataType::Int32, + false, + )])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![Arc::new(Int32Array::from(vec![1, 2, 3]))], + ) + .unwrap(); + let dataset = Arc::new( + Dataset::write( + RecordBatchIterator::new(vec![Ok(batch)], schema), + "memory://d2-unresolvable-covering", + None, + ) + .await + .unwrap(), + ); + + // A legal covering shape (`covering_fields` is the trailing slice of `fields`, + // which keeps at least one keyed entry -- see + // `IndexMetadata::validate_covering_fields`) whose carried id 9999 the schema + // cannot resolve. Keeping the shape legal is what makes this a test of the + // resolution failure rather than of metadata no commit can produce. + let index = IndexMetadata { + uuid: uuid::Uuid::new_v4(), + fields: vec![0, 9999], + name: "vector_idx".to_string(), + dataset_version: 1, + fragment_bitmap: Some(RoaringBitmap::new()), + index_details: None, + index_version: 0, + created_at: None, + base_id: None, + files: None, + covering_fields: vec![9999], + }; + index + .validate_covering_fields() + .expect("the fixture must be a legal covering declaration"); + + let result = ANNIvfSubIndexExec::try_new( + input, + dataset, + vec![index], + base_query(), + PreFilterSource::None, + ); + let err = result.expect_err("unresolvable covering field id must be rejected"); + assert!( + err.to_string().contains("covering field id 9999"), + "expected a manifest/schema inconsistency error, got: {err}" + ); + } + + /// The exec's declared output schema is the seam through which a query's covering + /// projection reaches the rest of the plan: `knn_combined` unions the flat paths + /// against this schema, so whatever is declared here is what every unindexed + /// fragment pays to read. + /// + /// All three states of [`Query::covering_projection`] must be visible in the + /// declaration, and the `Some(&[])` case must NOT collapse into the `None` case. + /// The two are indistinguishable in query results -- a covering column is + /// semantically transparent, so the values simply arrive via a take instead -- which + /// is why this is asserted on the declared schema rather than on any result. + #[tokio::test] + async fn test_ann_sub_index_declares_only_the_projected_covering_columns() { + use lance_datafusion::exec::OneShotExec; + + let input = Arc::new(OneShotExec::from_batch(RecordBatch::new_empty( + KNN_PARTITION_SCHEMA.clone(), + ))); + + let schema = Arc::new(ArrowSchema::new(vec![ + ArrowField::new("vec", DataType::Int32, false), + ArrowField::new("price", DataType::Int32, true), + ArrowField::new("payload", DataType::Utf8, true), + ])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(Int32Array::from(vec![1, 2, 3])), + Arc::new(Int32Array::from(vec![Some(10), Some(20), Some(30)])), + Arc::new(StringArray::from(vec!["a", "b", "c"])), + ], + ) + .unwrap(); + let dataset = Arc::new( + Dataset::write( + RecordBatchIterator::new(vec![Ok(batch)], schema), + "memory://d2-projected-covering", + None, + ) + .await + .unwrap(), + ); + let vec_id = dataset.schema().field("vec").unwrap().id; + let price_id = dataset.schema().field("price").unwrap().id; + let payload_id = dataset.schema().field("payload").unwrap().id; + + // Declared in the reverse of field-id order, so the assertions below distinguish + // declaration order from schema order. + let index = IndexMetadata { + uuid: uuid::Uuid::new_v4(), + fields: vec![vec_id, payload_id, price_id], + name: "vector_idx".to_string(), + dataset_version: 1, + fragment_bitmap: Some(RoaringBitmap::new()), + index_details: None, + index_version: 0, + created_at: None, + base_id: None, + files: None, + covering_fields: vec![payload_id, price_id], + }; + index + .validate_covering_fields() + .expect("the fixture must be a legal covering declaration"); + + let declared_for = |projection: Option>| { + let mut query = base_query(); + query.covering_projection = projection.map(Arc::from); + let exec = ANNIvfSubIndexExec::try_new_with_resolved_covering( + input.clone(), + dataset.clone(), + vec![index.clone()], + query, + PreFilterSource::None, + ) + .unwrap(); + exec.schema() + .fields() + .iter() + .map(|f| f.name().clone()) + .collect::>() + }; + + assert_eq!( + declared_for(None), + vec![DIST_COL, ROW_ID, "payload", "price"], + "no projection means no narrowing: every declared covering column" + ); + assert_eq!( + declared_for(Some(vec!["price".to_string()])), + vec![DIST_COL, ROW_ID, "price"], + "only the covering column this query reads is declared" + ); + assert_eq!( + declared_for(Some(vec!["price".to_string(), "payload".to_string()])), + vec![DIST_COL, ROW_ID, "payload", "price"], + "covering columns are declared in the index's order, not the request's" + ); + assert_eq!( + declared_for(Some(vec![])), + vec![DIST_COL, ROW_ID], + "a query needing no covering column must declare a plain-index schema, which \ + is what lets the flat paths skip the covering read entirely" + ); + + let mut unverified_query = base_query(); + unverified_query.covering_projection = Some(Arc::from(["price".to_string()])); + let unverified = ANNIvfSubIndexExec::try_new( + input, + dataset, + vec![index], + unverified_query, + PreFilterSource::None, + ) + .unwrap(); + assert_eq!( + unverified + .schema() + .fields() + .iter() + .map(|field| field.name().as_str()) + .collect::>(), + vec![DIST_COL, ROW_ID], + "the synchronous constructor cannot verify storage and must disable covering" + ); + } + + /// The exec declares one output schema but executes every delta segment, so delta + /// segments that disagree on covering fields must be rejected up front with a clear + /// error, not fail later with a stream/plan schema mismatch. (The commit boundary + /// rejects new mixed sets; this defends against already-committed metadata.) + #[tokio::test] + async fn test_ann_sub_index_rejects_mixed_covering_deltas() { + use lance_datafusion::exec::OneShotExec; + + let input = Arc::new(OneShotExec::from_batch(RecordBatch::new_empty( + KNN_PARTITION_SCHEMA.clone(), + ))); + + let schema = Arc::new(ArrowSchema::new(vec![ArrowField::new( + "id", + DataType::Int32, + false, + )])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![Arc::new(Int32Array::from(vec![1, 2, 3]))], + ) + .unwrap(); + let dataset = Arc::new( + Dataset::write( + RecordBatchIterator::new(vec![Ok(batch)], schema), + "memory://d2-mixed-covering-deltas", + None, + ) + .await + .unwrap(), + ); + let id_field_id = dataset.schema().field("id").unwrap().id; + // This exec path only resolves `covering_fields` against the live schema, so a + // key id that is not itself a real schema field is harmless here; it stands in + // for the vector key column a real covering index would also carry in `fields`. + let key_field_id = id_field_id + 1; + + // `covering_fields` is always the trailing entries of `fields`. + let covered_delta = IndexMetadata { + uuid: uuid::Uuid::new_v4(), + fields: vec![key_field_id, id_field_id], + name: "vector_idx".to_string(), + dataset_version: 1, + fragment_bitmap: Some(RoaringBitmap::new()), + index_details: None, + index_version: 0, + created_at: None, + base_id: None, + files: None, + covering_fields: vec![id_field_id], + }; + let plain_delta = IndexMetadata { + uuid: uuid::Uuid::new_v4(), + covering_fields: Vec::new(), + ..covered_delta.clone() + }; + + let result = ANNIvfSubIndexExec::try_new( + input, + dataset, + vec![covered_delta, plain_delta], + base_query(), + PreFilterSource::None, + ); + let err = result.expect_err("delta segments with mixed covering must be rejected"); + assert!( + err.to_string().contains("disagree on covering fields"), + "expected a delta covering mismatch error, got: {err}" + ); + } + #[test] fn test_effective_query_parallelism_clamps_to_cpu_pool() { let mut query = base_query(); @@ -2826,7 +3699,12 @@ mod tests { } fn find_partitions(&self, _query: &Query) -> Result<(UInt32Array, Float32Array)> { - unimplemented!() + // No additional partitions to reveal; the covered re-rank in + // late_search treats an empty result as "nothing more to search". + Ok(( + UInt32Array::from(Vec::::new()), + Float32Array::from(Vec::::new()), + )) } fn total_partitions(&self) -> usize { @@ -3194,6 +4072,69 @@ mod tests { ); } + /// A `FilterLoader` that yields a fixed allow-list of row addresses, so the + /// resulting prefilter mask is a bounded `AllowList` (`max_len()` is `Some`). + struct StaticAllowList(Vec); + + #[async_trait] + impl FilterLoader for StaticAllowList { + async fn load(self: Box) -> lance_core::Result { + Ok(lance_select::RowAddrMask::from_allowed( + lance_select::RowAddrTreeMap::from_iter(self.0.iter().copied()), + )) + } + } + + /// Like `empty_prefilter`, but the prefilter mask is a bounded allow-list of + /// the given row addresses, which drives the late-search no-rows shortcut. + async fn bounded_prefilter(allow: Vec) -> Arc { + static NEXT_PREFILTER_DATASET_ID: AtomicUsize = AtomicUsize::new(100_000); + let schema = Arc::new(ArrowSchema::new(vec![ArrowField::new( + "id", + DataType::Int32, + false, + )])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![Arc::new(Int32Array::from(vec![1, 2, 3]))], + ) + .unwrap(); + let uri = format!( + "memory://bounded-prefilter-{}", + NEXT_PREFILTER_DATASET_ID.fetch_add(1, Ordering::Relaxed) + ); + let dataset = Arc::new( + Dataset::write( + RecordBatchIterator::new(vec![Ok(batch)].into_iter(), schema), + &uri, + None, + ) + .await + .unwrap(), + ); + let mut indexed_fragments = RoaringBitmap::new(); + for fragment in dataset.manifest.fragments.iter() { + indexed_fragments.insert(fragment.id as u32); + } + let index = IndexMetadata { + uuid: uuid::Uuid::new_v4(), + fields: vec![], + name: "test".to_string(), + dataset_version: 1, + fragment_bitmap: Some(indexed_fragments), + index_details: None, + index_version: 0, + created_at: None, + base_id: None, + files: None, + covering_fields: Vec::new(), + }; + let loader: Box = Box::new(StaticAllowList(allow)); + let prefilter = Arc::new(DatasetPreFilter::new(dataset, &[index], Some(loader))); + prefilter.wait_for_ready().await.unwrap(); + prefilter + } + fn prepared_metrics() -> Arc { Arc::new(AnnIndexMetrics::new(&ExecutionPlanMetricsSet::new(), 0)) } @@ -3384,16 +4325,18 @@ mod tests { query.minimum_nprobes = 0; query.maximum_nprobes = Some(3); let state = Arc::new(ANNIvfEarlySearchResults::new(1, query.k)); - state.record_batch( - &RecordBatch::try_new( - KNN_INDEX_SCHEMA.clone(), - vec![ - Arc::new(Float32Array::from(vec![0.0])), - Arc::new(UInt64Array::from(vec![999])), - ], + state + .record_batch( + &RecordBatch::try_new( + KNN_INDEX_SCHEMA.clone(), + vec![ + Arc::new(Float32Array::from(vec![0.0])), + Arc::new(UInt64Array::from(vec![999])), + ], + ) + .unwrap(), ) - .unwrap(), - ); + .unwrap(); let prefilter = empty_prefilter().await; let batches = ANNIvfSubIndexExec::late_search( @@ -3407,6 +4350,7 @@ mod tests { state.clone(), usize::MAX, None, + false, ) .try_collect::>() .await @@ -3418,6 +4362,32 @@ mod tests { assert_eq!(state.num_results_found.load(Ordering::Relaxed), 2); } + /// `record_batch` must resolve `_rowid` by name, not by position: a covered ANN + /// batch appends covering columns after `_rowid`, but nothing here should assume + /// that specific layout. Puts another UInt64 column ahead of `_rowid` so a + /// positional read would silently record the decoy column's values instead. + #[test] + fn test_early_search_results_record_batch_resolves_row_id_by_name() { + let schema = Arc::new(ArrowSchema::new(vec![ + ArrowField::new(DIST_COL, DataType::Float32, true), + ArrowField::new("decoy_u64_col", DataType::UInt64, true), + ROW_ID_FIELD.clone(), + ])); + let batch = RecordBatch::try_new( + schema, + vec![ + Arc::new(Float32Array::from(vec![0.0, 1.0])), + Arc::new(UInt64Array::from(vec![111, 222])), + Arc::new(UInt64Array::from(vec![7, 8])), + ], + ) + .unwrap(); + + let state = ANNIvfEarlySearchResults::new(1, 2); + state.record_batch(&batch).unwrap(); + assert_eq!(*state.initial_ids.lock().unwrap(), vec![7, 8]); + } + fn row_ids_of(batches: &[RecordBatch]) -> Vec { batches .iter() @@ -3522,6 +4492,7 @@ mod tests { state.clone(), usize::MAX, Some(seg_mask), + false, ) .try_collect::>() .await @@ -3601,6 +4572,7 @@ mod tests { state.clone(), usize::MAX, Some(seg_mask), + false, ) .try_collect::>() .await @@ -3614,6 +4586,211 @@ mod tests { ); } + /// With a bounded selective prefilter, the late-search no-rows shortcut must fire + /// for covered indexes too -- but emit nothing: the shortcut's `[_distance, _rowid]` + /// batch would drop the covering columns, so the end-of-stream covered recovery + /// (`covered_not_found_batch`) emits the missing rows instead, with a take bounded + /// by the prefilter size. Searching instead of shortcutting is the failure mode this + /// pins down: a prefilter-matched row with no index entry (null vector) can never be + /// found, so the search degenerates to scoring every partition on every query. + #[tokio::test] + async fn test_late_search_covered_skips_search_on_selective_prefilter() { + // Bounded prefilter (one allowed row) + k=2 => max_results=1 <= k, and an + // empty initial state => found_so_far=0 < max_results: shortcut criteria. + async fn run_late_search(has_covered: bool) -> (Vec, Vec) { + let (index, _prepared, searched_partitions, _threads) = + prepared_index(vec![21, 22, 23]); + let mut query = base_query(); + query.k = 2; + query.minimum_nprobes = 0; + query.maximum_nprobes = Some(3); + let state = Arc::new(ANNIvfEarlySearchResults::new(1, query.k)); + let prefilter = bounded_prefilter(vec![0]).await; + + let batches = ANNIvfSubIndexExec::late_search( + index, + query, + Arc::new(UInt32Array::from(vec![0, 1, 2])), + Arc::new(Float32Array::from(vec![0.1, 0.2, 0.3])), + prefilter.clone(), + prefilter, + prepared_metrics(), + state, + usize::MAX, + None, + has_covered, + ) + .try_collect::>() + .await + .unwrap(); + let searched = searched_partitions.lock().unwrap().clone(); + (batches, searched) + } + + // Control: a non-covered index takes the shortcut, so no partition is + // searched and the only batch is the [_distance, _rowid] not-found batch. + let (batches, searched) = run_late_search(false).await; + assert!( + searched.is_empty(), + "non-covered query should take the shortcut and skip the search, searched={searched:?}" + ); + // The non-covered control for the covered/non-covered contrast below: guard + // against a batch count of zero, which would make `all(..)` vacuously true + // regardless of whether the shortcut actually emitted anything. + assert_eq!(batches.len(), 1, "the shortcut emits exactly one batch"); + assert!( + batches.iter().all(|b| b.num_columns() == 2), + "the shortcut emits only [_distance, _rowid]" + ); + + // Covered: the shortcut fires too (no partition searched), but emits nothing -- + // the end-of-stream covered recovery supplies the not-found rows with their + // covering columns. + let (batches, searched) = run_late_search(true).await; + assert!( + searched.is_empty(), + "covered query must take the shortcut and skip the search, searched={searched:?}" + ); + assert!( + batches.is_empty(), + "the covered shortcut emits nothing; recovery happens at end of stream" + ); + } + + /// Arming must record the arming delta's ownership scope, and the recovery must + /// emit only rows an armed scope selects. A sibling delta that skipped its late + /// search without arming never emits its not-found rows on the non-covered path + /// (its `seg_mask` filters them out of the shortcut batch), so recovering them + /// for a covered index would make the covered and non-covered result sets + /// diverge -- e.g. null-vector rows owned by the skipping delta surfacing at + /// INFINITY distance only when the index is covered. + #[tokio::test] + async fn test_covered_recovery_honors_arming_delta_ownership() { + // Delta B owns row 0 only; rows 1 and 2 belong to a sibling delta that will + // NOT arm (it skipped its late search entirely). + let seg_mask_b = Arc::new(RowAddrMask::from_allowed( + lance_select::RowAddrTreeMap::from_iter([0u64]), + )); + + // Delta B takes the covered no-rows shortcut: bounded prefilter of 3 rows, + // k = 4, nothing found so far. + let (index, _prepared, searched_partitions, _threads) = prepared_index(vec![21, 22, 23]); + let mut query = base_query(); + query.k = 4; + query.minimum_nprobes = 0; + query.maximum_nprobes = Some(3); + let state = Arc::new(ANNIvfEarlySearchResults::new(1, query.k)); + let prefilter = bounded_prefilter(vec![0, 1, 2]).await; + + let batches = ANNIvfSubIndexExec::late_search( + index, + query, + Arc::new(UInt32Array::from(vec![0, 1, 2])), + Arc::new(Float32Array::from(vec![0.1, 0.2, 0.3])), + prefilter.clone(), + prefilter.clone(), + prepared_metrics(), + state.clone(), + usize::MAX, + Some(seg_mask_b.clone()), + true, + ) + .try_collect::>() + .await + .unwrap(); + assert!(batches.is_empty(), "covered shortcut emits nothing inline"); + assert!( + searched_partitions.lock().unwrap().is_empty(), + "the shortcut must not search" + ); + + // Arming recorded delta B's scope (and only that one). + let scopes = std::mem::take(&mut *state.covered_recovery_scopes.lock().unwrap()); + assert_eq!(scopes.len(), 1, "exactly one delta armed"); + assert!(scopes[0].is_some(), "the armed scope is delta B's seg_mask"); + + // The recovery over a real dataset must return ONLY delta B's row 0 -- + // rows 1 and 2 are owned by the never-armed sibling. + let payload_schema = Arc::new(ArrowSchema::new(vec![ArrowField::new( + "payload", + DataType::Int32, + false, + )])); + let batch = RecordBatch::try_new( + payload_schema.clone(), + vec![Arc::new(Int32Array::from(vec![10, 11, 12]))], + ) + .unwrap(); + let dataset = Arc::new( + Dataset::write( + RecordBatchIterator::new(vec![Ok(batch)].into_iter(), payload_schema), + "memory://covered-recovery-ownership", + None, + ) + .await + .unwrap(), + ); + let mut fields = KNN_INDEX_SCHEMA.fields().to_vec(); + fields.push(Arc::new(ArrowField::new("payload", DataType::Int32, false))); + let output_schema = Arc::new(ArrowSchema::new(fields)); + + let emitted = Arc::new(std::sync::Mutex::new(roaring::RoaringTreemap::new())); + let recovered = ANNIvfSubIndexExec::covered_not_found_batch( + true, + prefilter.clone(), + dataset.clone(), + output_schema.clone(), + 4, + emitted.clone(), + scopes, + ) + .await + .unwrap() + .expect("delta B's unemitted row must be recovered"); + assert_eq!( + recovered + .column_by_name(ROW_ID) + .unwrap() + .as_primitive::() + .values(), + &[0], + "recovery must emit only rows the arming delta owns" + ); + assert_eq!( + recovered + .column_by_name("payload") + .unwrap() + .as_primitive::() + .values(), + &[10], + "the covering column value comes from the base table" + ); + + // A global (unrestricted) arming scope recovers every unemitted row -- the + // single-delta behaviour is unchanged. + let recovered = ANNIvfSubIndexExec::covered_not_found_batch( + true, + prefilter, + dataset, + output_schema, + 4, + emitted, + vec![None], + ) + .await + .unwrap() + .expect("a global scope recovers all unemitted rows"); + assert_eq!( + recovered + .column_by_name(ROW_ID) + .unwrap() + .as_primitive::() + .values(), + &[0, 1, 2], + "a global arming scope keeps the old recover-everything behaviour" + ); + } + #[tokio::test] async fn test_delta_skipping_late_search_releases_sibling() { let prefilter = empty_prefilter().await; @@ -3634,6 +4811,7 @@ mod tests { state.clone(), usize::MAX, None, + false, ) .try_collect::>(); @@ -3653,6 +4831,7 @@ mod tests { state, usize::MAX, None, + false, ) .try_collect::>(); @@ -4060,6 +5239,7 @@ mod tests { query_parallelism: DEFAULT_QUERY_PARALLELISM, dist_q_c: 0.0, approx_mode: Default::default(), + covering_projection: None, }; async fn multivector_scoring( @@ -4118,6 +5298,85 @@ mod tests { } } + /// Covered multivector scoring must carry each row's covering values through + /// the per-row-id reduction, aligned by row id -- and must FAIL LOUDLY (not + /// emit nulls) if an emitted row id ever has no recorded covering source. + #[tokio::test] + async fn test_multivector_score_carries_covering_values() { + let query = Query { + column: "vector".to_string(), + key: Arc::new(generate_random_array(1)), + k: 10, + lower_bound: None, + upper_bound: None, + minimum_nprobes: 1, + maximum_nprobes: None, + ef: None, + refine_factor: None, + metric_type: Some(DistanceType::Cosine), + use_index: true, + query_parallelism: DEFAULT_QUERY_PARALLELISM, + dist_q_c: 0.0, + approx_mode: Default::default(), + covering_projection: None, + }; + + let mut fields = KNN_INDEX_SCHEMA.fields().to_vec(); + fields.push(Arc::new(ArrowField::new("payload", DataType::Int32, false))); + let covered_schema = Arc::new(ArrowSchema::new(fields)); + + // Three sub-query streams over overlapping row ids 1..=4; the covering value + // is a pure function of the row id (100 + row_id), identical across streams, + // as the contract requires. + let inputs = (0..3u64) + .map(|i| { + let batch = RecordBatch::try_new( + covered_schema.clone(), + vec![ + Arc::new(Float32Array::from(vec![i as f32 + 1.0, i as f32 + 2.0])), + Arc::new(UInt64Array::from(vec![i + 1, i + 2])), + Arc::new(Int32Array::from(vec![ + 100 + (i + 1) as i32, + 100 + (i + 2) as i32, + ])), + ], + ) + .unwrap(); + let input: Arc = Arc::new(TestingExec::new(vec![batch])); + input + }) + .collect::>(); + + let ctx = Arc::new(datafusion::execution::context::TaskContext::default()); + let plan = MultivectorScoringExec::try_new(inputs, query).unwrap(); + let batches = plan + .execute(0, ctx) + .unwrap() + .try_collect::>() + .await + .unwrap(); + + let mut seen = 0; + for batch in batches { + assert_eq!(batch.num_columns(), 3, "covering column must flow through"); + let row_ids = batch[ROW_ID].as_primitive::(); + let payloads = batch + .column_by_name("payload") + .unwrap() + .as_primitive::(); + assert_eq!(payloads.null_count(), 0, "no fabricated nulls"); + for (row_id, payload) in row_ids.values().iter().zip(payloads.values().iter()) { + assert_eq!( + *payload, + 100 + *row_id as i32, + "covering value must stay aligned with its row id" + ); + seen += 1; + } + } + assert_eq!(seen, 4, "all four distinct row ids are emitted"); + } + /// A test dataset for testing the nprobes parameter. /// /// The dataset has 100 partitions and filterable columns setup to easily create