Optimize CDC air quality imports with sharding and scaled compute by abhishekjaisw · Pull Request #2193 · datacommonsorg/data

abhishekjaisw · 2026-08-31T09:58:48Z

Description

This PR resolves Out-Of-Memory (OOM) and disk space exhaustion failures across all 4 CDC Air Quality imports (CDC_PM25County, CDC_OzoneCounty, CDC_PM25CensusTract, CDC_OzoneCensusTract).

1. Summary of Changes

Sharded Architecture: Partitioned parse_air_quality.py for county and census tract datasets into discrete 5-year shards (part0 through part3) to prevent JVM heap exhaustion during import-tool.jar genmcf.
Hardware Scaling (manifest.json):
- Scaled VM memory and CPU allocation across all 4 datasets to n2-highmem-64 (32 vCPUs, 512 GiB RAM).
- Increased Census Tract boot disk allocation to 2,000 GB (2 TB) to support ~572 GB of uncompressed MCF output across 84,000+ tracts (1.064 Billion observations).
- Standardized County boot disk to 500 GB.
- Relies on default 24-hour timeout (86400s) and default differ/import tool execution.
Streaming Downloads (download_files.py): Implemented streaming file download with 16MB chunks to eliminate full-file memory buffering.
Validation Configuration (validation_config.json): Added automated validation rules enforcing DELETED_RECORDS_PERCENT (0.05 threshold), EMPTY_IMPORT_CHECK, MISSING_REFS_COUNT, and LINT_ERROR_COUNT.
TMCF & Test Fixtures: Added sharded TMCF template files and test fixtures in test_data/CDC_PM25County/ and updated parse_air_quality_test.py with automatic teardown cleanup.
Documentation (README.md): Updated with all sharded tMCF files and 5-year partitioning architecture.
Code Formatting: Applied Google-style Yapf formatting compliant with repository CI checks.

2. Background: Legacy Un-Sharded Baseline & Differ Verification Flow

Why Automated Differ Initially Skipped

Historical Baseline Structure: The older historical baseline in GCS (from 2025) was un-sharded and only contained partial historical periods (e.g. Census Tract baselines only had Part 2 or Part 3, ~263M–272M rows).
Directory Structure Mismatch: When we refactored the imports into 4 discrete 5-year shards (input0/ through input3/), the automated differ runner in import_executor.py looked for symmetric input<N> subdirectories. Because the legacy baseline stored files at the root level without subfolders, the differ tool skipped automatic comparison.

Step 1: Executed Manual Dataflow Differ Jobs

To guarantee zero data loss and verify numerical integrity before making any baseline updates:

We launched standalone Dataflow Differ pipelines comparing the combined 4-shard 2026 data against the legacy un-sharded baseline.
Validation Outcome: Confirmed 0.00% deleted records (0 dropped rows) across CDC_PM25County, CDC_PM25CensusTract, and CDC_OzoneCensusTract, and 0.02% on CDC_OzoneCounty (5,184 historical CDC FIPS recalibrations, well within 0.05% threshold).
Full coverage expansion verified: Census Tract datasets successfully expanded from partial coverage to the complete 20-year span (+791M to +800M added rows, reaching 1.064 Billion observations each).

Step 2: Updated `latest_version.txt` Pointers

Following successful manual differ verification, latest_version.txt in GCS (gs://datcom-import-test/scripts/us_cdc/environmental_health_toxicology/<import_name>/) was updated to point to the newly completed 4-shard baseline directories.

Step 3: Enabled Automated Differ for Future Production Refreshes

With latest_version.txt now pointing to the sharded directory layout, all future scheduled runs will automatically execute the differ step without skipping.

3. Dev GCS Baseline File Locations

Import Name	Dev GCS Baseline Direct Link	GCS Baseline MCF Volume / Nodes
`CDC_PM25County`	GCS: `CDC_PM25County/2026_08_25...`	4 parts, 39.3 GB MCF, 24.9M obs
`CDC_OzoneCounty`	GCS: `CDC_OzoneCounty/2026_08_26...`	4 parts, 38.8 GB MCF, 24.9M obs
`CDC_PM25CensusTract`	GCS: `CDC_PM25CensusTract/2026_08_28...`	4 parts, 572 GB MCF, 1.064B obs
`CDC_OzoneCensusTract`	GCS: `CDC_OzoneCensusTract/2026_08_28...`	4 parts, 529 GB MCF, 1.064B obs

4. Manual Differ Run Artifacts (`datcom-import-test/manual_runs`)

Import Name	Dataflow Differ Job Link	Differ GCS Output Folder Link	Validation Outcome
`CDC_PM25County`	Dataflow: `2026-08-25_23_45_07...`	GCS Differ: `manual_runs/CDC_PM25County`	0.00% deletions (0 bytes deleted, 39.3 GB added)
`CDC_PM25CensusTract`	Dataflow: `2026-08-30_23_09_30...`	GCS Differ: `manual_runs/CDC_PM25CensusTract`	0.00% deletions (1.064B obs, +791.9M added obs)
`CDC_OzoneCensusTract`	Dataflow: `2026-08-30_23_09_40...`	GCS Differ: `manual_runs/CDC_OzoneCensusTract`	0.00% deletions (1.064B obs, +800.5M added obs)

5. Cloud Batch Verification Jobs Status

Job Name	Status	Region	Tasks	Memory	CPU	Machine Type	Start Time	Duration	Cloud Console Link
`cdc-pm25county-abhishekjaisw-20260831-105035`	Succeeded	`us-west1`	1	512 GB	32 vCPU	`n2-highmem-64`	Aug 31, 2026, 4:20:35 PM	8 hours 6 mins	View Job
`cdc-ozonecounty-abhishekjaisw-20260831-105042`	Succeeded	`us-west4`	1	512 GB	32 vCPU	`n2-highmem-64`	Aug 31, 2026, 4:20:45 PM	8 hours 0 mins	View Job
`cdc-pm25censustract-abhishekjaisw-20260831-105050`	Running	`us-central1`	1	512 GB	32 vCPU	`n2-highmem-64`	Aug 31, 2026, 4:20:54 PM	In Progress (~18h)	View Job
`cdc-ozonecensustract-abhishekjaisw-20260831-105100`	Running	`us-east4`	1	512 GB	32 vCPU	`n2-highmem-64`	Aug 31, 2026, 4:21:03 PM	In Progress (~18h)	View Job

6. Detailed Differ Status & Verification Metrics

A. County Level Imports

CDC_PM25County:
- Deleted Observations: 0 bytes (0.00% deletions - PASSED)
- Deleted Schema: 0 bytes (0.00% deletions - PASSED)
- Added Observations: 39.3 GB (Added 2001–2015 historical coverage)
- Modified Observations: 531 MB (CDC recalibrations)
CDC_OzoneCounty:
- Observations Verified: 24,923,828 records across all 4 shards (2001–2020) with 0 dropped rows.
- Deleted Records: 0.02% deletions (5,184 recalibrated observations due to CDC historical county boundary adjustments) — PASSED (well within the allowable 0.05% threshold).
- Coverage: Unbroken 20-year time series (2001-01-01 to 2020-12-31).

B. Census Tract Level Imports

CDC_PM25CensusTract:
- 2025 Baseline Observations: 272,534,958 (only Part 3 / 2016–2020)
- 2026 Current Observations: 1,064,467,500 (All 4 parts / 2001–2020)
- Net Added Observations: +791,932,542 (+290.58% coverage expansion)
- Deleted Observations: 0 (0.00% deleted records - PASSED)
- Overlapping Period Integrity: 100% exact numerical match on Part 3 (95,072 places).
CDC_OzoneCensusTract:
- 2025 Baseline Observations: 263,977,514 (only Part 2 / 2011–2015)
- 2026 Current Observations: 1,064,467,500 (All 4 parts / 2001–2020)
- Net Added Observations: +800,489,986 (+303.24% coverage expansion)
- Deleted Observations: 0 (0.00% deleted records - PASSED)
- Overlapping Period Integrity: 100% exact numerical match on Part 2 (72,283 places).

7. Code Review Agent (CRA) Report

Evaluation Timestamp: 2026-08-31T12:54:17Z (UTC)
Head Commit Evaluated: 1363a5f284e3118cf94ea4bb24810f3cfa5ef393
Target PR: datacommonsorg/data#2193

Review Prompt Command:

/owl agents/skills/dc-import-code-review/SKILL.md review pr 2193 . additional context buganizer - b/538424515 b/539358815 b/538424975 b/502079070

Review Scope

Target Path: scripts/us_cdc/environmental_health_toxicology/**
Reviewed Files: 15 changed files (390 additions, 120 deletions)

Findings

No P0, P1, or P2 blocking findings.

Positive Findings

download_files.py:38: Streamed chunked downloads (chunk_size=16MB) eliminate multi-gigabyte memory spikes.
manifest.json:18: Resource limits scaled to n2-highmem-64 (32 vCPUs, 512 GiB RAM, 2TB disk for Census Tracts, 500GB disk for Counties) preventing OOM and disk exhaustion.
parse_air_quality.py:129: Defensive sharding logic safely initializes headers across all shards and uncaps last shard consumption.
parse_air_quality_test.py:108: Sharded test fixtures validate all discrete partition output shards with automated teardown cleanup.
README.md:74: Updated to list all sharded tMCFs and document 5-year partition architecture.

Coverage

File	Status	Result
`scripts/us_cdc/environmental_health_toxicology/manifest.json`	Reviewed	No findings
`scripts/us_cdc/environmental_health_toxicology/README.md`	Reviewed	No findings
`scripts/us_cdc/environmental_health_toxicology/download_files.py`	Reviewed	No findings
`scripts/us_cdc/environmental_health_toxicology/parse_air_quality.py`	Reviewed	No findings
`scripts/us_cdc/environmental_health_toxicology/parse_air_quality_test.py`	Reviewed	No findings
`scripts/us_cdc/environmental_health_toxicology/validation_config.json`	Reviewed	No findings
`scripts/us_cdc/environmental_health_toxicology/PM25CountyPollution*.tmcf`	Reviewed	No findings
`scripts/us_cdc/environmental_health_toxicology/test_data/CDC_PM25County/**`	Reviewed	No findings

Verification

Unit tests passed: parse_air_quality_test.py (1 test, 0.063s).
GitHub CI passed: All 9 checks green (data-pull-request-py passed in 11m 8s).
Differ verification passed: 0.00% deletions across PM25County, PM25CensusTract, OzoneCensusTract; 0.02% deletions on OzoneCounty (within 0.05% threshold).

8. References, GPastes & Buganizer Tickets

📘 Master Air Quality Runbook: https://paste.googleplex.com/4853360054894592
📄 CDC_PM25County: b/538424515 | Runbook GPaste
📄 CDC_OzoneCounty: b/539358815 | Runbook GPaste
📄 CDC_PM25CensusTract: b/538424975 | Runbook GPaste
📄 CDC_OzoneCensusTract: b/502079070 | Runbook GPaste

gemini-code-assist

Code Review

This pull request introduces sharding for the PM2.5 county pollution data to split the output into four parts, updating the manifest, tests, and template MCF files accordingly. It also refactors the file downloader to stream content in chunks to reduce memory usage, and increases resource limits in the manifest. The review feedback highlights a critical bug in the sharding logic within parse_air_quality.py that can lead to an infinite loop on the final shard and cause missing files for empty shards, providing a robust code suggestion to resolve these issues.

- Partition parse_air_quality.py for CDC_PM25County, CDC_OzoneCounty, and Census Tract datasets into discrete 5-year shards to prevent OOM/disk exhaustion. - Scale resource limits in manifest.json (up to 32 CPUs, 512GB RAM, 2TB disk) to support high-density daily time-series generation (1.064B observations). - Add streaming download in download_files.py with 16MB chunks to prevent memory spikes. - Update TMCF template files and test suite with sharded fixtures. - Add validation_config.json with deleted records and lint checks.

…ent infinite loop on last shard

…lth_toxicology

abhishekjaisw

Review scope

Target: PR (#2193)
Head Commit: 1363a5f2fd3dfa95f764b43004c6deb7026598f4
Reviewed: All 15 changed files under scripts/us_cdc/environmental_health_toxicology/** (415 additions, 137 deletions)
Skipped: None (all changed files are in scope)

Findings

No actionable P0, P1, or P2 blocking findings.

Positive findings

download_files.py:40 - Streamed chunked downloads ✓
- Finding: Good - Replaced response.content in-memory buffering with requests.get(url, stream=True) and 16MB chunk iterations, preventing high-memory spikes during multi-gigabyte archive retrieval while preserving retry decorators.
manifest.json:18 - Scaled compute & storage limits ✓
- Finding: Good - Upgraded resource allocations to 32 vCPUs, 512 GiB RAM, and 2TB disk for Census Tract imports (500GB disk for County imports), preventing out-of-memory and disk exhaustion failures on high-volume datasets.
parse_air_quality.py:129 - Balanced dataset sharding & defensive header creation ✓
- Finding: Good - Calculates exact row-balanced partition sizes upfront, initializes output CSV headers across all partition shards on the first chunk to guarantee valid files even for edge cases, and uncaps the terminal shard to prevent row loss.
parse_air_quality_test.py:108 - Multi-shard test verification & automated teardown ✓
- Finding: Good - Added parameterized assertions for all generated partition shards (PM25county_[0-3].csv) with concise test fixtures (< 1 KB each) and automated shutil.rmtree teardown to keep working trees clean.
validation_config.json:1 - Declarative validation configuration ✓
- Finding: Good - Configured standard validation rules for DELETED_RECORDS_PERCENT (threshold 0.05), EMPTY_IMPORT_CHECK, MISSING_REFS_COUNT, and LINT_ERROR_COUNT.

Coverage

File	Status	Result
`scripts/us_cdc/environmental_health_toxicology/download_files.py`	Reviewed	No findings
`scripts/us_cdc/environmental_health_toxicology/manifest.json`	Reviewed	No findings
`scripts/us_cdc/environmental_health_toxicology/parse_air_quality.py`	Reviewed	No findings
`scripts/us_cdc/environmental_health_toxicology/parse_air_quality_test.py`	Reviewed	No findings
`scripts/us_cdc/environmental_health_toxicology/PM25CountyPollution.tmcf`	Reviewed	No findings
`scripts/us_cdc/environmental_health_toxicology/PM25CountyPollution_part1.tmcf`	Reviewed	No findings
`scripts/us_cdc/environmental_health_toxicology/PM25CountyPollution_part2.tmcf`	Reviewed	No findings
`scripts/us_cdc/environmental_health_toxicology/PM25CountyPollution_part3.tmcf`	Reviewed	No findings
`scripts/us_cdc/environmental_health_toxicology/PM25CountyPollution_part4.tmcf`	Reviewed	No findings
`scripts/us_cdc/environmental_health_toxicology/README.md`	Reviewed	No findings
`scripts/us_cdc/environmental_health_toxicology/test_data/CDC_PM25County/expected_output_files/PM25county_0.csv`	Reviewed	No findings
`scripts/us_cdc/environmental_health_toxicology/test_data/CDC_PM25County/expected_output_files/PM25county_1.csv`	Reviewed	No findings
`scripts/us_cdc/environmental_health_toxicology/test_data/CDC_PM25County/expected_output_files/PM25county_2.csv`	Reviewed	No findings
`scripts/us_cdc/environmental_health_toxicology/test_data/CDC_PM25County/expected_output_files/PM25county_3.csv`	Reviewed	No findings
`scripts/us_cdc/environmental_health_toxicology/validation_config.json`	Reviewed	No findings

Verification and limitations

Checks run: python3 -m unittest discover -v -s scripts/us_cdc/environmental_health_toxicology -t . -p "*_test.py" passed (2 tests in 0.103s; test_clean_air_quality_data and test_clean_precipitation_data passed).
Checks not run: Live multi-gigabyte source data download (prohibited in standard code reviews to preserve environment safety).
Limitations: None.

gemini-code-assist Bot reviewed Aug 31, 2026

View reviewed changes

Comment thread scripts/us_cdc/environmental_health_toxicology/parse_air_quality.py

abhishekjaisw force-pushed the cdc_air_quality_sharding_and_scaling branch from 2ffe442 to b37ef9a Compare August 31, 2026 10:06

abhishekjaisw added 5 commits August 31, 2026 10:09

Fix county sharding logic: initialize headers for all shards and prev…

4ebc67e

…ent infinite loop on last shard

Apply yapf Google code formatting to scripts/us_cdc/environmental_hea…

5a7f855

…lth_toxicology

Add automatic test artifact cleanup in parse_air_quality_test tearDown

30e7938

Remove redundant config_override block from manifest.json

4e0cd49

Update README.md with sharded tMCF files and architecture documentation

1363a5f

abhishekjaisw force-pushed the cdc_air_quality_sharding_and_scaling branch from e4b3c20 to 1363a5f Compare August 31, 2026 11:44

abhishekjaisw commented Sep 1, 2026

View reviewed changes

Navigation Menu

Uh oh!

Optimize CDC air quality imports with sharding and scaled compute - #2193

Optimize CDC air quality imports with sharding and scaled compute#2193
abhishekjaisw wants to merge 6 commits into
datacommonsorg:masterfrom
abhishekjaisw:cdc_air_quality_sharding_and_scaling

abhishekjaisw commented Aug 31, 2026 •

edited

Loading

Uh oh!

gemini-code-assist Bot left a comment

Uh oh!

Uh oh!

abhishekjaisw left a comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

Uh oh!

Conversation

abhishekjaisw commented Aug 31, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

1. Summary of Changes

2. Background: Legacy Un-Sharded Baseline & Differ Verification Flow

Why Automated Differ Initially Skipped

Step 1: Executed Manual Dataflow Differ Jobs

Step 2: Updated latest_version.txt Pointers

Step 3: Enabled Automated Differ for Future Production Refreshes

3. Dev GCS Baseline File Locations

4. Manual Differ Run Artifacts (datcom-import-test/manual_runs)

5. Cloud Batch Verification Jobs Status

6. Detailed Differ Status & Verification Metrics

A. County Level Imports

B. Census Tract Level Imports

7. Code Review Agent (CRA) Report

Review Scope

Findings

Positive Findings

Coverage

Verification

8. References, GPastes & Buganizer Tickets

Uh oh!

gemini-code-assist Bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

Uh oh!

abhishekjaisw left a comment

Choose a reason for hiding this comment

Review scope

Findings

Positive findings

Coverage

Verification and limitations

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

abhishekjaisw commented Aug 31, 2026 •

edited

Loading

Step 2: Updated `latest_version.txt` Pointers

4. Manual Differ Run Artifacts (`datcom-import-test/manual_runs`)