fix: batch update (#133)

quotaclimat/data_processing/mediatree/api_import.py

@@ -25,7 +25,9 @@
 import modin.pandas as pd
 from modin.pandas import json_normalize
 from quotaclimat.utils.sentry import sentry_init
-
+logging.getLogger('modin.logger.default').setLevel(logging.ERROR)
+logging.getLogger('distributed.scheduler').setLevel(logging.ERROR)
+logging.getLogger('distributed').setLevel(logging.ERROR)
 sentry_init()
 
 #read whole file to a string

quotaclimat/data_processing/mediatree/detect_keywords.py

@@ -13,6 +13,8 @@
 import modin.pandas as pd
 import dask
 from quotaclimat.utils.logger import getLogger
+logging.getLogger('modin.logger.default').setLevel(logging.ERROR)
+logging.getLogger('distributed.scheduler').setLevel(logging.ERROR)
 dask.config.set({'dataframe.query-planning': True})
 
 def get_cts_in_ms_for_keywords(subtitle_duration: List[dict], keywords: List[str], theme: str) -> List[dict]:

quotaclimat/data_processing/mediatree/update_pg_keywords.py

@@ -1,20 +1,22 @@
 ### Library imports
 import requests
 import modin.pandas as pd
-import json
 
 import logging
+
 from sqlalchemy.orm import Session
 from postgres.schemas.models import Keywords
 from quotaclimat.data_processing.mediatree.detect_keywords import *
+from sqlalchemy import func, select
 
 def update_keywords(session: Session, batch_size: int = 50000) -> list:
-    saved_keywords = get_keywords_columns(session)
-    total_updates = len(saved_keywords)
-    logging.info(f"Updating {total_updates} saved keywords - keywords_with_timestamp and number_of_keywords columns")
+    total_updates = get_total_count_saved_keywords(session)
+    logging.info(f"Updating {total_updates} saved keywords")
+
     for i in range(0, total_updates, batch_size):
-        batch_updates = saved_keywords[i:i+batch_size]
-        for keyword_id, plaintext, keywords_with_timestamp, number_of_keywords, start, srt, theme in batch_updates:
+        current_batch_saved_keywords = get_keywords_columns(session, i, batch_size)
+        logging.info(f"Updating {len(current_batch_saved_keywords)} elements from {i} offsets - batch size {batch_size}")
+        for keyword_id, plaintext, keywords_with_timestamp, number_of_keywords, start, srt, theme in current_batch_saved_keywords:
             matching_themes, new_keywords_with_timestamp, new_number_of_keywords = get_themes_keywords_duration(plaintext, srt, start)
 
             if(number_of_keywords != new_number_of_keywords or keywords_with_timestamp != new_keywords_with_timestamp or theme != matching_themes):
@@ -32,7 +34,7 @@ def update_keywords(session: Session, batch_size: int = 50000) -> list:
     logging.info("updated all keywords")
 
 
-def get_keywords_columns(session: Session) -> list:
+def get_keywords_columns(session: Session, page: int = 0, batch_size: int = 50000) -> list:
     return (
         session.query(
             Keywords.id,
@@ -43,9 +45,16 @@ def get_keywords_columns(session: Session) -> list:
             Keywords.srt,
             Keywords.theme
         )
+        .offset(page)
+        .limit(batch_size)
         .all()
     )
 
+def get_total_count_saved_keywords(session: Session) -> int:
+        statement = select(func.count()).select_from(Keywords)
+        return session.execute(statement).scalar()
+
+
 def update_keyword_row(session: Session, keyword_id: int, new_number_of_keywords: int, new_keywords_with_timestamp: List[dict], matching_themes: List[str]):
     session.query(Keywords).filter(Keywords.id == keyword_id).update(
         {

test/sitemap/test_update_pg_keywords.py

@@ -123,7 +123,7 @@ def test_insert_data_in_sitemap_table():
     # check the value is well existing
     result_before_update = get_keyword(primary_key)
     session = get_db_session(conn)
-    update_keywords(session)
+    update_keywords(session, batch_size=50)
     result_after_update = get_keyword(primary_key)
 
     new_theme, new_keywords_with_timestamp, new_value = get_themes_keywords_duration(plaintext, srt, start)
@@ -153,4 +153,5 @@ def test_insert_data_in_sitemap_table():
     assert result_after_update.keywords_with_timestamp == new_keywords_with_timestamp
     assert expected_keywords_with_timestamp == new_keywords_with_timestamp
     # theme
-    assert result_after_update.theme == ["changement_climatique_constat", "adaptation_climatique_solutions_indirectes"]
+    assert result_after_update.theme == ["changement_climatique_constat", "adaptation_climatique_solutions_indirectes"]
+    assert new_theme == ["changement_climatique_constat", "adaptation_climatique_solutions_indirectes"]