deriva-ml 1.17.10__py3-none-any.whl

deriva_ml/model/catalog.py

@@ -0,0 +1,485 @@
+"""
+Model management for Deriva ML catalogs.
+
+This module provides the DerivaModel class which augments the standard Deriva model class with
+ML-specific functionality. It handles schema management, feature definitions, and asset tracking.
+"""
+
+from __future__ import annotations
+
+# Standard library imports
+from collections import Counter, defaultdict
+from graphlib import CycleError, TopologicalSorter
+from typing import Any, Callable, Final, Iterable, NewType, TypeAlias
+
+from deriva.core.ermrest_catalog import ErmrestCatalog
+
+# Deriva imports
+from deriva.core.ermrest_model import Column, FindAssociationResult, Model, Schema, Table
+
+# Third-party imports
+from pydantic import ConfigDict, validate_call
+
+from deriva_ml.core.definitions import (
+    ML_SCHEMA,
+    RID,
+    DerivaAssetColumns,
+    TableDefinition,
+)
+from deriva_ml.core.exceptions import DerivaMLException, DerivaMLTableTypeError
+
+# Local imports
+from deriva_ml.feature import Feature
+from deriva_ml.protocols.dataset import DatasetLike
+
+try:
+    from icecream import ic
+except ImportError:  # Graceful fallback if IceCream isn't installed.
+    ic = lambda *a: None if not a else (a[0] if len(a) == 1 else a)  # noqa
+
+
+# Define common types:
+TableInput: TypeAlias = str | Table
+SchemaDict: TypeAlias = dict[str, Schema]
+FeatureList: TypeAlias = Iterable[Feature]
+SchemaName = NewType("SchemaName", str)
+ColumnSet: TypeAlias = set[Column]
+AssociationResult: TypeAlias = FindAssociationResult
+TableSet: TypeAlias = set[Table]
+PathList: TypeAlias = list[list[Table]]
+
+# Define constants:
+VOCAB_COLUMNS: Final[set[str]] = {"NAME", "URI", "SYNONYMS", "DESCRIPTION", "ID"}
+ASSET_COLUMNS: Final[set[str]] = {"Filename", "URL", "Length", "MD5", "Description"}
+
+FilterPredicate = Callable[[Table], bool]
+
+
+class DerivaModel:
+    """Augmented interface to deriva model class.
+
+    This class provides a number of DerivaML specific methods that augment the interface in the deriva model class.
+
+    Attributes:
+        domain_schema: Schema name for domain-specific tables and relationships.
+        model: ERMRest model for the catalog.
+        catalog: ERMRest catalog for the model
+        hostname: ERMRest catalog for the model
+        ml_schema: The ML schema for the catalog.
+        domain_schema: The domain schema for the catalog.
+
+    """
+
+    def __init__(
+        self,
+        model: Model,
+        ml_schema: str = ML_SCHEMA,
+        domain_schema: str | None = None,
+    ):
+        """Create and initialize a DerivaML instance.
+
+        This method will connect to a catalog, and initialize local configuration for the ML execution.
+        This class is intended to be used as a base class on which domain-specific interfaces are built.
+
+        Args:
+            model: The ERMRest model for the catalog.
+            ml_schema: The ML schema name.
+            domain_schema: The domain schema name.
+        """
+        self.model = model
+        self.configuration = None
+        self.catalog: ErmrestCatalog = self.model.catalog
+        self.hostname = self.catalog.deriva_server.server if isinstance(self.catalog, ErmrestCatalog) else "localhost"
+
+        self.ml_schema = ml_schema
+        builtin_schemas = ("public", self.ml_schema, "www", "WWW")
+        if domain_schema:
+            self.domain_schema = domain_schema
+        else:
+            if len(user_schemas := {k for k in self.model.schemas.keys()} - set(builtin_schemas)) == 1:
+                self.domain_schema = user_schemas.pop()
+            else:
+                raise DerivaMLException(f"Ambiguous domain schema: {user_schemas}")
+
+    def refresh_model(self) -> None:
+        self.model = self.catalog.getCatalogModel()
+
+    @property
+    def schemas(self) -> dict[str, Schema]:
+        return self.model.schemas
+
+    @property
+    def chaise_config(self) -> dict[str, Any]:
+        """Return the chaise configuration."""
+        return self.model.chaise_config
+
+    def __getattr__(self, name: str) -> Any:
+        # Called only if `name` is not found in Manager.  Delegate attributes to model class.
+        return getattr(self.model, name)
+
+    def name_to_table(self, table: TableInput) -> Table:
+        """Return the table object corresponding to the given table name.
+
+        If the table name appears in more than one schema, return the first one you find.
+
+        Args:
+          table: A ERMRest table object or a string that is the name of the table.
+
+        Returns:
+          Table object.
+        """
+        if isinstance(table, Table):
+            return table
+        if table in (s := self.model.schemas[self.domain_schema].tables):
+            return s[table]
+        for s in [self.model.schemas[sname] for sname in [self.domain_schema, self.ml_schema, "WWW"]]:
+            if table in s.tables.keys():
+                return s.tables[table]
+        raise DerivaMLException(f"The table {table} doesn't exist.")
+
+    def is_vocabulary(self, table_name: TableInput) -> bool:
+        """Check if a given table is a controlled vocabulary table.
+
+        Args:
+          table_name: A ERMRest table object or the name of the table.
+
+        Returns:
+          Table object if the table is a controlled vocabulary, False otherwise.
+
+        Raises:
+          DerivaMLException: if the table doesn't exist.
+
+        """
+        vocab_columns = {"NAME", "URI", "SYNONYMS", "DESCRIPTION", "ID"}
+        table = self.name_to_table(table_name)
+        return vocab_columns.issubset({c.name.upper() for c in table.columns})
+
+    def is_association(
+        self,
+        table_name: str | Table,
+        unqualified: bool = True,
+        pure: bool = True,
+        min_arity: int = 2,
+        max_arity: int = 2,
+    ) -> bool | set[str] | int:
+        """Check the specified table to see if it is an association table.
+
+        Args:
+            table_name: param unqualified:
+            pure: return: (Default value = True)
+            table_name: str | Table:
+            unqualified:  (Default value = True)
+
+        Returns:
+
+
+        """
+        table = self.name_to_table(table_name)
+        return table.is_association(unqualified=unqualified, pure=pure, min_arity=min_arity, max_arity=max_arity)
+
+    def find_association(self, table1: Table | str, table2: Table | str) -> tuple[Table, Column, Column]:
+        """Given two tables, return an association table that connects the two and the two columns used to link them..
+
+        Raises:
+            DerivaML exception if there is either not an association table or more than one association table.
+        """
+        table1 = self.name_to_table(table1)
+        table2 = self.name_to_table(table2)
+
+        tables = [
+            (a.table, a.self_fkey.columns[0].name, other_key.columns[0].name)
+            for a in table1.find_associations(pure=False)
+            if len(a.other_fkeys) == 1 and (other_key := a.other_fkeys.pop()).pk_table == table2
+        ]
+
+        if len(tables) == 1:
+            return tables[0]
+        elif len(tables) == 0:
+            raise DerivaMLException(f"No association tables found between {table1.name} and {table2.name}.")
+        else:
+            raise DerivaMLException(
+                f"There are {len(tables)} association tables between {table1.name} and {table2.name}."
+            )
+
+    def is_asset(self, table_name: TableInput) -> bool:
+        """True if the specified table is an asset table.
+
+        Args:
+            table_name: str | Table:
+
+        Returns:
+            True if the specified table is an asset table, False otherwise.
+
+        """
+        asset_columns = {"Filename", "URL", "Length", "MD5", "Description"}
+        table = self.name_to_table(table_name)
+        return asset_columns.issubset({c.name for c in table.columns})
+
+    def find_assets(self, with_metadata: bool = False) -> list[Table]:
+        """Return the list of asset tables in the current model"""
+        return [t for s in self.model.schemas.values() for t in s.tables.values() if self.is_asset(t)]
+
+    def find_vocabularies(self) -> list[Table]:
+        """Return a list of all the controlled vocabulary tables in the domain schema."""
+        return [t for s in self.model.schemas.values() for t in s.tables.values() if self.is_vocabulary(t)]
+
+    @validate_call(config=ConfigDict(arbitrary_types_allowed=True))
+    def find_features(self, table: TableInput) -> Iterable[Feature]:
+        """List the names of the features in the specified table.
+
+        Args:
+            table: The table to find features for.
+            table: Table | str:
+
+        Returns:
+            An iterable of FeatureResult instances that describe the current features in the table.
+        """
+        table = self.name_to_table(table)
+
+        def is_feature(a: FindAssociationResult) -> bool:
+            """Check if association represents a feature.
+
+            Args:
+                a: Association result to check
+            Returns:
+                bool: True if association represents a feature
+            """
+            return {
+                "Feature_Name",
+                "Execution",
+                a.self_fkey.foreign_key_columns[0].name,
+            }.issubset({c.name for c in a.table.columns})
+
+        return [
+            Feature(a, self) for a in table.find_associations(min_arity=3, max_arity=3, pure=False) if is_feature(a)
+        ]
+
+    def lookup_feature(self, table: TableInput, feature_name: str) -> Feature:
+        """Lookup the named feature associated with the provided table.
+
+        Args:
+            table: param feature_name:
+            table: str | Table:
+            feature_name: str:
+
+        Returns:
+            A Feature class that represents the requested feature.
+
+        Raises:
+          DerivaMLException: If the feature cannot be found.
+        """
+        table = self.name_to_table(table)
+        try:
+            return [f for f in self.find_features(table) if f.feature_name == feature_name][0]
+        except IndexError:
+            raise DerivaMLException(f"Feature {table.name}:{feature_name} doesn't exist.")
+
+    def asset_metadata(self, table: str | Table) -> set[str]:
+        """Return the metadata columns for an asset table."""
+
+        table = self.name_to_table(table)
+
+        if not self.is_asset(table):
+            raise DerivaMLTableTypeError("asset table", table.name)
+        return {c.name for c in table.columns} - DerivaAssetColumns
+
+    def apply(self) -> None:
+        """Call ERMRestModel.apply"""
+        if self.catalog == "file-system":
+            raise DerivaMLException("Cannot apply() to non-catalog model.")
+        else:
+            self.model.apply()
+
+    def list_dataset_element_types(self) -> list[Table]:
+        """
+        Lists the data types of elements contained within a dataset.
+
+        This method analyzes the dataset and identifies the data types for all
+        elements within it. It is useful for understanding the structure and
+        content of the dataset and allows for better manipulation and usage of its
+        data.
+
+        Returns:
+            list[str]: A list of strings where each string represents a data type
+            of an element found in the dataset.
+
+        """
+
+        dataset_table = self.name_to_table("Dataset")
+
+        def domain_table(table: Table) -> bool:
+            return table.schema.name == self.domain_schema or table.name == dataset_table.name
+
+        return [t for a in dataset_table.find_associations() if domain_table(t := a.other_fkeys.pop().pk_table)]
+
+    def _prepare_wide_table(self,
+                            dataset,
+                            dataset_rid: RID,
+                            include_tables: list[str]) -> tuple[dict[str, Any], list[tuple]]:
+        """
+        Generates details of a wide table from the model
+
+        Args:
+            include_tables (list[str] | None): List of table names to include in the denormalized dataset. If None,
+                all tables from the dataset will be included.
+
+        Returns:
+            str: SQL query string that represents the process of denormalization.
+        """
+
+        # Skip over tables that we don't want to include in the denormalized dataset.
+        # Also, strip off the Dataset/Dataset_X part of the path so we don't include dataset columns in the denormalized
+        # table.
+        include_tables = set(include_tables)
+        for t in include_tables:
+            # Check to make sure the table is in the catalog.
+            _ = self.name_to_table(t)
+
+        table_paths = [
+            path
+            for path in self._schema_to_paths()
+            if path[-1].name in include_tables and include_tables.intersection({p.name for p in path})
+        ]
+        paths_by_element = defaultdict(list)
+        for p in table_paths:
+            paths_by_element[p[2].name].append(p)
+
+        # Get the names of all of the tables that can be dataset elements.
+        dataset_element_tables = {
+            e.name for e in self.list_dataset_element_types() if e.schema.name == self.domain_schema
+        }
+
+        skip_columns = {"RCT", "RMT", "RCB", "RMB"}
+        element_tables = {}
+        for element_table, paths in paths_by_element.items():
+            graph = {}
+            for path in paths:
+                for left, right in zip(path[0:], path[1:]):
+                    graph.setdefault(left.name, set()).add(right.name)
+
+            # New lets remove any cycles that we may have in the graph.
+            # We will use a topological sort to find the order in which we need to join the tables.
+            # If we find a cycle, we will remove the table from the graph and splice in an additional ON clause.
+            # We will then repeat the process until there are no cycles.
+            graph_has_cycles = True
+            element_join_tables = []
+            element_join_conditions = {}
+            while graph_has_cycles:
+                try:
+                    ts = TopologicalSorter(graph)
+                    element_join_tables = list(reversed(list(ts.static_order())))
+                    graph_has_cycles = False
+                except CycleError as e:
+                    cycle_nodes = e.args[1]
+                    if len(cycle_nodes) > 3:
+                        raise DerivaMLException(f"Unexpected cycle found when normalizing dataset {cycle_nodes}")
+                    # Remove cycle from graph and splice in additional ON constraint.
+                    graph[cycle_nodes[1]].remove(cycle_nodes[0])
+
+            # The Dataset_Version table is a special case as it points to dataset and dataset to version.
+            if "Dataset_Version" in element_join_tables:
+                element_join_tables.remove("Dataset_Version")
+
+            for path in paths:
+                for left, right in zip(path[0:], path[1:]):
+                    if right.name == "Dataset_Version":
+                        # The Dataset_Version table is a special case as it points to dataset and dataset to version.
+                        continue
+                    if element_join_tables.index(right.name) < element_join_tables.index(left.name):
+                        continue
+                    table_relationship = self._table_relationship(left, right)
+                    element_join_conditions.setdefault(right.name, set()).add(
+                        (table_relationship[0], table_relationship[1])
+                    )
+            element_tables[element_table] = (element_join_tables, element_join_conditions)
+        # Get the list of columns that will appear in the final denormalized dataset.
+        denormalized_columns = [
+            (table_name, c.name)
+            for table_name in include_tables
+            if not self.is_association(table_name)  # Don't include association columns in the denormalized view.'
+            for c in self.name_to_table(table_name).columns
+            if (not include_tables or table_name in include_tables) and (c.name not in skip_columns)
+        ]
+        return element_tables, denormalized_columns
+
+    def _table_relationship(
+        self,
+        table1: TableInput,
+        table2: TableInput,
+    ) -> tuple[Column, Column]:
+        """Return columns used to relate two tables."""
+        table1 = self.name_to_table(table1)
+        table2 = self.name_to_table(table2)
+        relationships = [
+            (fk.foreign_key_columns[0], fk.referenced_columns[0]) for fk in table1.foreign_keys if fk.pk_table == table2
+        ]
+        relationships.extend(
+            [(fk.referenced_columns[0], fk.foreign_key_columns[0]) for fk in table1.referenced_by if fk.table == table2]
+        )
+        if len(relationships) != 1:
+            raise DerivaMLException(
+                f"Ambiguous linkage between {table1.name} and {table2.name}: {[(r[0].name, r[1].name) for r in relationships]}"
+            )
+        return relationships[0]
+
+    def _schema_to_paths(
+        self,
+        root: Table | None = None,
+        path: list[Table] | None = None,
+    ) -> list[list[Table]]:
+        """Return a list of paths through the schema graph.
+
+        Args:
+            root: The root table to start from.
+            path: The current path being built.
+
+        Returns:
+            A list of paths through the schema graph.
+        """
+        path = path or []
+
+        root = root or self.model.schemas[self.ml_schema].tables["Dataset"]
+        path = path.copy() if path else []
+        parent = path[-1] if path else None  # Table that we are coming from.
+        path.append(root)
+        paths = [path]
+
+        def find_arcs(table: Table) -> set[Table]:
+            """Given a path through the model, return the FKs that link the tables"""
+            arc_list = [fk.pk_table for fk in table.foreign_keys] + [fk.table for fk in table.referenced_by]
+            arc_list = [t for t in arc_list if t.schema.name in {self.domain_schema, self.ml_schema}]
+            domain_tables = [t for t in arc_list if t.schema.name == self.domain_schema]
+            if multiple_columns := [c for c, cnt in Counter(domain_tables).items() if cnt > 1]:
+                raise DerivaMLException(f"Ambiguous relationship in {table.name} {multiple_columns}")
+            return set(arc_list)
+
+        def is_nested_dataset_loopback(n1: Table, n2: Table) -> bool:
+            """Test to see if node is an association table used to link elements to datasets."""
+            # If we have node_name <- node_name_dataset-> Dataset then we are looping
+            # back around to a new dataset element
+            dataset_table = self.model.schemas[self.ml_schema].tables["Dataset"]
+            assoc_table = [a for a in dataset_table.find_associations() if a.table == n2]
+            return len(assoc_table) == 1 and n1 != dataset_table
+
+        # Don't follow vocabulary terms back to their use.
+        if self.is_vocabulary(root):
+            return paths
+
+        for child in find_arcs(root):
+            if child.name in {"Dataset_Execution", "Dataset_Dataset", "Execution"}:
+                continue
+            if child == parent:
+                # Don't loop back via referred_by
+                continue
+            if is_nested_dataset_loopback(root, child):
+                continue
+            if child in path:
+                raise DerivaMLException(f"Cycle in schema path: {child.name} path:{[p.name for p in path]}")
+
+            paths.extend(self._schema_to_paths(child, path))
+        return paths
+
+    @validate_call(config=ConfigDict(arbitrary_types_allowed=True))
+    def create_table(self, table_def: TableDefinition) -> Table:
+        """Create a new table from TableDefinition."""
+        return self.model.schemas[self.domain_schema].create_table(table_def.model_dump())