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

19
slice.py
View file

@ -3,8 +3,6 @@ import duckdb
import geohash import geohash
import geopandas as gpd import geopandas as gpd
import pandas as pd import pandas as pd
import pyarrow.parquet as pq
import pyarrow as pa
import json import json
import cbor2 import cbor2
from pathlib import Path 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}'...") print(f"📂 Processing '{input_file}', output will be saved in '{output_dir}'...")
# Connect to DuckDB and load spatial extension
con = duckdb.connect() con = duckdb.connect()
con.execute("INSTALL spatial; LOAD spatial;") 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'") 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(""" bbox_results = con.execute("""
SELECT feature_id, rowid, ST_XMin(ST_Envelope(geometry)), ST_YMin(ST_Envelope(geometry)), SELECT feature_id, rowid, ST_XMin(ST_Envelope(geometry)), ST_YMin(ST_Envelope(geometry)),
ST_XMax(ST_Envelope(geometry)), ST_YMax(ST_Envelope(geometry)) ST_XMax(ST_Envelope(geometry)), ST_YMax(ST_Envelope(geometry))
FROM geoparquet; FROM geoparquet;
""").fetchall() """).fetchall()
# Map geohash codes to rowids
geohash_mapping: Dict[str, List[int]] = {} geohash_mapping: Dict[str, List[int]] = {}
for feature_id, rowid, min_x, min_y, max_x, max_y in bbox_results: 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) intersecting_geohashes = get_geohashes_from_bbox(min_x, min_y, max_x, max_y)
for geohash_code in intersecting_geohashes: 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] = []
geohash_mapping[geohash_code].append(rowid) geohash_mapping[geohash_code].append(rowid)
geohash_data: Dict[str, gpd.GeoDataFrame] = {} # Process each geohash tile
for geohash_code, rowids in geohash_mapping.items(): for geohash_code, rowids in geohash_mapping.items():
geohash_bbox = geohash.bbox(geohash_code) geohash_bbox = geohash.bbox(geohash_code)
# Fetch and clip geometries for this geohash tile
filtered_data = con.execute(f""" filtered_data = con.execute(f"""
SELECT feature_id, tags, SELECT feature_id, tags,
ST_AsWKB(ST_Intersection(geometry, ST_AsWKB(ST_Intersection(geometry,
@ -112,21 +113,19 @@ def slice_and_split_geoparquet(input_file: str, output_dir: str, export_parquet:
""").fetchdf() """).fetchdf()
if not filtered_data.empty: 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
if not filtered_data.empty:
filtered_data["geometry"] = filtered_data["clipped_geom"].apply( filtered_data["geometry"] = filtered_data["clipped_geom"].apply(
lambda x: wkb.loads(bytes(x)) if isinstance(x, (bytes, bytearray)) else None lambda x: wkb.loads(bytes(x)) if isinstance(x, (bytes, bytearray)) else None
) )
filtered_data.drop(columns=["clipped_geom"], inplace=True) filtered_data.drop(columns=["clipped_geom"], inplace=True)
# Drop rows with invalid geometries
filtered_data = filtered_data.dropna(subset=["geometry"]) filtered_data = filtered_data.dropna(subset=["geometry"])
if not filtered_data.empty: if not filtered_data.empty:
gdf = gpd.GeoDataFrame(filtered_data, geometry="geometry", crs="EPSG:4326") gdf = gpd.GeoDataFrame(filtered_data, geometry="geometry", crs="EPSG:4326")
geohash_data[geohash_code] = gdf
for geohash_code, gdf in geohash_data.items(): # Export to desired formats
if export_parquet: if export_parquet:
geohash_file = output_path / f"{geohash_code}.parquet" geohash_file = output_path / f"{geohash_code}.parquet"
append_to_parquet(geohash_file, gdf) append_to_parquet(geohash_file, gdf)