import argparse import json import duckdb import geohash import geopandas as gpd import pandas as pd from pathlib import Path from shapely import wkb from shapely.geometry import box from tqdm import tqdm # For progress bar from typing import List, Dict, Optional # Configure logging import logging logging.basicConfig(format="%(levelname)s: %(message)s", level=logging.INFO) logger = logging.getLogger(__name__) 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, add_water: bool): """Appends new data to an existing Parquet file while ensuring correct merging and avoiding duplicates.""" if file_path.exists(): try: existing_data = gpd.read_parquet(file_path) # Remove sea polygons to avoid duplicates (only if we're adding water polygons) if add_water: existing_data = existing_data[~existing_data["tags"].apply(lambda tags: isinstance(tags, dict) and tags.get("natural") == "water")] # Ensure no duplicates based on feature_id (for land geometries) if "feature_id" in existing_data.columns and "feature_id" in new_data.columns: new_unique_data = new_data[~new_data["feature_id"].isin(existing_data["feature_id"])] combined_data = pd.concat([existing_data, new_unique_data], ignore_index=True) else: combined_data = pd.concat([existing_data, new_data], ignore_index=True) combined_data.to_parquet(file_path, index=False) except Exception as e: logger.warning(f"Error while merging {file_path}: {e}") else: new_data.to_parquet(file_path, index=False) 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") logger.info(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).""" import cbor2 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) logger.info(f"✅ Saved CBOR: {cbor_file}") def geohash_bbox(geohash_code): """Returns a bounding box polygon for a geohash.""" bbox = geohash.bbox(geohash_code) return box(bbox["w"], bbox["s"], bbox["e"], bbox["n"]) def ensure_crs_consistency(gdf): """Ensures the GeoDataFrame is in EPSG:4326 to prevent CRS mismatches.""" if gdf.crs is None or gdf.crs.to_string() != "EPSG:4326": gdf = gdf.to_crs("EPSG:4326") return gdf def add_water_to_geohash(geohash_code: str, geohash_file: Path, water_gdf: gpd.GeoDataFrame): """Adds water polygons to a geohash tile.""" try: existing_gdf = gpd.read_parquet(geohash_file) except Exception as e: logger.warning(f"Failed to read {geohash_file}. Skipping. Error: {e}") return False existing_gdf = ensure_crs_consistency(existing_gdf) # Get the geohash bounding box bbox = geohash_bbox(geohash_code) # Extract water polygons that overlap with this geohash water_in_tile = water_gdf[water_gdf.intersects(bbox)].copy() if water_in_tile.empty: return False # No water polygons added # Clip water polygons to the geohash tile boundary water_in_tile["geometry"] = water_in_tile.intersection(bbox) # Assign "water" metadata and a unique identifier for water polygons water_in_tile["feature_id"] = [f"water_{geohash_code}_{i}" for i in range(len(water_in_tile))] water_in_tile["tags"] = [{"natural": "water", "water": "sea"}] * len(water_in_tile) # Ensure columns match before merging for col in ["feature_id", "tags"]: if col not in existing_gdf.columns: existing_gdf[col] = None if col not in water_in_tile.columns: water_in_tile[col] = None # Merge updated water polygons into geohash file updated_gdf = gpd.GeoDataFrame(pd.concat([existing_gdf, water_in_tile], ignore_index=True), crs="EPSG:4326") # Save back to parquet updated_gdf.to_parquet(geohash_file, index=False) return True # Water polygons added def slice_and_split_geoparquet(input_file: str, output_dir: str, water_file: Optional[str], export_parquet: bool, export_geojson: bool, export_cbor: bool): """Splits the GeoParquet file into multiple files based on 4-character Geohash tiles and optionally adds water polygons.""" input_path, output_path = Path(input_file), Path(output_dir) if not input_path.exists(): raise FileNotFoundError(f"File '{input_file}' not found.") output_path.mkdir(parents=True, exist_ok=True) # Load water dataset if provided water_gdf = None if water_file: water_path = Path(water_file) if not water_path.exists(): raise FileNotFoundError(f"File '{water_file}' not found.") water_gdf = ensure_crs_consistency(gpd.read_parquet(water_path)) # Connect to DuckDB and load spatial extension con = duckdb.connect() con.execute("INSTALL spatial; LOAD spatial;") # Load the GeoParquet file into DuckDB 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'") # Precompute bounding boxes for all features bbox_results = con.execute(""" SELECT feature_id, rowid, ST_XMin(ST_Envelope(geometry)), ST_YMin(ST_Envelope(geometry)), ST_XMax(ST_Envelope(geometry)), ST_YMax(ST_Envelope(geometry)) FROM geoparquet; """).fetchall() # Map geohash codes to rowids geohash_mapping: Dict[str, List[int]] = {} for feature_id, 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) for geohash_code in intersecting_geohashes: if geohash_code not in geohash_mapping: geohash_mapping[geohash_code] = [] geohash_mapping[geohash_code].append(rowid) # Process each geohash tile total_tiles = len(geohash_mapping) tiles_with_water = 0 with tqdm(total=total_tiles, desc="Processing geohash tiles") as pbar: for geohash_code, rowids in geohash_mapping.items(): geohash_bbox = geohash.bbox(geohash_code) # Fetch and clip geometries for this geohash tile filtered_data = con.execute(f""" SELECT feature_id, tags, 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: # Convert WKB geometries to Shapely geometries 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) # Drop rows with invalid geometries filtered_data = filtered_data.dropna(subset=["geometry"]) if not filtered_data.empty: gdf = gpd.GeoDataFrame(filtered_data, geometry="geometry", crs="EPSG:4326") # Export to GeoParquet if export_parquet: geohash_file = output_path / f"{geohash_code}.parquet" append_to_parquet(geohash_file, gdf, add_water=(water_gdf is not None)) if water_gdf is not None and add_water_to_geohash(geohash_code, geohash_file, water_gdf): tiles_with_water += 1 # Export to GeoJSON if export_geojson: save_as_geojson(geohash_code, gdf, output_path) # Export to CBOR if export_cbor: save_as_cbor(geohash_code, gdf, output_path) pbar.update(1) con.close() print(f"🎉 Processing complete! Generated {total_tiles} tiles, {tiles_with_water} of which had water polygons added.") def main(): parser = argparse.ArgumentParser(description="Split a GeoParquet file into Geohash tiles, optionally add water polygons, and export in desired formats.") parser.add_argument("-i", "--input", required=True, help="Path to the input GeoParquet file.") parser.add_argument("-o", "--output", required=True, help="Directory to save output files.") parser.add_argument("-w", "--water", required=False, help="Path to the water.parquet file.") 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.output, args.water, args.parquet, args.geojson, args.cbor) except Exception as e: logger.error(f"❌ Error: {e}") if __name__ == "__main__": main()