compress files for export with gzip #451

Author: Hagellach37

mapswipe_workers/mapswipe_workers/generate_stats/project_stats.py

@@ -1,6 +1,8 @@
 import datetime
 import os
 from typing import List
+import gzip
+import io
 
 import pandas as pd
 from psycopg2 import sql
@@ -31,10 +33,18 @@ def write_sql_to_csv(filename: str, sql_query: sql.SQL):
     Use the copy statement to write data from postgres to a csv file.
     """
 
+    temp_file = "temp.csv"
     pg_db = auth.postgresDB()
-    with open(filename, "w") as f:
+    with open(temp_file, "w") as f:
         pg_db.copy_expert(sql_query, f)
-    logger.info(f"wrote csv file from sql: {filename}")
+
+    with open(temp_file, 'rb') as f_in, gzip.open(filename, 'wb') as f_out:
+        f_out.writelines(f_in)
+
+    # remove temp file
+    os.remove(temp_file)
+
+    logger.info(f"wrote gzipped csv file from sql: {filename}")
 
 
 def load_df_from_csv(filename: str) -> pd.DataFrame:
@@ -44,7 +54,11 @@ def load_df_from_csv(filename: str) -> pd.DataFrame:
     """
     dtype_dict = {"project_id": str, "group_id": str, "task_id": str}
 
-    df = pd.read_csv(filename, dtype=dtype_dict)
+    df = pd.read_csv(
+        filename,
+        dtype=dtype_dict,
+        compression="gzip"
+    )
     logger.info(f"loaded pandas df from {filename}")
     return df
 
@@ -322,11 +336,11 @@ def get_per_project_statistics(project_id: str, project_info: pd.Series) -> dict
     """
 
     # set filenames
-    results_filename = f"{DATA_PATH}/api/results/results_{project_id}.csv"
-    tasks_filename = f"{DATA_PATH}/api/tasks/tasks_{project_id}.csv"
-    groups_filename = f"{DATA_PATH}/api/groups/groups_{project_id}.csv"
-    agg_results_filename = f"{DATA_PATH}/api/agg_results/agg_results_{project_id}.csv"
-    agg_results_by_user_id_filename = f"{DATA_PATH}/api/users/users_{project_id}.csv"
+    results_filename = f"{DATA_PATH}/api/results/results_{project_id}.csv.gz"
+    tasks_filename = f"{DATA_PATH}/api/tasks/tasks_{project_id}.csv.gz"
+    groups_filename = f"{DATA_PATH}/api/groups/groups_{project_id}.csv.gz"
+    agg_results_filename = f"{DATA_PATH}/api/agg_results/agg_results_{project_id}.csv.gz"
+    agg_results_by_user_id_filename = f"{DATA_PATH}/api/users/users_{project_id}.csv.gz"
     project_stats_by_date_filename = f"{DATA_PATH}/api/history/history_{project_id}.csv"
 
     # load data from postgres or local storage if already downloaded
@@ -339,11 +353,23 @@ def get_per_project_statistics(project_id: str, project_info: pd.Series) -> dict
         groups_df = get_groups(groups_filename, project_id)
         tasks_df = get_tasks(tasks_filename, project_id)
 
+        if any("maxar" in s for s in project_info["tile_server_names"]):
+            add_metadata = True
+
         # aggregate results by task id
         agg_results_df = get_agg_results_by_task_id(results_df, tasks_df)
-        agg_results_df.to_csv(agg_results_filename, index_label="idx")
+        agg_results_df.to_csv(
+            agg_results_filename,
+            index_label="idx",
+            compression="gzip"
+        )
+
+        geojson_functions.gzipped_csv_to_gzipped_geojson(
+            filename=agg_results_filename,
+            geometry_field="geom",
+            add_metadata=add_metadata
+        )
         logger.info(f"saved agg results for {project_id}: {agg_results_filename}")
-        geojson_functions.csv_to_geojson(agg_results_filename, "geom")
 
         # aggregate results by user id
         # TODO: solve memory issue for agg results by user id
@@ -361,10 +387,6 @@ def get_per_project_statistics(project_id: str, project_info: pd.Series) -> dict
             sentry.capture_exception()
             logger.info(f"failed to agg results by user id for {project_id}")
 
-        if any("maxar" in s for s in project_info["tile_server_names"]):
-            add_metadata_to_csv(agg_results_filename)
-            geojson_functions.add_metadata_to_geojson(agg_results_filename)
-
         project_stats_by_date_df = project_stats_by_date.get_project_history(
             results_df, groups_df
         )
@@ -380,7 +402,10 @@ def get_per_project_statistics(project_id: str, project_info: pd.Series) -> dict
             # do not do this for ArbitraryGeometry / BuildingFootprint projects
             logger.info(f"do NOT generate tasking manager geometries for {project_id}")
         else:
-            tasking_manager_geometries.generate_tasking_manager_geometries(project_id)
+            tasking_manager_geometries.generate_tasking_manager_geometries(
+                project_id=project_id,
+                agg_results_filename=agg_results_filename
+            )
 
         # prepare output of function
         project_stats_dict = {

mapswipe_workers/mapswipe_workers/generate_stats/tasking_manager_geometries.py

@@ -1,4 +1,5 @@
 import csv
+import gzip
 import threading
 from queue import Queue
 
@@ -8,15 +9,15 @@
 from mapswipe_workers.utils import geojson_functions, tile_functions
 
 
-def load_data(project_id: str, csv_file: str) -> list:
+def load_data(project_id: str, gzipped_csv_file: str) -> list:
     """
     This will load the aggregated results csv file into a list of dictionaries.
     For further steps we currently rely on task_x, task_y, task_z and yes_share and
     maybe_share and wkt
     """
 
     project_data = []
-    with open(csv_file, "r") as f:
+    with gzip.open(gzipped_csv_file, mode="rt") as f:
         reader = csv.reader(f, delimiter=",")
 
         for i, row in enumerate(reader):
@@ -416,7 +417,7 @@ def dissolve_project_data(project_data):
     return dissolved_geometry
 
 
-def generate_tasking_manager_geometries(project_id: str):
+def generate_tasking_manager_geometries(project_id: str, agg_results_filename):
     """
     This functions runs the workflow to create a GeoJSON file ready to be used in the
     HOT Tasking Manager.
@@ -428,14 +429,13 @@ def generate_tasking_manager_geometries(project_id: str):
     Finally, both data sets are saved into GeoJSON files.
     """
 
-    raw_data_filename = f"{DATA_PATH}/api/agg_results/agg_results_{project_id}.csv"
     filtered_data_filename = f"{DATA_PATH}/api/yes_maybe/yes_maybe_{project_id}.geojson"
     tasking_manager_data_filename = (
         f"{DATA_PATH}/api/hot_tm/hot_tm_{project_id}.geojson"
     )
 
     # load project data from existing files
-    results = load_data(project_id, raw_data_filename)
+    results = load_data(project_id, agg_results_filename)
 
     # filter yes and maybe results
     filtered_results = filter_data(results)

mapswipe_workers/mapswipe_workers/utils/geojson_functions.py

@@ -1,12 +1,67 @@
 import json
 import os
+import gzip
+import shutil
 import subprocess
 
 from osgeo import ogr, osr
 
 from mapswipe_workers.definitions import logger
 
 
+def gzipped_csv_to_gzipped_geojson(
+        filename: str,
+        geometry_field: str = "geom",
+        add_metadata: bool = False
+):
+    """Use ogr2ogr to convert csv file to GeoJSON.
+
+    Check if file is compressed.
+    """
+    csv_file = "temp.csv"
+    geojson_file = "temp.geojson"
+    outfile = filename.replace(".csv", f"_{geometry_field}.geojson")
+    filename_without_path = csv_file.split("/")[-1].replace(".csv", "")
+
+    with gzip.open(filename, 'rb') as f_in:
+        with open(csv_file, 'wb') as f_out:
+            shutil.copyfileobj(f_in, f_out)
+
+    # need to remove file here because ogr2ogr can't overwrite when choosing GeoJSON
+    if os.path.isfile(geojson_file):
+        os.remove(geojson_file)
+
+    # TODO: remove geom column from normal attributes in sql query
+    subprocess.run(
+        [
+            "ogr2ogr",
+            "-f",
+            "GeoJSON",
+            geojson_file,
+            csv_file,
+            "-sql",
+            f'SELECT *, CAST({geometry_field} as geometry) FROM "{filename_without_path}"',  # noqa E501
+        ],
+        check=True,
+    )
+    logger.info(f"converted {filename} to {outfile}.")
+
+    if add_metadata:
+        add_metadata_to_geojson(geojson_file)
+
+    cast_datatypes_for_geojson(geojson_file)
+
+    with open(geojson_file, "r") as f:
+        json_data = json.load(f)
+
+    with gzip.open(outfile, 'wt') as fout:
+        json.dump(json_data, fout)
+
+    # remove temp files
+    os.remove(csv_file)
+    os.remove(geojson_file)
+
+
 def csv_to_geojson(filename: str, geometry_field: str = "geom"):
     """
     Use ogr2ogr to convert csv file to GeoJSON