diff --git a/python/python/tests/test_index.py b/python/python/tests/test_index.py index 1cf2c733c..94268a53e 100644 --- a/python/python/tests/test_index.py +++ b/python/python/tests/test_index.py @@ -372,6 +372,31 @@ async def test_create_vector_index(some_table: AsyncTable): assert stats.num_indices == 1 +@pytest.mark.asyncio +async def test_create_ivf_index_reports_unsplittable_partitions(db_async): + dim = 8 + num_partitions = 300 # More than 256 selects hierarchical k-means. + base_vectors = [[float(row == column) for column in range(dim)] for row in range(5)] + vectors = pa.array(base_vectors * 200, pa.list_(pa.float32(), dim)) + table = await db_async.create_table( + "unsplittable_partitions", + pa.table({"vector": vectors}), + ) + + error_pattern = ( + rf"Cannot create {num_partitions} IVF partitions: k-means could only form" + ) + with pytest.raises(RuntimeError, match=error_pattern): + await table.create_index( + "vector", + config=IvfFlat( + distance_type="dot", + num_partitions=num_partitions, + max_iterations=10, + ), + ) + + @pytest.mark.asyncio async def test_create_4bit_ivfpq_index(some_table: AsyncTable): # Can create