mirror of
https://github.com/PlaneQuery/OpenAirframes.git
synced 2026-08-09 09:56:01 +02:00
FEATURE: Add contributions framework. Fix and improve daily adsb release using Github actions for map reduce.
This commit is contained in:
ggman12
@@ -0,0 +1,205 @@
|
||||
"""
|
||||
Combines chunk parquet files and compresses to final aircraft CSV.
|
||||
This is the reduce phase of the map-reduce pipeline.
|
||||
|
||||
Memory-efficient: processes each chunk separately, compresses, then combines.
|
||||
|
||||
Usage:
|
||||
python -m src.adsb.combine_chunks_to_csv --chunks-dir data/output/adsb_chunks
|
||||
"""
|
||||
import gc
|
||||
import os
|
||||
import sys
|
||||
import glob
|
||||
import argparse
|
||||
from datetime import datetime, timedelta
|
||||
|
||||
import polars as pl
|
||||
|
||||
from src.adsb.download_adsb_data_to_parquet import OUTPUT_DIR, get_resource_usage
|
||||
from src.adsb.compress_adsb_to_aircraft_data import compress_multi_icao_df, COLUMNS
|
||||
|
||||
|
||||
DEFAULT_CHUNK_DIR = os.path.join(OUTPUT_DIR, "adsb_chunks")
|
||||
FINAL_OUTPUT_DIR = "./data/planequery_aircraft"
|
||||
os.makedirs(FINAL_OUTPUT_DIR, exist_ok=True)
|
||||
|
||||
|
||||
def get_target_day() -> datetime:
|
||||
"""Get yesterday's date (the day we're processing)."""
|
||||
return datetime.utcnow() - timedelta(days=1)
|
||||
|
||||
|
||||
def process_single_chunk(chunk_path: str) -> pl.DataFrame:
|
||||
"""Load and compress a single chunk parquet file."""
|
||||
print(f"Processing {os.path.basename(chunk_path)}... | {get_resource_usage()}")
|
||||
|
||||
# Load chunk - only columns we need
|
||||
needed_columns = ['time', 'icao'] + COLUMNS
|
||||
df = pl.read_parquet(chunk_path, columns=needed_columns)
|
||||
print(f" Loaded {len(df)} rows")
|
||||
|
||||
# Compress to aircraft records (one per ICAO) using shared function
|
||||
compressed = compress_multi_icao_df(df, verbose=True)
|
||||
print(f" Compressed to {len(compressed)} aircraft records")
|
||||
|
||||
del df
|
||||
gc.collect()
|
||||
|
||||
return compressed
|
||||
|
||||
|
||||
def combine_compressed_chunks(compressed_dfs: list[pl.DataFrame]) -> pl.DataFrame:
|
||||
"""Combine multiple compressed DataFrames.
|
||||
|
||||
Since chunks are partitioned by ICAO hash, each ICAO only appears in one chunk.
|
||||
No deduplication needed here - just concatenate.
|
||||
"""
|
||||
print(f"Combining {len(compressed_dfs)} compressed chunks... | {get_resource_usage()}")
|
||||
|
||||
# Concat all
|
||||
combined = pl.concat(compressed_dfs)
|
||||
print(f"Combined: {len(combined)} records")
|
||||
|
||||
return combined
|
||||
|
||||
|
||||
def download_and_merge_base_release(compressed_df: pl.DataFrame) -> pl.DataFrame:
|
||||
"""Download base release and merge with new data."""
|
||||
from src.get_latest_planequery_aircraft_release import download_latest_aircraft_adsb_csv
|
||||
|
||||
print("Downloading base ADS-B release...")
|
||||
try:
|
||||
base_path = download_latest_aircraft_adsb_csv(
|
||||
output_dir="./data/planequery_aircraft_base"
|
||||
)
|
||||
print(f"Download returned: {base_path}")
|
||||
|
||||
if base_path and os.path.exists(str(base_path)):
|
||||
print(f"Loading base release from {base_path}")
|
||||
base_df = pl.read_csv(base_path)
|
||||
print(f"Base release has {len(base_df)} records")
|
||||
|
||||
# Ensure columns match
|
||||
base_cols = set(base_df.columns)
|
||||
new_cols = set(compressed_df.columns)
|
||||
print(f"Base columns: {sorted(base_cols)}")
|
||||
print(f"New columns: {sorted(new_cols)}")
|
||||
|
||||
# Add missing columns
|
||||
for col in new_cols - base_cols:
|
||||
base_df = base_df.with_columns(pl.lit(None).alias(col))
|
||||
for col in base_cols - new_cols:
|
||||
compressed_df = compressed_df.with_columns(pl.lit(None).alias(col))
|
||||
|
||||
# Reorder columns to match
|
||||
compressed_df = compressed_df.select(base_df.columns)
|
||||
|
||||
# Concat and deduplicate by icao (keep new data - it comes last)
|
||||
combined = pl.concat([base_df, compressed_df])
|
||||
print(f"After concat: {len(combined)} records")
|
||||
|
||||
deduplicated = combined.unique(subset=["icao"], keep="last")
|
||||
|
||||
print(f"Combined with base: {len(combined)} -> {len(deduplicated)} after dedup")
|
||||
|
||||
del base_df, combined
|
||||
gc.collect()
|
||||
|
||||
return deduplicated
|
||||
else:
|
||||
print(f"No base release found at {base_path}, using only new data")
|
||||
return compressed_df
|
||||
except Exception as e:
|
||||
import traceback
|
||||
print(f"Failed to download base release: {e}")
|
||||
traceback.print_exc()
|
||||
return compressed_df
|
||||
|
||||
|
||||
def cleanup_chunks(date_str: str, chunks_dir: str):
|
||||
"""Delete chunk parquet files after successful merge."""
|
||||
pattern = os.path.join(chunks_dir, f"chunk_*_{date_str}.parquet")
|
||||
chunk_files = glob.glob(pattern)
|
||||
for f in chunk_files:
|
||||
try:
|
||||
os.remove(f)
|
||||
print(f"Deleted {f}")
|
||||
except Exception as e:
|
||||
print(f"Failed to delete {f}: {e}")
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description="Combine chunk parquets to final CSV")
|
||||
parser.add_argument("--date", type=str, help="Date in YYYY-MM-DD format (default: yesterday)")
|
||||
parser.add_argument("--chunks-dir", type=str, default=DEFAULT_CHUNK_DIR, help="Directory containing chunk parquet files")
|
||||
parser.add_argument("--skip-base", action="store_true", help="Skip downloading and merging base release")
|
||||
parser.add_argument("--keep-chunks", action="store_true", help="Keep chunk files after merging")
|
||||
args = parser.parse_args()
|
||||
|
||||
if args.date:
|
||||
target_day = datetime.strptime(args.date, "%Y-%m-%d")
|
||||
else:
|
||||
target_day = get_target_day()
|
||||
|
||||
date_str = target_day.strftime("%Y-%m-%d")
|
||||
chunks_dir = args.chunks_dir
|
||||
|
||||
print(f"Combining chunks for {date_str}")
|
||||
print(f"Chunks directory: {chunks_dir}")
|
||||
print(f"Resource usage at start: {get_resource_usage()}")
|
||||
|
||||
# Find chunk files
|
||||
pattern = os.path.join(chunks_dir, f"chunk_*_{date_str}.parquet")
|
||||
chunk_files = sorted(glob.glob(pattern))
|
||||
|
||||
if not chunk_files:
|
||||
print(f"No chunk files found matching: {pattern}")
|
||||
sys.exit(1)
|
||||
|
||||
print(f"Found {len(chunk_files)} chunk files")
|
||||
|
||||
# Process each chunk separately to save memory
|
||||
compressed_chunks = []
|
||||
for chunk_path in chunk_files:
|
||||
compressed = process_single_chunk(chunk_path)
|
||||
compressed_chunks.append(compressed)
|
||||
gc.collect()
|
||||
|
||||
# Combine all compressed chunks
|
||||
combined = combine_compressed_chunks(compressed_chunks)
|
||||
|
||||
# Free memory from individual chunks
|
||||
del compressed_chunks
|
||||
gc.collect()
|
||||
print(f"After combining: {get_resource_usage()}")
|
||||
|
||||
# Merge with base release
|
||||
if not args.skip_base:
|
||||
combined = download_and_merge_base_release(combined)
|
||||
|
||||
# Convert list columns to strings for CSV compatibility
|
||||
for col in combined.columns:
|
||||
if combined[col].dtype == pl.List:
|
||||
combined = combined.with_columns(
|
||||
pl.col(col).list.join(",").alias(col)
|
||||
)
|
||||
|
||||
# Sort by time for consistent output
|
||||
if 'time' in combined.columns:
|
||||
combined = combined.sort('time')
|
||||
|
||||
# Write final CSV
|
||||
output_path = os.path.join(FINAL_OUTPUT_DIR, f"planequery_aircraft_adsb_{date_str}.csv")
|
||||
combined.write_csv(output_path)
|
||||
print(f"Wrote {len(combined)} records to {output_path}")
|
||||
|
||||
# Cleanup
|
||||
if not args.keep_chunks:
|
||||
cleanup_chunks(date_str, chunks_dir)
|
||||
|
||||
print(f"Done! | {get_resource_usage()}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user