add python API

Signed-off-by: BubbleCal <[email protected]>

python/python/lance/dataset.py

@@ -3849,6 +3849,7 @@ def __init__(self, dataset: LanceDataset, index_name: str):
         stats = dataset.stats.index_stats(index_name)
 
         self.dataset = dataset
+        self.index_name = index_name
         self.stats = stats
 
     def num_partitions(self) -> int:
@@ -3874,7 +3875,7 @@ def centroids(self) -> np.ndarray:
         )
 
     def partition_reader(
-        self, partition_id: int, with_vector: bool = False
+        self, partition_id: int, *, with_vector: bool = False
     ) -> pa.RecordBatchReader:
         """
         Returns a reader for the given IVF partition
@@ -3893,6 +3894,6 @@ def partition_reader(
             A record batch reader for the given partition
         """
 
-        return self.dataset._ds.partition_reader(
-            self.stats["indices"][0]["partitions"][partition_id]
+        return self.dataset._ds.read_index_partition(
+            self.index_name, partition_id, with_vector
         )

python/python/tests/test_vector_index.py

@@ -13,6 +13,7 @@
 import pyarrow.compute as pc
 import pytest
 from lance import LanceFragment
+from lance.dataset import VectorIndexReader
 
 torch = pytest.importorskip("torch")
 from lance.util import validate_vector_index  # noqa: E402
@@ -1129,3 +1130,16 @@ def test_drop_indices(indexed_dataset):
     )
 
     assert len(results) == 15
+
+
+def test_read_partition(indexed_dataset):
+    idx_name = indexed_dataset.list_indices()[0]["name"]
+    reader = VectorIndexReader(indexed_dataset, idx_name)
+
+    num_rows = indexed_dataset.count_rows()
+    row_sum = 0
+    for part_id in range(reader.num_partitions()):
+        part_reader = reader.partition_reader(part_id)
+        for batch in part_reader:
+            row_sum += batch.num_rows
+    assert row_sum == num_rows

python/src/dataset.rs

@@ -31,7 +31,7 @@ use arrow_array::Array;
 use futures::{StreamExt, TryFutureExt};
 use lance::dataset::builder::DatasetBuilder;
 use lance::dataset::refs::{Ref, TagContents};
-use lance::dataset::scanner::MaterializationStyle;
+use lance::dataset::scanner::{DatasetRecordBatchStream, MaterializationStyle};
 use lance::dataset::statistics::{DataStatistics, DatasetStatisticsExt};
 use lance::dataset::{
     fragment::FileFragment as LanceFileFragment,
@@ -1558,6 +1558,27 @@ impl Dataset {
 
         Ok(())
     }
+
+    #[pyo3(signature = (index_name,partition_id, with_vector=false))]
+    fn read_index_partition(
+        &self,
+        index_name: String,
+        partition_id: usize,
+        with_vector: bool,
+    ) -> PyResult<PyArrowType<Box<dyn RecordBatchReader + Send>>> {
+        let stream = RT
+            .block_on(
+                None,
+                self.ds
+                    .read_index_partition(&index_name, partition_id, with_vector),
+            )?
+            .map_err(|err| PyValueError::new_err(err.to_string()))?;
+
+        let reader = Box::new(LanceReader::from_stream(DatasetRecordBatchStream::new(
+            stream,
+        )));
+        Ok(PyArrowType(reader))
+    }
 }
 
 impl Dataset {

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

@@ -143,7 +143,7 @@ mod test {
             Ok(())
         }
 
-        async fn to_batch_stream(&self, with_vector: bool) -> Result<SendableRecordBatchStream> {
+        async fn to_batch_stream(&self, _with_vector: bool) -> Result<SendableRecordBatchStream> {
             unimplemented!("only for SubIndex")
         }