diff --git a/spatialyze/database.py b/spatialyze/database.py index 1bb69429..a1f1aea9 100644 --- a/spatialyze/database.py +++ b/spatialyze/database.py @@ -3,7 +3,6 @@ from typing import TYPE_CHECKING, Callable, NamedTuple import duckdb -import pandas as pd import shapely.geometry from .data_types.camera_key import CameraKey @@ -103,10 +102,10 @@ class Database: def __init__(self, connection: "duckdb.DuckDBPyConnection"): self.connection = connection + self.connection.install_extension("spatial") + self.connection.load_extension("spatial") self.cursor = self.connection.cursor() - self.cursor.execute("INSTALL spatial;") - self.cursor.execute("LOAD spatial;") self.cursor.commit() self.connection.commit() @@ -287,7 +286,8 @@ def insert_camera(self, camera: list[CameraConfig]): cursor = self.connection.cursor() cursor.executemany( "INSERT INTO Camera VALUES " - "(?, ?, ?, ?, ST_GeomFromWKB(?), ?, ?, ST_GeomFromWKB(?), ?, ?, ?, ?)", + "(?, ?, ?, ?, ST_GeomFromWKB(?), ?," + " ?, ST_GeomFromWKB(?), ?, ?, ?, ?)", map(_config, camera), ) @@ -333,12 +333,9 @@ def predicate(self, predicate: "PredicateNode", temporal: bool = True): for frame_number, camera_id, filename, *item_ids in self.execute(sql_str) ] - def sql(self, query: str) -> pd.DataFrame: - results, cursor = self.execute_and_cursor(query) - description = cursor.description - assert description is not None - cursor.close() - return pd.DataFrame(results, columns=[d.name for d in description]) + def sql(self, query: str) -> duckdb.DuckDBPyRelation: + with self.connection.cursor() as cursor: + return cursor.sql(query) def _join_table(temporal: bool):