cognite-neat 0.103.1__py3-none-any.whl → 0.105.0__py3-none-any.whl

cognite/neat/_graph/extractors/_classic_cdf/_sequences.py

@@ -1,37 +1,251 @@
-from collections.abc import Iterable
+import itertools
+import json
+from collections.abc import Callable, Iterable, Set
 from pathlib import Path
+from typing import Any
 
 from cognite.client import CogniteClient
-from cognite.client.data_classes import Sequence, SequenceFilter, SequenceList
+from cognite.client.data_classes import Sequence, SequenceFilter
+from rdflib import RDF, XSD, Literal, Namespace, URIRef
 
-from ._base import ClassicCDFBaseExtractor, InstanceIdPrefix
+from cognite.neat._client.data_classes.neat_sequence import NeatSequence, NeatSequenceList
+from cognite.neat._shared import Triple
 
+from ._base import DEFAULT_SKIP_METADATA_VALUES, ClassicCDFBaseExtractor, InstanceIdPrefix
 
-class SequencesExtractor(ClassicCDFBaseExtractor[Sequence]):
-    """Extract data from Cognite Data Fusions Sequences into Neat."""
+
+class SequencesExtractor(ClassicCDFBaseExtractor[NeatSequence]):
+    """Extract data from Cognite Data Fusions Sequences into Neat.
+
+    Args:
+        items (Iterable[T_CogniteResource]): An iterable of classic resource.
+        namespace (Namespace, optional): The namespace to use. Defaults to DEFAULT_NAMESPACE.
+        to_type (Callable[[T_CogniteResource], str | None], optional): A function to convert an item to a type.
+            Defaults to None. If None or if the function returns None, the asset will be set to the default type.
+        total (int, optional): The total number of items to load. If passed, you will get a progress bar if rich
+            is installed. Defaults to None.
+        limit (int, optional): The maximal number of items to load. Defaults to None. This is typically used for
+            testing setup of the extractor. For example, if you are extracting 100 000 assets, you might want to
+            limit the extraction to 1000 assets to test the setup.
+        unpack_metadata (bool, optional): Whether to unpack metadata. Defaults to False, which yields the metadata as
+            a JSON string.
+        skip_metadata_values (set[str] | frozenset[str] | None, optional): If you are unpacking metadata, then
+           values in this set will be skipped.
+        camel_case (bool, optional): Whether to use camelCase instead of snake_case for property names.
+            Defaults to True.
+        as_write (bool, optional): Whether to use the write/request format of the items. Defaults to False.
+        unpack_columns (bool, optional): Whether to unpack columns. Defaults to False.
+    """
 
     _default_rdf_type = "Sequence"
+    _column_rdf_type = "ColumnClass"
     _instance_id_prefix = InstanceIdPrefix.sequence
 
+    def __init__(
+        self,
+        items: Iterable[NeatSequence],
+        namespace: Namespace | None = None,
+        to_type: Callable[[NeatSequence], str | None] | None = None,
+        total: int | None = None,
+        limit: int | None = None,
+        unpack_metadata: bool = True,
+        skip_metadata_values: Set[str] | None = DEFAULT_SKIP_METADATA_VALUES,
+        camel_case: bool = True,
+        as_write: bool = False,
+        unpack_columns: bool = False,
+    ):
+        super().__init__(
+            items, namespace, to_type, total, limit, unpack_metadata, skip_metadata_values, camel_case, as_write
+        )
+        self.unpack_columns = unpack_columns
+
     @classmethod
-    def _from_dataset(cls, client: CogniteClient, data_set_external_id: str) -> tuple[int | None, Iterable[Sequence]]:
+    def from_dataset(
+        cls,
+        client: CogniteClient,
+        data_set_external_id: str,
+        namespace: Namespace | None = None,
+        to_type: Callable[[NeatSequence], str | None] | None = None,
+        limit: int | None = None,
+        unpack_metadata: bool = True,
+        skip_metadata_values: Set[str] | None = DEFAULT_SKIP_METADATA_VALUES,
+        camel_case: bool = True,
+        as_write: bool = False,
+        unpack_columns: bool = False,
+    ):
+        total, items = cls._from_dataset(client, data_set_external_id)
+        return cls(
+            items,
+            namespace,
+            to_type,
+            total,
+            limit,
+            unpack_metadata,
+            skip_metadata_values,
+            camel_case,
+            as_write,
+            unpack_columns,
+        )
+
+    @classmethod
+    def from_hierarchy(
+        cls,
+        client: CogniteClient,
+        root_asset_external_id: str,
+        namespace: Namespace | None = None,
+        to_type: Callable[[NeatSequence], str | None] | None = None,
+        limit: int | None = None,
+        unpack_metadata: bool = True,
+        skip_metadata_values: Set[str] | None = DEFAULT_SKIP_METADATA_VALUES,
+        camel_case: bool = True,
+        as_write: bool = False,
+        unpack_columns: bool = False,
+    ):
+        total, items = cls._from_hierarchy(client, root_asset_external_id)
+        return cls(
+            items,
+            namespace,
+            to_type,
+            total,
+            limit,
+            unpack_metadata,
+            skip_metadata_values,
+            camel_case,
+            as_write,
+            unpack_columns,
+        )
+
+    @classmethod
+    def from_file(
+        cls,
+        file_path: str | Path,
+        namespace: Namespace | None = None,
+        to_type: Callable[[NeatSequence], str | None] | None = None,
+        limit: int | None = None,
+        unpack_metadata: bool = True,
+        skip_metadata_values: Set[str] | None = DEFAULT_SKIP_METADATA_VALUES,
+        camel_case: bool = True,
+        as_write: bool = False,
+        unpack_columns: bool = False,
+    ):
+        total, items = cls._from_file(file_path)
+        return cls(
+            items,
+            namespace,
+            to_type,
+            total,
+            limit,
+            unpack_metadata,
+            skip_metadata_values,
+            camel_case,
+            as_write,
+            unpack_columns,
+        )
+
+    @classmethod
+    def _from_dataset(
+        cls, client: CogniteClient, data_set_external_id: str
+    ) -> tuple[int | None, Iterable[NeatSequence]]:
         total = client.sequences.aggregate_count(
             filter=SequenceFilter(data_set_ids=[{"externalId": data_set_external_id}])
         )
         items = client.sequences(data_set_external_ids=data_set_external_id)
-        return total, items
+        return total, cls._lookup_rows(items, client)
 
     @classmethod
     def _from_hierarchy(
         cls, client: CogniteClient, root_asset_external_id: str
-    ) -> tuple[int | None, Iterable[Sequence]]:
+    ) -> tuple[int | None, Iterable[NeatSequence]]:
         total = client.sequences.aggregate_count(
             filter=SequenceFilter(asset_subtree_ids=[{"externalId": root_asset_external_id}])
         )
         items = client.sequences(asset_subtree_external_ids=[root_asset_external_id])
-        return total, items
+        return total, cls._lookup_rows(items, client)
 
     @classmethod
-    def _from_file(cls, file_path: str | Path) -> tuple[int | None, Iterable[Sequence]]:
-        sequences = SequenceList.load(Path(file_path).read_text())
+    def _from_file(cls, file_path: str | Path) -> tuple[int | None, Iterable[NeatSequence]]:
+        sequences = NeatSequenceList.load(Path(file_path).read_text())
         return len(sequences), sequences
+
+    @classmethod
+    def _lookup_rows(cls, sequence_iterable: Iterable[Sequence], client: CogniteClient) -> Iterable[NeatSequence]:
+        iterator = iter(sequence_iterable)
+        for sequences in iter(lambda: list(itertools.islice(iterator, client.config.max_workers)), []):
+            # The PySDK uses max_workers to limit the number of requests made in parallel.
+            # We can only get one set of sequence rows per request, so we chunk the sequences up into groups of
+            # max_workers and then make a request to get all the rows for those sequences in one go.
+            sequence_list = list(sequences)
+            row_list = client.sequences.rows.retrieve(id=[seq.id for seq in sequence_list])
+            rows_by_sequence_id = {row.id: row.rows for row in row_list}
+            for seq in sequence_list:
+                yield NeatSequence.from_cognite_sequence(seq, rows_by_sequence_id.get(seq.id))
+
+    def _item2triples_special_cases(self, id_: URIRef, dumped: dict[str, Any]) -> list[Triple]:
+        """For sequences, columns and rows are special cases.'"""
+        if self.unpack_columns:
+            return self._unpack_columns(id_, dumped)
+        else:
+            return self._default_columns_and_rows(id_, dumped)
+
+    def _default_columns_and_rows(self, id_: URIRef, dumped: dict[str, Any]) -> list[Triple]:
+        triples: list[Triple] = []
+        if "columns" in dumped:
+            columns = dumped.pop("columns")
+            triples.extend(
+                [
+                    (
+                        id_,
+                        self.namespace.columns,
+                        # Rows have a rowNumber, so we introduce colNumber here to be consistent.
+                        Literal(json.dumps({"colNumber": no, **col}), datatype=XSD._NS["json"]),
+                    )
+                    for no, col in enumerate(columns, 1)
+                ]
+            )
+        if "rows" in dumped:
+            rows = dumped.pop("rows")
+            triples.extend(
+                [(id_, self.namespace.rows, Literal(json.dumps(row), datatype=XSD._NS["json"])) for row in rows]
+            )
+        return triples
+
+    def _unpack_columns(self, id_: URIRef, dumped: dict[str, Any]) -> list[Triple]:
+        triples: list[Triple] = []
+        columnValueTypes: list[str] = []
+        column_order: list[str] = []
+        if columns := dumped.pop("columns", None):
+            for col in columns:
+                external_id = col.pop("externalId")
+                column_order.append(external_id)
+                value_type = col.pop("valueType")
+                columnValueTypes.append(value_type)
+
+                col_id = self.namespace[f"Column_{external_id}"]
+                triples.append((id_, self.namespace[external_id], col_id))
+                type_ = self.namespace[self._column_rdf_type]
+                triples.append((col_id, RDF.type, type_))
+                if metadata := col.pop("metadata", None):
+                    triples.extend(self._metadata_to_triples(col_id, metadata))
+                # Should only be name and description left in col
+                for key, value in col.items():
+                    if value is None:
+                        continue
+                    triples.append((col_id, self.namespace[key], Literal(value, datatype=XSD.string)))
+
+            triples.append(
+                (id_, self.namespace.columnOrder, Literal(json.dumps(column_order), datatype=XSD._NS["json"]))
+            )
+            triples.append(
+                (id_, self.namespace.columnValueTypes, Literal(json.dumps(columnValueTypes), datatype=XSD._NS["json"]))
+            )
+        if rows := dumped.pop("rows", None):
+            values_by_column: list[list[Any]] = [[] for _ in column_order]
+            for row in rows:
+                for i, value in enumerate(row["values"]):
+                    values_by_column[i].append(value)
+            for col_name, values in zip(column_order, values_by_column, strict=False):
+                triples.append(
+                    (id_, self.namespace[f"{col_name}Values"], Literal(json.dumps(values), datatype=XSD._NS["json"]))
+                )
+
+        return triples

cognite/neat/_graph/extractors/_mock_graph_generator.py

@@ -46,7 +46,7 @@ class MockGraphGenerator(BaseExtractor):
             # fixes potential issues with circular dependencies
             from cognite.neat._rules.transformers import DMSToInformation
 
-            self.rules = DMSToInformation().transform(rules).rules
+            self.rules = DMSToInformation().transform(rules)
         elif isinstance(rules, InformationRules):
             self.rules = rules
         else:
@@ -183,7 +183,7 @@ def _get_generation_order(
     parent_col: str = "source_class",
     child_col: str = "target_class",
 ) -> dict:
-    parent_child_list: list[list[str]] = class_linkage[[parent_col, child_col]].values.tolist()
+    parent_child_list: list[list[str]] = class_linkage[[parent_col, child_col]].values.tolist()  # type: ignore[assignment]
     # Build a directed graph and a list of all names that have no parent
     graph: dict[str, set] = {name: set() for tup in parent_child_list for name in tup}
     has_parent: dict[str, bool] = {name: False for tup in parent_child_list for name in tup}

cognite/neat/_graph/loaders/_rdf2dms.py

@@ -19,6 +19,7 @@ from pydantic import BaseModel, ValidationInfo, create_model, field_validator
 from rdflib import RDF, URIRef
 
 from cognite.neat._client import NeatClient
+from cognite.neat._constants import is_readonly_property
 from cognite.neat._graph._tracking import LogTracker, Tracker
 from cognite.neat._issues import IssueList, NeatIssue, NeatIssueList
 from cognite.neat._issues.errors import (
@@ -303,6 +304,9 @@ class DMSLoader(CDFLoader[dm.InstanceApply]):
             if isinstance(prop, dm.EdgeConnection):
                 edge_by_property[prop_id] = prop_id, prop
             if isinstance(prop, dm.MappedProperty):
+                if is_readonly_property(prop.container, prop.container_property_identifier):
+                    continue
+
                 if isinstance(prop.type, dm.DirectRelation):
                     if prop.container == dm.ContainerId("cdf_cdm", "CogniteTimeSeries") and prop_id == "unit":
                         unit_properties.append(prop_id)
@@ -343,9 +347,14 @@ class DMSLoader(CDFLoader[dm.InstanceApply]):
 
             return value
 
-        def parse_json_string(cls, value: Any, info: ValidationInfo) -> dict:
+        def parse_json_string(cls, value: Any, info: ValidationInfo) -> dict | list:
             if isinstance(value, dict):
                 return value
+            elif isinstance(value, list):
+                try:
+                    return [json.loads(v) if isinstance(v, str) else v for v in value]
+                except json.JSONDecodeError as error:
+                    raise ValueError(f"Not valid JSON string for {info.field_name}: {value}, error {error}") from error
             elif isinstance(value, str):
                 try:
                     return json.loads(value)
@@ -401,7 +410,9 @@ class DMSLoader(CDFLoader[dm.InstanceApply]):
             space=self.instance_space,
             external_id=identifier,
             type=(dm.DirectRelationReference(view_id.space, view_id.external_id) if type_ is not None else None),
-            sources=[dm.NodeOrEdgeData(source=view_id, properties=dict(created.model_dump().items()))],
+            sources=[
+                dm.NodeOrEdgeData(source=view_id, properties=dict(created.model_dump(exclude_unset=True).items()))
+            ],
         )
 
     def _create_edges(

cognite/neat/_graph/transformers/__init__.py

@@ -15,7 +15,7 @@ from ._prune_graph import (
     PruneTypes,
 )
 from ._rdfpath import AddSelfReferenceProperty, MakeConnectionOnExactMatch
-from ._value_type import ConvertLiteral, LiteralToEntity, SplitMultiValueProperty
+from ._value_type import ConnectionToLiteral, ConvertLiteral, LiteralToEntity, SplitMultiValueProperty
 
 __all__ = [
     "AddAssetDepth",
@@ -26,6 +26,7 @@ __all__ = [
     "AssetSequenceConnector",
     "AssetTimeSeriesConnector",
     "AttachPropertyFromTargetToSource",
+    "ConnectionToLiteral",
     "ConvertLiteral",
     "LiteralToEntity",
     "MakeConnectionOnExactMatch",
@@ -55,4 +56,5 @@ Transformers = (
     | PruneInstancesOfUnknownType
     | ConvertLiteral
     | LiteralToEntity
+    | ConnectionToLiteral
 )

cognite/neat/_graph/transformers/_base.py

@@ -1,7 +1,27 @@
+import dataclasses
+import warnings
 from abc import ABC, abstractmethod
-from typing import ClassVar
+from typing import ClassVar, TypeAlias, cast
 
 from rdflib import Graph
+from rdflib.query import ResultRow
+
+from cognite.neat._issues.warnings import NeatValueWarning
+from cognite.neat._shared import Triple
+from cognite.neat._utils.collection_ import iterate_progress_bar
+from cognite.neat._utils.graph_transformations_report import GraphTransformationResult
+
+To_Add_Triples: TypeAlias = list[Triple]
+To_Remove_Triples: TypeAlias = list[Triple]
+
+
+@dataclasses.dataclass
+class RowTransformationOutput:
+    remove_triples: To_Remove_Triples = dataclasses.field(default_factory=list)
+    add_triples: To_Add_Triples = dataclasses.field(default_factory=list)
+    instances_removed_count: int = 0
+    instances_added_count: int = 0
+    instances_modified_count: int = 0
 
 
 class BaseTransformer(ABC):
@@ -12,3 +32,91 @@ class BaseTransformer(ABC):
     @abstractmethod
     def transform(self, graph: Graph) -> None:
         raise NotImplementedError()
+
+
+class BaseTransformerStandardised(ABC):
+    """Standardised base transformer to use in case a transformer is adding or removing triples from a graph. If you
+    are doing more specialised operations, please overwrite the .transform() method.
+    """
+
+    description: str
+    _use_only_once: bool = False
+    _need_changes: ClassVar[frozenset[str]] = frozenset()
+    _use_iterate_bar_threshold: int = 500
+
+    @abstractmethod
+    def operation(self, query_result_row: ResultRow) -> RowTransformationOutput:
+        """The operations to perform on each row resulting from the ._iterate_query() method.
+        The operation should return a list of triples to add and to remove.
+        """
+        raise NotImplementedError()
+
+    @abstractmethod
+    def _count_query(self) -> str:
+        """
+        Overwrite to fetch all affected properties in the graph as a result of the transformation.
+        Returns:
+            A query string.
+        """
+        raise NotImplementedError()
+
+    @abstractmethod
+    def _iterate_query(self) -> str:
+        """
+        The query to use for extracting target triples from the graph and performing the transformation.
+        Returns:
+            A query string.
+        """
+        raise NotImplementedError()
+
+    def _skip_count_query(self) -> str:
+        """
+        The query to use for extracting target triples from the graph and performing the transformation.
+        Returns:
+            A query string.
+        """
+        return ""
+
+    def transform(self, graph: Graph) -> GraphTransformationResult:
+        outcome = GraphTransformationResult(self.__class__.__name__)
+        outcome.added = outcome.modified = outcome.removed = 0
+
+        iteration_count_res = list(graph.query(self._count_query()))
+        iteration_count = int(iteration_count_res[0][0])  # type: ignore [index, arg-type]
+
+        outcome.affected_nodes_count = iteration_count
+
+        if self._skip_count_query():
+            skipped_count_res = list(graph.query(self._skip_count_query()))
+            skipped_count = int(skipped_count_res[0][0])  # type: ignore [index, arg-type]
+            warnings.warn(
+                NeatValueWarning(f"Skipping {skipped_count} properties in transformation {self.__class__.__name__}"),
+                stacklevel=2,
+            )
+            outcome.skipped = skipped_count
+
+        if iteration_count == 0:
+            return outcome
+
+        result_iterable = graph.query(self._iterate_query())
+        if iteration_count > self._use_iterate_bar_threshold:
+            result_iterable = iterate_progress_bar(  # type: ignore[misc, assignment]
+                result_iterable,
+                total=iteration_count,
+                description=self.description,
+            )
+
+        for row in result_iterable:
+            row = cast(ResultRow, row)
+            row_output = self.operation(row)
+
+            outcome.added += row_output.instances_added_count
+            outcome.removed += row_output.instances_removed_count
+            outcome.modified += row_output.instances_modified_count
+
+            for triple in row_output.add_triples:
+                graph.add(triple)
+            for triple in row_output.remove_triples:
+                graph.remove(triple)
+
+        return outcome

cognite/neat/_graph/transformers/_classic_cdf.py

@@ -21,6 +21,7 @@ from cognite.neat._utils.rdf_ import (
 from ._base import BaseTransformer
 
 
+# TODO: standardise
 class AddAssetDepth(BaseTransformer):
     description: str = "Adds depth of asset in the asset hierarchy to the graph"
     _use_only_once: bool = True
@@ -84,6 +85,7 @@ class AddAssetDepth(BaseTransformer):
             return None
 
 
+# TODO: standardise
 class BaseAssetConnector(BaseTransformer, ABC):
     _asset_type: URIRef = DEFAULT_NAMESPACE.Asset
     _item_type: URIRef
@@ -166,6 +168,7 @@ class AssetEventConnector(BaseAssetConnector):
     _connection_type = DEFAULT_NAMESPACE.event
 
 
+# TODO: standardise
 class AssetRelationshipConnector(BaseTransformer):
     description: str = "Connects assets via relationships"
     _use_only_once: bool = True
@@ -242,6 +245,7 @@ class AssetRelationshipConnector(BaseTransformer):
                     graph.remove((relationship_id, self.relationship_target_xid_prop, None))
 
 
+# TODO: standardise
 class RelationshipAsEdgeTransformer(BaseTransformer):
     """Converts relationships into edges in the graph.
 
@@ -371,7 +375,8 @@ WHERE {{
     ) -> list[Triple]:
         relationship_triples = cast(list[Triple], list(graph.query(f"DESCRIBE <{relationship_id}>")))
         object_by_predicates = cast(
-            dict[str, URIRef | Literal], {remove_namespace_from_uri(row[1]): row[2] for row in relationship_triples}
+            dict[str, URIRef | Literal],
+            {remove_namespace_from_uri(row[1]): row[2] for row in relationship_triples if row[1] != RDF.type},
         )
         source_external_id = cast(URIRef, object_by_predicates["sourceExternalId"])
         target_source_id = cast(URIRef, object_by_predicates["targetExternalId"])