feat: introduce IndexFileVersion for Python API (#4040)

Signed-off-by: BubbleCal <bubble-cal@outlook.com>

Author: BubbleCal

python/python/lance/dataset.py

@@ -1947,6 +1947,8 @@ def create_index(
             - num_bits
                 The number of bits for PQ (Product Quantization). Default is 8.
                 Only 4, 8 are supported.
+            - index_file_version
+                The version of the index file. Default is "V3".
 
         Optional parameters for `IVF_HNSW_*`:
             max_level

python/python/lance/indices.py

@@ -2,6 +2,7 @@
 # SPDX-FileCopyrightText: Copyright The Lance Authors
 import math
 import warnings
+from enum import Enum
 from typing import TYPE_CHECKING, Optional, Union
 
 import numpy as np
@@ -16,6 +17,11 @@
     from .dependencies import torch
 
 
+class IndexFileVersion(str, Enum):
+    LEGACY = "Legacy"
+    V3 = "V3"
+
+
 class PqModel:
     """
     A class that represents a trained PQ model

python/python/tests/test_vector_index.py

@@ -14,6 +14,7 @@
 import pytest
 from lance import LanceFragment
 from lance.dataset import VectorIndexReader
+from lance.indices import IndexFileVersion
 from lance.util import validate_vector_index  # noqa: E402
 from lance.vector import vec_to_table  # noqa: E402
 
@@ -238,7 +239,10 @@ def test_f16_cuda(tmp_path):
     validate_vector_index(dataset, "vector")
 
 
-def test_index_with_nans(tmp_path):
+@pytest.mark.parametrize(
+    "index_file_version", [IndexFileVersion.V3, IndexFileVersion.LEGACY]
+)
+def test_index_with_nans(tmp_path, index_file_version):
     # 1024 rows, the entire table should be sampled
     tbl = create_table(nvec=1000, nans=24)
 
@@ -248,11 +252,17 @@ def test_index_with_nans(tmp_path):
         index_type="IVF_PQ",
         num_partitions=4,
         num_sub_vectors=16,
+        index_file_version=index_file_version,
     )
+    idx_stats = dataset.stats.index_stats("vector_idx")
+    assert idx_stats["indices"][0]["index_file_version"] == index_file_version
     validate_vector_index(dataset, "vector")
 
 
-def test_torch_index_with_nans(tmp_path):
+@pytest.mark.parametrize(
+    "index_file_version", [IndexFileVersion.V3, IndexFileVersion.LEGACY]
+)
+def test_torch_index_with_nans(tmp_path, index_file_version):
     torch = pytest.importorskip("torch")
 
     # 1024 rows, the entire table should be sampled
@@ -266,7 +276,10 @@ def test_torch_index_with_nans(tmp_path):
         num_sub_vectors=16,
         accelerator=torch.device("cpu"),
         one_pass_ivfpq=True,
+        index_file_version=index_file_version,
     )
+    idx_stats = dataset.stats.index_stats("vector_idx")
+    assert idx_stats["indices"][0]["index_file_version"] == index_file_version
     validate_vector_index(dataset, "vector")
 
 
@@ -658,6 +671,7 @@ def test_pre_populated_ivf_centroids(dataset, tmp_path: Path):
             "num_sub_vectors": 8,
             "transposed": True,
         },
+        "index_file_version": IndexFileVersion.V3,
     }
 
     with pytest.raises(KeyError, match='Index "non-existent_idx" not found'):

python/src/dataset.rs

@@ -31,7 +31,6 @@ use pyo3::{
 use pyo3::{prelude::*, IntoPyObjectExt};
 use snafu::location;
 
-use lance::dataset::builder::DatasetBuilder;
 use lance::dataset::refs::{Ref, TagContents};
 use lance::dataset::scanner::{
     DatasetRecordBatchStream, ExecutionStatsCallback, MaterializationStyle,
@@ -54,6 +53,7 @@ use lance::dataset::{
 use lance::dataset::{ColumnAlteration, ProjectionRequest};
 use lance::index::vector::utils::get_vector_type;
 use lance::index::{vector::VectorIndexParams, DatasetIndexInternalExt};
+use lance::{dataset::builder::DatasetBuilder, index::vector::IndexFileVersion};
 use lance_arrow::as_fixed_size_list_array;
 use lance_index::scalar::inverted::query::{
     BooleanQuery, BoostQuery, FtsQuery, MatchQuery, MultiMatchQuery, Operator, PhraseQuery,
@@ -2027,6 +2027,7 @@ fn prepare_vector_index_params(
     let mut hnsw_params = HnswBuildParams::default();
     let mut pq_params = PQBuildParams::default();
     let mut sq_params = SQBuildParams::default();
+    let mut index_file_version = IndexFileVersion::V3;
 
     if let Some(kwargs) = kwargs {
         // Parse metric type
@@ -2143,9 +2144,15 @@ fn prepare_vector_index_params(
             let codebook = as_fixed_size_list_array(batch.column(0));
             pq_params.codebook = Some(codebook.values().clone())
         };
+
+        if let Some(version) = kwargs.get_item("index_file_version")? {
+            let version: String = version.extract()?;
+            index_file_version = IndexFileVersion::try_from(&version)
+                .map_err(|e| PyValueError::new_err(format!("Invalid index_file_version: {e}")))?;
+        }
     }
 
-    match index_type {
+    let mut params = match index_type {
         "IVF_FLAT" => Ok(Box::new(VectorIndexParams::ivf_flat(
             ivf_params.num_partitions,
             m_type,
@@ -2178,7 +2185,9 @@ fn prepare_vector_index_params(
         _ => Err(PyValueError::new_err(format!(
             "Index type '{index_type}' is not supported."
         ))),
-    }
+    }?;
+    params.version(index_file_version);
+    Ok(params)
 }
 
 #[pyclass(name = "_FragmentWriteProgress", module = "_lib")]

rust/lance/src/index/vector.rs

@@ -41,6 +41,7 @@ use lance_io::traits::Reader;
 use lance_linalg::distance::*;
 use lance_table::format::Index as IndexMetadata;
 use object_store::path::Path;
+use serde::Serialize;
 use snafu::location;
 use tempfile::tempdir;
 use tracing::instrument;
@@ -64,12 +65,25 @@ pub enum StageParams {
 // The version of the index file.
 // `Legacy` is used for only IVF_PQ index, and is the default value.
 // The other index types are using `V3`.
-#[derive(Debug, Clone)]
+#[derive(Debug, Clone, Serialize)]
 pub enum IndexFileVersion {
     Legacy,
     V3,
 }
 
+impl IndexFileVersion {
+    pub fn try_from(version: &str) -> Result<Self> {
+        match version.to_lowercase().as_str() {
+            "legacy" => Ok(Self::Legacy),
+            "v3" => Ok(Self::V3),
+            _ => Err(Error::Index {
+                message: format!("Invalid index file version: {}", version),
+                location: location!(),
+            }),
+        }
+    }
+}
+
 /// The parameters to build vector index.
 #[derive(Debug, Clone)]
 pub struct VectorIndexParams {

rust/lance/src/index/vector/ivf.rs

@@ -14,9 +14,9 @@ use super::{
     pq::{build_pq_model, PQIndex},
     utils::maybe_sample_training_data,
 };
-use crate::dataset::builder::DatasetBuilder;
 use crate::index::vector::utils::{get_vector_dim, get_vector_type};
 use crate::index::DatasetIndexInternalExt;
+use crate::{dataset::builder::DatasetBuilder, index::vector::IndexFileVersion};
 use crate::{
     dataset::Dataset,
     index::{pb, prefilter::PreFilter, vector::ivf::io::write_pq_partitions, INDEX_FILE_NAME},
@@ -775,6 +775,7 @@ pub struct IvfIndexStatistics {
     partitions: Vec<IvfIndexPartitionStatistics>,
     centroids: Vec<Vec<f32>>,
     loss: Option<f64>,
+    index_file_version: IndexFileVersion,
 }
 
 fn centroids_to_vectors(centroids: &FixedSizeListArray) -> Result<Vec<Vec<f32>>> {
@@ -880,6 +881,7 @@ impl Index for IVFIndex {
             partitions: partitions_statistics,
             centroids: centroid_vecs,
             loss: self.ivf.loss(),
+            index_file_version: IndexFileVersion::Legacy,
         })?)
     }
 

rust/lance/src/index/vector/ivf/v2.rs

@@ -10,7 +10,10 @@ use std::{
     sync::{Arc, Weak},
 };
 
-use crate::index::vector::builder::{index_type_string, IvfIndexBuilder};
+use crate::index::vector::{
+    builder::{index_type_string, IvfIndexBuilder},
+    IndexFileVersion,
+};
 use crate::{
     index::{
         vector::{utils::PartitionLoadLock, VectorIndex},
@@ -412,6 +415,7 @@ impl<S: IvfSubIndex + 'static, Q: Quantization + 'static> Index for IVFIndex<S,
             partitions: partitions_statistics,
             centroids: centroid_vecs,
             loss: self.ivf.loss(),
+            index_file_version: IndexFileVersion::V3,
         })?)
     }