Merge pull request pinecone-io#137 from pinecone-io/upsert_df_sync

Upsert dataframe - support non-async option

Author: igiloh-pinecone

pinecone/core/grpc/index_grpc.py

@@ -450,28 +450,45 @@ def _upsert_batch(self,
         request = UpsertRequest(vectors=vectors, **args_dict)
         return self._wrap_grpc_call(self.stub.Upsert, request, timeout=timeout, **kwargs)
 
-    def upsert_dataframe(self,
+    def upsert_from_dataframe(self,
                          df,
-                         namespase: str = None,
+                         namespace: str = None,
                          batch_size: int = 500,
+                         use_async_requests: bool = True,
                          show_progress: bool = True) -> None:
         """Upserts a dataframe into the index.
 
         Args:
             df: A pandas dataframe with the following columns: id, vector, and metadata.
             namespace: The namespace to upsert into.
             batch_size: The number of rows to upsert in a single batch.
+            use_async_requests: Whether to upsert multiple requests at the same time using asynchronous request mechanism.
+                                Set to `False`
             show_progress: Whether to show a progress bar.
         """
-        if find_spec("pandas") is None:
-            raise ImportError("pandas not found. Please install pandas to use this method.")
-
-        async_results = [
-            self.upsert(vectors=chunk, namespace=namespase, async_req=True)
-            for chunk in tqdm(self._iter_dataframe(df, batch_size=batch_size),
-                              total=len(df) // batch_size, disable=not show_progress)
-        ]
-        res = [async_result.result() for async_result in async_results]
+        try:
+            import pandas as pd
+        except ImportError:
+            raise RuntimeError("The `pandas` package is not installed. Please install pandas to use `upsert_from_dataframe()`")
+
+        if not isinstance(df, pd.DataFrame):
+            raise ValueError(f"Only pandas dataframes are supported. Found: {type(df)}")
+
+        pbar = tqdm(total=len(df), disable=not show_progress, desc="sending upsert requests")
+        results = []
+        for chunk in self._iter_dataframe(df, batch_size=batch_size):
+            res = self.upsert(vectors=chunk, namespace=namespace, async_req=use_async_requests)
+            pbar.update(len(chunk))
+            results.append(res)
+
+        if use_async_requests:
+            results = [async_result.result() for async_result in tqdm(results, desc="collecting async responses")]
+
+        upserted_count = 0
+        for res in results:
+            upserted_count += res.upserted_count
+
+        return UpsertResponse(upserted_count=upserted_count)
 
     @staticmethod
     def _iter_dataframe(df, batch_size):

pinecone/index.py

@@ -229,9 +229,9 @@ def _vector_transform(item: Union[Vector, Tuple]):
             **{k: v for k, v in kwargs.items() if k in _OPENAPI_ENDPOINT_PARAMS}
         )
 
-    def upsert_dataframe(self,
+    def upsert_from_dataframe(self,
                          df,
-                         namespase: str = None,
+                         namespace: str = None,
                          batch_size: int = 500,
                          show_progress: bool = True) -> None:
         """Upserts a dataframe into the index.
@@ -242,15 +242,24 @@ def upsert_dataframe(self,
             batch_size: The number of rows to upsert in a single batch.
             show_progress: Whether to show a progress bar.
         """
-        if find_spec("pandas") is None:
-            raise ImportError("pandas not found. Please install pandas to use this method.")
+        try:
+            import pandas as pd
+        except ImportError:
+            raise RuntimeError("The `pandas` package is not installed. Please install pandas to use `upsert_from_dataframe()`")
 
+        if not isinstance(df, pd.DataFrame):
+            raise ValueError(f"Only pandas dataframes are supported. Found: {type(df)}")
+
+        upserted_count = 0
         pbar = tqdm(total=len(df), disable=not show_progress)
         for i in range(0, len(df), batch_size):
             batch = df.iloc[i:i + batch_size].to_dict(orient="records")
-            self.upsert(batch, namespace=namespase)
+            res = self.upsert(batch, namespace=namespace)
+            upserted_count += res.upserted_count
             pbar.update(len(batch))
 
+        return UpsertResponse(upserted_count=upserted_count)
+
     @validate_and_convert_errors
     def delete(self,
                ids: Optional[List[str]] = None,