From 0d89d1519fc6b8ddd05a2588138e2e85f5a921b1 Mon Sep 17 00:00:00 2001 From: Tom Steenbergen <41334387+TomSteenbergen@users.noreply.github.com> Date: Mon, 8 Jul 2024 12:27:39 +0200 Subject: [PATCH] fix: Remove redundant batching in PostgreSQLOnlineStore.online_write_batch and fix progress bar (#4331) * Remove batching and fix tqdm progress bar Signed-off-by: TomSteenbergen * Comment Signed-off-by: TomSteenbergen * Remove test changes Signed-off-by: TomSteenbergen * Update comment Signed-off-by: TomSteenbergen --------- Signed-off-by: TomSteenbergen --- .../feast/infra/online_stores/contrib/postgres.py | 13 +++++-------- 1 file changed, 5 insertions(+), 8 deletions(-) diff --git a/sdk/python/feast/infra/online_stores/contrib/postgres.py b/sdk/python/feast/infra/online_stores/contrib/postgres.py index 8715f0f65b..330b50bc78 100644 --- 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,