fix: one more finite vector fix (#3673)

one more place that needed valid part id filter

Author: chebbyChefNEQ

python/python/lance/vector.py

@@ -516,7 +516,7 @@ def _partition_assignment() -> Iterable[pa.RecordBatch]:
 
                 split_columns = []
                 if num_sub_vectors is not None:
-                    residual_vecs = vecs - kmeans.centroids[partitions]
+                    residual_vecs = vecs[mask_gpu] - kmeans.centroids[partitions]
                     for i in range(num_sub_vectors):
                         subvector_tensor = residual_vecs[
                             :, i * subvector_size : (i + 1) * subvector_size
@@ -685,7 +685,7 @@ def _partition_and_pq_codes_assignment() -> Iterable[pa.RecordBatch]:
                 assert vecs.shape[0] == ids.shape[0]
 
                 # Ignore any invalid vectors.
-                mask_gpu = partitions.isfinite()
+                mask_gpu = partitions.isfinite() & (partitions >= 0)
                 ids = ids.to(ivf_kmeans.device)[mask_gpu].cpu().reshape(-1)
                 partitions = partitions[mask_gpu].cpu()
                 vecs = vecs[mask_gpu]
@@ -746,8 +746,6 @@ def _partition_and_pq_codes_assignment() -> Iterable[pa.RecordBatch]:
     LOGGER.info("Saved precomputed pq_codes to %s", dst_dataset_uri)
 
     shuffle_buffers = [
-        data_file.path()
-        for frag in ds.get_fragments()
-        for data_file in frag.data_files()
+        data_file.path for frag in ds.get_fragments() for data_file in frag.data_files()
     ]
     return dst_dataset_uri, shuffle_buffers

python/python/tests/test_vector_index.py

@@ -254,6 +254,22 @@ def test_index_with_nans(tmp_path):
     validate_vector_index(dataset, "vector")
 
 
+def test_torch_index_with_nans(tmp_path):
+    # 1024 rows, the entire table should be sampled
+    tbl = create_table(nvec=1000, nans=24)
+
+    dataset = lance.write_dataset(tbl, tmp_path)
+    dataset = dataset.create_index(
+        "vector",
+        index_type="IVF_PQ",
+        num_partitions=4,
+        num_sub_vectors=16,
+        accelerator=torch.device("cpu"),
+        one_pass_ivfpq=True,
+    )
+    validate_vector_index(dataset, "vector")
+
+
 def test_index_with_no_centroid_movement(tmp_path):
     # this test makes the centroids essentially [1..]
     # this makes sure the early stop condition in the index building code