pbf/navsea_build_kyushu_core_sqlite.py

from __future__ import annotations

import datetime as dt
import sqlite3
from pathlib import Path


SOURCE_DB = Path("out/navsea_kyushu_semantic.sqlite")
OUT = Path("out/navsea_kyushu_core.sqlite")
KEEP_GROUNDING = ("干潮帯", "0-2m", "0-5m", "0-10m", "2-5m", "2-10m", "5-10m")


def open_db() -> sqlite3.Connection:
    if OUT.exists():
        OUT.unlink()
    OUT.parent.mkdir(parents=True, exist_ok=True)
    conn = sqlite3.connect(OUT)
    conn.execute("PRAGMA journal_mode=OFF")
    conn.execute("PRAGMA synchronous=OFF")
    conn.execute("PRAGMA temp_store=MEMORY")
    conn.execute("PRAGMA cache_size=150000")
    conn.execute("PRAGMA locking_mode=EXCLUSIVE")
    return conn


def create_schema(conn: sqlite3.Connection) -> None:
    conn.executescript(
        """
        CREATE TABLE meta (
          key TEXT PRIMARY KEY,
          value TEXT NOT NULL
        );

        CREATE TABLE semantic_object (
          feature_id INTEGER PRIMARY KEY,
          fid TEXT,
          source_layer TEXT,
          canonical_object_type TEXT NOT NULL,
          object_family TEXT NOT NULL,
          detection_class TEXT NOT NULL,
          capability TEXT,
          z INTEGER,
          x INTEGER,
          y INTEGER,
          geom_type TEXT,
          is_land INTEGER NOT NULL DEFAULT 0,
          is_hole INTEGER NOT NULL DEFAULT 0,
          is_collision INTEGER NOT NULL DEFAULT 0,
          is_grounding INTEGER NOT NULL DEFAULT 0
        );

        CREATE INDEX idx_semantic_object_fid ON semantic_object(fid);
        CREATE INDEX idx_semantic_object_tile ON semantic_object(z, x, y);
        CREATE INDEX idx_semantic_object_capability ON semantic_object(capability, canonical_object_type);

        CREATE TABLE grid_cell (
          tile_key TEXT PRIMARY KEY,
          z INTEGER NOT NULL,
          x INTEGER NOT NULL,
          y INTEGER NOT NULL,
          land_count INTEGER NOT NULL DEFAULT 0,
          hole_count INTEGER NOT NULL DEFAULT 0,
          collision_count INTEGER NOT NULL DEFAULT 0,
          grounding_count INTEGER NOT NULL DEFAULT 0,
          blocked INTEGER NOT NULL DEFAULT 0,
          blocked_reason TEXT
        );

        CREATE INDEX idx_grid_cell_zxy ON grid_cell(z, x, y);
        CREATE INDEX idx_grid_cell_blocked ON grid_cell(blocked, z);
        """
    )
    conn.executemany(
        "INSERT INTO meta(key, value) VALUES (?, ?)",
        [
            ("source_db", str(SOURCE_DB)),
            ("source_scope", "九州"),
            ("build_time", dt.datetime.now().isoformat(timespec="seconds")),
            ("notes", "九州核心版语义索引库，只保留陆地、碰撞物、浅水搁浅层"),
            ("keep_grounding", ",".join(KEEP_GROUNDING)),
        ],
    )
    conn.commit()


def main() -> None:
    if not SOURCE_DB.exists():
        raise SystemExit(f"missing source db: {SOURCE_DB}")

    src = sqlite3.connect(SOURCE_DB)
    src.row_factory = sqlite3.Row
    conn = open_db()
    create_schema(conn)

    placeholders = ",".join("?" for _ in KEEP_GROUNDING)
    sql = f"""
        SELECT
          feature_id, fid, source_layer, canonical_object_type, object_family, detection_class,
          capability, z, x, y, geom_type, is_land, is_hole, is_collision, is_grounding
        FROM semantic_object
        WHERE is_land = 1
           OR is_hole = 1
           OR is_collision = 1
           OR (is_grounding = 1 AND canonical_object_type IN ({placeholders}))
        ORDER BY z, x, y, feature_id
    """

    batch = []
    grid = {}
    count = 0
    batch_size = 4000

    for row in src.execute(sql, KEEP_GROUNDING):
        row = dict(row)
        batch.append(
            (
                row["feature_id"],
                row["fid"],
                row["source_layer"],
                row["canonical_object_type"],
                row["object_family"],
                row["detection_class"],
                row["capability"],
                row["z"],
                row["x"],
                row["y"],
                row["geom_type"],
                row["is_land"],
                row["is_hole"],
                row["is_collision"],
                row["is_grounding"],
            )
        )
        tile_key = f'{row["z"]}/{row["x"]}/{row["y"]}'
        c = grid.get(tile_key)
        if c is None:
            c = {"z": row["z"], "x": row["x"], "y": row["y"], "land": 0, "hole": 0, "collision": 0, "grounding": 0}
            grid[tile_key] = c
        c["land"] += int(row["is_land"])
        c["hole"] += int(row["is_hole"])
        c["collision"] += int(row["is_collision"])
        c["grounding"] += int(row["is_grounding"])
        count += 1

        if len(batch) >= batch_size:
            conn.executemany(
                """
                INSERT INTO semantic_object(
                  feature_id, fid, source_layer, canonical_object_type, object_family, detection_class,
                  capability, z, x, y, geom_type, is_land, is_hole, is_collision, is_grounding
                ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
                """,
                batch,
            )
            conn.commit()
            batch.clear()
            print(f"已写入 {count} 条核心对象")

    if batch:
        conn.executemany(
            """
            INSERT INTO semantic_object(
              feature_id, fid, source_layer, canonical_object_type, object_family, detection_class,
              capability, z, x, y, geom_type, is_land, is_hole, is_collision, is_grounding
            ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
            """,
            batch,
        )
        conn.commit()

    cell_rows = []
    for tile_key, c in grid.items():
        blocked = 0
        reason = None
        if c["land"] > c["hole"] and c["land"] > 0:
            blocked = 1
            reason = "land"
        elif c["collision"] > 0:
            blocked = 1
            reason = "collision"
        elif c["grounding"] > 0:
            blocked = 1
            reason = "grounding"
        cell_rows.append(
            (
                tile_key,
                c["z"],
                c["x"],
                c["y"],
                c["land"],
                c["hole"],
                c["collision"],
                c["grounding"],
                blocked,
                reason,
            )
        )

    conn.executemany(
        """
        INSERT INTO grid_cell(
          tile_key, z, x, y, land_count, hole_count, collision_count, grounding_count,
          blocked, blocked_reason
        ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
        """,
        cell_rows,
    )
    conn.commit()
    conn.execute("ANALYZE")
    conn.commit()

    obj_count = conn.execute("SELECT COUNT(*) FROM semantic_object").fetchone()[0]
    cell_count = conn.execute("SELECT COUNT(*) FROM grid_cell").fetchone()[0]
    conn.close()
    src.close()

    print(f"完成: objects={obj_count}, cells={cell_count}")
    print(f"输出文件: {OUT.resolve()}")


if __name__ == "__main__":
    main()