From 71389d1f624b5cd8f0d1abb4aa4ee2f250b06f14 Mon Sep 17 00:00:00 2001
From: didier amyot <didier.amyot@all-dressed-programming.com>
Date: Thu, 2 Jan 2025 15:19:55 -0500
Subject: [PATCH] Added duckdb generator

---
 flake.nix                  |   3 +
 scripts/generate_duckdb.py | 116 +++++++++++++++++++++++++++++++++++++
 2 files changed, 119 insertions(+)
 create mode 100644 scripts/generate_duckdb.py

diff --git a/flake.nix b/flake.nix
index ad37df5..b4c6e1f 100644
--- a/flake.nix
+++ b/flake.nix
@@ -64,10 +64,13 @@
         python-virtual-env = pkgs.python312.withPackages (ps: [
           ps.matplotlib
           ps.jupyter
+          ps.tqdm
           ps.duckdb
           ps.pandas
           ps.unidecode
           ps.geopy
+          ps.click
+          ps.jinja2
         ]);
       in {
         packages = {
diff --git a/scripts/generate_duckdb.py b/scripts/generate_duckdb.py
new file mode 100644
index 0000000..716e327
--- /dev/null
+++ b/scripts/generate_duckdb.py
@@ -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()