cognite-neat 0.91.0__py3-none-any.whl → 0.92.1__py3-none-any.whl

cognite/neat/_version.py

@@ -1 +1 @@
-__version__ = "0.91.0"
+__version__ = "0.92.1"

cognite/neat/graph/extractors/__init__.py

@@ -1,5 +1,6 @@
 from ._base import BaseExtractor
 from ._classic_cdf._assets import AssetsExtractor
+from ._classic_cdf._classic import ClassicGraphExtractor
 from ._classic_cdf._events import EventsExtractor
 from ._classic_cdf._files import FilesExtractor
 from ._classic_cdf._labels import LabelsExtractor
@@ -15,6 +16,7 @@ from ._rdf_file import RdfFileExtractor
 __all__ = [
     "BaseExtractor",
     "AssetsExtractor",
+    "ClassicGraphExtractor",
     "MockGraphGenerator",
     "RelationshipsExtractor",
     "TimeSeriesExtractor",

cognite/neat/graph/extractors/_classic_cdf/_assets.py

@@ -8,12 +8,12 @@ from cognite.client.data_classes import Asset, AssetFilter, AssetList
 from rdflib import RDF, Literal, Namespace
 
 from cognite.neat.graph.models import Triple
-from cognite.neat.utils.auxiliary import create_sha256_hash
 
-from ._base import DEFAULT_SKIP_METADATA_VALUES, ClassicCDFExtractor
+from ._base import DEFAULT_SKIP_METADATA_VALUES, ClassicCDFBaseExtractor, InstanceIdPrefix
+from ._labels import LabelsExtractor
 
 
-class AssetsExtractor(ClassicCDFExtractor[Asset]):
+class AssetsExtractor(ClassicCDFBaseExtractor[Asset]):
     """Extract data from Cognite Data Fusions Assets into Neat.
 
     Args:
@@ -108,7 +108,7 @@ class AssetsExtractor(ClassicCDFExtractor[Asset]):
 
     def _item2triples(self, asset: Asset) -> list[Triple]:
         """Converts an asset to triples."""
-        id_ = self.namespace[f"Asset_{asset.id}"]
+        id_ = self.namespace[f"{InstanceIdPrefix.asset}{asset.id}"]
 
         type_ = self._get_rdf_type(asset)
 
@@ -151,7 +151,7 @@ class AssetsExtractor(ClassicCDFExtractor[Asset]):
                     (
                         id_,
                         self.namespace.label,
-                        self.namespace[f"Label_{create_sha256_hash(label.dump()['externalId'])}"],
+                        self.namespace[f"{InstanceIdPrefix.label}{LabelsExtractor._label_id(label)}"],
                     )
                 )
 
@@ -160,17 +160,17 @@ class AssetsExtractor(ClassicCDFExtractor[Asset]):
 
         # Create connections:
         if asset.parent_id:
-            triples.append((id_, self.namespace.parent, self.namespace[f"Asset_{asset.parent_id}"]))
+            triples.append((id_, self.namespace.parent, self.namespace[f"{InstanceIdPrefix.asset}{asset.parent_id}"]))
 
         if asset.root_id:
-            triples.append((id_, self.namespace.root, self.namespace[f"Asset_{asset.root_id}"]))
+            triples.append((id_, self.namespace.root, self.namespace[f"{InstanceIdPrefix.asset}{asset.root_id}"]))
 
         if asset.data_set_id:
             triples.append(
                 (
                     id_,
                     self.namespace.dataset,
-                    self.namespace[f"Dataset_{asset.data_set_id}"],
+                    self.namespace[f"{InstanceIdPrefix.data_set}{asset.data_set_id}"],
                 )
             )
 

cognite/neat/graph/extractors/_classic_cdf/_base.py

@@ -1,5 +1,6 @@
 import json
 import re
+import sys
 from abc import ABC, abstractmethod
 from collections.abc import Callable, Iterable, Set
 from typing import Generic, TypeVar
@@ -16,8 +17,32 @@ T_CogniteResource = TypeVar("T_CogniteResource", bound=CogniteResource)
 
 DEFAULT_SKIP_METADATA_VALUES = frozenset({"nan", "null", "none", ""})
 
+if sys.version_info >= (3, 11):
+    from enum import StrEnum
+else:
+    from backports.strenum import StrEnum
 
-class ClassicCDFExtractor(BaseExtractor, ABC, Generic[T_CogniteResource]):
+
+class InstanceIdPrefix(StrEnum):
+    asset = "Asset_"
+    label = "Label_"
+    relationship = "Relationship_"
+    sequence = "Sequence_"
+    file = "File_"
+    time_series = "TimeSeries_"
+    event = "Event_"
+    data_set = "DataSet_"
+
+    @classmethod
+    def from_str(cls, raw: str) -> "InstanceIdPrefix":
+        raw = raw.title() + "_"
+        if raw == "Timeseries_":
+            return cls.time_series
+        else:
+            return cls(raw)
+
+
+class ClassicCDFBaseExtractor(BaseExtractor, ABC, Generic[T_CogniteResource]):
     """This is the Base Extractor for all classic CDF resources.
 
     A classic resource is recognized in that it has a metadata attribute of type dict[str, str].

cognite/neat/graph/extractors/_classic_cdf/_classic.py

@@ -0,0 +1,208 @@
+from collections import defaultdict
+from collections.abc import Iterable, Sequence
+from typing import ClassVar, NamedTuple
+
+from cognite.client import CogniteClient
+from rdflib import Namespace
+
+from cognite.neat.constants import DEFAULT_NAMESPACE
+from cognite.neat.graph.extractors._base import BaseExtractor
+from cognite.neat.graph.models import Triple
+from cognite.neat.utils.collection_ import chunker
+from cognite.neat.utils.rdf_ import remove_namespace_from_uri
+
+from ._assets import AssetsExtractor
+from ._base import ClassicCDFBaseExtractor, InstanceIdPrefix
+from ._data_sets import DataSetExtractor
+from ._events import EventsExtractor
+from ._files import FilesExtractor
+from ._labels import LabelsExtractor
+from ._relationships import RelationshipsExtractor
+from ._sequences import SequencesExtractor
+from ._timeseries import TimeSeriesExtractor
+
+
+class _ClassicCoreType(NamedTuple):
+    extractor_cls: (
+        type[AssetsExtractor]
+        | type[TimeSeriesExtractor]
+        | type[SequencesExtractor]
+        | type[EventsExtractor]
+        | type[FilesExtractor]
+    )
+    resource_type: InstanceIdPrefix
+    api_name: str
+
+
+class ClassicGraphExtractor(BaseExtractor):
+    """This extractor extracts all classic CDF Resources.
+
+    The Classic Graph consists of the following core resource type.
+
+    Classic Node CDF Resources:
+     - Assets
+     - TimeSeries
+     - Sequences
+     - Events
+     - Files
+
+    All the classic node CDF resources can have one or more connections to one or more assets. This
+    will match a direct relationship in the data modeling of CDF.
+
+    In addition, you have relationships between the classic node CDF resources. This matches an edge
+    in the data modeling of CDF.
+
+    Finally, you have labels and data sets that to organize the graph. In which data sets have a similar,
+    but different, role as a space in data modeling. While labels can be compared to node types in data modeling,
+    used to quickly filter and find nodes/edges.
+
+    This extractor will extract the classic CDF graph into Neat starting from either a data set or a root asset.
+
+    It works as follows:
+
+    1. Extract all core nodes (assets, time series, sequences, events, files) filtered by the given data set or
+       root asset.
+    2. Extract all relationships starting from any of the extracted core nodes.
+    3. Extract all core nodes that are targets of the relationships that are not already extracted.
+    4. Extract all labels that are connected to the extracted core nodes/relationships.
+    5. Extract all data sets that are connected to the extracted core nodes/relationships.
+
+    Args:
+        client (CogniteClient): The Cognite client to use.
+        data_set_external_id (str, optional): The data set external id to extract from. Defaults to None.
+        root_asset_external_id (str, optional): The root asset external id to extract from. Defaults to None.
+        namespace (Namespace, optional): The namespace to use. Defaults to DEFAULT_NAMESPACE.
+    """
+
+    # These are the core resource types in the classic CDF.
+    _classic_node_types: ClassVar[tuple[_ClassicCoreType, ...]] = (
+        _ClassicCoreType(AssetsExtractor, InstanceIdPrefix.asset, "assets"),
+        _ClassicCoreType(TimeSeriesExtractor, InstanceIdPrefix.time_series, "time_series"),
+        _ClassicCoreType(SequencesExtractor, InstanceIdPrefix.sequence, "sequences"),
+        _ClassicCoreType(EventsExtractor, InstanceIdPrefix.event, "events"),
+        _ClassicCoreType(FilesExtractor, InstanceIdPrefix.file, "files"),
+    )
+
+    def __init__(
+        self,
+        client: CogniteClient,
+        data_set_external_id: str | None = None,
+        root_asset_external_id: str | None = None,
+        namespace: Namespace | None = None,
+    ):
+        self._client = client
+        if sum([bool(data_set_external_id), bool(root_asset_external_id)]) != 1:
+            raise ValueError("Exactly one of data_set_external_id or root_asset_external_id must be set.")
+        self._root_asset_external_id = root_asset_external_id
+        self._data_set_external_id = data_set_external_id
+        self._namespace = namespace or DEFAULT_NAMESPACE
+
+        self._source_external_ids_by_type: dict[InstanceIdPrefix, set[str]] = defaultdict(set)
+        self._target_external_ids_by_type: dict[InstanceIdPrefix, set[str]] = defaultdict(set)
+        self._labels: set[str] = set()
+        self._data_set_ids: set[int] = set()
+
+    def extract(self) -> Iterable[Triple]:
+        """Extracts all classic CDF Resources."""
+        yield from self._extract_core_start_nodes()
+
+        yield from self._extract_start_node_relationships()
+
+        yield from self._extract_core_end_nodes()
+
+        yield from self._extract_labels()
+        yield from self._extract_data_sets()
+
+    def _extract_core_start_nodes(self):
+        for core_node in self._classic_node_types:
+            if self._data_set_external_id:
+                extractor = core_node.extractor_cls.from_dataset(
+                    self._client, self._data_set_external_id, self._namespace, unpack_metadata=False
+                )
+            elif self._root_asset_external_id:
+                extractor = core_node.extractor_cls.from_hierarchy(
+                    self._client, self._root_asset_external_id, self._namespace, unpack_metadata=False
+                )
+            else:
+                raise ValueError("Exactly one of data_set_external_id or root_asset_external_id must be set.")
+
+            yield from self._extract_with_logging_label_dataset(extractor, core_node.resource_type)
+
+    def _extract_start_node_relationships(self):
+        for start_resource_type, source_external_ids in self._source_external_ids_by_type.items():
+            start_type = start_resource_type.removesuffix("_")
+            for chunk in self._chunk(list(source_external_ids), description=f"Extracting {start_type} relationships"):
+                relationship_iterator = self._client.relationships(
+                    source_external_ids=list(chunk), source_types=[start_type]
+                )
+                extractor = RelationshipsExtractor(relationship_iterator, self._namespace, unpack_metadata=False)
+                # This is a private attribute, but we need to set it to log the target nodes.
+                extractor._log_target_nodes = True
+
+                yield from extractor.extract()
+
+                # After the extraction is done, we need to update all the new target nodes so
+                # we can extract them in the next step.
+                for end_type, target_external_ids in extractor._target_external_ids_by_type.items():
+                    for external_id in target_external_ids:
+                        # We only want to extract the target nodes that are not already extracted.
+                        # Even though _source_external_ids_by_type is a defaultdict, we have to check if the key exists.
+                        # This is because we might not have extracted any nodes of that type yet, and looking up
+                        # a key that does not exist will create it. We are iterating of this dictionary, and
+                        # we do not want to create new keys while iterating.
+                        if (
+                            end_type not in self._source_external_ids_by_type
+                            or external_id not in self._source_external_ids_by_type[end_type]
+                        ):
+                            self._target_external_ids_by_type[end_type].add(external_id)
+
+    def _extract_core_end_nodes(self):
+        for core_node in self._classic_node_types:
+            target_external_ids = self._target_external_ids_by_type[core_node.resource_type]
+            api = getattr(self._client, core_node.api_name)
+            for chunk in self._chunk(
+                list(target_external_ids),
+                description=f"Extracting end nodes {core_node.resource_type.removesuffix('_')}",
+            ):
+                resource_iterator = api.retrieve_multiple(external_ids=list(chunk), ignore_unknown_ids=True)
+                extractor = core_node.extractor_cls(resource_iterator, self._namespace, unpack_metadata=False)
+                yield from self._extract_with_logging_label_dataset(extractor)
+
+    def _extract_labels(self):
+        for chunk in self._chunk(list(self._labels), description="Extracting labels"):
+            label_iterator = self._client.labels.retrieve(external_id=list(chunk), ignore_unknown_ids=True)
+            yield from LabelsExtractor(label_iterator, self._namespace).extract()
+
+    def _extract_data_sets(self):
+        for chunk in self._chunk(list(self._data_set_ids), description="Extracting data sets"):
+            data_set_iterator = self._client.data_sets.retrieve_multiple(ids=list(chunk), ignore_unknown_ids=True)
+            yield from DataSetExtractor(data_set_iterator, self._namespace, unpack_metadata=False).extract()
+
+    def _extract_with_logging_label_dataset(
+        self, extractor: ClassicCDFBaseExtractor, resource_type: InstanceIdPrefix | None = None
+    ) -> Iterable[Triple]:
+        for triple in extractor.extract():
+            if triple[1] == self._namespace.external_id and resource_type is not None:
+                self._source_external_ids_by_type[resource_type].add(remove_namespace_from_uri(triple[2]))
+            elif triple[1] == self._namespace.label:
+                self._labels.add(remove_namespace_from_uri(triple[2]).removeprefix(InstanceIdPrefix.label))
+            elif triple[1] == self._namespace.dataset:
+                self._data_set_ids.add(
+                    int(remove_namespace_from_uri(triple[2]).removeprefix(InstanceIdPrefix.data_set))
+                )
+            yield triple
+
+    @staticmethod
+    def _chunk(items: Sequence, description: str) -> Iterable:
+        to_iterate: Iterable = chunker(items, chunk_size=1000)
+        try:
+            from rich.progress import track
+        except ModuleNotFoundError:
+            ...
+        else:
+            to_iterate = track(
+                to_iterate,
+                total=(len(items) // 1000) + 1,
+                description=description,
+            )
+        return to_iterate

cognite/neat/graph/extractors/_classic_cdf/_data_sets.py

@@ -0,0 +1,110 @@
+from collections.abc import Set
+from datetime import datetime, timezone
+from pathlib import Path
+
+from cognite.client import CogniteClient
+from cognite.client.data_classes import DataSet, DataSetList
+from cognite.client.utils.useful_types import SequenceNotStr
+from rdflib import RDF, Literal, Namespace
+
+from cognite.neat.graph.models import Triple
+
+from ._base import DEFAULT_SKIP_METADATA_VALUES, ClassicCDFBaseExtractor, InstanceIdPrefix
+
+
+class DataSetExtractor(ClassicCDFBaseExtractor[DataSet]):
+    """Extract DataSets from Cognite Data Fusions into Neat.
+
+    Args:
+        items (Iterable[Asset]): An iterable of assets.
+        namespace (Namespace, optional): The namespace to use. Defaults to DEFAULT_NAMESPACE.
+        to_type (Callable[[Asset], str | None], optional): A function to convert an asset to a type. Defaults to None.
+            If None or if the function returns None, the asset will be set to the default type "Asset".
+        total (int, optional): The total number of assets to load. If passed, you will get a progress bar if rich
+            is installed. Defaults to None.
+        limit (int, optional): The maximal number of assets 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): A set of values to skip when unpacking
+            metadata. Defaults to frozenset({"nan", "null", "none", ""}).
+    """
+
+    _default_rdf_type = "DataSet"
+
+    @classmethod
+    def from_dataset(
+        cls,
+        client: CogniteClient,
+        data_set_external_id: SequenceNotStr[str],
+        namespace: Namespace | None = None,
+        unpack_metadata: bool = True,
+        skip_metadata_values: Set[str] | None = DEFAULT_SKIP_METADATA_VALUES,
+    ):
+        return cls(
+            client.data_sets.retrieve_multiple(external_ids=data_set_external_id),
+            namespace=namespace,
+            total=len(data_set_external_id),
+            unpack_metadata=unpack_metadata,
+            skip_metadata_values=skip_metadata_values,
+        )
+
+    @classmethod
+    def from_file(
+        cls,
+        file_path: str,
+        namespace: Namespace | None = None,
+        unpack_metadata: bool = True,
+        skip_metadata_values: Set[str] | None = DEFAULT_SKIP_METADATA_VALUES,
+    ):
+        data_sets = DataSetList.load(Path(file_path).read_text())
+        return cls(
+            data_sets,
+            namespace=namespace,
+            total=len(data_sets),
+            unpack_metadata=unpack_metadata,
+            skip_metadata_values=skip_metadata_values,
+        )
+
+    def _item2triples(self, item: DataSet) -> list[Triple]:
+        """Converts an asset to triples."""
+        id_ = self.namespace[f"{InstanceIdPrefix.data_set}{item.id}"]
+
+        type_ = self._get_rdf_type(item)
+
+        triples: list[Triple] = [(id_, RDF.type, self.namespace[type_])]
+
+        # Create attributes
+        if item.name:
+            triples.append((id_, self.namespace.name, Literal(item.name)))
+
+        if item.description:
+            triples.append((id_, self.namespace.description, Literal(item.description)))
+
+        if item.external_id:
+            triples.append((id_, self.namespace.external_id, Literal(item.external_id)))
+
+        # properties' ref creation and update
+        triples.append(
+            (
+                id_,
+                self.namespace.created_time,
+                Literal(datetime.fromtimestamp(item.created_time / 1000, timezone.utc)),
+            )
+        )
+        triples.append(
+            (
+                id_,
+                self.namespace.last_updated_time,
+                Literal(datetime.fromtimestamp(item.last_updated_time / 1000, timezone.utc)),
+            )
+        )
+
+        if item.write_protected:
+            triples.append((id_, self.namespace.write_protected, Literal(item.write_protected)))
+
+        if item.metadata:
+            triples.extend(self._metadata_to_triples(id_, item.metadata))
+
+        return triples

cognite/neat/graph/extractors/_classic_cdf/_events.py

@@ -8,10 +8,10 @@ from rdflib import RDF, Literal, Namespace
 
 from cognite.neat.graph.models import Triple
 
-from ._base import DEFAULT_SKIP_METADATA_VALUES, ClassicCDFExtractor
+from ._base import DEFAULT_SKIP_METADATA_VALUES, ClassicCDFBaseExtractor, InstanceIdPrefix
 
 
-class EventsExtractor(ClassicCDFExtractor[Event]):
+class EventsExtractor(ClassicCDFBaseExtractor[Event]):
     """Extract data from Cognite Data Fusions Events into Neat.
 
     Args:
@@ -55,6 +55,31 @@ class EventsExtractor(ClassicCDFExtractor[Event]):
             skip_metadata_values=skip_metadata_values,
         )
 
+    @classmethod
+    def from_hierarchy(
+        cls,
+        client: CogniteClient,
+        root_asset_external_id: str,
+        namespace: Namespace | None = None,
+        to_type: Callable[[Event], str | None] | None = None,
+        limit: int | None = None,
+        unpack_metadata: bool = True,
+        skip_metadata_values: Set[str] | None = DEFAULT_SKIP_METADATA_VALUES,
+    ):
+        total = client.events.aggregate_count(
+            filter=EventFilter(asset_subtree_ids=[{"externalId": root_asset_external_id}])
+        )
+
+        return cls(
+            client.events(asset_subtree_external_ids=[root_asset_external_id]),
+            namespace,
+            to_type,
+            total,
+            limit,
+            unpack_metadata=unpack_metadata,
+            skip_metadata_values=skip_metadata_values,
+        )
+
     @classmethod
     def from_file(
         cls,
@@ -78,7 +103,7 @@ class EventsExtractor(ClassicCDFExtractor[Event]):
         )
 
     def _item2triples(self, event: Event) -> list[Triple]:
-        id_ = self.namespace[f"Event_{event.id}"]
+        id_ = self.namespace[f"{InstanceIdPrefix.event}{event.id}"]
 
         type_ = self._get_rdf_type(event)
 
@@ -146,12 +171,12 @@ class EventsExtractor(ClassicCDFExtractor[Event]):
                 (
                     id_,
                     self.namespace.data_set_id,
-                    self.namespace[f"Dataset_{event.data_set_id}"],
+                    self.namespace[f"{InstanceIdPrefix.data_set}{event.data_set_id}"],
                 )
             )
 
         if event.asset_ids:
             for asset_id in event.asset_ids:
-                triples.append((id_, self.namespace.asset, self.namespace[f"Asset_{asset_id}"]))
+                triples.append((id_, self.namespace.asset, self.namespace[f"{InstanceIdPrefix.asset}{asset_id}"]))
 
         return triples

cognite/neat/graph/extractors/_classic_cdf/_files.py

@@ -1,18 +1,18 @@
 from collections.abc import Callable, Set
 from datetime import datetime, timezone
 from pathlib import Path
-from urllib.parse import quote
 
 from cognite.client import CogniteClient
-from cognite.client.data_classes import FileMetadata, FileMetadataList
+from cognite.client.data_classes import FileMetadata, FileMetadataFilter, FileMetadataList
 from rdflib import RDF, Literal, Namespace
 
 from cognite.neat.graph.models import Triple
 
-from ._base import DEFAULT_SKIP_METADATA_VALUES, ClassicCDFExtractor
+from ._base import DEFAULT_SKIP_METADATA_VALUES, ClassicCDFBaseExtractor, InstanceIdPrefix
+from ._labels import LabelsExtractor
 
 
-class FilesExtractor(ClassicCDFExtractor[FileMetadata]):
+class FilesExtractor(ClassicCDFBaseExtractor[FileMetadata]):
     """Extract data from Cognite Data Fusions files metadata into Neat.
 
     Args:
@@ -53,6 +53,31 @@ class FilesExtractor(ClassicCDFExtractor[FileMetadata]):
             skip_metadata_values=skip_metadata_values,
         )
 
+    @classmethod
+    def from_hierarchy(
+        cls,
+        client: CogniteClient,
+        root_asset_external_id: str,
+        namespace: Namespace | None = None,
+        to_type: Callable[[FileMetadata], str | None] | None = None,
+        limit: int | None = None,
+        unpack_metadata: bool = True,
+        skip_metadata_values: Set[str] | None = DEFAULT_SKIP_METADATA_VALUES,
+    ):
+        total = client.files.aggregate(
+            filter=FileMetadataFilter(asset_subtree_ids=[{"externalId": root_asset_external_id}])
+        )[0].count
+
+        return cls(
+            client.files(asset_subtree_external_ids=[root_asset_external_id]),
+            namespace,
+            to_type,
+            total,
+            limit,
+            unpack_metadata=unpack_metadata,
+            skip_metadata_values=skip_metadata_values,
+        )
+
     @classmethod
     def from_file(
         cls,
@@ -75,7 +100,7 @@ class FilesExtractor(ClassicCDFExtractor[FileMetadata]):
         )
 
     def _item2triples(self, file: FileMetadata) -> list[Triple]:
-        id_ = self.namespace[f"File_{file.id}"]
+        id_ = self.namespace[f"{InstanceIdPrefix.file}{file.id}"]
 
         type_ = self._get_rdf_type(file)
 
@@ -153,7 +178,7 @@ class FilesExtractor(ClassicCDFExtractor[FileMetadata]):
                     (
                         id_,
                         self.namespace.label,
-                        self.namespace[f"Label_{quote(label.dump()['externalId'])}"],
+                        self.namespace[f"{InstanceIdPrefix.label}{LabelsExtractor._label_id(label)}"],
                     )
                 )
 
@@ -166,12 +191,12 @@ class FilesExtractor(ClassicCDFExtractor[FileMetadata]):
                 (
                     id_,
                     self.namespace.data_set_id,
-                    self.namespace[f"Dataset_{file.data_set_id}"],
+                    self.namespace[f"{InstanceIdPrefix.data_set}{file.data_set_id}"],
                 )
             )
 
         if file.asset_ids:
             for asset_id in file.asset_ids:
-                triples.append((id_, self.namespace.asset, self.namespace[f"Asset_{asset_id}"]))
+                triples.append((id_, self.namespace.asset, self.namespace[f"{InstanceIdPrefix.asset}{asset_id}"]))
 
         return triples

cognite/neat/graph/extractors/_classic_cdf/_labels.py

@@ -1,18 +1,18 @@
 from collections.abc import Callable, Set
 from datetime import datetime, timezone
 from pathlib import Path
+from urllib.parse import quote
 
 from cognite.client import CogniteClient
-from cognite.client.data_classes import LabelDefinition, LabelDefinitionList
+from cognite.client.data_classes import Label, LabelDefinition, LabelDefinitionList
 from rdflib import RDF, Literal, Namespace
 
 from cognite.neat.graph.models import Triple
-from cognite.neat.utils.auxiliary import create_sha256_hash
 
-from ._base import DEFAULT_SKIP_METADATA_VALUES, ClassicCDFExtractor
+from ._base import DEFAULT_SKIP_METADATA_VALUES, ClassicCDFBaseExtractor, InstanceIdPrefix
 
 
-class LabelsExtractor(ClassicCDFExtractor[LabelDefinition]):
+class LabelsExtractor(ClassicCDFBaseExtractor[LabelDefinition]):
     """Extract data from Cognite Data Fusions Labels into Neat.
 
     Args:
@@ -78,7 +78,7 @@ class LabelsExtractor(ClassicCDFExtractor[LabelDefinition]):
         if not label.external_id:
             return []
 
-        id_ = self.namespace[f"Label_{create_sha256_hash(label.external_id)}"]
+        id_ = self.namespace[f"{InstanceIdPrefix.label}{self._label_id(label)}"]
 
         type_ = self._get_rdf_type(label)
         # Set rdf type
@@ -107,8 +107,16 @@ class LabelsExtractor(ClassicCDFExtractor[LabelDefinition]):
                 (
                     id_,
                     self.namespace.data_set_id,
-                    self.namespace[f"Dataset_{label.data_set_id}"],
+                    self.namespace[f"{InstanceIdPrefix.data_set}{label.data_set_id}"],
                 )
             )
 
         return triples
+
+    @staticmethod
+    def _label_id(label: Label | LabelDefinition) -> str:
+        # external_id can create ill-formed URIs, so we create websafe URIs
+        # since labels do not have internal ids, we use the external_id as the id
+        if label.external_id is None:
+            raise ValueError("External id must be set of the label")
+        return quote(label.external_id)

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

@@ -1,7 +1,7 @@
-from collections.abc import Callable, Set
+from collections import defaultdict
+from collections.abc import Callable, Iterable, Set
 from datetime import datetime, timezone
 from pathlib import Path
-from urllib.parse import quote
 
 from cognite.client import CogniteClient
 from cognite.client.data_classes import Relationship, RelationshipList
@@ -10,10 +10,11 @@ from rdflib import RDF, Literal, Namespace
 from cognite.neat.graph.models import Triple
 from cognite.neat.utils.auxiliary import create_sha256_hash
 
-from ._base import DEFAULT_SKIP_METADATA_VALUES, ClassicCDFExtractor
+from ._base import DEFAULT_SKIP_METADATA_VALUES, ClassicCDFBaseExtractor, InstanceIdPrefix
+from ._labels import LabelsExtractor
 
 
-class RelationshipsExtractor(ClassicCDFExtractor[Relationship]):
+class RelationshipsExtractor(ClassicCDFBaseExtractor[Relationship]):
     """Extract data from Cognite Data Fusions Relationships into Neat.
 
     Args:
@@ -34,6 +35,31 @@ class RelationshipsExtractor(ClassicCDFExtractor[Relationship]):
 
     _default_rdf_type = "Relationship"
 
+    def __init__(
+        self,
+        items: Iterable[Relationship],
+        namespace: Namespace | None = None,
+        to_type: Callable[[Relationship], 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,
+    ):
+        super().__init__(
+            items,
+            namespace=namespace,
+            to_type=to_type,
+            total=total,
+            limit=limit,
+            unpack_metadata=unpack_metadata,
+            skip_metadata_values=skip_metadata_values,
+        )
+        # 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)
+
     @classmethod
     def from_dataset(
         cls,
@@ -79,8 +105,13 @@ class RelationshipsExtractor(ClassicCDFExtractor[Relationship]):
         """Converts an asset to triples."""
 
         if relationship.external_id and relationship.source_external_id and relationship.target_external_id:
+            if self._log_target_nodes and relationship.target_type and relationship.target_external_id:
+                self._target_external_ids_by_type[InstanceIdPrefix.from_str(relationship.target_type)].add(
+                    relationship.target_external_id
+                )
+
             # relationships do not have an internal id, so we generate one
-            id_ = self.namespace[f"Relationship_{create_sha256_hash(relationship.external_id)}"]
+            id_ = self.namespace[f"{InstanceIdPrefix.relationship}{create_sha256_hash(relationship.external_id)}"]
 
             type_ = self._get_rdf_type(relationship)
             # Set rdf type
@@ -178,7 +209,7 @@ class RelationshipsExtractor(ClassicCDFExtractor[Relationship]):
                         (
                             id_,
                             self.namespace.label,
-                            self.namespace[f"Label_{quote(label.dump()['externalId'])}"],
+                            self.namespace[f"{InstanceIdPrefix.label}{LabelsExtractor._label_id(label)}"],
                         )
                     )
 
@@ -188,7 +219,7 @@ class RelationshipsExtractor(ClassicCDFExtractor[Relationship]):
                     (
                         id_,
                         self.namespace.dataset,
-                        self.namespace[f"Dataset_{relationship.data_set_id}"],
+                        self.namespace[f"{InstanceIdPrefix.data_set}{relationship.data_set_id}"],
                     )
                 )
 

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

@@ -8,10 +8,10 @@ from rdflib import RDF, Literal, Namespace
 
 from cognite.neat.graph.models import Triple
 
-from ._base import DEFAULT_SKIP_METADATA_VALUES, ClassicCDFExtractor
+from ._base import DEFAULT_SKIP_METADATA_VALUES, ClassicCDFBaseExtractor, InstanceIdPrefix
 
 
-class SequencesExtractor(ClassicCDFExtractor[Sequence]):
+class SequencesExtractor(ClassicCDFBaseExtractor[Sequence]):
     """Extract data from Cognite Data Fusions Sequences into Neat.
 
     Args:
@@ -56,6 +56,31 @@ class SequencesExtractor(ClassicCDFExtractor[Sequence]):
             skip_metadata_values=skip_metadata_values,
         )
 
+    @classmethod
+    def from_hierarchy(
+        cls,
+        client: CogniteClient,
+        root_asset_external_id: str,
+        namespace: Namespace | None = None,
+        to_type: Callable[[Sequence], str | None] | None = None,
+        limit: int | None = None,
+        unpack_metadata: bool = True,
+        skip_metadata_values: Set[str] | None = DEFAULT_SKIP_METADATA_VALUES,
+    ):
+        total = client.sequences.aggregate_count(
+            filter=SequenceFilter(asset_subtree_ids=[{"externalId": root_asset_external_id}])
+        )
+
+        return cls(
+            client.sequences(asset_subtree_external_ids=[root_asset_external_id]),
+            namespace,
+            to_type,
+            total,
+            limit,
+            unpack_metadata=unpack_metadata,
+            skip_metadata_values=skip_metadata_values,
+        )
+
     @classmethod
     def from_file(
         cls,
@@ -78,7 +103,7 @@ class SequencesExtractor(ClassicCDFExtractor[Sequence]):
         )
 
     def _item2triples(self, sequence: Sequence) -> list[Triple]:
-        id_ = self.namespace[f"Sequence_{sequence.id}"]
+        id_ = self.namespace[f"{InstanceIdPrefix.sequence}{sequence.id}"]
 
         type_ = self._get_rdf_type(sequence)
         # Set rdf type
@@ -121,7 +146,7 @@ class SequencesExtractor(ClassicCDFExtractor[Sequence]):
                 (
                     id_,
                     self.namespace.data_set_id,
-                    self.namespace[f"Dataset_{sequence.data_set_id}"],
+                    self.namespace[f"{InstanceIdPrefix.data_set}{sequence.data_set_id}"],
                 )
             )
 
@@ -130,7 +155,7 @@ class SequencesExtractor(ClassicCDFExtractor[Sequence]):
                 (
                     id_,
                     self.namespace.asset,
-                    self.namespace[f"Asset_{sequence.asset_id}"],
+                    self.namespace[f"{InstanceIdPrefix.asset}{sequence.asset_id}"],
                 )
             )
 

cognite/neat/graph/extractors/_classic_cdf/_timeseries.py

@@ -9,10 +9,10 @@ from rdflib import RDF, Literal, Namespace, URIRef
 
 from cognite.neat.graph.models import Triple
 
-from ._base import DEFAULT_SKIP_METADATA_VALUES, ClassicCDFExtractor
+from ._base import DEFAULT_SKIP_METADATA_VALUES, ClassicCDFBaseExtractor, InstanceIdPrefix
 
 
-class TimeSeriesExtractor(ClassicCDFExtractor[TimeSeries]):
+class TimeSeriesExtractor(ClassicCDFBaseExtractor[TimeSeries]):
     """Extract data from Cognite Data Fusions TimeSeries into Neat.
 
     Args:
@@ -58,6 +58,31 @@ class TimeSeriesExtractor(ClassicCDFExtractor[TimeSeries]):
             skip_metadata_values=skip_metadata_values,
         )
 
+    @classmethod
+    def from_hierarchy(
+        cls,
+        client: CogniteClient,
+        root_asset_external_id: str,
+        namespace: Namespace | None = None,
+        to_type: Callable[[TimeSeries], str | None] | None = None,
+        limit: int | None = None,
+        unpack_metadata: bool = True,
+        skip_metadata_values: Set[str] | None = DEFAULT_SKIP_METADATA_VALUES,
+    ):
+        total = client.time_series.aggregate_count(
+            filter=TimeSeriesFilter(asset_subtree_ids=[{"externalId": root_asset_external_id}])
+        )
+
+        return cls(
+            client.time_series(asset_external_ids=[root_asset_external_id]),
+            namespace,
+            to_type,
+            total,
+            limit,
+            unpack_metadata=unpack_metadata,
+            skip_metadata_values=skip_metadata_values,
+        )
+
     @classmethod
     def from_file(
         cls,
@@ -80,7 +105,7 @@ class TimeSeriesExtractor(ClassicCDFExtractor[TimeSeries]):
         )
 
     def _item2triples(self, timeseries: TimeSeries) -> list[Triple]:
-        id_ = self.namespace[f"TimeSeries_{timeseries.id}"]
+        id_ = self.namespace[f"{InstanceIdPrefix.time_series}{timeseries.id}"]
 
         # Set rdf type
         type_ = self._get_rdf_type(timeseries)
@@ -158,7 +183,7 @@ class TimeSeriesExtractor(ClassicCDFExtractor[TimeSeries]):
                 (
                     id_,
                     self.namespace.dataset,
-                    self.namespace[f"Dataset_{timeseries.data_set_id}"],
+                    self.namespace[f"{InstanceIdPrefix.data_set}{timeseries.data_set_id}"],
                 )
             )
 
@@ -167,7 +192,7 @@ class TimeSeriesExtractor(ClassicCDFExtractor[TimeSeries]):
                 (
                     id_,
                     self.namespace.asset,
-                    self.namespace[f"Asset_{timeseries.asset_id}"],
+                    self.namespace[f"{InstanceIdPrefix.asset}{timeseries.asset_id}"],
                 )
             )
 

cognite/neat/utils/collection_.py

@@ -1,5 +1,5 @@
 from collections import Counter
-from collections.abc import Sequence
+from collections.abc import Iterable, Sequence
 from typing import TypeVar
 
 T_Element = TypeVar("T_Element")
@@ -14,5 +14,6 @@ def most_occurring_element(list_of_elements: list[T_Element]) -> T_Element:
     return counts.most_common(1)[0][0]
 
 
-def chunker(sequence: Sequence[T_Element], chunk_size: int) -> list[Sequence[T_Element]]:
-    return [sequence[i : i + chunk_size] for i in range(0, len(sequence), chunk_size)]
+def chunker(sequence: Sequence[T_Element], chunk_size: int) -> Iterable[Sequence[T_Element]]:
+    for i in range(0, len(sequence), chunk_size):
+        yield sequence[i : i + chunk_size]

{cognite_neat-0.91.0.dist-info → cognite_neat-0.92.1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: cognite-neat
-Version: 0.91.0
+Version: 0.92.1
 Summary: Knowledge graph transformation
 Home-page: https://cognite-neat.readthedocs-hosted.com/
 License: Apache-2.0
@@ -59,7 +59,7 @@ Description-Content-Type: text/markdown
 
 # kNowlEdge grAph Transformer (NEAT)
 
-[![release](https://img.shields.io/github/actions/workflow/status/cognitedata/neat/release_pypi.yaml?style=for-the-badge)](https://github.com/cognitedata/neat/actions/workflows/release_pypi.yaml)
+[![release](https://img.shields.io/github/actions/workflow/status/cognitedata/neat/release.yaml?style=for-the-badge)](https://github.com/cognitedata/neat/actions/workflows/release.yaml)
 [![Documentation Status](https://readthedocs.com/projects/cognite-neat/badge/?version=latest&style=for-the-badge)](https://cognite-neat.readthedocs-hosted.com/en/latest/?badge=latest)
 [![Github](https://shields.io/badge/github-cognite/neat-green?logo=github&style=for-the-badge)](https://github.com/cognitedata/neat)
 [![PyPI](https://img.shields.io/pypi/v/cognite-neat?style=for-the-badge)](https://pypi.org/project/cognite-neat/)

{cognite_neat-0.91.0.dist-info → cognite_neat-0.92.1.dist-info}/RECORD

@@ -1,6 +1,6 @@
 cognite/neat/__init__.py,sha256=AiexNcHdAHFbrrbo9c65gtil1dqx_SGraDH1PSsXjKE,126
 cognite/neat/_shared.py,sha256=RSaHm2eJceTlvb-hMMe4nHgoHdPYDfN3XcxDXo24k3A,1530
-cognite/neat/_version.py,sha256=RBon26SFk4Bs08h0XAmTYW8lrYs3N9ihifF9r6VT4Vc,23
+cognite/neat/_version.py,sha256=UAHzlES5E7AnCRpSS4aAUrJQvopOZoD9QiA0Ws0h2t0,23
 cognite/neat/app/api/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 cognite/neat/app/api/asgi/metrics.py,sha256=nxFy7L5cChTI0a-zkCiJ59Aq8yLuIJp5c9Dg0wRXtV0,152
 cognite/neat/app/api/configuration.py,sha256=L1DCtLZ1HZku8I2z-JWd5RDsXhIsboFsKwAMhkrm-bY,3600
@@ -49,17 +49,19 @@ cognite/neat/graph/examples/Knowledge-Graph-Nordic44-dirty.xml,sha256=ujJip6XBs5
 cognite/neat/graph/examples/Knowledge-Graph-Nordic44.xml,sha256=U2Ns-M4LRjT1fBkhmRj63ur7jDzlRtHK9yOLf_npZ_g,1437996
 cognite/neat/graph/examples/__init__.py,sha256=yAjHVY3b5jOjmbW-iLbhvu7BG014TpGi3K4igkDqW5I,368
 cognite/neat/graph/examples/skos-capturing-sheet-wind-topics.xlsx,sha256=CV_yK5ZSbYS_ktfIZUPD8Sevs47zpswLXQUDFkGE4Gw,45798
-cognite/neat/graph/extractors/__init__.py,sha256=Mytvccgz6unU0uAh7ny9btlI1Ha1dEBrjlW7JCJG_lM,1772
+cognite/neat/graph/extractors/__init__.py,sha256=9A_AsBwdbnYDo-20c18wcu__S-h9PHCeQB_fGteK5mo,1858
 cognite/neat/graph/extractors/_base.py,sha256=8IWygpkQTwo0UOmbbwWVI7540_klTVdUVX2JjVPFRIs,498
 cognite/neat/graph/extractors/_classic_cdf/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-cognite/neat/graph/extractors/_classic_cdf/_assets.py,sha256=xVqWC_lyqX5RKGfTxX0Gb96A0GAyM13d6YIjd1xEbpI,6524
-cognite/neat/graph/extractors/_classic_cdf/_base.py,sha256=oLqYxNkAgYCZlAGMqfOVTDbRnouyf72m54d1BfcpeAA,4494
-cognite/neat/graph/extractors/_classic_cdf/_events.py,sha256=8shnrQD9LZCrKPIt7xeU6ac0-lSkIwGe9KB0DdJ7Eik,5600
-cognite/neat/graph/extractors/_classic_cdf/_files.py,sha256=wa2jzbpOj4BssVY7PbgdGymFZGSrvd1GZaRVN-JjOXQ,6588
-cognite/neat/graph/extractors/_classic_cdf/_labels.py,sha256=Xvyv22jebq4s94Sl1h2SMEeBvOkoTIPujB9tczRgcDo,4359
-cognite/neat/graph/extractors/_classic_cdf/_relationships.py,sha256=VxmwnZ_vN_P14lXDuEYfR68OK-ZnLrvGqSFW9n0KHeY,7429
-cognite/neat/graph/extractors/_classic_cdf/_sequences.py,sha256=sTtVPngffQw2fyvoRG1ABC2yFkepA6skFT1XdCzg4cU,5086
-cognite/neat/graph/extractors/_classic_cdf/_timeseries.py,sha256=zmVf5pO-G4ro3elKt0sCTfTtLulkEFoEahLgP9J5kKE,6581
+cognite/neat/graph/extractors/_classic_cdf/_assets.py,sha256=3zpwp4LIXf_dQgfXEQCmp0CNprBYGLC-VjNmhMXc-4k,6604
+cognite/neat/graph/extractors/_classic_cdf/_base.py,sha256=5jCegqeqQr_Bxwskt0ns7aiivzz1KyEIxqD5Xojoa7Y,5071
+cognite/neat/graph/extractors/_classic_cdf/_classic.py,sha256=Al4lO-vydyrqFjAt_OanmuIjd2fM4SS1TtukZj3YBM8,10186
+cognite/neat/graph/extractors/_classic_cdf/_data_sets.py,sha256=_jn2YM_87flrJ2pdLUKApQ3I1M02zjnhDSGKhl0vdDU,4297
+cognite/neat/graph/extractors/_classic_cdf/_events.py,sha256=08c1IavgRWdBem8YbSQnzXEYa41QYojmTilzZBbfAR0,6508
+cognite/neat/graph/extractors/_classic_cdf/_files.py,sha256=hi5Zf_-JWbRN1wEoGc-JIT81YOcWGxWwPAh162Urs4s,7554
+cognite/neat/graph/extractors/_classic_cdf/_labels.py,sha256=fNi_eqG2rfxMHyPUzAqRSARp_QRY3drYHH_WcoLzz0M,4766
+cognite/neat/graph/extractors/_classic_cdf/_relationships.py,sha256=wSYQK1QiRD9aYnjOJWXOCjpEnpH6s0LqPcqi3gC9ckU,8792
+cognite/neat/graph/extractors/_classic_cdf/_sequences.py,sha256=HMC-tmGE1Wgc7c1KeDyx0BlnYOxGbyhUHkoOHl1Yz60,6006
+cognite/neat/graph/extractors/_classic_cdf/_timeseries.py,sha256=ZO7-dgIq7EavvmVfo0gAWHHM_WoNcfPJe7zRArt4W6w,7502
 cognite/neat/graph/extractors/_dexpi.py,sha256=s3Ff3s3CMy0Xalixtfvnu-3yrEsCF26-9WVhIIZBZQI,9432
 cognite/neat/graph/extractors/_dms.py,sha256=NbqY1nrisn-dJq8_qeCatpgJwm6-cKFsTISDR5afdbU,6688
 cognite/neat/graph/extractors/_iodd.py,sha256=RiUxhVsV99YxrWjMxPQcLQzYG0xLYqZCEHduCoG6P8Q,6759
@@ -179,7 +181,7 @@ cognite/neat/utils/cdf/loaders/__init__.py,sha256=s2aPR5XLo6WZ0ybstAJlcGFYkA7CyH
 cognite/neat/utils/cdf/loaders/_base.py,sha256=ryNC_AMXIESWXuTVJ-02L-HSVSpD6V49XdLTRYeFg70,1764
 cognite/neat/utils/cdf/loaders/_data_modeling.py,sha256=0jynnwZ7utnG2KIkVRJBNvDshUCzwX29LZuGSlm6qUM,13216
 cognite/neat/utils/cdf/loaders/_ingestion.py,sha256=QbF4ntaGUum6yTRTq9WNg8n49TQpfzF1T7H-Bx740ws,6326
-cognite/neat/utils/collection_.py,sha256=fX4eAMau5AsUs1Dm5PlzINxR64kUjncEEHeoL_JT9rU,582
+cognite/neat/utils/collection_.py,sha256=m5pn78NRQ9vx8l-NdMTAdYDLCTB8CQLX-rqxAWt7vyg,602
 cognite/neat/utils/rdf_.py,sha256=VXDBQUt86vRntiGhejK35PlsbvKCUkuQQa1cMYz4SIc,5656
 cognite/neat/utils/regex_patterns.py,sha256=rXOBW9d-n_eAffm9zVRIPFB7ElLS8VDdsvFDsr01Q-M,2155
 cognite/neat/utils/spreadsheet.py,sha256=LI0c7dlW0zXHkHw0NvB-gg6Df6cDcE3FbiaHBYLXdzQ,2714
@@ -214,8 +216,8 @@ cognite/neat/workflows/steps_registry.py,sha256=FjMsFBlFFy82ABUzDnWoFidYODV3pp3c
 cognite/neat/workflows/tasks.py,sha256=dqlJwKAb0jlkl7abbY8RRz3m7MT4SK8-7cntMWkOYjw,788
 cognite/neat/workflows/triggers.py,sha256=_BLNplzoz0iic367u1mhHMHiUrCwP-SLK6_CZzfODX0,7071
 cognite/neat/workflows/utils.py,sha256=gKdy3RLG7ctRhbCRwaDIWpL9Mi98zm56-d4jfHDqP1E,453
-cognite_neat-0.91.0.dist-info/LICENSE,sha256=W8VmvFia4WHa3Gqxq1Ygrq85McUNqIGDVgtdvzT-XqA,11351
-cognite_neat-0.91.0.dist-info/METADATA,sha256=9QzJcHigRT9OudOBhVmpmitBZepeNs9wNm9emotHwUE,9441
-cognite_neat-0.91.0.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
-cognite_neat-0.91.0.dist-info/entry_points.txt,sha256=61FPqiWb25vbqB0KI7znG8nsg_ibLHBvTjYnkPvNFso,50
-cognite_neat-0.91.0.dist-info/RECORD,,
+cognite_neat-0.92.1.dist-info/LICENSE,sha256=W8VmvFia4WHa3Gqxq1Ygrq85McUNqIGDVgtdvzT-XqA,11351
+cognite_neat-0.92.1.dist-info/METADATA,sha256=YHC8hU6I2-w_BCUtS64Nz2X41ypo6W6dzei3NIhzOUw,9431
+cognite_neat-0.92.1.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
+cognite_neat-0.92.1.dist-info/entry_points.txt,sha256=61FPqiWb25vbqB0KI7znG8nsg_ibLHBvTjYnkPvNFso,50
+cognite_neat-0.92.1.dist-info/RECORD,,