linkml-store 0.3.0__py3-none-any.whl

linkml_store/api/queries.py

@@ -0,0 +1,69 @@
+from collections import namedtuple
+from typing import Any, Dict, List, Optional, Tuple, Union
+
+import pandas as pd
+from pydantic import BaseModel
+
+# defined a named tuple called between with two values (start, end):
+# This is used in the Query class to represent a range of values
+# This is used in the Query class to represent a range of values
+Between = namedtuple("Between", "min max")
+
+FACET_GROUP_ATOM = Union[str, int, float, Between]
+FACET_GROUP = Union[FACET_GROUP_ATOM, Tuple[FACET_GROUP_ATOM, ...]]
+
+
+class Query(BaseModel):
+    """
+    A query object.
+
+    - In SQL this would be a SQL query string
+    """
+
+    from_table: Optional[str] = None
+    select_cols: Optional[List[str]] = None
+    where_clause: Optional[Union[str, List[str], Dict[str, Any]]] = None
+    sort_by: Optional[List[str]] = None
+    limit: Optional[int] = None
+    offset: Optional[int] = None
+    include_facet_counts: bool = False
+    facet_slots: Optional[List[str]] = None
+
+
+class FacetCountResult(BaseModel):
+    """
+    A facet count result
+    """
+
+    as_dict: Dict[FACET_GROUP, List[Tuple[FACET_GROUP, int]]]
+
+
+class QueryResult(BaseModel):
+    """
+    A query result.
+
+    TODO: make this a subclass of Collection
+    """
+
+    query: Optional[Query] = None
+    search_term: Optional[str] = None
+    num_rows: int
+    offset: Optional[int] = 0
+    rows: Optional[List[Dict[str, Any]]] = None
+    ranked_rows: Optional[List[Tuple[float, Dict[str, Any]]]] = None
+    _rows_dataframe: Optional[pd.DataFrame] = None
+    facet_counts: Optional[Dict[str, List[Tuple[FACET_GROUP, int]]]] = None
+
+    @property
+    def rows_dataframe(self) -> pd.DataFrame:
+        if self.ranked_rows is not None:
+            self._rows_dataframe = pd.DataFrame([{"score": score, **row} for score, row in self.ranked_rows])
+        if self._rows_dataframe is None and self.rows:
+            self._rows_dataframe = pd.DataFrame(self.rows)
+        return self._rows_dataframe
+
+    def set_rows(self, rows: pd.DataFrame):
+        self._rows_dataframe = rows
+
+    class Config:
+        arbitrary_types_allowed = True

linkml_store/api/stores/chromadb/__init__.py

@@ -0,0 +1,7 @@
+"""
+Adapter for ChromaDB vector database.
+
+.. warning::
+
+    Support for ChromaDB is experimental and may change in the future.
+"""

linkml_store/api/stores/chromadb/chromadb_collection.py

@@ -0,0 +1,121 @@
+"""
+ChromaDB Collection
+"""
+
+import logging
+from typing import Any, Dict, List, Optional, Tuple, Union
+
+from chromadb.api.models.Collection import Collection as ChromaCollection
+from linkml_runtime.linkml_model import SlotDefinition
+
+from linkml_store.api import Collection
+from linkml_store.api.collection import DEFAULT_FACET_LIMIT, OBJECT
+from linkml_store.api.queries import Query, QueryResult
+from linkml_store.index import Indexer
+
+logger = logging.getLogger(__name__)
+
+
+class ChromaDBCollection(Collection):
+    """
+    A wrapper for ChromaDB collections.
+    """
+
+    @property
+    def native_collection(self) -> ChromaCollection:
+        return self.parent.client.get_collection(self.name)
+
+    def insert(self, objs: Union[OBJECT, List[OBJECT]], **kwargs):
+        if not isinstance(objs, list):
+            objs = [objs]
+
+        documents = []
+        metadatas = []
+        ids = []
+        indexer = Indexer()
+
+        for obj in objs:
+            obj_id = self.object_identifier(obj)
+            ids.append(obj_id)
+            doc_text = indexer.object_to_text(obj)
+            documents.append(doc_text)
+            # TODO: handle nesting
+            metadata = {k: v for k, v in obj.items()}
+            metadatas.append(metadata)
+
+        self.native_collection.add(
+            documents=documents,
+            metadatas=metadatas,
+            ids=ids,
+        )
+
+    def delete(self, objs: Union[OBJECT, List[OBJECT]], **kwargs) -> int:
+        if not isinstance(objs, list):
+            objs = [objs]
+        ids = [obj["id"] for obj in objs]
+        self.native_collection.delete(ids=ids)
+        return len(ids)
+
+    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.native_collection.get(where=where)
+        ids = [result["id"] for result in results]
+        self.native_collection.delete(ids=ids)
+        return len(ids)
+
+    def query(self, query: Query, **kwargs) -> QueryResult:
+        chroma_filter = self._build_chroma_filter(query.where_clause)
+        if query.limit:
+            results = self.native_collection.get(where=chroma_filter, limit=query.limit)
+        else:
+            results = self.native_collection.get(where=chroma_filter)
+
+        count = len(results)
+        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 = {}
+        cd = self.class_definition()
+        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]
+
+            if sd.multivalued:
+                facet_results = self.native_collection.aggregate(
+                    aggregation=[
+                        {"$match": where} if where else {"$match": {}},
+                        {"$unwind": f"${col}"},
+                        {"$group": {"_id": f"${col}", "count": {"$sum": 1}}},
+                        {"$sort": {"count": -1}},
+                        {"$limit": facet_limit},
+                    ]
+                )
+            else:
+                facet_results = self.native_collection.aggregate(
+                    aggregation=[
+                        {"$match": where} if where else {"$match": {}},
+                        {"$group": {"_id": f"${col}", "count": {"$sum": 1}}},
+                        {"$sort": {"count": -1}},
+                        {"$limit": facet_limit},
+                    ]
+                )
+
+            results[col] = [(result["_id"], result["count"]) for result in facet_results]
+
+        return results
+
+    def _build_chroma_filter(self, where_clause: Dict[str, Any]) -> Dict[str, Any]:
+        chroma_filter = {}
+        for field, value in where_clause.items():
+            chroma_filter[field] = value
+        return chroma_filter

linkml_store/api/stores/chromadb/chromadb_database.py

@@ -0,0 +1,89 @@
+# chromadb_database.py
+
+import logging
+from typing import Optional
+
+import chromadb
+from chromadb.config import Settings
+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.chromadb.chromadb_collection import ChromaDBCollection
+
+logger = logging.getLogger(__name__)
+
+
+class ChromaDBDatabase(Database):
+    _client: chromadb.Client = None
+    collection_class = ChromaDBCollection
+
+    def __init__(self, handle: Optional[str] = None, **kwargs):
+        if handle is None:
+            handle = ".chromadb"
+        super().__init__(handle=handle, **kwargs)
+
+    @property
+    def client(self) -> chromadb.Client:
+        if self._client is None:
+            self._client = chromadb.Client(
+                Settings(
+                    chroma_db_impl="duckdb+parquet",
+                    persist_directory=self.handle,
+                )
+            )
+        return self._client
+
+    def commit(self, **kwargs):
+        pass
+
+    def close(self, **kwargs):
+        if self._client:
+            self._client.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.client.list_collections():
+            if collection_name not in self._collections:
+                collection = ChromaDBCollection(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.client.list_collections():
+            sb.add_class(collection_name)
+            chroma_collection = self.client.get_collection(collection_name)
+            sample_doc = chroma_collection.peek(1)
+            if sample_doc:
+                for field, value in sample_doc[0].items():
+                    if field == "_id":
+                        continue
+                    sd = SlotDefinition(field)
+                    if isinstance(value, list):
+                        sd.multivalued = True
+                    if isinstance(value, dict):
+                        sd.inlined = 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/dremio/__init__.py

@@ -0,0 +1,10 @@
+"""Dremio database adapter for linkml-store.
+
+This module provides a Dremio adapter that uses Arrow Flight SQL for high-performance
+data access to Dremio data lakehouse.
+"""
+
+from linkml_store.api.stores.dremio.dremio_collection import DremioCollection
+from linkml_store.api.stores.dremio.dremio_database import DremioDatabase
+
+__all__ = ["DremioDatabase", "DremioCollection"]