fix(python): accept numpy ivf_centroids without num_partitions

python/python/lance/dataset.py

@@ -4188,14 +4188,26 @@ def _create_index_impl(
                 if _check_for_numpy(ivf_centroids) and isinstance(
                     ivf_centroids, np.ndarray
                 ):
-                    if (
-                        len(ivf_centroids.shape) != 2
-                        or ivf_centroids.shape[0] != num_partitions
-                    ):
+                    if len(ivf_centroids.shape) != 2:
                         raise ValueError(
                             f"Ivf centroids must be 2D array: (clusters, dim), "
                             f"got {ivf_centroids.shape}"
                         )
+                    if ivf_centroids.shape[0] == 0:
+                        # num_partitions was derived from shape[0] above, and
+                        # zero partitions panics in the Rust residual step.
+                        raise ValueError(
+                            "Ivf centroids must have at least one cluster, "
+                            f"got {ivf_centroids.shape}"
+                        )
+                    if (
+                        num_partitions is not None
+                        and ivf_centroids.shape[0] != num_partitions
+                    ):
+                        raise ValueError(
+                            f"Ivf centroids has {ivf_centroids.shape[0]} clusters, "
+                            f"but num_partitions={num_partitions}"
+                        )
                     if ivf_centroids.dtype not in [np.float16, np.float32, np.float64]:
                         raise TypeError(
                             "IVF centroids must be floating number"
@@ -4348,8 +4360,9 @@ def create_index(
             It can be either :py:class:`np.ndarray`,
             :py:class:`pyarrow.FixedSizeListArray` or
             :py:class:`pyarrow.FixedShapeTensorArray`.
-            A ``num_partitions x dimension`` array of existing K-mean centroids
-            for IVF clustering. If not provided, a new KMeans model will be trained.
+            A ``num_clusters x dimension`` array of existing K-mean centroids
+            for IVF clustering. The row count determines the number of IVF
+            partitions. If not provided, a new KMeans model will be trained.
         pq_codebook : optional,
             It can be :py:class:`np.ndarray`, :py:class:`pyarrow.FixedSizeListArray`,
             or :py:class:`pyarrow.FixedShapeTensorArray`.

python/python/tests/test_vector_index.py

@@ -1569,6 +1569,49 @@ def test_pre_populated_ivf_centroids(dataset, tmp_path: Path):
     partition_keys = {"size"}
     assert all([partition_keys == set(p.keys()) for p in partitions])
 
+    # num_partitions is deprecated in favor of target_partition_size, so
+    # centroids supplied without it must not be rejected. Seven clusters, so the
+    # assertion below tells the new index apart from the five-cluster one above
+    # and from the four target_partition_size would have picked.
+    new_centroids = np.random.randn(7, 128).astype(np.float32)
+    dataset_with_index = dataset.create_index(
+        ["vector"],
+        index_type="IVF_PQ",
+        metric="cosine",
+        ivf_centroids=new_centroids,
+        # 1000 rows / 250 = 4, so this diverges from the centroid count.
+        target_partition_size=250,
+        num_sub_vectors=8,
+        replace=True,
+    )
+    stats = dataset_with_index.stats.index_stats("vector_idx")
+    assert stats["indices"][0]["num_partitions"] == 7
+
+    # A count that disagrees with an explicitly passed num_partitions is still
+    # rejected, and the message now names both numbers.
+    with pytest.raises(ValueError, match="but num_partitions=4"):
+        dataset.create_index(
+            ["vector"],
+            index_type="IVF_PQ",
+            metric="cosine",
+            ivf_centroids=new_centroids,
+            num_partitions=4,
+            num_sub_vectors=8,
+        )
+
+    # A zero-row array passes the 2D check, and the Rust residual step panics on
+    # the empty centroid buffer.
+    with pytest.raises(ValueError, match="at least one cluster"):
+        dataset.create_index(
+            ["vector"],
+            index_type="IVF_PQ",
+            metric="cosine",
+            ivf_centroids=np.empty((0, 128), dtype=np.float32),
+            num_sub_vectors=8,
+            # Otherwise the duplicate-name check intercepts first.
+            replace=True,
+        )
+
 
 def test_create_ivf_pq_skip_transpose(dataset, tmp_path: Path):
     ds = lance.write_dataset(