linkml-store 0.0.0__py3-none-any.whl → 0.1.6__py3-none-any.whl

linkml_store/api/stores/duckdb/duckdb_collection.py

@@ -1,22 +1,24 @@
-from dataclasses import dataclass
+import logging
 from typing import Any, Dict, List, Optional, Union
 
 import sqlalchemy as sqla
 from linkml_runtime.linkml_model import ClassDefinition, SlotDefinition
-from sqlalchemy import Column, Table, delete, insert, text
+from sqlalchemy import Column, Table, delete, insert, inspect, text
 from sqlalchemy.sql.ddl import CreateTable
 
 from linkml_store.api import Collection
-from linkml_store.api.collection import OBJECT
+from linkml_store.api.collection import DEFAULT_FACET_LIMIT, OBJECT
+from linkml_store.api.queries import Query
 from linkml_store.api.stores.duckdb.mappings import TMAP
 from linkml_store.utils.sql_utils import facet_count_sql
 
+logger = logging.getLogger(__name__)
+
 
-@dataclass
 class DuckDBCollection(Collection):
     _table_created: bool = None
 
-    def add(self, objs: Union[OBJECT, List[OBJECT]], **kwargs):
+    def insert(self, objs: Union[OBJECT, List[OBJECT]], **kwargs):
         if not isinstance(objs, list):
             objs = [objs]
         if not objs:
@@ -26,6 +28,7 @@ class DuckDBCollection(Collection):
             cd = self.induce_class_definition_from_objects(objs)
         self._create_table(cd)
         table = self._sqla_table(cd)
+        logger.info(f"Inserting into: {self._alias} // T={table.name}")
         engine = self.parent.engine
         col_names = [c.name for c in table.columns]
         objs = [{k: obj.get(k, None) for k in col_names} for obj in objs]
@@ -51,39 +54,58 @@ class DuckDBCollection(Collection):
                 conn.commit()
         return len(objs)
 
-    def delete_where(self, where: Optional[Dict[str, Any]] = None, **kwargs) -> int:
+    def delete_where(self, where: Optional[Dict[str, Any]] = None, missing_ok=True, **kwargs) -> int:
+        logger.info(f"Deleting from {self._target_class_name} where: {where}")
+        if where is None:
+            where = {}
         cd = self.class_definition()
+        if not cd:
+            logger.info(f"No class definition found for {self._target_class_name}, assuming not prepopulated")
+            return 0
         table = self._sqla_table(cd)
         engine = self.parent.engine
+        inspector = inspect(engine)
+        table_exists = table.name in inspector.get_table_names()
+        if not table_exists:
+            logger.info(f"Table {table.name} does not exist, assuming no data")
+            return 0
         with engine.connect() as conn:
             conditions = [table.c[k] == v for k, v in where.items()]
             stmt = delete(table).where(*conditions)
             stmt = stmt.compile(engine)
-            conn.execute(stmt)
+            result = conn.execute(stmt)
+            deleted_rows_count = result.rowcount
+            if deleted_rows_count == 0 and not missing_ok:
+                raise ValueError(f"No rows found for {where}")
             conn.commit()
-        return 0
+            return deleted_rows_count
 
-    def query_facets(self, where: Dict = None, facet_columns: List[str] = None) -> Dict[str, Dict[str, int]]:
+    def query_facets(
+        self, where: Dict = None, facet_columns: List[str] = None, facet_limit=DEFAULT_FACET_LIMIT, **kwargs
+    ) -> Dict[str, Dict[str, int]]:
         results = {}
         cd = self.class_definition()
         with self.parent.engine.connect() as conn:
             if not facet_columns:
                 facet_columns = list(self.class_definition().attributes.keys())
             for col in facet_columns:
+                logger.debug(f"Faceting on {col}")
                 if isinstance(col, tuple):
                     sd = SlotDefinition(name="PLACEHOLDER")
                 else:
                     sd = cd.attributes[col]
                 facet_query = self._create_query(where_clause=where)
                 facet_query_str = facet_count_sql(facet_query, col, multivalued=sd.multivalued)
+                logger.debug(f"Facet query: {facet_query_str}")
                 rows = list(conn.execute(text(facet_query_str)))
                 results[col] = rows
             return results
 
     def _sqla_table(self, cd: ClassDefinition) -> Table:
+        schema_view = self.parent.schema_view
         metadata_obj = sqla.MetaData()
         cols = []
-        for att in cd.attributes.values():
+        for att in schema_view.class_induced_slots(cd.name):
             typ = TMAP.get(att.range, sqla.String)
             if att.inlined:
                 typ = sqla.JSON
@@ -93,17 +115,28 @@ class DuckDBCollection(Collection):
                 typ = sqla.ARRAY(typ, dimensions=1)
             col = Column(att.name, typ)
             cols.append(col)
-        t = Table(self.name, metadata_obj, *cols)
+        t = Table(self._alias, metadata_obj, *cols)
         return t
 
     def _create_table(self, cd: ClassDefinition):
-        if self._table_created:
+        if self._table_created or self.metadata.is_prepopulated:
+            logger.info(f"Already have table for: {cd.name}")
             return
+        query = Query(
+            from_table="information_schema.tables", where_clause={"table_type": "BASE TABLE", "table_name": self._alias}
+        )
+        qr = self.parent.query(query)
+        if qr.num_rows > 0:
+            logger.info(f"Table already exists for {cd.name}")
+            self._table_created = True
+            self.metadata.is_prepopulated = True
+            return
+        logger.info(f"Creating table for {cd.name}")
         t = self._sqla_table(cd)
         ct = CreateTable(t)
         ddl = str(ct.compile(self.parent.engine))
         with self.parent.engine.connect() as conn:
             conn.execute(text(ddl))
             conn.commit()
-        if not self._table_created:
-            self._table_created = True
+        self._table_created = True
+        self.metadata.is_prepopulated = True

linkml_store/api/stores/duckdb/duckdb_database.py

@@ -1,15 +1,14 @@
 import json
 import logging
-from dataclasses import dataclass
 from typing import Optional
 
 import pandas as pd
 import sqlalchemy
 from duckdb import DuckDBPyConnection
 from linkml_runtime import SchemaView
-from linkml_runtime.linkml_model import SlotDefinition
+from linkml_runtime.linkml_model import ClassDefinition, SlotDefinition
 from linkml_runtime.utils.schema_builder import SchemaBuilder
-from sqlalchemy import text
+from sqlalchemy import NullPool, text
 
 from linkml_store.api import Database
 from linkml_store.api.queries import Query, QueryResult
@@ -20,54 +19,36 @@ TYPE_MAP = {
     "VARCHAR": "string",
     "BIGINT": "integer",
     "BOOLEAN": "boolean",
+    "DATE": "date",
+    "DOUBLE": "float",
+    "INTEGER": "integer",
 }
 
 
 logger = logging.getLogger(__name__)
 
 
-def run_query(con: DuckDBPyConnection, query: Query, **kwargs):
-    """
-    Run a query and return the result.
-
-    >>> import duckdb
-    >>> con = duckdb.connect("db/mgi.db")
-    >>> query = Query(from_table="gaf_association", limit=5)
-    >>> result = run_query(con, query)
-    >>> print(result.num_rows)
-    532233
-
-    :param con:
-    :param query:
-    :return:
-    """
-    count_query_str = query_to_sql(query, count=True)
-    num_rows = con.execute(count_query_str).fetchall()[0][0]
-    logger.debug(f"num_rows: {num_rows}")
-    query_str = query_to_sql(query, **kwargs)
-    logger.debug(f"query_str: {query_str}")
-    rows = con.execute(query_str).fetchdf()
-    qr = QueryResult(query=query, num_rows=num_rows)
-    qr.set_rows(rows)
-    return qr
-
-
-@dataclass
 class DuckDBDatabase(Database):
     _connection: DuckDBPyConnection = None
     _engine: sqlalchemy.Engine = None
+    collection_class = DuckDBCollection
 
-    def __post_init__(self):
-        if not self.handle:
-            self.handle = "duckdb:///:memory:"
+    def __init__(self, handle: Optional[str] = None, **kwargs):
+        if handle is None:
+            handle = "duckdb:///:memory:"
+        super().__init__(handle=handle, **kwargs)
 
     @property
     def engine(self) -> sqlalchemy.Engine:
         if not self._engine:
             handle = self.handle
             if not handle.startswith("duckdb://") and not handle.startswith(":"):
-                handle = f"duckdb://{handle}"
-            self._engine = sqlalchemy.create_engine(handle)
+                handle = f"duckdb:///{handle}"
+            if ":memory:" not in handle:
+                # TODO: investigate this; duckdb appears to be prematurely caching
+                self._engine = sqlalchemy.create_engine(handle, poolclass=NullPool)
+            else:
+                self._engine = sqlalchemy.create_engine(handle)
         return self._engine
 
     def commit(self, **kwargs):
@@ -80,6 +61,14 @@ class DuckDBDatabase(Database):
     def query(self, query: Query, **kwargs) -> QueryResult:
         json_encoded_cols = []
         if query.from_table:
+            if not query.from_table.startswith("information_schema"):
+                meta_query = Query(
+                    from_table="information_schema.tables", where_clause={"table_name": query.from_table}
+                )
+                qr = self.query(meta_query)
+                if qr.num_rows == 0:
+                    logger.debug(f"Table {query.from_table} not created yet")
+                    return QueryResult(query=query, num_rows=0, rows=[])
             sv = self._schema_view
             if sv:
                 cd = None
@@ -127,21 +116,15 @@ class DuckDBDatabase(Database):
                 collection = DuckDBCollection(name=table_name, parent=self)
                 self._collections[table_name] = collection
 
-    def create_collection(self, name: str, alias: Optional[str] = None, **kwargs) -> DuckDBCollection:
-        collection = DuckDBCollection(name=name, parent=self)
-        if not self._collections:
-            self._collections = {}
-        if not alias:
-            alias = name
-        self._collections[alias] = collection
-        return collection
-
     def induce_schema_view(self) -> SchemaView:
         # TODO: unify schema introspection
+        # TODO: handle case where schema is provided in advance
+        logger.info(f"Inducing schema view for {self.metadata.handle}")
         sb = SchemaBuilder()
         schema = sb.schema
         query = Query(from_table="information_schema.tables", where_clause={"table_type": "BASE TABLE"})
         qr = self.query(query)
+        logger.info(f"Found {qr.num_rows} information_schema.tables // {qr.rows}")
         if qr.num_rows:
             table_names = [row["table_name"] for row in qr.rows]
             for tbl in table_names:
@@ -162,5 +145,13 @@ class DuckDBDatabase(Database):
                 row["column_name"], required=row["is_nullable"] == "NO", multivalued=multivalued, range=rng
             )
             sb.schema.classes[tbl_name].attributes[sd.name] = sd
+            logger.info(f"Introspected slot: {tbl_name}.{sd.name}: {sd.range}")
         sb.add_defaults()
+        for cls_name in schema.classes:
+            if cls_name in self.metadata.collections:
+                collection_metadata = self.metadata.collections[cls_name]
+                if collection_metadata.attributes:
+                    del schema.classes[cls_name]
+                    cls = ClassDefinition(name=collection_metadata.type, attributes=collection_metadata.attributes)
+                    schema.classes[cls.name] = cls
         return SchemaView(schema)

linkml_store/api/stores/hdf5/hdf5_collection.py

@@ -0,0 +1,104 @@
+import json
+import logging
+from typing import Any, Dict, List, Optional, Tuple, Union
+
+import h5py
+
+from linkml_store.api import Collection
+from linkml_store.api.collection import DEFAULT_FACET_LIMIT, OBJECT
+from linkml_store.api.queries import Query, QueryResult
+
+logger = logging.getLogger(__name__)
+
+
+class HDF5Collection(Collection):
+
+    @property
+    def hdf5_group(self) -> h5py.Group:
+        return self.parent.file[self.name]
+
+    def insert(self, objs: Union[OBJECT, List[OBJECT]], **kwargs):
+        if not isinstance(objs, list):
+            objs = [objs]
+
+        for obj in objs:
+            if "id" not in obj:
+                raise ValueError("Each object must have an 'id' field.")
+            obj_id = str(obj["id"])
+            for key, value in obj.items():
+                if key == "id":
+                    continue
+                if isinstance(value, (dict, list)):
+                    value = json.dumps(value)
+                self.hdf5_group.create_dataset(f"{obj_id}/{key}", data=value)
+
+    def delete(self, objs: Union[OBJECT, List[OBJECT]], **kwargs) -> int:
+        if not isinstance(objs, list):
+            objs = [objs]
+        count = 0
+        for obj in objs:
+            if "id" not in obj:
+                raise ValueError("Each object must have an 'id' field.")
+            obj_id = str(obj["id"])
+            if obj_id in self.hdf5_group:
+                del self.hdf5_group[obj_id]
+                count += 1
+        return count
+
+    def delete_where(self, where: Optional[Dict[str, Any]] = None, missing_ok=True, **kwargs) -> int:
+        logger.info(f"Deleting from {self._target_class_name} where: {where}")
+        if where is None:
+            where = {}
+        results = self.query(Query(where_clause=where)).rows
+        count = self.delete(results)
+        return count
+
+    def query(self, query: Query, **kwargs) -> QueryResult:
+        results = []
+        for obj_id in self.hdf5_group:
+            obj = {"id": obj_id}
+            for key, value in self.hdf5_group[obj_id].items():
+                try:
+                    obj[key] = json.loads(value[()])
+                except json.JSONDecodeError:
+                    obj[key] = value[()]
+            if self._match_where_clause(obj, query.where_clause):
+                results.append(obj)
+
+        count = len(results)
+        if query.limit:
+            results = results[: query.limit]
+        return QueryResult(query=query, num_rows=count, rows=results)
+
+    def query_facets(
+        self, where: Dict = None, facet_columns: List[str] = None, facet_limit=DEFAULT_FACET_LIMIT, **kwargs
+    ) -> Dict[str, List[Tuple[Any, int]]]:
+        results = {}
+        if not facet_columns:
+            facet_columns = list(self.class_definition().attributes.keys())
+
+        for col in facet_columns:
+            logger.debug(f"Faceting on {col}")
+            facet_counts = {}
+            for obj in self.query(Query(where_clause=where)).rows:
+                if col in obj:
+                    value = obj[col]
+                    if isinstance(value, list):
+                        for v in value:
+                            facet_counts[v] = facet_counts.get(v, 0) + 1
+                    else:
+                        facet_counts[value] = facet_counts.get(value, 0) + 1
+            facet_counts = sorted(facet_counts.items(), key=lambda x: x[1], reverse=True)[:facet_limit]
+            results[col] = facet_counts
+
+        return results
+
+    def _match_where_clause(self, obj: Dict[str, Any], where_clause: Optional[Dict[str, Any]]) -> bool:
+        if where_clause is None:
+            return True
+        for key, value in where_clause.items():
+            if key not in obj:
+                return False
+            if obj[key] != value:
+                return False
+        return True

linkml_store/api/stores/hdf5/hdf5_database.py

@@ -0,0 +1,79 @@
+# hdf5_database.py
+
+import logging
+from typing import Optional
+
+import h5py
+from linkml_runtime import SchemaView
+from linkml_runtime.linkml_model import ClassDefinition, SlotDefinition
+from linkml_runtime.utils.schema_builder import SchemaBuilder
+
+from linkml_store.api import Database
+from linkml_store.api.queries import Query, QueryResult
+from linkml_store.api.stores.hdf5.hdf5_collection import HDF5Collection
+
+logger = logging.getLogger(__name__)
+
+
+class HDF5Database(Database):
+    _file: h5py.File = None
+    collection_class = HDF5Collection
+
+    def __init__(self, handle: Optional[str] = None, **kwargs):
+        if handle is None:
+            handle = "linkml_store.h5"
+        super().__init__(handle=handle, **kwargs)
+
+    @property
+    def file(self) -> h5py.File:
+        if self._file is None:
+            self._file = h5py.File(self.handle, "a")
+        return self._file
+
+    def commit(self, **kwargs):
+        self.file.flush()
+
+    def close(self, **kwargs):
+        if self._file:
+            self._file.close()
+
+    def query(self, query: Query, **kwargs) -> QueryResult:
+        if query.from_table:
+            collection = self.get_collection(query.from_table)
+            return collection.query(query, **kwargs)
+
+    def init_collections(self):
+        if self._collections is None:
+            self._collections = {}
+
+        for collection_name in self.file:
+            if collection_name not in self._collections:
+                collection = HDF5Collection(name=collection_name, parent=self)
+                self._collections[collection_name] = collection
+
+    def induce_schema_view(self) -> SchemaView:
+        logger.info(f"Inducing schema view for {self.handle}")
+        sb = SchemaBuilder()
+        schema = sb.schema
+
+        for collection_name in self.file:
+            sb.add_class(collection_name)
+            hdf5_group = self.file[collection_name]
+            for field in hdf5_group:
+                if field == "_id":
+                    continue
+                sd = SlotDefinition(field)
+                if isinstance(hdf5_group[field][()], list):
+                    sd.multivalued = True
+                sb.schema.classes[collection_name].attributes[sd.name] = sd
+
+        sb.add_defaults()
+        for cls_name in schema.classes:
+            if cls_name in self.metadata.collections:
+                collection_metadata = self.metadata.collections[cls_name]
+                if collection_metadata.attributes:
+                    del schema.classes[cls_name]
+                    cls = ClassDefinition(name=collection_metadata.type, attributes=collection_metadata.attributes)
+                    schema.classes[cls.name] = cls
+
+        return SchemaView(schema)

linkml_store/api/stores/mongodb/mongodb_collection.py

@@ -1,56 +1,102 @@
-from dataclasses import dataclass
-from typing import Any, Dict, List, Optional, Union
+import logging
+from typing import Any, Dict, List, Optional, Tuple, Union
+
+from linkml_runtime.linkml_model import SlotDefinition
+from pymongo.collection import Collection as MongoCollection
 
 from linkml_store.api import Collection
-from linkml_store.api.collection import OBJECT
+from linkml_store.api.collection import DEFAULT_FACET_LIMIT, OBJECT
+from linkml_store.api.queries import Query, QueryResult
+
+logger = logging.getLogger(__name__)
 
 
-@dataclass
 class MongoDBCollection(Collection):
-    """
-    A wrapper around a MongoDB collection
-    """
 
-    def add(self, objs: Union[OBJECT, List[OBJECT]], **kwargs):
-        if not isinstance(objs, list):
-            objs = [objs]
-        if not objs:
-            return
-        cd = self.class_definition()
-        if not cd:
-            cd = self.induce_class_definition_from_objects(objs)
-        collection = self.parent.database[self.name]
-        collection.insert_many(objs)
+    @property
+    def mongo_collection(self) -> MongoCollection:
+        if not self.name:
+            raise ValueError("Collection name not set")
+        return self.parent.native_db[self.name]
 
-    def delete(self, objs: Union[OBJECT, List[OBJECT]], **kwargs) -> int:
+    def insert(self, objs: Union[OBJECT, List[OBJECT]], **kwargs):
         if not isinstance(objs, list):
             objs = [objs]
-        cd = self.class_definition()
-        if not cd:
-            cd = self.induce_class_definition_from_objects(objs)
-        collection = self.parent.database[self.name]
-        deleted_count = 0
-        for obj in objs:
-            result = collection.delete_one(obj)
-            deleted_count += result.deleted_count
-        return deleted_count
+        self.mongo_collection.insert_many(objs)
 
-    def delete_where(self, where: Optional[Dict[str, Any]] = None, **kwargs) -> int:
-        collection = self.parent.database[self.name]
-        result = collection.delete_many(where)
-        return result.deleted_count
+    def query(self, query: Query, **kwargs) -> QueryResult:
+        mongo_filter = self._build_mongo_filter(query.where_clause)
+        if query.limit:
+            cursor = self.mongo_collection.find(mongo_filter).limit(query.limit)
+        else:
+            cursor = self.mongo_collection.find(mongo_filter)
+
+        rows = list(cursor)
+        count = self.mongo_collection.count_documents(mongo_filter)
+
+        return QueryResult(query=query, num_rows=count, rows=rows)
+
+    def _build_mongo_filter(self, where_clause: Dict[str, Any]) -> Dict[str, Any]:
+        mongo_filter = {}
+        if where_clause:
+            for field, value in where_clause.items():
+                mongo_filter[field] = value
+        return mongo_filter
 
-    def query_facets(self, where: Dict = None, facet_columns: List[str] = None) -> Dict[str, Dict[str, int]]:
+    def query_facets(
+        self, where: Dict = None, facet_columns: List[str] = None, facet_limit=DEFAULT_FACET_LIMIT, **kwargs
+    ) -> Dict[str, List[Tuple[Any, int]]]:
         results = {}
-        _cd = self.class_definition()
-        collection = self.parent.database[self.name]
+        cd = self.class_definition()
         if not facet_columns:
             facet_columns = list(self.class_definition().attributes.keys())
+
         for col in facet_columns:
-            facet_pipeline = [
-                {"$match": where} if where else {"$match": {}},
-                {"$group": {"_id": f"${col}", "count": {"$sum": 1}}},
-            ]
-            facet_results = list(collection.aggregate(facet_pipeline))
-            results[col] = [(row["_id"], row["count"]) for row in facet_results]
+            logger.debug(f"Faceting on {col}")
+            if isinstance(col, tuple):
+                sd = SlotDefinition(name="PLACEHOLDER")
+            else:
+                sd = cd.attributes[col]
+
+            if sd.multivalued:
+                facet_pipeline = [
+                    {"$match": where} if where else {"$match": {}},
+                    {"$unwind": f"${col}"},
+                    {"$group": {"_id": f"${col}", "count": {"$sum": 1}}},
+                    {"$sort": {"count": -1}},
+                    {"$limit": facet_limit},
+                ]
+            else:
+                facet_pipeline = [
+                    {"$match": where} if where else {"$match": {}},
+                    {"$group": {"_id": f"${col}", "count": {"$sum": 1}}},
+                    {"$sort": {"count": -1}},
+                    {"$limit": facet_limit},
+                ]
+
+            facet_results = list(self.mongo_collection.aggregate(facet_pipeline))
+            results[col] = [(result["_id"], result["count"]) for result in facet_results]
+
         return results
+
+    def delete(self, objs: Union[OBJECT, List[OBJECT]], **kwargs) -> int:
+        if not isinstance(objs, list):
+            objs = [objs]
+        filter_conditions = []
+        for obj in objs:
+            filter_condition = {}
+            for key, value in obj.items():
+                filter_condition[key] = value
+            filter_conditions.append(filter_condition)
+        result = self.mongo_collection.delete_many({"$or": filter_conditions})
+        return result.deleted_count
+
+    def delete_where(self, where: Optional[Dict[str, Any]] = None, missing_ok=True, **kwargs) -> int:
+        logger.info(f"Deleting from {self._target_class_name} where: {where}")
+        if where is None:
+            where = {}
+        result = self.mongo_collection.delete_many(where)
+        deleted_rows_count = result.deleted_count
+        if deleted_rows_count == 0 and not missing_ok:
+            raise ValueError(f"No rows found for {where}")
+        return deleted_rows_count