129 lines
3.7 KiB
Python
129 lines
3.7 KiB
Python
import duckdb
|
|
import jinja2
|
|
from pathlib import Path
|
|
import click
|
|
from tqdm import tqdm
|
|
from posix import mkdir
|
|
from datetime import datetime, date
|
|
|
|
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,
|
|
extract_date Date,
|
|
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
|
|
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, extract_date, relation_id, points, tags, highway, length) (
|
|
SELECT '{{ city }}' AS city_name,
|
|
'{{extract_date}}' AS extract_date,
|
|
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, extract_date: date) -> None:
|
|
city = extract_city(osm_pbf_path)
|
|
sql = sql_tmpl.render(city=sql_safe(city),
|
|
file_name=sql_safe(osm_pbf_path),
|
|
extract_date=str(extract_date))
|
|
conn.execute(sql)
|
|
|
|
def append_data_to_duckdb_file(duckdb_file: Path, osm_pbf_path: Path, extract_date: date) -> None:
|
|
with duckdb.connect(str(duckdb_file)) as conn:
|
|
append_data_to_duckdb(conn, osm_pbf_path, extract_date)
|
|
|
|
@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))
|
|
@click.option('--date', '--date-end', type=click.DateTime(formats=["%Y-%m-%d"]))
|
|
def create_duckdb_file(output_file: Path, input_dir: Path, date: datetime) -> 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, date.date())
|
|
except Exception as e:
|
|
click.echo(f"Could not process {osm_file}")
|
|
|
|
|
|
|
|
if __name__ == '__main__':
|
|
create_duckdb_file()
|