diff --git a/python/python/lance/dataset.py b/python/python/lance/dataset.py index 662e9307f59..048df3b7ac3 100644 --- a/python/python/lance/dataset.py +++ b/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`. diff --git a/python/python/tests/test_vector_index.py b/python/python/tests/test_vector_index.py index ab5f26379b6..eb2c036f76b 100644 --- a/python/python/tests/test_vector_index.py +++ b/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(