randogoth/dopecarpet
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

faster slice

Browse source
This commit is contained in:
randogoth 2025-03-04 15:15:11 +00:00
parent cdabeb730b
commit 2edd6ca38c
1 changed files with 26 additions and 27 deletions
Download patch file Download diff file Expand all files Collapse all files

53
slice.py
View file

@ -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()