Added duckdb generator

This commit is contained in:
didier committed 2025-01-02 15:19:55 -05:00
1 parent f17e7c70fc
commit 71389d1f62
2 files changed
+119

No files matched your search

+116
View File
@@ -0,0 +1,116 @@
import duckdb
import jinja2
from pathlib import Path
import click
from tqdm import tqdm
from posix import mkdir
sql_tmpl = jinja2.Template('''
LOAD spatial;
CREATE TABLE IF NOT EXISTS city_road(
city_name TEXT,
relation_id BIGINT,
points DOUBLE[][],
tags MAP(VARCHAR, VARCHAR),
highway TEXT,
length DOUBLE
);
CREATE TEMPORARY VIEW osm_quebec AS (
SELECT * FROM ST_READOSM ('{{ file_name }}')
);
CREATE TEMPORARY VIEW interresting_ways AS (
SELECT * FROM osm_quebec
WHERE tags['highway'][1] IS NOT NULL
);
CREATE TEMPORARY VIEW relations AS (
SELECT id,
tags,
UNNEST(refs) AS refs,
UNNEST(RANGE(LEN(refs))) as index
FROM interresting_ways
);
CREATE TEMPORARY VIEW road_points AS (
SELECT osm_quebec.id AS id,
osm_quebec.lon AS lon,
osm_quebec.lat AS lat,
relations.tags AS tags,
relations.index AS relation_index,
relations.id AS relation_id
FROM relations
JOIN osm_quebec
ON osm_quebec.id = relations.refs
WHERE
osm_quebec.kind = 'node'
ORDER BY relation_id, relation_index
);
CREATE TEMPORARY VIEW road_points_with_distance AS (
SELECT *,
st_point(lat, lon) AS point,
lag(st_point(lat, lon)) OVER (PARTITION BY relation_id ORDER BY relation_index) AS lag_point
FROM road_points
);
CREATE TEMPORARY VIEW data AS (
SELECT
'{{ city }}' AS city_name,
relation_id AS relation_id,
LIST([road_points_with_distance.lat, road_points_with_distance.lon]) AS points,
ANY_VALUE(tags) AS tags,
ANY_VALUE(tags)['highway'][1] AS highway,
SUM(ST_Distance_Spheroid(point, lag_point)) AS length
FROM road_points_with_distance
GROUP BY relation_id
);
INSERT INTO city_road(city_name, relation_id, points, tags, highway, length) (
SELECT city_name, relation_id, points, tags, highway, length FROM data
);
''')
def extract_city(file_name: str) -> str:
return file_name.name[:file_name.name.index(".")].replace("_quebec", "")
def sql_safe(w: str) -> str:
return str(w).replace("'", r"\\'")
def append_data_to_duckdb(conn: duckdb.DuckDBPyConnection, osm_pbf_path: Path) -> None:
city = extract_city(osm_pbf_path)
sql = sql_tmpl.render(city=sql_safe(city), file_name=sql_safe(osm_pbf_path))
conn.execute(sql)
def append_data_to_duckdb_file(duckdb_file: Path, osm_pbf_path: Path) -> None:
with duckdb.connect(str(duckdb_file)) as conn:
append_data_to_duckdb(conn, osm_pbf_path)
@click.command()
@click.option('--output-file', type=click.Path(exists=False, file_okay=True, dir_okay=False, writable=True, path_type=Path))
@click.option('--input-dir', type=click.Path(exists=False, file_okay=False, dir_okay=True, readable=True, path_type=Path))
def create_duckdb_file(output_file: Path, input_dir: Path) -> None:
output_file.parent.mkdir(parents=True, exist_ok=True)
files = [f for f in input_dir.glob('**/*osm.pbf') if f.is_file()]
for osm_file in tqdm(files):
if "'" in str(osm_file):
continue
try:
append_data_to_duckdb_file(output_file, osm_file)
except Exception as e:
click.echo(f"Could not process {osm_file}")
if __name__ == '__main__':
create_duckdb_file()