feat: support transforming selected fragments in vector transform sta…

…ge for ivf_pq index (#2657)

The `transform_vectors()` api transforms the ivf-pq indices to get an
output table with the row_id, pq_code, and partition_id. This PR
modifies this API so that this transformation is done for selected
fragments that the user specifies as a parameter. If `fragments=None`,
the default will use all fragments of the dataset.

python/python/lance/indices.py

@@ -7,6 +7,7 @@
 import numpy as np
 import pyarrow as pa
 
+from lance import LanceFragment
 from lance.file import LanceFileReader, LanceFileWriter
 
 from .lance import indices
@@ -379,6 +380,7 @@ def transform_vectors(
         ivf: IvfModel,
         pq: PqModel,
         dest_uri: str,
+        fragments: Optional[list[LanceFragment]],
         partition_ds_uri: Optional[str] = None,
     ):
         """
@@ -396,6 +398,9 @@ def transform_vectors(
         dest_uri: str
             The URI to save the transformed vectors to.  The URI can be a local file
             path or a cloud storage path.
+        fragments: list[LanceFragment]
+            The list of data fragments to use when computing the transformed vectors.
+            This is an optional parameter (the default uses all fragments).
         partition_ds_uri: str
             The URI of a precomputed partitions dataset.  This allows the partition
             transform to be skipped, using the precomputed value instead.  This is
@@ -404,6 +409,13 @@ def transform_vectors(
         dimension = self.dataset.schema.field(self.column[0]).type.list_size
         num_subvectors = pq.num_subvectors
         distance_type = ivf.distance_type
+        if fragments is None:
+            fragments = [f._fragment for f in self.dataset.get_fragments()]
+        elif len(fragments) == 0:
+            raise ValueError("fragments must be a non-empty list or None")
+        else:
+            fragments = [f._fragment for f in fragments]
+
         indices.transform_vectors(
             self.dataset._ds,
             self.column[0],
@@ -413,6 +425,7 @@ def transform_vectors(
             ivf.centroids,
             pq.codebook,
             dest_uri,
+            fragments,
         )
 
     def _determine_num_partitions(self, num_partitions: Optional[int], num_rows: int):

python/python/tests/test_indices.py

@@ -7,10 +7,12 @@
 from lance.file import LanceFileReader
 from lance.indices import IndicesBuilder, IvfModel, PqModel
 
-NUM_ROWS = 10000
+NUM_ROWS_PER_FRAGMENT = 10000
 DIMENSION = 128
 NUM_SUBVECTORS = 8
-NUM_PARTITIONS = 100
+NUM_FRAGMENTS = 3
+NUM_ROWS = NUM_ROWS_PER_FRAGMENT * NUM_FRAGMENTS
+NUM_PARTITIONS = round(np.sqrt(NUM_ROWS))
 
 
 @pytest.fixture(params=[np.float16, np.float32, np.float64], ids=["f16", "f32", "f64"])
@@ -19,7 +21,9 @@ def rand_dataset(tmpdir, request):
     vectors.shape = -1
     vectors = pa.FixedSizeListArray.from_arrays(vectors, DIMENSION)
     table = pa.Table.from_arrays([vectors], names=["vectors"])
-    ds = lance.write_dataset(table, str(tmpdir / "dataset"))
+    uri = str(tmpdir / "dataset")
+
+    ds = lance.write_dataset(table, uri, max_rows_per_file=NUM_ROWS_PER_FRAGMENT)
 
     return ds
 
@@ -115,12 +119,15 @@ def rand_pq(rand_dataset, rand_ivf):
 
 
 def test_vector_transform(tmpdir, rand_dataset, rand_ivf, rand_pq):
+    fragments = list(rand_dataset.get_fragments())
+
     builder = IndicesBuilder(rand_dataset, "vectors")
-    builder.transform_vectors(rand_ivf, rand_pq, str(tmpdir / "transformed"))
+    uri = str(tmpdir / "transformed")
+    builder.transform_vectors(rand_ivf, rand_pq, uri, fragments=fragments)
 
-    reader = LanceFileReader(str(tmpdir / "transformed"))
+    reader = LanceFileReader(uri)
 
-    assert reader.metadata().num_rows == 10000
+    assert reader.metadata().num_rows == (NUM_ROWS_PER_FRAGMENT * len(fragments))
     data = next(reader.read_all(batch_size=10000).to_batches())
 
     row_id = data.column("_rowid")
@@ -131,3 +138,9 @@ def test_vector_transform(tmpdir, rand_dataset, rand_ivf, rand_pq):
 
     part_id = data.column("__ivf_part_id")
     assert part_id.type == pa.uint32()
+
+    # test when fragments = None
+    builder.transform_vectors(rand_ivf, rand_pq, uri, fragments=None)
+    reader = LanceFileReader(uri)
+
+    assert reader.metadata().num_rows == (NUM_ROWS_PER_FRAGMENT * NUM_FRAGMENTS)

python/src/fragment.rs

@@ -124,7 +124,7 @@ impl FileFragment {
         self.fragment.id()
     }
 
-    fn metadata(&self) -> FragmentMetadata {
+    pub fn metadata(&self) -> FragmentMetadata {
         FragmentMetadata::new(self.fragment.metadata().clone())
     }
 

python/src/indices.rs

@@ -12,6 +12,7 @@ use lance_index::vector::{
 use lance_linalg::distance::DistanceType;
 use pyo3::{pyfunction, types::PyModule, wrap_pyfunction, PyObject, PyResult, Python};
 
+use crate::fragment::FileFragment;
 use crate::{dataset::Dataset, error::PythonErrorExt, file::object_store_from_uri_or_path, RT};
 
 async fn do_train_ivf_model(
@@ -147,11 +148,14 @@ async fn do_transform_vectors(
     ivf_centroids: FixedSizeListArray,
     pq_model: ProductQuantizer,
     dst_uri: &str,
+    fragments: Vec<FileFragment>,
 ) -> PyResult<()> {
     let num_rows = dataset.ds.count_rows(None).await.infer_error()?;
+    let fragments = fragments.iter().map(|item| item.metadata().inner).collect();
     let transform_input = dataset
         .ds
         .scan()
+        .with_fragments(fragments)
         .project(&[column])
         .infer_error()?
         .with_row_id()
@@ -188,6 +192,7 @@ pub fn transform_vectors(
     ivf_centroids: PyArrowType<ArrayData>,
     pq_codebook: PyArrowType<ArrayData>,
     dst_uri: &str,
+    fragments: Vec<FileFragment>,
 ) -> PyResult<()> {
     let ivf_centroids = ivf_centroids.0;
     let ivf_centroids = FixedSizeListArray::from(ivf_centroids);
@@ -203,7 +208,15 @@ pub fn transform_vectors(
     );
     RT.block_on(
         Some(py),
-        do_transform_vectors(dataset, column, distance_type, ivf_centroids, pq, dst_uri),
+        do_transform_vectors(
+            dataset,
+            column,
+            distance_type,
+            ivf_centroids,
+            pq,
+            dst_uri,
+            fragments,
+        ),
     )?
 }