From 727f21d745a25316884732b6fffabdef0f89424b Mon Sep 17 00:00:00 2001 From: yangjie01 Date: Wed, 2 Sep 2026 14:29:32 +0800 Subject: [PATCH 1/2] fix(python): accept numpy ivf_centroids without num_partitions The numpy branch compared the centroid count against num_partitions unconditionally, but num_partitions defaults to None and is documented as deprecated in favor of target_partition_size. A valid 2D centroid array supplied without num_partitions therefore failed the shape check and was reported as "must be 2D array". Split the two checks: reject a non-2D array, and compare against num_partitions only when the caller set it. The Rust core already resolves the partition count from the centroids when num_partitions is absent. --- python/python/lance/dataset.py | 18 +++++++++----- python/python/tests/test_vector_index.py | 30 ++++++++++++++++++++++++ 2 files changed, 42 insertions(+), 6 deletions(-) diff --git a/python/python/lance/dataset.py b/python/python/lance/dataset.py index 616c5f31004..fe38a2beab0 100644 --- a/python/python/lance/dataset.py +++ b/python/python/lance/dataset.py @@ -4105,14 +4105,19 @@ 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 ( + 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" @@ -4265,8 +4270,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`. diff --git a/python/python/tests/test_vector_index.py b/python/python/tests/test_vector_index.py index 253592dd35b..cf8a38aaecd 100644 --- a/python/python/tests/test_vector_index.py +++ b/python/python/tests/test_vector_index.py @@ -1529,6 +1529,36 @@ 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, + ) + def test_create_ivf_pq_skip_transpose(dataset, tmp_path: Path): ds = lance.write_dataset( From 859ce5fb6cc33228fbb5efe49cecbb8717533808 Mon Sep 17 00:00:00 2001 From: yangjie01 Date: Sat, 5 Sep 2026 10:43:01 +0800 Subject: [PATCH 2/2] fix: reject zero-row ivf_centroids before deriving the partition count --- python/python/lance/dataset.py | 7 +++++++ python/python/tests/test_vector_index.py | 13 +++++++++++++ 2 files changed, 20 insertions(+) diff --git a/python/python/lance/dataset.py b/python/python/lance/dataset.py index fe38a2beab0..df5b41cb324 100644 --- a/python/python/lance/dataset.py +++ b/python/python/lance/dataset.py @@ -4110,6 +4110,13 @@ def _create_index_impl( 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 diff --git a/python/python/tests/test_vector_index.py b/python/python/tests/test_vector_index.py index cf8a38aaecd..3272b49cf79 100644 --- a/python/python/tests/test_vector_index.py +++ b/python/python/tests/test_vector_index.py @@ -1559,6 +1559,19 @@ def test_pre_populated_ivf_centroids(dataset, tmp_path: Path): 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(