Fix pylance typo and float32 array conversion. (#463)

lancedb · Jan 26, 2023 · 4dc6706 · 4dc6706
1 parent b043fd0
commit 4dc6706
Showing 1 changed file with 4 additions and 4 deletions.
diff --git a/rust/pylance/python/lance/dataset.py b/rust/pylance/python/lance/dataset.py
@@ -48,7 +48,7 @@ def scanner(self, columns: Optional[list[str]] = None,
         nearest: dict, default None
             Get the rows corresponding to the K most similar vectors
             nearest should look like {
-              "columns": <embedding col name>,
+              "column": <embedding col name>,
               "q": <query vector as pa.Float32Array>,
               "k": 10
             }
@@ -85,7 +85,7 @@ def to_table(self, columns: Optional[list[str]] = None,
         nearest: dict, default None
             Get the rows corresponding to the K most similar vectors
             nearest should look like {
-              "columns": <embedding col name>,
+              "column": <embedding col name>,
               "q": <query vector as pa.Float32Array>,
               "k": 10
             }
@@ -272,7 +272,7 @@ def columns(self, cols: Optional[list[str]] = None) -> ScannerBuilder:
         return self
 
     def nearest(self, column: Optional[str] = None,
-                q: Optional[pa.Float32Array] = None,
+                q: Optional[pa.FloatingPointArray] = None,
                 k: Optional[int] = None) -> ScannerBuilder:
         if column is None or q is None:
             self._nearest = None
@@ -281,7 +281,7 @@ def nearest(self, column: Optional[str] = None,
         if self.ds.schema.get_field_index(column) < 0:
             raise ValueError(f"Embedding column {column} not in dataset")
         if isinstance(q, (np.ndarray, list, tuple)):
-            q = pa.Float32Array.from_pandas(q)
+            q = pa.FloatingPointArray.from_pandas(q, type=pa.float32())
         if k is not None and int(k) <= 0:
             raise ValueError(f"Nearest-K must be > 0 but got {k}")
         self._nearest = {