perf: use modin instead of pandas

docker-compose.yml

@@ -17,6 +17,7 @@ services:
       POSTGRES_PASSWORD: password
       POSTGRES_HOST: postgres_db
       POSTGRES_PORT: 5432
+      MODIN_ENGINE: dask #TODO test me
     tty: true # colorize terminal
     volumes:
       - ./quotaclimat/:/app/quotaclimat/
@@ -142,6 +143,7 @@ services:
       MEDIATREE_PASSWORD:  /run/secrets/pwd_api
       MEDIATREE_AUTH_URL: https://keywords.mediatree.fr/api/auth/token/
       KEYWORDS_URL: https://keywords.mediatree.fr/api/subtitle/ # https://keywords.mediatree.fr/docs/#api-Subtitle-SubtitleList
+      MODIN_ENGINE: dask #TODO test me
     volumes:
       - ./quotaclimat/:/app/quotaclimat/
       - ./postgres/:/app/postgres/

pyproject.toml

@@ -44,6 +44,7 @@ tenacity = "^8.2.3"
 sentry-sdk = "^1.40.5"
 coverage = "^7.4.2"
 modin = {extras = ["all"], version = "^0.27.0"}
+dask-expr = "^0.4.2"
 
 [build-system]
 requires = ["poetry-core>=1.1"]

quotaclimat/data_processing/mediatree/api_import.py

@@ -90,7 +90,8 @@ async def get_and_save_api_data(exit_event):
             try:
                 df = extract_api_sub(token, channel, type_sub, date_epoch)
                 if(df is not None):
-                    save_to_pg(df, keywords_table, conn)
+                    # must ._to_pandas() because modin to_sql is not working
+                    save_to_pg(df._to_pandas(), keywords_table, conn)
                 else: 
                     logging.info("Nothing to save to Postgresql")
             except Exception as err:
@@ -230,7 +231,7 @@ def parse_reponse_subtitle(response_sub, channel = None) -> Optional[pd.DataFram
         log_dataframe_size(new_df, channel)
 
         logging.debug("Parsed %s" % (new_df.head(1).to_string()))
-        logging.info("Parsed Schema\n%s", new_df.dtypes)
+        logging.debug("Parsed Schema\n%s", new_df.dtypes)
 
         return new_df
     else:

quotaclimat/data_processing/mediatree/detect_keywords.py

@@ -10,6 +10,9 @@
 import swifter
 from itertools import groupby
 import modin.pandas as pd
+import dask
+
+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]:
     result = []

test/sitemap/test_main_import_api.py

@@ -0,0 +1,32 @@
+import logging
+
+from quotaclimat.data_processing.mediatree.update_pg_keywords import *
+
+from postgres.insert_data import (clean_data,
+                                  insert_data_in_sitemap_table)
+
+from postgres.schemas.models import create_tables, get_db_session, get_keyword, connect_to_db
+from postgres.insert_data import save_to_pg
+from quotaclimat.data_processing.mediatree.detect_keywords import *
+from quotaclimat.data_processing.mediatree.api_import import *
+
+import time as t
+
+
+def test_main_api_import():
+    create_tables()
+    conn = connect_to_db()
+    json_file_path = 'test/sitemap/mediatree.json'
+    with open(json_file_path, 'r') as file:
+        json_response = json.load(file)
+        start_time = t.time()
+        df = parse_reponse_subtitle(json_response)
+        df = filter_and_tag_by_theme(df)
+        df["id"] = add_primary_key(df)
+        end_time = t.time()
+        logging.info(f"Elapsed time for api import {end_time - start_time}")
+        # must df._to_pandas() because to_sql does not handle modin dataframe
+        save_to_pg(df._to_pandas(), keywords_table, conn)
+        session = get_db_session(conn)
+        saved_keywords = get_keywords_columns(session)
+        assert len(saved_keywords) == len(df)

test/sitemap/test_mediatree.py

@@ -101,7 +101,7 @@ def test_parse_reponse_subtitle():
     expected_result['start'] = pd.to_datetime(expected_result['start'], unit='s').dt.tz_localize('UTC').dt.tz_convert('Europe/Paris')
     df = parse_reponse_subtitle(json_response)
     debug_df(df)
-    # df = df.to_pandas()
+
     pd.testing.assert_frame_equal(df._to_pandas().reset_index(drop=True), expected_result.reset_index(drop=True))
 
 def test_get_includes_or_query():

test/sitemap/test_update_pg_keywords.py

@@ -60,7 +60,7 @@ def test_insert_data_in_sitemap_table():
     }]) 
     df['start'] = pd.to_datetime(df['start'], unit='ms').dt.tz_localize('UTC').dt.tz_convert('Europe/Paris')
 
-    assert save_to_pg(df, keywords_table, conn) == 1
+    assert save_to_pg(df._to_pandas(), keywords_table, conn) == 1
 
     # check the value is well existing
     result_before_update = get_keyword(primary_key)