cognite-neat 0.106.0__py3-none-any.whl → 0.108.0__py3-none-any.whl

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

@@ -1,11 +1,16 @@
+import typing
+import warnings
 from collections import defaultdict
 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 Relationship, RelationshipList
-from rdflib import Namespace
+from rdflib import Namespace, URIRef
 
+from cognite.neat._issues.warnings import NeatValueWarning
+from cognite.neat._shared import Triple
 from cognite.neat._utils.auxiliary import create_sha256_hash
 
 from ._base import DEFAULT_SKIP_METADATA_VALUES, ClassicCDFBaseExtractor, InstanceIdPrefix, T_CogniteResource
@@ -28,9 +33,18 @@ class RelationshipsExtractor(ClassicCDFBaseExtractor[Relationship]):
         skip_metadata_values: Set[str] | None = DEFAULT_SKIP_METADATA_VALUES,
         camel_case: bool = True,
         as_write: bool = False,
+        prefix: str | None = None,
+        identifier: typing.Literal["id", "externalId"] = "id",
     ):
+        # This is used by the ClassicExtractor to log the target nodes, such
+        # that it can extract them.
+        # It is private to avoid exposing it to the user.
+        self._target_external_ids_by_type: dict[InstanceIdPrefix, set[str]] = defaultdict(set)
+        self._log_target_nodes = False
+        # Ensure that this becomes an iterator, even if it is a list.
+        to_iterate = (self._log_target_nodes_if_set(item) for item in items)
         super().__init__(
-            items,
+            to_iterate,
             namespace=namespace,
             to_type=to_type,
             total=total,
@@ -39,12 +53,40 @@ class RelationshipsExtractor(ClassicCDFBaseExtractor[Relationship]):
             skip_metadata_values=skip_metadata_values,
             camel_case=camel_case,
             as_write=as_write,
+            prefix=prefix,
+            identifier=identifier,
         )
-        # This is used by the ClassicExtractor to log the target nodes, such
-        # that it can extract them.
-        # It is private to avoid exposing it to the user.
-        self._log_target_nodes = False
-        self._target_external_ids_by_type: dict[InstanceIdPrefix, set[str]] = defaultdict(set)
+        self._uri_by_external_id_by_by_type: dict[InstanceIdPrefix, dict[str, URIRef]] = defaultdict(dict)
+        self._target_triples: list[tuple[URIRef, URIRef, str, str]] = []
+
+    def _log_target_nodes_if_set(self, item: Relationship) -> Relationship:
+        if not self._log_target_nodes:
+            return item
+        if item.target_type and item.target_external_id:
+            self._target_external_ids_by_type[InstanceIdPrefix.from_str(item.target_type)].add(item.target_external_id)
+        return item
+
+    def _item2triples_special_cases(self, id_: URIRef, dumped: dict[str, Any]) -> list[Triple]:
+        if self.identifier == "externalId":
+            return []
+        triples: list[Triple] = []
+        if (source_external_id := dumped.pop("sourceExternalId")) and "sourceType" in dumped:
+            source_type = dumped["sourceType"]
+            try:
+                source_uri = self._uri_by_external_id_by_by_type[InstanceIdPrefix.from_str(source_type)][
+                    source_external_id
+                ]
+            except KeyError:
+                warnings.warn(
+                    NeatValueWarning(f"Missing externalId {source_external_id} for {source_type}"), stacklevel=2
+                )
+            else:
+                triples.append((id_, self.namespace["sourceExternalId"], source_uri))
+        if (target_external_id := dumped.pop("targetExternalId")) and "targetType" in dumped:
+            target_type = dumped["targetType"]
+            # We do not yet have the target nodes, so we log them for later extraction.
+            self._target_triples.append((id_, self.namespace["targetExternalId"], target_type, target_external_id))
+        return triples
 
     @classmethod
     def _from_dataset(
@@ -67,10 +109,6 @@ class RelationshipsExtractor(ClassicCDFBaseExtractor[Relationship]):
         return len(relationships), relationships
 
     def _fallback_id(self, item: Relationship) -> str | None:
-        if item.external_id and item.source_external_id and item.target_external_id:
-            if self._log_target_nodes and item.target_type and item.target_external_id:
-                self._target_external_ids_by_type[InstanceIdPrefix.from_str(item.target_type)].add(
-                    item.target_external_id
-                )
+        if item.external_id:
             return create_sha256_hash(item.external_id)
         return None

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

@@ -1,5 +1,6 @@
 import itertools
 import json
+import typing
 from collections.abc import Callable, Iterable, Set
 from pathlib import Path
 from typing import Any
@@ -52,10 +53,22 @@ class SequencesExtractor(ClassicCDFBaseExtractor[NeatSequence]):
         skip_metadata_values: Set[str] | None = DEFAULT_SKIP_METADATA_VALUES,
         camel_case: bool = True,
         as_write: bool = False,
+        prefix: str | None = None,
+        identifier: typing.Literal["id", "externalId"] = "id",
         unpack_columns: bool = False,
     ):
         super().__init__(
-            items, namespace, to_type, total, limit, unpack_metadata, skip_metadata_values, camel_case, as_write
+            items,
+            namespace,
+            to_type,
+            total,
+            limit,
+            unpack_metadata,
+            skip_metadata_values,
+            camel_case,
+            as_write,
+            prefix,
+            identifier,
         )
         self.unpack_columns = unpack_columns
 
@@ -71,6 +84,8 @@ class SequencesExtractor(ClassicCDFBaseExtractor[NeatSequence]):
         skip_metadata_values: Set[str] | None = DEFAULT_SKIP_METADATA_VALUES,
         camel_case: bool = True,
         as_write: bool = False,
+        prefix: str | None = None,
+        identifier: typing.Literal["id", "externalId"] = "id",
         unpack_columns: bool = False,
     ):
         total, items = cls._handle_no_access(lambda: cls._from_dataset(client, data_set_external_id))
@@ -84,6 +99,8 @@ class SequencesExtractor(ClassicCDFBaseExtractor[NeatSequence]):
             skip_metadata_values,
             camel_case,
             as_write,
+            prefix,
+            identifier,
             unpack_columns,
         )
 
@@ -99,6 +116,8 @@ class SequencesExtractor(ClassicCDFBaseExtractor[NeatSequence]):
         skip_metadata_values: Set[str] | None = DEFAULT_SKIP_METADATA_VALUES,
         camel_case: bool = True,
         as_write: bool = False,
+        prefix: str | None = None,
+        identifier: typing.Literal["id", "externalId"] = "id",
         unpack_columns: bool = False,
     ):
         total, items = cls._handle_no_access(lambda: cls._from_hierarchy(client, root_asset_external_id))
@@ -112,6 +131,8 @@ class SequencesExtractor(ClassicCDFBaseExtractor[NeatSequence]):
             skip_metadata_values,
             camel_case,
             as_write,
+            prefix,
+            identifier,
             unpack_columns,
         )
 
@@ -126,6 +147,8 @@ class SequencesExtractor(ClassicCDFBaseExtractor[NeatSequence]):
         skip_metadata_values: Set[str] | None = DEFAULT_SKIP_METADATA_VALUES,
         camel_case: bool = True,
         as_write: bool = False,
+        prefix: str | None = None,
+        identifier: typing.Literal["id", "externalId"] = "id",
         unpack_columns: bool = False,
     ):
         total, items = cls._from_file(file_path)
@@ -139,6 +162,8 @@ class SequencesExtractor(ClassicCDFBaseExtractor[NeatSequence]):
             skip_metadata_values,
             camel_case,
             as_write,
+            prefix,
+            identifier,
             unpack_columns,
         )
 

cognite/neat/_graph/extractors/_dms.py

@@ -1,44 +1,68 @@
-from collections.abc import Iterable, Iterator
+import urllib.parse
+from collections.abc import Iterable, Iterator, Set
+from functools import cached_property
 from typing import cast
 
 from cognite.client import CogniteClient
 from cognite.client import data_modeling as dm
 from cognite.client.data_classes.data_modeling import DataModelIdentifier
 from cognite.client.data_classes.data_modeling.instances import Instance, PropertyValue
+from cognite.client.utils.useful_types import SequenceNotStr
 from rdflib import RDF, XSD, Literal, Namespace, URIRef
 
-from cognite.neat._constants import DEFAULT_SPACE_URI
+from cognite.neat._config import GLOBAL_CONFIG
+from cognite.neat._constants import DEFAULT_SPACE_URI, is_readonly_property
 from cognite.neat._issues.errors import ResourceRetrievalError
 from cognite.neat._shared import Triple
+from cognite.neat._utils.auxiliary import string_to_ideal_type
+from cognite.neat._utils.collection_ import iterate_progress_bar
 
 from ._base import BaseExtractor
 
+DEFAULT_EMPTY_VALUES = frozenset({"nan", "null", "none", "", " ", "nil", "n/a", "na", "unknown", "undefined"})
+
 
 class DMSExtractor(BaseExtractor):
     """Extract data from Cognite Data Fusion DMS instances into Neat.
 
     Args:
-        items: The items to extract.
-        total: The total number of items to extract. If provided, this will be used to estimate the progress.
+        total_instances_pair_by_view: A dictionary where the key is the view id and the value is a tuple with the total
+            number of instances and an iterable of instances.
         limit: The maximum number of items to extract.
         overwrite_namespace: If provided, this will overwrite the space of the extracted items.
+        unpack_json: If True, JSON objects will be unpacked into RDF literals.
+        empty_values: If unpack_json is True, when unpacking JSON objects, if a key has a value in this set, it will be
+            considered as an empty value and skipped.
+        str_to_ideal_type: If unpack_json is True, when unpacking JSON objects, if the value is a string, the extractor
+            will try to convert it to the ideal type.
     """
 
     def __init__(
         self,
-        items: Iterable[Instance],
-        total: int | None = None,
+        total_instances_pair_by_view: dict[dm.ViewId, tuple[int | None, Iterable[Instance]]],
         limit: int | None = None,
         overwrite_namespace: Namespace | None = None,
+        unpack_json: bool = False,
+        empty_values: Set[str] = DEFAULT_EMPTY_VALUES,
+        str_to_ideal_type: bool = False,
     ) -> None:
-        self.items = items
-        self.total = total
+        self.total_instances_pair_by_view = total_instances_pair_by_view
         self.limit = limit
         self.overwrite_namespace = overwrite_namespace
+        self.unpack_json = unpack_json
+        self.empty_values = empty_values
+        self.str_to_ideal_type = str_to_ideal_type
 
     @classmethod
     def from_data_model(
-        cls, client: CogniteClient, data_model: DataModelIdentifier, limit: int | None = None
+        cls,
+        client: CogniteClient,
+        data_model: DataModelIdentifier,
+        limit: int | None = None,
+        overwrite_namespace: Namespace | None = None,
+        instance_space: str | SequenceNotStr[str] | None = None,
+        unpack_json: bool = False,
+        str_to_ideal_type: bool = False,
     ) -> "DMSExtractor":
         """Create an extractor from a data model.
 
@@ -46,28 +70,79 @@ class DMSExtractor(BaseExtractor):
             client: The Cognite client to use.
             data_model: The data model to extract.
             limit: The maximum number of instances to extract.
+            overwrite_namespace: If provided, this will overwrite the space of the extracted items.
+            instance_space: The space to extract instances from.
+            unpack_json: If True, JSON objects will be unpacked into RDF literals.
         """
         retrieved = client.data_modeling.data_models.retrieve(data_model, inline_views=True)
         if not retrieved:
             raise ResourceRetrievalError(dm.DataModelId.load(data_model), "data model", "Data Model is missing in CDF")
-        return cls.from_views(client, retrieved.latest_version().views, limit)
+        return cls.from_views(
+            client,
+            retrieved.latest_version().views,
+            limit,
+            overwrite_namespace,
+            instance_space,
+            unpack_json,
+            str_to_ideal_type,
+        )
 
     @classmethod
-    def from_views(cls, client: CogniteClient, views: Iterable[dm.View], limit: int | None = None) -> "DMSExtractor":
+    def from_views(
+        cls,
+        client: CogniteClient,
+        views: Iterable[dm.View],
+        limit: int | None = None,
+        overwrite_namespace: Namespace | None = None,
+        instance_space: str | SequenceNotStr[str] | None = None,
+        unpack_json: bool = False,
+        str_to_ideal_type: bool = False,
+    ) -> "DMSExtractor":
         """Create an extractor from a set of views.
 
         Args:
             client: The Cognite client to use.
             views: The views to extract.
             limit: The maximum number of instances to extract.
+            overwrite_namespace: If provided, this will overwrite the space of the extracted items.
+            instance_space: The space to extract instances from.
+            unpack_json: If True, JSON objects will be unpacked into RDF literals.
+            str_to_ideal_type: If True, when unpacking JSON objects, if the value is a string, the extractor will try to
+                convert it to the ideal type.
         """
-        return cls(_InstanceIterator(client, views), total=None, limit=limit)
+        total_instances_pair_by_view: dict[dm.ViewId, tuple[int | None, Iterable[Instance]]] = {}
+        for view in views:
+            instance_iterator = _ViewInstanceIterator(client, view, instance_space)
+            total_instances_pair_by_view[view.as_id()] = (instance_iterator.count, instance_iterator)
+
+        return cls(
+            total_instances_pair_by_view=total_instances_pair_by_view,
+            limit=limit,
+            overwrite_namespace=overwrite_namespace,
+            unpack_json=unpack_json,
+            str_to_ideal_type=str_to_ideal_type,
+        )
 
     def extract(self) -> Iterable[Triple]:
-        for count, item in enumerate(self.items, 1):
-            if self.limit and count > self.limit:
-                break
-            yield from self._extract_instance(item)
+        total_instances = sum(total for total, _ in self.total_instances_pair_by_view.values() if total is not None)
+        use_progress_bar = (
+            GLOBAL_CONFIG.use_iterate_bar_threshold and total_instances > GLOBAL_CONFIG.use_iterate_bar_threshold
+        )
+
+        for view_id, (total, instances) in self.total_instances_pair_by_view.items():
+            if total == 0:
+                continue
+            if use_progress_bar and total is not None:
+                instances = iterate_progress_bar(
+                    instances,
+                    total,
+                    f"Extracting instances from {view_id.space}:{view_id.external_id}(version={view_id.version})",
+                )
+
+            for count, item in enumerate(instances, 1):
+                if self.limit and count > self.limit:
+                    break
+                yield from self._extract_instance(item)
 
     def _extract_instance(self, instance: Instance) -> Iterable[Triple]:
         if isinstance(instance, dm.Edge):
@@ -82,7 +157,6 @@ class DMSExtractor(BaseExtractor):
                 # If the edge has properties, we create a node for the edge and connect it to the start and end nodes.
                 id_ = self._as_uri_ref(instance)
                 yield id_, RDF.type, self._as_uri_ref(instance.type)
-                yield id_, RDF.type, self._get_namespace(instance.space).Edge
                 yield (
                     id_,
                     self._as_uri_ref(dm.DirectRelationReference(instance.space, "startNode")),
@@ -98,6 +172,9 @@ class DMSExtractor(BaseExtractor):
             id_ = self._as_uri_ref(instance)
             if instance.type:
                 type_ = self._as_uri_ref(cast(dm.DirectRelationReference, instance.type))
+            elif len(instance.properties) == 1:
+                view_id = next(iter(instance.properties.keys()))
+                type_ = self._get_namespace(view_id.space)[urllib.parse.quote(view_id.external_id)]
             else:
                 type_ = self._get_namespace(instance.space).Node
 
@@ -105,54 +182,126 @@ class DMSExtractor(BaseExtractor):
         else:
             raise NotImplementedError(f"Unknown instance type {type(instance)}")
 
+        if self.overwrite_namespace:
+            # If the namespace is overwritten, keep the original space as a property to avoid losing information.
+            yield id_, self._get_namespace(instance.space)["space"], Literal(instance.space)
+
         for view_id, properties in instance.properties.items():
             namespace = self._get_namespace(view_id.space)
             for key, value in properties.items():
-                for object_ in self._get_objects(value):
-                    yield id_, namespace[key], object_
+                for predicate_str, object_ in self._get_predicate_objects_pair(key, value):
+                    yield id_, namespace[urllib.parse.quote(predicate_str)], object_
 
-    def _get_objects(self, value: PropertyValue) -> Iterable[Literal | URIRef]:
+    def _get_predicate_objects_pair(self, key: str, value: PropertyValue) -> Iterable[tuple[str, Literal | URIRef]]:
         if isinstance(value, str | float | bool | int):
-            yield Literal(value)
+            yield key, Literal(value)
         elif isinstance(value, dict) and "space" in value and "externalId" in value:
-            yield self._as_uri_ref(dm.DirectRelationReference.load(value))
+            yield key, self._as_uri_ref(dm.DirectRelationReference.load(value))
+        elif isinstance(value, dict) and self.unpack_json:
+            for sub_key, sub_value in value.items():
+                if isinstance(sub_value, str):
+                    if sub_value.casefold() in self.empty_values:
+                        continue
+                    if self.str_to_ideal_type:
+                        yield sub_key, Literal(string_to_ideal_type(sub_value))
+                    else:
+                        yield sub_key, Literal(sub_value)
+                elif isinstance(sub_value, int | float | bool):
+                    yield sub_key, Literal(sub_value)
+                elif isinstance(sub_value, dict):
+                    yield from self._get_predicate_objects_pair(f"{key}_{sub_key}", sub_value)
+                elif isinstance(sub_value, list):
+                    for item in sub_value:
+                        yield from self._get_predicate_objects_pair(f"{key}_{sub_key}", item)
+                else:
+                    yield sub_key, Literal(str(sub_value))
         elif isinstance(value, dict):
             # This object is a json object.
-            yield Literal(str(value), datatype=XSD._NS["json"])
+            yield key, Literal(str(value), datatype=XSD._NS["json"])
         elif isinstance(value, list):
             for item in value:
-                yield from self._get_objects(item)
+                yield from self._get_predicate_objects_pair(key, item)
 
     def _as_uri_ref(self, instance: Instance | dm.DirectRelationReference) -> URIRef:
-        return self._get_namespace(instance.space)[instance.external_id]
+        return self._get_namespace(instance.space)[urllib.parse.quote(instance.external_id)]
 
     def _get_namespace(self, space: str) -> Namespace:
         if self.overwrite_namespace:
             return self.overwrite_namespace
-        return Namespace(DEFAULT_SPACE_URI.format(space=space))
+        return Namespace(DEFAULT_SPACE_URI.format(space=urllib.parse.quote(space)))
 
 
-class _InstanceIterator(Iterator[Instance]):
-    def __init__(self, client: CogniteClient, views: Iterable[dm.View]):
+class _ViewInstanceIterator(Iterable[Instance]):
+    def __init__(self, client: CogniteClient, view: dm.View, instance_space: str | SequenceNotStr[str] | None = None):
         self.client = client
-        self.views = views
+        self.view = view
+        self.instance_space = instance_space
+
+    @cached_property
+    def count(self) -> int:
+        node_count = edge_count = 0
+        if self.view.used_for in ("node", "all"):
+            node_count = int(
+                self.client.data_modeling.instances.aggregate(
+                    view=self.view.as_id(),
+                    aggregates=dm.aggregations.Count("externalId"),
+                    instance_type="node",
+                    space=self.instance_space,
+                ).value
+            )
+        if self.view.used_for in ("edge", "all"):
+            edge_count = int(
+                self.client.data_modeling.instances.aggregate(
+                    view=self.view.as_id(),
+                    aggregates=dm.aggregations.Count("externalId"),
+                    instance_type="edge",
+                    space=self.instance_space,
+                ).value
+            )
+        return node_count + edge_count
 
     def __iter__(self) -> Iterator[Instance]:
-        return self
-
-    def __next__(self) -> Instance:  # type: ignore[misc]
-        for view in self.views:
-            # All nodes and edges with properties
-            yield from self.client.data_modeling.instances(chunk_size=None, instance_type="node", sources=[view])
-            yield from self.client.data_modeling.instances(chunk_size=None, instance_type="edge", sources=[view])
-
-            for prop in view.properties.values():
-                if isinstance(prop, dm.EdgeConnection):
-                    # Get all edges with properties
-                    yield from self.client.data_modeling.instances(
-                        chunk_size=None,
-                        instance_type="edge",
-                        filter=dm.filters.Equals(
-                            ["edge", "type"], {"space": prop.type.space, "externalId": prop.type.external_id}
-                        ),
-                    )
+        view_id = self.view.as_id()
+        read_only_properties = {
+            prop_id
+            for prop_id, prop in self.view.properties.items()
+            if isinstance(prop, dm.MappedProperty)
+            and is_readonly_property(prop.container, prop.container_property_identifier)
+        }
+        # All nodes and edges with properties
+        if self.view.used_for in ("node", "all"):
+            node_iterable: Iterable[Instance] = self.client.data_modeling.instances(
+                chunk_size=None, instance_type="node", sources=[view_id], space=self.instance_space
+            )
+            if read_only_properties:
+                node_iterable = self._remove_read_only_properties(node_iterable, read_only_properties, view_id)
+            yield from node_iterable
+
+        if self.view.used_for in ("edge", "all"):
+            yield from self.client.data_modeling.instances(
+                chunk_size=None, instance_type="edge", sources=[view_id], space=self.instance_space
+            )
+
+        for prop in self.view.properties.values():
+            if isinstance(prop, dm.EdgeConnection):
+                if prop.edge_source:
+                    # All edges with properties are extracted from the edge source
+                    continue
+                yield from self.client.data_modeling.instances(
+                    chunk_size=None,
+                    instance_type="edge",
+                    filter=dm.filters.Equals(
+                        ["edge", "type"], {"space": prop.type.space, "externalId": prop.type.external_id}
+                    ),
+                    space=self.instance_space,
+                )
+
+    @staticmethod
+    def _remove_read_only_properties(
+        nodes: Iterable[Instance], read_only_properties: Set[str], view_id: dm.ViewId
+    ) -> Iterable[Instance]:
+        for node in nodes:
+            if properties := node.properties.get(view_id):
+                for read_only in read_only_properties:
+                    properties.pop(read_only, None)
+            yield node