[PULL REQUEST] Optimize data insertion for ASE data

pyproject.toml

@@ -3,11 +3,11 @@ name = "Estimates-Program"
 version = "0.0.0" # Not used, refer to config.yml instead
 requires-python = ">=3.11"
 dependencies = [
+    "black>=25.12.0,<26.0.0",
     "cerberus>=1.3.8,<2.0.0",
     "mssql>=1.0.1,<2.0.0",
     "numpy>=2.4.0,<3.0.0",
     "pandas>=2.3.3,<3.0.0",
-    "polars>=1.36.1,<2.0.0",
     "pyodbc>=5.3.0,<6.0.0",
     "pyyaml>=6.0.3,<7.0.0",
     "sqlalchemy>=2.0.45,<3.0.0",

python/ase.py

@@ -2,12 +2,12 @@
 # wiki page for more details:
 # https://github.com/SANDAG/Estimates-Program/wiki/Population-by-Age-Sex-Ethnicity
 
+import csv
 import functools
 import logging
 
 import numpy as np
 import pandas as pd
-import polars as pl
 import sqlalchemy as sql
 
 import python.tests as tests
@@ -78,7 +78,7 @@ def run_ase(year: int) -> None:
     ase_outputs = _create_ase(year, ase_inputs)
     _validate_ase_outputs(ase_outputs)
 
-    _insert_ase(ase_outputs)
+    _insert_ase(year, ase_outputs)
 
 
 @functools.lru_cache(maxsize=1)
@@ -915,15 +915,15 @@ def _validate_ase_outputs(ase_outputs: dict[str, pd.DataFrame]) -> None:
         )
 
 
-def _insert_ase(ase_outputs: dict[str, pd.DataFrame]) -> None:
+def _insert_ase(year: int, ase_outputs: dict[str, pd.DataFrame]) -> None:
     """Insert age/sex/ethnicity population by type to database."""
     for pop_type, output in ase_outputs.items():
         logger.info("Loading Estimates for " + pop_type)
 
-        # Convert the DataFrame to a Polars DataFrame
-        # Polars used solely for write to CSV performance
-        pl_df = pl.from_pandas(
-            output[
+        # Write the DataFrame to a CSV file
+        csv_temp_location = utils.BULK_INSERT_STAGING / (pop_type + ".txt")
+        (
+            output.loc[lambda df: df["value"] != 0][
                 [
                     "run_id",
                     "year",
@@ -934,29 +934,25 @@ def _insert_ase(ase_outputs: dict[str, pd.DataFrame]) -> None:
                     "ethnicity",
                     "value",
                 ]
-            ],
-            include_index=False,
-        )
-
-        # Write the DataFrame to a CSV file
-        pl_df.write_csv(
-            utils.BULK_INSERT_STAGING / (pop_type + ".txt"),
-            include_header=False,
-            separator="|",
-            quote_style="never",
+            ].to_csv(
+                csv_temp_location,
+                header=False,
+                index=False,
+                sep="|",
+                quoting=csv.QUOTE_NONE,
+            )
         )
 
         # Bulk insert the CSV file into the production database
         with utils.ESTIMATES_ENGINE.connect() as con:
-            fp = (utils.BULK_INSERT_STAGING / (pop_type + ".txt")).as_posix()
             query = sql.text(
                 f"""
                     BULK INSERT [outputs].[ase]
-                    FROM '{fp}'
+                    FROM '{csv_temp_location.as_posix()}'
                     WITH (
                         TABLOCK,
                         MAXERRORS=0,
-                        FIELDTERMINATOR = '|',
+                        FIELDTERMINATOR = '|',  
                         ROWTERMINATOR = '0x0A',
                         CHECK_CONSTRAINTS
                     )
@@ -966,4 +962,4 @@ def _insert_ase(ase_outputs: dict[str, pd.DataFrame]) -> None:
             con.commit()
 
         # Remove the temporary CSV file
-        (utils.BULK_INSERT_STAGING / (pop_type + ".txt")).unlink()
+        csv_temp_location.unlink()

python/utils.py

@@ -105,6 +105,7 @@
 RUN_INSTRUCTIONS = input_parser.run_instructions
 RUN_ID = input_parser.run_id
 MGRA_VERSION = input_parser.mgra_version
+MGRA_SERIES = int(MGRA_VERSION.replace("mgra", ""))
 
 logger.info(
     f"RUN_ID: {RUN_ID}, MGRA_VERSION: {MGRA_VERSION}, YEARS: {RUN_INSTRUCTIONS["years"]}"

sql/ase/insert_ase_zeros.sql

@@ -0,0 +1,55 @@
+DECLARE @run_id INTEGER = :run_id;
+DECLARE @year INTEGER = :year;
+DECLARE @series INTEGER = :series;
+
+DROP TABLE IF EXISTS [#shell];
+SELECT *, 0 AS [value]
+INTO [#shell]
+FROM (VALUES(@run_id)) AS [run_id]([run_id])
+CROSS JOIN (VALUES(@year)) AS [year]([year])
+CROSS JOIN (
+        SELECT DISTINCT [mgra]
+        FROM [demographic_warehouse].[dim].[mgra]
+        WHERE [series] = @series
+) AS [mgra]
+CROSS JOIN (
+        SELECT DISTINCT [long_name] AS [pop_type]
+        FROM [demographic_warehouse].[dim].[housing_type]
+        WHERE [housing_type_id] IN (1, 2, 3, 4, 6) -- Skip 'Group Quarters - Civilian'
+) AS [pop_type]
+CROSS JOIN (
+        SELECT DISTINCT [name] AS [age_group]
+        FROM [demographic_warehouse].[dim].[age_group]
+        WHERE [age_group_id] BETWEEN 1 AND 20
+    ) AS [age_group]
+CROSS JOIN (
+        SELECT DISTINCT [sex]
+        FROM [demographic_warehouse].[dim].[sex]
+        WHERE [sex_id] IN (1, 2)
+    ) AS [sex]
+CROSS JOIN (
+        SELECT DISTINCT [long_name] AS [ethnicity]
+        FROM [demographic_warehouse].[dim].[ethnicity]
+        WHERE [ethnicity_id] IN (1, 2, 3, 4, 5, 6, 8) -- Skip 'Non-Hispanic, Other'
+    ) AS [ethnicity]
+
+INSERT INTO [EstimatesProgram].[outputs].[ase]
+SELECT
+    [#shell].[run_id],
+    [#shell].[year],
+    [#shell].[mgra],
+    [#shell].[pop_type],
+    [#shell].[age_group],
+    [#shell].[sex],
+    [#shell].[ethnicity],
+    [#shell].[value]
+FROM [#shell]
+LEFT JOIN [EstimatesProgram].[outputs].[ase]
+    ON [#shell].[run_id] = [ase].[run_id]
+    AND [#shell].[year] = [ase].[year]
+    AND [#shell].[mgra] = [ase].[mgra]
+    AND [#shell].[pop_type] = [ase].[pop_type]
+    AND [#shell].[age_group] = [ase].[age_group]
+    AND [#shell].[sex] = [ase].[sex]
+    AND [#shell].[ethnicity] = [ase].[ethnicity]
+WHERE [ase].[value] IS NULL

sql/create_objects.sql

@@ -131,6 +131,87 @@ CREATE TABLE [outputs].[ase] (
     CONSTRAINT [chk_non_negative_outputs_ase] CHECK ([value] >= 0)
 )
 
+-- For purposes of data insertion speed, only non-zero ASE data is inserted into
+-- [outputs].[ase]. In case you want the full table with zeros, you can use the below
+-- function
+CREATE FUNCTION [outputs].[ase_with_zeros](
+    @run_id INTEGER
+)
+RETURNS @ase_with_zeros TABLE (
+        [run_id] INTEGER NOT NULL,
+        [year] INTEGER NOT NULL,
+        [mgra] INTEGER NOT NULL,
+        [pop_type] NVARCHAR(75) NOT NULL,
+        [age_group] NVARCHAR(15) NOT NULL,
+        [sex] NVARCHAR(6) NOT NULL,
+        [ethnicity] NVARCHAR(50) NOT NULL,
+        [value] INTEGER NOT NULL
+    )
+AS
+BEGIN
+    INSERT INTO @ase_with_zeros
+    DECLARE @run_id INTEGER = 134;
+    SELECT
+        [shell].[run_id],
+        [shell].[year],
+        [shell].[mgra],
+        [shell].[pop_type],
+        [shell].[age_group],
+        [shell].[sex],
+        [shell].[ethnicity],
+        ISNULL([ase].[value], 0) AS [value]
+    FROM (
+            SELECT
+                [run_id].[run_id],
+                [year].[year],
+                [mgra].[mgra],
+                [pop_type].[pop_type],
+                [age_group].[age_group],
+                [sex].[sex],
+                [ethnicity].[ethnicity]
+            FROM (VALUES(@run_id)) AS [run_id]([run_id])
+            CROSS JOIN (
+                SELECT DISTINCT [year]
+                FROM [outputs].[ase]
+                WHERE [run_id] = @run_id
+            ) AS [year]
+            CROSS JOIN (
+                SELECT DISTINCT [mgra]
+                FROM [inputs].[mgra]
+                WHERE [run_id] = @run_id
+            ) AS [mgra]
+            CROSS JOIN (
+                SELECT DISTINCT [pop_type]
+                FROM [outputs].[ase]
+                WHERE [run_id] = @run_id
+            ) AS [pop_type]
+            CROSS JOIN (
+                SELECT DISTINCT [age_group]
+                FROM [outputs].[ase]
+                WHERE [run_id] = @run_id
+            ) AS [age_group]
+            CROSS JOIN (
+                SELECT DISTINCT [sex]
+                FROM [outputs].[ase]
+                WHERE [run_id] = @run_id
+            ) AS [sex]
+            CROSS JOIN (
+                SELECT DISTINCT [ethnicity]
+                FROM [outputs].[ase]
+                WHERE [run_id] = @run_id
+            ) AS [ethnicity]
+        ) AS [shell]
+    LEFT JOIN [outputs].[ase]
+        ON [shell].[run_id] = [ase].[run_id]
+        AND [shell].[year] = [ase].[year]
+        AND [shell].[mgra] = [ase].[mgra]
+        AND [shell].[pop_type] = [ase].[pop_type]
+        AND [shell].[age_group] = [ase].[age_group]
+        AND [shell].[sex] = [ase].[sex]
+        AND [shell].[ethnicity] = [ase].[ethnicity]
+    RETURN;
+END
+
 CREATE TABLE [outputs].[gq] (
     [run_id] INT NOT NULL,
     [year] INT NOT NULL,