industrial-model 0.1.10__py3-none-any.whl → 0.1.11__py3-none-any.whl

industrial_model/__init__.py

@@ -1,6 +1,7 @@
 from .config import DataModelId
 from .engines import AsyncEngine, Engine
 from .models import (
+    AggregationResult,
     InstanceId,
     PaginatedResult,
     TViewInstance,
@@ -10,9 +11,11 @@ from .models import (
     ViewInstanceConfig,
     WritableViewInstance,
 )
-from .statements import and_, col, not_, or_, select
+from .statements import aggregate, and_, col, not_, or_, select
 
 __all__ = [
+    "aggregate",
+    "AggregationResult",
     "and_",
     "or_",
     "col",

industrial_model/cognite_adapters/__init__.py

@@ -11,9 +11,14 @@ from cognite.client.data_classes.data_modeling.query import (
 )
 
 from industrial_model.config import DataModelId
-from industrial_model.models import TViewInstance, TWritableViewInstance
-from industrial_model.statements import Statement
+from industrial_model.models import (
+    AggregationResult,
+    TViewInstance,
+    TWritableViewInstance,
+)
+from industrial_model.statements import AggregationStatement, Statement
 
+from .aggregation_mapper import AggregationMapper
 from .optimizer import QueryOptimizer
 from .query_mapper import QueryMapper
 from .query_result_mapper import (
@@ -34,15 +39,12 @@ class CogniteAdapter:
     ):
         self._cognite_client = cognite_client
 
-        dm = cognite_client.data_modeling.data_models.retrieve(
-            ids=data_model_id.as_tuple(),
-            inline_views=True,
-        ).latest_version()
-        view_mapper = ViewMapper(dm.views)
+        view_mapper = ViewMapper(cognite_client, data_model_id)
+        self._optmizer = QueryOptimizer(cognite_client, data_model_id)
         self._query_mapper = QueryMapper(view_mapper)
         self._result_mapper = QueryResultMapper(view_mapper)
         self._upsert_mapper = UpsertMapper(view_mapper)
-        self._optmizer = QueryOptimizer(cognite_client)
+        self._aggregation_mapper = AggregationMapper(view_mapper)
 
     def query(
         self, statement: Statement[TViewInstance], all_pages: bool
@@ -80,6 +82,29 @@ class CogniteAdapter:
             if not all_pages or last_page:
                 return data, next_cursor_
 
+    def aggregate(
+        self, statement: AggregationStatement[TViewInstance]
+    ) -> list[AggregationResult]:
+        query = self._aggregation_mapper.map(statement)
+
+        result = self._cognite_client.data_modeling.instances.aggregate(
+            view=query.view.as_id(),
+            aggregates=query.metric_aggregation,
+            filter=query.filters,
+            group_by=query.group_by_columns,
+            limit=query.limit,
+        )
+
+        return [
+            AggregationResult(
+                group=item.group,
+                value=item.aggregates[0].value,
+                aggregate=statement.aggregate_,
+            )
+            for item in result
+            if item.aggregates and item.aggregates[0].value is not None
+        ]
+
     def upsert(
         self, entries: list[TWritableViewInstance], replace: bool = False
     ) -> None:

industrial_model/cognite_adapters/aggregation_mapper.py

@@ -0,0 +1,64 @@
+from dataclasses import dataclass
+
+import cognite.client.data_classes.filters as filters
+from cognite.client.data_classes.aggregations import Count, MetricAggregation
+from cognite.client.data_classes.data_modeling import (
+    View,
+)
+
+from industrial_model.models import TViewInstance
+from industrial_model.statements import AggregationStatement
+
+from .filter_mapper import (
+    FilterMapper,
+)
+from .view_mapper import ViewMapper
+
+
+@dataclass
+class AggregationQuery:
+    view: View
+    metric_aggregation: MetricAggregation
+    filters: filters.Filter | None
+    group_by_columns: list[str]
+    limit: int
+
+
+class AggregationMapper:
+    def __init__(self, view_mapper: ViewMapper):
+        self._view_mapper = view_mapper
+        self._filter_mapper = FilterMapper(view_mapper)
+
+    def map(
+        self, statement: AggregationStatement[TViewInstance]
+    ) -> AggregationQuery:
+        root_node = statement.entity.get_view_external_id()
+
+        root_view = self._view_mapper.get_view(root_node)
+
+        filters_ = (
+            filters.And(
+                *self._filter_mapper.map(statement.where_clauses, root_view)
+            )
+            if statement.where_clauses
+            else None
+        )
+
+        metric_aggregation = (
+            Count(statement.aggregation_property.property)
+            if statement.aggregate_ == "count"
+            else None
+        )
+        if metric_aggregation is None:
+            raise ValueError(
+                f"Unsupported aggregate function: {statement.aggregate_}"
+            )
+        return AggregationQuery(
+            view=root_view,
+            metric_aggregation=metric_aggregation,
+            filters=filters_,
+            group_by_columns=[
+                column.property for column in statement.group_by_columns
+            ],
+            limit=statement.limit_,
+        )

industrial_model/cognite_adapters/optimizer.py

@@ -1,5 +1,8 @@
+from threading import Lock
+
 from cognite.client import CogniteClient
 
+from industrial_model.config import DataModelId
 from industrial_model.models import TViewInstance
 from industrial_model.statements import (
     BoolExpression,
@@ -14,11 +17,11 @@ SPACE_PROPERTY = "space"
 
 class QueryOptimizer:
     def __init__(
-        self,
-        cognite_client: CogniteClient,
+        self, cognite_client: CogniteClient, data_model_id: DataModelId
     ):
-        self._all_spaces: list[str] | None = None
+        self._all_spaces = data_model_id.instance_spaces
         self._cognite_client = cognite_client
+        self._lock = Lock()
 
     def optimize(self, statement: Statement[TViewInstance]) -> None:
         instance_spaces = statement.entity.view_config.get("instance_spaces")
@@ -58,7 +61,7 @@ class QueryOptimizer:
         return False
 
     def _find_spaces(self, instance_spaces_prefix: str) -> list[str]:
-        all_spaces = self._get_all_spaces()
+        all_spaces = self._load_spaces()
 
         return [
             space
@@ -66,12 +69,18 @@ class QueryOptimizer:
             if space.startswith(instance_spaces_prefix)
         ]
 
-    def _get_all_spaces(self) -> list[str]:
+    def _load_spaces(self) -> list[str]:
         all_spaces = self._all_spaces
-        if all_spaces is None:
+        if all_spaces:
+            return all_spaces
+
+        with self._lock:
+            if self._all_spaces:
+                return self._all_spaces
+
             all_spaces = self._cognite_client.data_modeling.spaces.list(
                 limit=-1
             ).as_ids()
 
-        self._all_spaces = all_spaces
-        return all_spaces
+            self._all_spaces = all_spaces
+            return all_spaces

industrial_model/cognite_adapters/view_mapper.py

@@ -1,16 +1,50 @@
+from threading import Lock
+
+from cognite.client import CogniteClient
 from cognite.client.data_classes.data_modeling import (
     View,
 )
 
+from industrial_model.config import DataModelId
+
 
 class ViewMapper:
-    def __init__(self, views: list[View]):
-        self._views_as_dict = {view.external_id: view for view in views}
+    def __init__(
+        self, cognite_client: CogniteClient, data_model_id: DataModelId
+    ):
+        self._cognite_client = cognite_client
+        self._data_model_id = data_model_id
+        self._views_as_dict: dict[str, View] | None = None
+
+        if data_model_id.views:
+            self._views_as_dict = {
+                view.external_id: view for view in data_model_id.views
+            }
+
+        self._lock = Lock()
 
     def get_view(self, view_external_id: str) -> View:
-        if view_external_id not in self._views_as_dict:
+        views = self._load_views()
+        if view_external_id not in views:
             raise ValueError(
                 f"View {view_external_id} is not available in data model"
             )
 
-        return self._views_as_dict[view_external_id]
+        return views[view_external_id]
+
+    def _load_views(self) -> dict[str, View]:
+        if self._views_as_dict:
+            return self._views_as_dict
+
+        with self._lock:
+            if self._views_as_dict:
+                return self._views_as_dict
+
+            dm = self._cognite_client.data_modeling.data_models.retrieve(
+                ids=self._data_model_id.as_tuple(),
+                inline_views=True,
+            ).latest_version()
+
+            views = {view.external_id: view for view in dm.views}
+            self._views_as_dict = views
+            return views

industrial_model/config.py

@@ -1,10 +1,19 @@
-from industrial_model.models import RootModel
+from dataclasses import dataclass, field
 
+from cognite.client.data_classes.data_modeling import View
+from pydantic import BaseModel, ConfigDict
+
+
+@dataclass
+class DataModelId(BaseModel):
+    model_config = ConfigDict(arbitrary_types_allowed=True)
 
-class DataModelId(RootModel):
     external_id: str
     space: str
     version: str
 
+    views: list[View] | None = field(default=None)
+    instance_spaces: list[str] | None = field(default=None)
+
     def as_tuple(self) -> tuple[str, str, str]:
         return self.space, self.external_id, self.version

industrial_model/engines/async_engine.py

@@ -2,12 +2,13 @@ from cognite.client import CogniteClient
 
 from industrial_model.config import DataModelId
 from industrial_model.models import (
+    AggregationResult,
     PaginatedResult,
     TViewInstance,
+    TWritableViewInstance,
     ValidationMode,
 )
-from industrial_model.models.entities import TWritableViewInstance
-from industrial_model.statements import Statement
+from industrial_model.statements import AggregationStatement, Statement
 from industrial_model.utils import run_async
 
 from .engine import Engine
@@ -37,6 +38,11 @@ class AsyncEngine:
             self._engine.query_all_pages, statement, validation_mode
         )
 
+    async def aggregate_async(
+        self, statement: AggregationStatement[TViewInstance]
+    ) -> list[AggregationResult]:
+        return await run_async(self._engine.aggregate, statement)
+
     async def upsert_async(
         self, entries: list[TWritableViewInstance], replace: bool = False
     ) -> None:

industrial_model/engines/engine.py

@@ -5,15 +5,14 @@ from cognite.client import CogniteClient
 from industrial_model.cognite_adapters import CogniteAdapter
 from industrial_model.config import DataModelId
 from industrial_model.models import (
+    AggregationResult,
+    EdgeContainer,
     PaginatedResult,
     TViewInstance,
-    ValidationMode,
-)
-from industrial_model.models.entities import (
-    EdgeContainer,
     TWritableViewInstance,
+    ValidationMode,
 )
-from industrial_model.statements import Statement
+from industrial_model.statements import AggregationStatement, Statement
 
 
 class Engine:
@@ -49,6 +48,11 @@ class Engine:
 
         return self._validate_data(statement.entity, data, validation_mode)
 
+    def aggregate(
+        self, statement: AggregationStatement[TViewInstance]
+    ) -> list[AggregationResult]:
+        return self._cognite_adapter.aggregate(statement)
+
     def upsert(
         self, entries: list[TWritableViewInstance], replace: bool = False
     ) -> None:

industrial_model/models/__init__.py

@@ -1,5 +1,6 @@
 from .base import RootModel
 from .entities import (
+    AggregationResult,
     EdgeContainer,
     InstanceId,
     PaginatedResult,
@@ -13,6 +14,7 @@ from .entities import (
 from .schemas import get_parent_and_children_nodes, get_schema_properties
 
 __all__ = [
+    "AggregationResult",
     "RootModel",
     "EdgeContainer",
     "InstanceId",

industrial_model/models/entities.py

@@ -90,4 +90,10 @@ class PaginatedResult(RootModel, Generic[TViewInstance]):
     next_cursor: str | None
 
 
+class AggregationResult(RootModel):
+    group: dict[str, str | int | float | bool] | None
+    value: float
+    aggregate: str
+
+
 ValidationMode = Literal["raiseOnError", "ignoreOnError"]

industrial_model/statements/__init__.py

@@ -1,5 +1,5 @@
 from dataclasses import dataclass, field
-from typing import Any, Generic, Self, TypeVar
+from typing import Any, Generic, Literal, Self, TypeVar
 
 from industrial_model.constants import DEFAULT_LIMIT, SORT_DIRECTION
 
@@ -17,6 +17,10 @@ from .expressions import (
 T = TypeVar("T")
 
 
+def _create_column(property: str | Column | Any) -> Column:
+    return property if isinstance(property, Column) else Column(property)
+
+
 @dataclass
 class Statement(Generic[T]):
     entity: type[T] = field(init=True)
@@ -33,16 +37,21 @@ class Statement(Generic[T]):
             self.where_clauses.append(expression)
         return self
 
-    def asc(self, property: Any) -> Self:
-        self.sort_clauses.append((Column(property), "ascending"))
-        return self
+    def asc(self, property: str | Column | Any) -> Self:
+        return self.sort(property, "ascending")
 
-    def desc(self, property: Any) -> Self:
-        self.sort_clauses.append((Column(property), "descending"))
-        return self
+    def desc(self, property: str | Column | Any) -> Self:
+        return self.sort(property, "descending")
 
-    def sort(self, property: Any, direction: SORT_DIRECTION) -> Self:
-        self.sort_clauses.append((Column(property), direction))
+    def sort(
+        self, property: str | Column | Any, direction: SORT_DIRECTION
+    ) -> Self:
+        self.sort_clauses.append(
+            (
+                _create_column(property),
+                direction,
+            )
+        )
         return self
 
     def limit(self, limit: int) -> Self:
@@ -54,11 +63,52 @@ class Statement(Generic[T]):
         return self
 
 
+@dataclass
+class AggregationStatement(Generic[T]):
+    aggregate_: Literal["count"] = field(init=False, default="count")
+    aggregation_property: Column = field(
+        init=False, default=Column("externalId")
+    )
+
+    entity: type[T] = field(init=True)
+    group_by_columns: list[Column] = field(init=False, default_factory=list)
+    where_clauses: list[Expression] = field(init=False, default_factory=list)
+    limit_: int = field(init=False, default=-1)
+
+    def aggregate(self, aggregates: Literal["count"]) -> Self:
+        self.aggregate_ = aggregates
+        return self
+
+    def aggregate_by(self, property: str | Column | Any) -> Self:
+        self.aggregation_property = _create_column(property)
+        return self
+
+    def group_by(self, *property: str | Column | Any) -> Self:
+        self.group_by_columns.extend(_create_column(p) for p in property)
+        return self
+
+    def where(self, *expressions: bool | Expression) -> Self:
+        for expression in expressions:
+            assert isinstance(expression, Expression)
+            self.where_clauses.append(expression)
+        return self
+
+    def limit(self, limit: int) -> Self:
+        self.limit_ = limit
+        return self
+
+
 def select(entity: type[T]) -> Statement[T]:
     return Statement(entity)
 
 
+def aggregate(entity: type[T]) -> AggregationStatement[T]:
+    return AggregationStatement(entity)
+
+
 __all__ = [
+    "aggregate",
+    "AggregationStatement",
     "Statement",
     "select",
     "Column",

{industrial_model-0.1.10.dist-info → industrial_model-0.1.11.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: industrial-model
-Version: 0.1.10
+Version: 0.1.11
 Summary: Industrial Model ORM
 Author-email: Lucas Alves <lucasrosaalves@gmail.com>
 Classifier: Programming Language :: Python

{industrial_model-0.1.10.dist-info → industrial_model-0.1.11.dist-info}/RECORD

@@ -1,30 +1,31 @@
-industrial_model/__init__.py,sha256=LJRHjaQvFZC-rvDubXvI9VaM9me0fqi4DG2mC0pEzC0,635
-industrial_model/config.py,sha256=wzSKVKHIdGlxCRtu0PIeL3SLYvnQBR4Py46CcETxa8U,238
+industrial_model/__init__.py,sha256=13TekQQlw7euJr1pIncV1R1Z8cWtvCPPx5e6s0NY_C0,711
+industrial_model/config.py,sha256=bNEfjgD2vth4RPGKDwjGmh0edleQ95g__UJGHV4vix0,520
 industrial_model/constants.py,sha256=wtFdxkR9gojqekwXeyVCof6VUkJYuqjjgAgVVKrcxaw,450
 industrial_model/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 industrial_model/utils.py,sha256=oh4AxwxXaWgIC2uolkCbvkgo0ququHB6yAPVIXy45Ts,663
-industrial_model/cognite_adapters/__init__.py,sha256=xRLmfPctoEvjRiXqD9Nd_e0uwDD_2VrsvjG-aJXmuR0,4658
+industrial_model/cognite_adapters/__init__.py,sha256=a7xPiSkuWkVxvsmVfFD-lOzdOc-tTNo8JmbyL5tC-68,5461
+industrial_model/cognite_adapters/aggregation_mapper.py,sha256=wG8ottGilTAJcpf2kQvhvyRmDBpSKiT4esx0P5kmZBU,1845
 industrial_model/cognite_adapters/filter_mapper.py,sha256=NqH-OW7_iKFY9POCG8W3KjkwXUgrZP1d_yxDx1J0fXM,3859
 industrial_model/cognite_adapters/models.py,sha256=2j2IS01uPkQEp9WdVk8seYzEqGcDdWFnpzXhusHB2zk,945
-industrial_model/cognite_adapters/optimizer.py,sha256=G8I07jJ9tarE5GZXXUSpTMKUei6ptV-cudJSLsoykX4,2223
+industrial_model/cognite_adapters/optimizer.py,sha256=BbbsA3r9ZKJt6mN1wVRZc046ov-vbml9_ESfurTNRUg,2477
 industrial_model/cognite_adapters/query_mapper.py,sha256=3fEcaLsGjLKIh-g1BbMcffQ6rp99JeCW555iJo8JW44,6260
 industrial_model/cognite_adapters/query_result_mapper.py,sha256=lsKz0wiqMOo54MVEw_YqLcHXa_HWYgAbSLnBbbdVupw,8451
 industrial_model/cognite_adapters/sort_mapper.py,sha256=RJUAYlZGXoYzK0PwX63cibRF_L-MUq9g2ZsC2EeNIF4,696
 industrial_model/cognite_adapters/upsert_mapper.py,sha256=tWEiBJQeeNz1HDu0AoBIfCw_LL156Zg4h7ORKlZ__uw,4870
 industrial_model/cognite_adapters/utils.py,sha256=rztCtS10ZTQmXfBv0nLgDiQIMWemhdSFK-SwrbRjVxM,4792
-industrial_model/cognite_adapters/view_mapper.py,sha256=lnv64KezSQTAr6XdcExa8d92GU5Ll9K9HfMcQGzhX6k,488
+industrial_model/cognite_adapters/view_mapper.py,sha256=qvT25gwIgPdAqZf-oQEh0GylUi5KEPwjLNmIl5ISQGA,1474
 industrial_model/engines/__init__.py,sha256=7aGHrUm2MxIq39vR8h0xu3i1zNOuT9H9U-q4lV3nErQ,102
-industrial_model/engines/async_engine.py,sha256=jnNGBVoD0xL52UEOZ2SzTU3Eogx_EJQXklPJu22iN0I,1325
-industrial_model/engines/engine.py,sha256=J7JlTN0V8H259Aidh7RkBc34UK04J5iTEsWLSHV-hzg,2834
-industrial_model/models/__init__.py,sha256=TVlMUoK7IThiZmH9CYe0pSqfJfVJ4-DmG9vFdUILKtU,624
+industrial_model/engines/async_engine.py,sha256=qILOHTBnq8DOwHoGnBQ6qrRh2zv2-0rjBd36hqWXM0w,1523
+industrial_model/engines/engine.py,sha256=fSJL0i7TwJsVP_-uNrN_WoxToCpwidIik58KREEw-oQ,3003
+industrial_model/models/__init__.py,sha256=klwPCssjGx-0rGi8aBLzP6GWO9kgR7zDAo93el9hQ1I,672
 industrial_model/models/base.py,sha256=jbiaICJ0R1mmxXtDqxxlVdq-tTX4RLdqnLTgs9HLm_4,1279
-industrial_model/models/entities.py,sha256=iLZ4gYGPlNCQJ33s0d6TenjMFuCR8TsvQPEpRAhzvlQ,2378
+industrial_model/models/entities.py,sha256=x4c-ch-9cRVVk039dx60kXSDFhsxQ3lIY96tJ4vPjYI,2506
 industrial_model/models/schemas.py,sha256=LKNPDnUy1jtMyOHDf28En9vThhdzOSswewIcjC_y-6U,4393
 industrial_model/queries/__init__.py,sha256=7aheTE5qs03rxWm9fmGWptbz_p9OIXXYD8if56cqs18,227
 industrial_model/queries/models.py,sha256=iiHQ7-cfg0nukEv5PoCx9QPF-w1gVSnoNbXBOK9Mzeo,1185
 industrial_model/queries/params.py,sha256=ehgCoR5n6E-tkEuoymZ2lkLcSzMaBAx_HnyJ7sWpqz0,964
-industrial_model/statements/__init__.py,sha256=9fD-qpNXIkrjoahxC_R6hS4DKSVelehimvRPKbpYfA0,1775
+industrial_model/statements/__init__.py,sha256=Y57G1TJlYxCHg3Zz7ojQoAMVnCMaUCRRWXEm2IcbAWo,3332
 industrial_model/statements/expressions.py,sha256=Sar1cIvy3sYi7tkWJN3ylHlZ252oN2mZJpZ1TX9jN3s,4940
-industrial_model-0.1.10.dist-info/METADATA,sha256=o0eF-e97Jrv_Z5e55btU-7jmfeZ5q-025Yy4H-isGVM,5698
-industrial_model-0.1.10.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-industrial_model-0.1.10.dist-info/RECORD,,
+industrial_model-0.1.11.dist-info/METADATA,sha256=IXdJyysssgu2JDxbSe6IN0Wg3US7Ax3XYu4NXUXU838,5698
+industrial_model-0.1.11.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+industrial_model-0.1.11.dist-info/RECORD,,