fix: Remove redundant batching in PostgreSQLOnlineStore.online_write_…

…batch and fix progress bar (feast-dev#4331) * Remove batching and fix tqdm progress bar Signed-off-by: TomSteenbergen <[email protected]> * Comment Signed-off-by: TomSteenbergen <[email protected]> * Remove test changes Signed-off-by: TomSteenbergen <[email protected]> * Update comment Signed-off-by: TomSteenbergen <[email protected]> --------- Signed-off-by: TomSteenbergen <[email protected]>
shuchu · Jul 8, 2024 · 0d89d15 · 0d89d15
1 parent b0dc683
commit 0d89d15
Showing 1 changed file with 5 additions and 8 deletions.
diff --git a/sdk/python/feast/infra/online_stores/contrib/postgres.py b/sdk/python/feast/infra/online_stores/contrib/postgres.py
@@ -75,7 +75,6 @@ def online_write_batch(
             Tuple[EntityKeyProto, Dict[str, ValueProto], datetime, Optional[datetime]]
         ],
         progress: Optional[Callable[[int], Any]],
-        batch_size: int = 5000,
     ) -> None:
         # Format insert values
         insert_values = []
@@ -118,15 +117,13 @@ def online_write_batch(
         """
         ).format(sql.Identifier(_table_id(config.project, table)))
 
-        # Push data in batches to online store
+        # Push data into the online store
         with self._get_conn(config) as conn, conn.cursor() as cur:
-            for i in range(0, len(insert_values), batch_size):
-                cur_batch = insert_values[i : i + batch_size]
-                cur.executemany(sql_query, cur_batch)
-                conn.commit()
+            cur.executemany(sql_query, insert_values)
+            conn.commit()
 
-                if progress:
-                    progress(len(cur_batch))
+        if progress:
+            progress(len(data))
 
     def online_read(
         self,