feat(index): covering columns for IVF vector indexes

java/lance-jni/src/utils.rs

@@ -515,6 +515,7 @@ pub fn get_vector_index_params(
                 version: IndexFileVersion::V3,
                 skip_transpose: false,
                 runtime_hints: Default::default(),
+                covering_columns: Default::default(),
             })
         },
     )?;

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.
+   *
+   * <p>This is the index's <em>keyed</em> 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<Integer> getFieldIds() {
     return fieldIds;
   }

python/src/indices.rs

@@ -674,9 +674,11 @@ 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<u32>,
-    /// 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<String>,
     /// The number of rows indexed by the index

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

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,
 };
@@ -66,6 +67,7 @@ use crate::vector::pq::storage::transpose;
 use crate::vector::quantizer::{QuantizerMetadata, QuantizerStorage};
 use crate::vector::storage::{
     DistCalculator, DistanceCalculatorOptions, QueryResidual, RabitRawQueryContext, VectorStore,
+    covering_field_indices_excluding,
 };
 
 pub const RABIT_METADATA_KEY: &str = "lance:rabit";
@@ -2048,6 +2050,29 @@ impl VectorStore for RabitQuantizationStorage {
         self.batch.schema_ref()
     }
 
+    /// RQ storage carries the row id plus many internal columns (the binary codes, the
+    /// extended-bit codes, and the per-row factor columns, some present only for higher
+    /// `num_bits`); the covering ("included") columns are everything else.
+    fn covering_field_indices(&self) -> Vec<usize> {
+        // Every RaBitQ-internal column plus the legacy `__ivf_part_id` column
+        // (left in pre-#3606 single-partition builds; unlike PQ, RaBitQ storage
+        // does not drop it at construction, so it is excluded here). Anything
+        // else is a user-declared covering column.
+        const INTERNAL: &[&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,
+        ];
+        covering_field_indices_excluding(self.schema().as_ref(), INTERNAL)
+    }
+
     fn to_batches(&self) -> Result<impl Iterator<Item = RecordBatch> + Send> {
         Ok(std::iter::once(self.batch.clone()))
     }
@@ -3035,6 +3060,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();

rust/lance-index/src/vector/flat/storage.rs

@@ -6,8 +6,11 @@ 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, covering_field_indices_excluding, remap_row_ids_by_name,
+};
 use crate::vector::utils::do_prefetch;
 use arrow::array::AsArray;
 use arrow::compute::concat_batches;
@@ -85,7 +88,7 @@ impl QuantizerStorage for FlatFloatStorage {
         frag_reuse_index: Option<Arc<dyn RowIdRemapper>>,
     ) -> Result<Self> {
         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
         };
@@ -194,6 +197,17 @@ impl VectorStore for FlatFloatStorage {
         self.batch.schema_ref()
     }
 
+    /// Flat storage is `[_rowid, flat, <included cols...>]`, so the covering columns
+    /// are every field except 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, so it is excluded here).
+    fn covering_field_indices(&self) -> Vec<usize> {
+        covering_field_indices_excluding(
+            self.schema().as_ref(),
+            &[ROW_ID, FLAT_COLUMN, PART_ID_COLUMN],
+        )
+    }
+
     fn as_any(&self) -> &dyn std::any::Any {
         self
     }
@@ -265,7 +279,7 @@ impl QuantizerStorage for FlatBinStorage {
         frag_reuse_index: Option<Arc<dyn RowIdRemapper>>,
     ) -> Result<Self> {
         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
         };
@@ -369,6 +383,17 @@ impl VectorStore for FlatBinStorage {
         self.batch.schema_ref()
     }
 
+    /// Flat storage is `[_rowid, flat, <included cols...>]`, so the covering columns
+    /// are every field except 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, so it is excluded here).
+    fn covering_field_indices(&self) -> Vec<usize> {
+        covering_field_indices_excluding(
+            self.schema().as_ref(),
+            &[ROW_ID, FLAT_COLUMN, PART_ID_COLUMN],
+        )
+    }
+
     fn as_any(&self) -> &dyn std::any::Any {
         self
     }
@@ -685,6 +710,69 @@ 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_float_covering_excludes_legacy_part_id_column() {
+        use crate::vector::PART_ID_COLUMN;
+        const DIM: i32 = 4;
+        const N: usize = 8;
+        let row_ids = UInt64Array::from_iter_values(0..N as u64);
+        let values =
+            arrow_array::Float32Array::from_iter_values((0..N * DIM as usize).map(|v| v as f32));
+        let vectors = FixedSizeListArray::try_new_from_values(values, DIM).unwrap();
+        let part_ids = arrow_array::UInt32Array::from_iter_values((0..N).map(|_| 0));
+        let batch = RecordBatch::try_from_iter_with_nullable(vec![
+            (ROW_ID, Arc::new(row_ids) as ArrayRef, true),
+            (FLAT_COLUMN, Arc::new(vectors) as ArrayRef, true),
+            (PART_ID_COLUMN, Arc::new(part_ids) as ArrayRef, true),
+        ])
+        .unwrap();
+        let storage = FlatFloatStorage::try_from_batch(
+            batch,
+            &FlatMetadata { dim: DIM as usize },
+            DistanceType::L2,
+            None,
+        )
+        .unwrap();
+        assert!(
+            storage.covering_field_indices().is_empty(),
+            "legacy __ivf_part_id must not be classified as a covering column"
+        );
+    }
+
+    #[test]
+    fn test_flat_bin_covering_excludes_legacy_part_id_column() {
+        use crate::vector::PART_ID_COLUMN;
+        const DIM: i32 = 8;
+        const N: usize = 8;
+        let row_ids = UInt64Array::from_iter_values(0..N as u64);
+        let codes =
+            arrow_array::UInt8Array::from_iter_values((0..N * DIM as usize).map(|v| v as u8));
+        let vectors = FixedSizeListArray::try_new_from_values(codes, DIM).unwrap();
+        let part_ids = arrow_array::UInt32Array::from_iter_values((0..N).map(|_| 0));
+        let batch = RecordBatch::try_from_iter_with_nullable(vec![
+            (ROW_ID, Arc::new(row_ids) as ArrayRef, true),
+            (FLAT_COLUMN, Arc::new(vectors) as ArrayRef, true),
+            (PART_ID_COLUMN, Arc::new(part_ids) as ArrayRef, true),
+        ])
+        .unwrap();
+        let storage = FlatBinStorage::try_from_batch(
+            batch,
+            &FlatMetadata { dim: DIM as usize },
+            DistanceType::L2,
+            None,
+        )
+        .unwrap();
+        assert!(
+            storage.covering_field_indices().is_empty(),
+            "legacy __ivf_part_id must not be classified as a covering column"
+        );
+    }
+
     #[test]
     fn test_flat_float_storage_distance_f16() {
         let storage = make_f16_storage();