Skip to content
Merged
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
17 changes: 7 additions & 10 deletions spatialyze/database.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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()

Expand Down Expand Up @@ -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),
)

Expand Down Expand Up @@ -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):
Expand Down