industrial-model 0.1.10__tar.gz → 0.1.12__tar.gz

{industrial_model-0.1.10 → industrial_model-0.1.12}/PKG-INFO

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

{industrial_model-0.1.10 → industrial_model-0.1.12}/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-0.1.10 → industrial_model-0.1.12}/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-0.1.12/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-0.1.10 → industrial_model-0.1.12}/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-0.1.12/industrial_model/cognite_adapters/view_mapper.py

@@ -0,0 +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, 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:
+        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 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-0.1.12/industrial_model/config.py

@@ -0,0 +1,19 @@
+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)
+
+    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-0.1.10 → industrial_model-0.1.12}/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-0.1.10 → industrial_model-0.1.12}/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-0.1.10 → industrial_model-0.1.12}/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-0.1.10 → industrial_model-0.1.12}/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 | InstanceId] | None
+    value: float
+    aggregate: str
+
+
 ValidationMode = Literal["raiseOnError", "ignoreOnError"]

industrial_model-0.1.12/industrial_model/statements/__init__.py

@@ -0,0 +1,122 @@
+from dataclasses import dataclass, field
+from typing import Any, Generic, Literal, Self, TypeVar
+
+from industrial_model.constants import DEFAULT_LIMIT, SORT_DIRECTION
+
+from .expressions import (
+    BoolExpression,
+    Column,
+    Expression,
+    LeafExpression,
+    and_,
+    col,
+    not_,
+    or_,
+)
+
+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)
+    where_clauses: list[Expression] = field(init=False, default_factory=list)
+    sort_clauses: list[tuple[Column, SORT_DIRECTION]] = field(
+        init=False, default_factory=list
+    )
+    limit_: int = field(init=False, default=DEFAULT_LIMIT)
+    cursor_: str | None = field(init=False, default=None)
+
+    def where(self, *expressions: bool | Expression) -> Self:
+        for expression in expressions:
+            assert isinstance(expression, Expression)
+            self.where_clauses.append(expression)
+        return self
+
+    def asc(self, property: str | Column | Any) -> Self:
+        return self.sort(property, "ascending")
+
+    def desc(self, property: str | Column | Any) -> Self:
+        return self.sort(property, "descending")
+
+    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:
+        self.limit_ = limit
+        return self
+
+    def cursor(self, cursor: str | None) -> Self:
+        self.cursor_ = cursor
+        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",
+    "col",
+    "Expression",
+    "LeafExpression",
+    "BoolExpression",
+    "and_",
+    "not_",
+    "or_",
+]

{industrial_model-0.1.10 → industrial_model-0.1.12}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "industrial-model"
-version = "0.1.10"
+version = "0.1.12"
 description = "Industrial Model ORM"
 readme = "README.md"
 requires-python = ">=3.11"

{industrial_model-0.1.10 → industrial_model-0.1.12}/tests/tests_adapter.py

@@ -1,7 +1,7 @@
 import datetime
 import json
 
-from industrial_model import col, select
+from industrial_model import aggregate, col, select
 
 from .hubs import generate_engine
 from .models import DescribableEntity, Event
@@ -17,7 +17,9 @@ if __name__ == "__main__":
         & (col(Event.start_date_time) < datetime.datetime(2025, 6, 1))
     )
 
-    statement = select(Event).limit(50).where(filter)
+    statement = (
+        select(Event).limit(2500).where(filter).asc(Event.start_date_time)
+    )
 
     result = [
         item.model_dump(mode="json")
@@ -25,3 +27,5 @@ if __name__ == "__main__":
     ]
     print(len(result))
     json.dump(result, open("entities.json", "w"), indent=2)
+
+    print(adapter.aggregate(aggregate(Event).group_by(Event.space)))

{industrial_model-0.1.10 → industrial_model-0.1.12}/uv.lock

@@ -215,7 +215,7 @@ wheels = [
 
 [[package]]
 name = "industrial-model"
-version = "0.1.10"
+version = "0.1.12"
 source = { editable = "." }
 dependencies = [
     { name = "anyio" },

industrial_model-0.1.10/.vscode/settings.json

@@ -1,8 +0,0 @@
-{
-    "editor.defaultFormatter": "charliermarsh.ruff",
-    "editor.formatOnSave": true,
-    "editor.codeActionsOnSave": {
-        "source.fixAll": "always",
-        "source.organizeImports": "always"
-    }
-}

industrial_model-0.1.10/industrial_model/cognite_adapters/view_mapper.py

@@ -1,16 +0,0 @@
-from cognite.client.data_classes.data_modeling import (
-    View,
-)
-
-
-class ViewMapper:
-    def __init__(self, views: list[View]):
-        self._views_as_dict = {view.external_id: view for view in views}
-
-    def get_view(self, view_external_id: str) -> View:
-        if view_external_id not in self._views_as_dict:
-            raise ValueError(
-                f"View {view_external_id} is not available in data model"
-            )
-
-        return self._views_as_dict[view_external_id]

industrial_model-0.1.10/industrial_model/config.py

@@ -1,10 +0,0 @@
-from industrial_model.models import RootModel
-
-
-class DataModelId(RootModel):
-    external_id: str
-    space: str
-    version: str
-
-    def as_tuple(self) -> tuple[str, str, str]:
-        return self.space, self.external_id, self.version

industrial_model-0.1.10/industrial_model/statements/__init__.py

@@ -1,72 +0,0 @@
-from dataclasses import dataclass, field
-from typing import Any, Generic, Self, TypeVar
-
-from industrial_model.constants import DEFAULT_LIMIT, SORT_DIRECTION
-
-from .expressions import (
-    BoolExpression,
-    Column,
-    Expression,
-    LeafExpression,
-    and_,
-    col,
-    not_,
-    or_,
-)
-
-T = TypeVar("T")
-
-
-@dataclass
-class Statement(Generic[T]):
-    entity: type[T] = field(init=True)
-    where_clauses: list[Expression] = field(init=False, default_factory=list)
-    sort_clauses: list[tuple[Column, SORT_DIRECTION]] = field(
-        init=False, default_factory=list
-    )
-    limit_: int = field(init=False, default=DEFAULT_LIMIT)
-    cursor_: str | None = field(init=False, default=None)
-
-    def where(self, *expressions: bool | Expression) -> Self:
-        for expression in expressions:
-            assert isinstance(expression, Expression)
-            self.where_clauses.append(expression)
-        return self
-
-    def asc(self, property: Any) -> Self:
-        self.sort_clauses.append((Column(property), "ascending"))
-        return self
-
-    def desc(self, property: Any) -> Self:
-        self.sort_clauses.append((Column(property), "descending"))
-        return self
-
-    def sort(self, property: Any, direction: SORT_DIRECTION) -> Self:
-        self.sort_clauses.append((Column(property), direction))
-        return self
-
-    def limit(self, limit: int) -> Self:
-        self.limit_ = limit
-        return self
-
-    def cursor(self, cursor: str | None) -> Self:
-        self.cursor_ = cursor
-        return self
-
-
-def select(entity: type[T]) -> Statement[T]:
-    return Statement(entity)
-
-
-__all__ = [
-    "Statement",
-    "select",
-    "Column",
-    "col",
-    "Expression",
-    "LeafExpression",
-    "BoolExpression",
-    "and_",
-    "not_",
-    "or_",
-]