Fix up call to compute_partitions in create_index

lancedb · Jul 3, 2024 · f9e7a41 · f9e7a41
1 parent 401617a
commit f9e7a41
Show file tree

Hide file tree

Showing 3 changed files with 3 additions and 3 deletions.
diff --git a/python/python/lance/dataset.py b/python/python/lance/dataset.py
@@ -1490,7 +1490,7 @@ def create_index(
                     accelerator,
                 )
                 partitions_file = compute_partitions(
-                    self, column, kmeans, batch_size=20480
+                    self, column[0], kmeans, batch_size=20480
                 )
                 kwargs["precomputed_partitions_file"] = partitions_file
 

diff --git a/python/python/lance/indices.py b/python/python/lance/indices.py
@@ -168,7 +168,7 @@ def train_ivf(
             # Use accelerator to train ivf centroids
             from .vector import train_ivf_centroids_on_accelerator
 
-            ivf_centroids, kmeans = train_ivf_centroids_on_accelerator(
+            ivf_centroids, _ = train_ivf_centroids_on_accelerator(
                 self.dataset,
                 column[0],
                 num_partitions,

diff --git a/python/python/lance/vector.py b/python/python/lance/vector.py
@@ -183,7 +183,7 @@ def train_ivf_centroids_on_accelerator(
         np.save(f, centroids)
     logging.info("Saved centroids to %s", f.name)
 
-    return centroids, compute_partitions(dataset, column, kmeans, batch_size=20480)
+    return centroids, kmeans
 
 
 def compute_partitions(