Reset DuckDB connection between operations to prevent OOM in building partitioning (#313)

Author: andersy005

ocr/pipeline/partition.py

@@ -18,55 +18,59 @@ def partition_buildings_by_geography(config: OCRConfig):
     install_load_extensions(aws=needs_s3, spatial=True, httpfs=True, con=connection)
     apply_s3_creds(region='us-west-2', con=connection)
 
-    if config.debug:
-        console.log(f'Loading buildings data from: {path}')
-
-    connection.execute(f"""
-        SET preserve_insertion_order=false;
-        CREATE TEMP TABLE buildings_temp AS
-        SELECT
-            *,
-            SUBSTRING(GEOID, 1, 2) AS state_fips,
-            SUBSTRING(GEOID, 3, 3) AS county_fips
-        FROM '{path}';
-    """)
+    consolidated_buildings_parquet = (
+        f'{config.vector.building_geoparquet_uri.parent / "consolidated-buildings.parquet"}'
+    )
 
     if config.debug:
-        console.log(f'Partitioning geoparquet regions to: {output_path}')
+        console.log(f'Creating a consolidated parquet file at: {consolidated_buildings_parquet}')
 
     connection.execute(f"""
-        COPY buildings_temp
-        TO '{output_path}' (
+        SET preserve_insertion_order=false;
+        COPY (
+            SELECT *
+            FROM '{path}'
+        )
+        TO '{consolidated_buildings_parquet}'
+        (
             FORMAT 'parquet',
-            PARTITION_BY (state_fips, county_fips),
             COMPRESSION 'zstd',
             OVERWRITE_OR_IGNORE true
         );""")
 
     if config.debug:
-        console.log(f'Partitioned buildings written to: {output_path}')
+        console.log(f'Consolidated buildings written to: {consolidated_buildings_parquet}')
 
-    consolidated_buildings_parquet = (
-        f'{config.vector.building_geoparquet_uri.parent / "consolidated-buildings.parquet"}'
-    )
+    connection.close()
 
     if config.debug:
-        console.log(f'Creating a consolidated parquet file at: {consolidated_buildings_parquet}')
+        console.log('Reconnecting to database for partitioned parquet creation')
+
+    connection = duckdb.connect(database=':memory:')
+    install_load_extensions(aws=needs_s3, spatial=True, httpfs=True, con=connection)
+    apply_s3_creds(region='us-west-2', con=connection)
+
+    if config.debug:
+        console.log(f'Partitioning geoparquet regions from: {path}')
 
     connection.execute(f"""
+        SET preserve_insertion_order=false;
         COPY (
-            SELECT * EXCLUDE (state_fips, county_fips)
-            FROM buildings_temp
+        SELECT *
+        FROM (
+            SELECT
+                *,
+                SUBSTRING(GEOID, 1, 2) AS state_fips,
+                SUBSTRING(GEOID, 3, 3) AS county_fips
+            FROM '{path}'
         )
-        TO '{consolidated_buildings_parquet}'
-        (
+        )
+        TO '{output_path}' (
             FORMAT 'parquet',
+            PARTITION_BY (state_fips, county_fips),
             COMPRESSION 'zstd',
-            OVERWRITE_OR_IGNORE true,
-            ROW_GROUP_SIZE 10000000
+            OVERWRITE_OR_IGNORE true
         );""")
 
     if config.debug:
-        console.log(f'Consolidated buildings written to: {consolidated_buildings_parquet}')
-
-    connection.execute('DROP TABLE buildings_temp')
+        console.log(f'Partitioned buildings written to: {output_path}')