import argparse import duckdb import geohash import geopandas as gpd import pandas as pd import json import cbor2 from pathlib import Path from shapely import wkb from typing import List, Dict def get_geohashes_from_bbox(min_x, min_y, max_x, max_y) -> List[str]: """Generates all 4-character Geohashes that intersect a bounding box.""" geohashes = set() step_size = 0.25 # 4-character Geohash grid size lat = min_y while lat <= max_y: lon = min_x while lon <= max_x: gh = geohash.encode(lat, lon, precision=4) geohashes.add(gh) lon += step_size lat += step_size return sorted(geohashes) def append_to_parquet(file_path: Path, new_data: gpd.GeoDataFrame): """Appends new data to an existing Parquet file while ensuring correct merging and removing duplicates.""" if file_path.exists(): try: existing_data = gpd.read_parquet(file_path) combined_data = pd.concat([existing_data, new_data], ignore_index=True) combined_data = combined_data.drop_duplicates(subset=["geometry"], keep="last").reset_index(drop=True) combined_data.to_parquet(file_path, index=False) print(f"✅ Appended new data to {file_path}, duplicates removed") except Exception as e: print(f"❌ Error while merging {file_path}: {e}") else: new_data.to_parquet(file_path, index=False) print(f"✅ Created new file: {file_path}") def save_as_geojson(geohash_code: str, gdf: gpd.GeoDataFrame, output_path: Path): """Saves a GeoDataFrame as a GeoJSON file.""" geojson_file = output_path / f"{geohash_code}.geojson" gdf.to_file(geojson_file, driver="GeoJSON") print(f"✅ Saved GeoJSON: {geojson_file}") def save_as_cbor(geohash_code: str, gdf: gpd.GeoDataFrame, output_path: Path): """Saves a GeoDataFrame as a CBOR file (compact binary format).""" cbor_file = output_path / f"{geohash_code}.cbor" geojson_dict = json.loads(gdf.to_json()) with open(cbor_file, "wb") as f: cbor2.dump(geojson_dict, f) print(f"✅ Saved CBOR: {cbor_file}") def slice_and_split_geoparquet(input_file: str, coastline_file: str, output_dir: str, export_parquet: bool, export_geojson: bool, export_cbor: bool): """Splits the GeoParquet file into multiple files based on 4-character Geohash tiles, keeping only coast-adjacent areas.""" input_path, coastline_path, output_path = Path(input_file), Path(coastline_file), Path(output_dir) if not input_path.exists() or not coastline_path.exists(): raise FileNotFoundError(f"File '{input_file}' or '{coastline_file}' not found.") output_path.mkdir(parents=True, exist_ok=True) print(f"📂 Processing '{input_file}', output will be saved in '{output_dir}'...") con = duckdb.connect() con.execute("INSTALL spatial; LOAD spatial;") # Load coastline data and create a 30km buffer con.execute(f""" CREATE TEMP TABLE coastline AS SELECT ST_Buffer(geometry, 30000) AS buffer_geom FROM read_parquet('{coastline_path}'); """) con.execute(f"CREATE TEMP TABLE geoparquet AS SELECT * FROM read_parquet('{input_path}') WHERE ST_GeometryType(geometry) IS NOT NULL AND ST_GeometryType(geometry) != 'POINT'") bbox_results = con.execute(""" SELECT rowid, ST_XMin(ST_Envelope(geometry)), ST_YMin(ST_Envelope(geometry)), ST_XMax(ST_Envelope(geometry)), ST_YMax(ST_Envelope(geometry)) FROM geoparquet; """).fetchall() geohash_mapping: Dict[str, List[int]] = {} for rowid, min_x, min_y, max_x, max_y in bbox_results: intersecting_geohashes = get_geohashes_from_bbox(min_x, min_y, max_x, max_y) # Check if geohash is near coastline buffer valid_geohashes = [] for gh in intersecting_geohashes: gh_bbox = geohash.bbox(gh) intersects = con.execute(f""" SELECT COUNT(*) FROM coastline WHERE ST_Intersects(buffer_geom, ST_MakeEnvelope({gh_bbox['w']}, {gh_bbox['s']}, {gh_bbox['e']}, {gh_bbox['n']})); """).fetchone()[0] > 0 if intersects: valid_geohashes.append(gh) for geohash_code in valid_geohashes: if geohash_code not in geohash_mapping: geohash_mapping[geohash_code] = [] geohash_mapping[geohash_code].append(rowid) for geohash_code, rowids in geohash_mapping.items(): geohash_bbox = geohash.bbox(geohash_code) filtered_data = con.execute(f""" SELECT ST_AsWKB(ST_Intersection(geometry, ST_MakeEnvelope({geohash_bbox['w']}, {geohash_bbox['s']}, {geohash_bbox['e']}, {geohash_bbox['n']}))) AS clipped_geom FROM geoparquet WHERE rowid IN ({','.join(map(str, rowids))}) AND ST_Intersects(geometry, ST_MakeEnvelope({geohash_bbox['w']}, {geohash_bbox['s']}, {geohash_bbox['e']}, {geohash_bbox['n']})); """).fetchdf() if not filtered_data.empty: filtered_data = filtered_data[filtered_data["clipped_geom"].apply(lambda x: isinstance(x, (bytes, bytearray)))] if not filtered_data.empty: filtered_data["geometry"] = filtered_data["clipped_geom"].apply(lambda x: wkb.loads(bytes(x)) if isinstance(x, (bytes, bytearray)) else None) filtered_data.drop(columns=["clipped_geom"], inplace=True) filtered_data = filtered_data.dropna(subset=["geometry"]) if not filtered_data.empty: filtered_data["natural"] = "water" filtered_data["water"] = "sea" gdf = gpd.GeoDataFrame(filtered_data, geometry="geometry", crs="EPSG:4326") if export_parquet: append_to_parquet(output_path / f"{geohash_code}.parquet", gdf) if export_geojson: save_as_geojson(geohash_code, gdf, output_path) if export_cbor: save_as_cbor(geohash_code, gdf, output_path) con.close() print("🎉 Processing complete!") def main(): parser = argparse.ArgumentParser(description="Split a GeoParquet file into coastal Geohash tiles with export format options.") parser.add_argument("-i", "--input", required=True, help="Path to the input GeoParquet file.") parser.add_argument("-c", "--coastline", required=True, help="Path to the coastline Parquet file.") parser.add_argument("-o", "--output", required=True, help="Directory to save output files.") parser.add_argument("--parquet", action="store_true", help="Export as GeoParquet") parser.add_argument("--geojson", action="store_true", help="Export as GeoJSON") parser.add_argument("--cbor", action="store_true", help="Export as CBOR") args = parser.parse_args() if not (args.parquet or args.geojson or args.cbor): args.parquet = True # Default to GeoParquet if no options are given try: slice_and_split_geoparquet(args.input, args.coastline, args.output, args.parquet, args.geojson, args.cbor) except Exception as e: print(f"❌ Error: {e}") if __name__ == "__main__": main()