fix: Remove redundant batching in PostgreSQLOnlineStore.online_write_batch and fix progress bar (feast-dev#4331)

TomSteenbergen · web-flow · commit 0d89d1519fc6 · 2024-07-08T06:27:39.000-04:00
* Remove batching and fix tqdm progress bar

Signed-off-by: TomSteenbergen &lt;tomsteenbergen1995@gmail.com&gt;

* Comment

Signed-off-by: TomSteenbergen &lt;tomsteenbergen1995@gmail.com&gt;

* Remove test changes

Signed-off-by: TomSteenbergen &lt;tomsteenbergen1995@gmail.com&gt;

* Update comment

Signed-off-by: TomSteenbergen &lt;tomsteenbergen1995@gmail.com&gt;

---------

Signed-off-by: TomSteenbergen &lt;tomsteenbergen1995@gmail.com&gt;
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,
