import argparse 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 import json # Configure logging import logging logging.basicConfig(format="%(levelname)s: %(message)s", level=logging.INFO) logger = logging.getLogger(__name__) def load_tags_config(tags_file: Path) -> Dict: """Loads the tags configuration from a JSON file.""" with open(tags_file, "r") as f: return json.load(f) def filter_tags(tags: Dict, tags_config: Dict) -> Dict: """Filters tags based on the rules defined in the tags configuration.""" filtered_tags = {} for key, value in tags.items(): if key in tags_config: if tags_config[key] is True: # Keep all values for this key filtered_tags[key] = value elif isinstance(tags_config[key], list) and value in tags_config[key]: # Keep only specific values for this key filtered_tags[key] = value return filtered_tags 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], tags_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, optionally adds water polygons, and filters tags.""" 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 tags configuration if provided tags_config = {} if tags_file: tags_path = Path(tags_file) if not tags_path.exists(): raise FileNotFoundError(f"File '{tags_file}' not found.") tags_config = load_tags_config(tags_path) # 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: # Filter tags if tags configuration is provided if tags_config: filtered_data["tags"] = filtered_data["tags"].apply( lambda tags: filter_tags(tags, tags_config) if isinstance(tags, dict) else {} ) 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, filter tags, 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("-t", "--tags", required=False, help="Path to the tags.json 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.tags, args.parquet, args.geojson, args.cbor) except Exception as e: logger.error(f"❌ Error: {e}") if __name__ == "__main__": main()