ras2cng.duckdb_session

ras2cng.duckdb_session

DuckDB helpers for querying GeoParquet (including GeoParquet geometry columns).

DuckSession

DuckDB session with spatial extension pre-loaded when available.

Source code in ras2cng/duckdb_session.py

class DuckSession:
    """DuckDB session with spatial extension pre-loaded when available."""

    def __init__(self, db_path: str = ":memory:"):
        self.con = duckdb.connect(db_path)
        self._load_spatial_extension()

    def _load_spatial_extension(self):
        # Prefer LOAD first (works if extension already installed), then INSTALL.
        try:
            self.con.execute("LOAD spatial;")
            return
        except Exception:
            pass

        try:
            self.con.execute("INSTALL spatial;")
            self.con.execute("LOAD spatial;")
        except Exception as e:
            raise RuntimeError(
                "DuckDB spatial extension could not be loaded. "
                "If you're offline, pre-install extensions or skip spatial queries. "
                f"Underlying error: {e}"
            )

    def register_parquet(self, path: Path, name: str = "_"):
        """Register a (Geo)Parquet file as a view.

        If a GeoParquet 'geometry' column exists (WKB), it will be converted into DuckDB's
        GEOMETRY type in the view so that ST_* functions work.
        """

        p = str(Path(path))

        # Try to wrap geometry column into GEOMETRY type.
        try:
            self.con.execute(
                f"""
                CREATE OR REPLACE VIEW {name} AS
                SELECT * EXCLUDE (geometry),
                       ST_GeomFromWKB(geometry) AS geometry
                FROM read_parquet('{p}');
                """
            )
        except Exception:
            # Fallback for non-GeoParquet or files without geometry column.
            self.con.execute(
                f"CREATE OR REPLACE VIEW {name} AS SELECT * FROM read_parquet('{p}');"
            )
        return self

    def query(self, sql: str) -> pd.DataFrame:
        return self.con.execute(sql).df()

    def close(self):
        self.con.close()

register_parquet(path, name='_')

Register a (Geo)Parquet file as a view.

If a GeoParquet 'geometry' column exists (WKB), it will be converted into DuckDB's GEOMETRY type in the view so that ST_* functions work.

Source code in ras2cng/duckdb_session.py

def register_parquet(self, path: Path, name: str = "_"):
    """Register a (Geo)Parquet file as a view.

    If a GeoParquet 'geometry' column exists (WKB), it will be converted into DuckDB's
    GEOMETRY type in the view so that ST_* functions work.
    """

    p = str(Path(path))

    # Try to wrap geometry column into GEOMETRY type.
    try:
        self.con.execute(
            f"""
            CREATE OR REPLACE VIEW {name} AS
            SELECT * EXCLUDE (geometry),
                   ST_GeomFromWKB(geometry) AS geometry
            FROM read_parquet('{p}');
            """
        )
    except Exception:
        # Fallback for non-GeoParquet or files without geometry column.
        self.con.execute(
            f"CREATE OR REPLACE VIEW {name} AS SELECT * FROM read_parquet('{p}');"
        )
    return self