Added duckdb generator
This commit is contained in:
1 parent
f17e7c70fc
commit
71389d1f62
2 files changed
+119
No files matched your search
@@ -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()
|
||||
Reference in new issue
Block a user