fix: index stats reports incorrect partition size (#4847)

BubbleCal · web-flow · commit 00703be7a2e6 · 2025-09-30T15:44:12.000+08:00
fix #4843 Signed-off-by: BubbleCal <bubble-cal@outlook.com>
diff --git a/python/python/tests/test_vector_index.py b/python/python/tests/test_vector_index.py
@@ -609,6 +609,30 @@ def test_create_ivf_pq_with_target_partition_size(dataset, tmp_path):
     assert ann_ds.stats.index_stats("vector_idx")["indices"][0]["num_partitions"] == 2
 
 
+def test_index_size_stats(tmp_path: Path):
+    num_rows = 512
+    dims = 32
+    schema = pa.schema([pa.field("a", pa.list_(pa.float32(), dims), False)])
+    values = pc.random(num_rows * dims).cast("float32")
+    table = pa.Table.from_pydict(
+        {"a": pa.FixedSizeListArray.from_arrays(values, dims)}, schema=schema
+    )
+
+    base_dir = tmp_path / "test"
+
+    dataset = lance.write_dataset(table, base_dir)
+
+    index_name = "vec_idx"
+    dataset.create_index(
+        "a", "IVF_PQ", name=index_name, num_partitions=2, num_sub_vectors=1
+    )
+
+    # Expect to see non-zero sizes here but all sizes are zero
+    stats = dataset.stats.index_stats(index_name)
+    stats = stats["indices"][0]
+    assert stats["partitions"][0]["size"] + stats["partitions"][1]["size"] == num_rows
+
+
 def test_ivf_flat_over_binary_vector(tmp_path):
     dim = 128
     nvec = 1000
diff --git a/rust/lance-index/src/vector/storage.rs b/rust/lance-index/src/vector/storage.rs
@@ -248,6 +248,10 @@ impl<Q: Quantization> IvfQuantizationStorage<Q> {
         self.reader.num_rows()
     }
 
+    pub fn partition_size(&self, part_id: usize) -> usize {
+        self.ivf.partition_size(part_id)
+    }
+
     pub fn quantizer(&self) -> Result<Quantizer> {
         let metadata = self.metadata();
         Q::from_metadata(metadata, self.distance_type)
diff --git a/rust/lance/src/index/vector/ivf/v2.rs b/rust/lance/src/index/vector/ivf/v2.rs
@@ -376,7 +376,7 @@ impl<S: IvfSubIndex + 'static, Q: Quantization + 'static> Index for IVFIndex<S,
     fn statistics(&self) -> Result<serde_json::Value> {
         let partitions_statistics = (0..self.ivf.num_partitions())
             .map(|part_id| IvfIndexPartitionStatistics {
-                size: self.ivf.partition_size(part_id) as u32,
+                size: self.storage.partition_size(part_id) as u32,
             })
             .collect::<Vec<_>>();
 
