diff --git a/slice.py b/slice.py index eb5db5d..28ed8bb 100644 --- a/slice.py +++ b/slice.py @@ -3,8 +3,6 @@ import duckdb import geohash import geopandas as gpd import pandas as pd -import pyarrow.parquet as pq -import pyarrow as pa import json import cbor2 from pathlib import Path @@ -73,19 +71,22 @@ def slice_and_split_geoparquet(input_file: str, output_dir: str, export_parquet: print(f"📂 Processing '{input_file}', output will be saved in '{output_dir}'...") + # 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: @@ -93,11 +94,11 @@ def slice_and_split_geoparquet(input_file: str, output_dir: str, export_parquet: geohash_mapping[geohash_code] = [] geohash_mapping[geohash_code].append(rowid) - geohash_data: Dict[str, gpd.GeoDataFrame] = {} - + # Process each geohash tile 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, @@ -112,30 +113,28 @@ def slice_and_split_geoparquet(input_file: str, output_dir: str, export_parquet: """).fetchdf() if not filtered_data.empty: - filtered_data = filtered_data[filtered_data["clipped_geom"].apply(lambda x: isinstance(x, (bytes, bytearray)))] + # 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: - 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) + gdf = gpd.GeoDataFrame(filtered_data, geometry="geometry", crs="EPSG:4326") - filtered_data = filtered_data.dropna(subset=["geometry"]) - - if not filtered_data.empty: - gdf = gpd.GeoDataFrame(filtered_data, geometry="geometry", crs="EPSG:4326") - geohash_data[geohash_code] = gdf - - for geohash_code, gdf in geohash_data.items(): - if export_parquet: - geohash_file = output_path / f"{geohash_code}.parquet" - append_to_parquet(geohash_file, gdf) - - if export_geojson: - save_as_geojson(geohash_code, gdf, output_path) - - if export_cbor: - save_as_cbor(geohash_code, gdf, output_path) + # Export to desired formats + if export_parquet: + geohash_file = output_path / f"{geohash_code}.parquet" + append_to_parquet(geohash_file, 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!") @@ -159,4 +158,4 @@ def main(): print(f"❌ Error: {e}") if __name__ == "__main__": - main() + main() \ No newline at end of file