wedata-feature-engineering 0.1.3__py3-none-any.whl → 0.1.5__py3-none-any.whl

wedata/feature_store/utils/feature_spec_utils.py

@@ -0,0 +1,286 @@
+import logging
+from dataclasses import dataclass
+from functools import reduce
+from typing import Dict, List, Set, Tuple, Type, Union
+
+import yaml
+from mlflow.utils.file_utils import YamlSafeDumper
+
+from feature_store.entities.column_info import ColumnInfo
+from feature_store.entities.feature_column_info import FeatureColumnInfo
+from feature_store.entities.feature_spec import FeatureSpec
+from feature_store.entities.on_demand_column_info import OnDemandColumnInfo
+from feature_store.entities.source_data_column_info import SourceDataColumnInfo
+from feature_store.utils.topological_sort import topological_sort
+
+DEFAULT_GRAPH_DEPTH_LIMIT = 5
+
+COLUMN_INFO_TYPE_SOURCE = "SOURCE"
+COLUMN_INFO_TYPE_ON_DEMAND = "ON_DEMAND"
+COLUMN_INFO_TYPE_FEATURE = "FEATURE"
+
+_logger = logging.getLogger(__name__)
+
+
+@dataclass
+class FeatureExecutionGroup:
+    type: str  # could be FEATURE, ON_DEMAND, SOURCE
+    features: Union[
+        List[FeatureColumnInfo], List[OnDemandColumnInfo], List[SourceDataColumnInfo]
+    ]
+
+
+# Small number has high priority. Besides SOURCE, preferring FEATURE over ON_DEMAND in topological
+# sorting to make sure ON_DEMAND columns after FEATURE in simple cases to align with previous
+# assumption before implementing TLT.
+# NOTE: changing this priority may cause performance regression, proceed with caution.
+COLUMN_TYPE_PRIORITY = {
+    COLUMN_INFO_TYPE_SOURCE: 0,
+    COLUMN_INFO_TYPE_ON_DEMAND: 1,
+    COLUMN_INFO_TYPE_FEATURE: 2,
+}
+
+
+class _GraphNode:
+    def __init__(self, column_info: ColumnInfo):
+        info = column_info.info
+        self.column_info = column_info
+        self.output_name = info.output_name
+
+        if isinstance(column_info.info, SourceDataColumnInfo):
+            self.input_names = set()
+            self.type = COLUMN_INFO_TYPE_SOURCE
+        elif isinstance(column_info.info, FeatureColumnInfo):
+            self.input_names = set(info.lookup_key)
+            self.type = COLUMN_INFO_TYPE_FEATURE
+        elif isinstance(column_info.info, OnDemandColumnInfo):
+            self.input_names = set(info.input_bindings.values())
+            self.type = COLUMN_INFO_TYPE_ON_DEMAND
+        else:
+            raise ValueError("unknown column info type")
+
+    def __str__(self):
+        return "node<" + self.output_name + ">"
+
+    def __repr__(self):
+        return str(self)
+
+
+def _column_info_sort_key(node: _GraphNode) -> Tuple[int, str]:
+    """
+    Returns a tuple of an int and a str as the sorting key for _GraphNode. Priority is determined by
+    the first element and then use the second element to break ties.
+    """
+    return COLUMN_TYPE_PRIORITY[node.type], node.output_name
+
+
+def _should_be_grouped(node: _GraphNode) -> bool:
+    """
+    Returns True if the given node is of type that should be grouped together as much as possible.
+    """
+    return node.type == COLUMN_INFO_TYPE_FEATURE
+
+
+def _validate_graph_depth(nodes: List[_GraphNode], depth_limit: int):
+    name_to_node = {node.output_name: node for node in nodes}
+    visited_depth = {}
+
+    def dfs(node: _GraphNode, depth: int):
+        if depth > depth_limit:
+            raise ValueError(
+                f"The given graph contains a dependency path longer than the limit {depth_limit}"
+            )
+        if (
+            node.output_name in visited_depth
+            and depth <= visited_depth[node.output_name]
+        ):
+            return
+        visited_depth[node.output_name] = depth
+        for column_name in node.input_names:
+            dependency = name_to_node[column_name]
+            dfs(dependency, depth + 1)
+
+    for node in nodes:
+        dfs(node, 1)
+
+
+def get_encoded_graph_map(column_infos: List[ColumnInfo]) -> Dict[str, List[str]]:
+    """
+    Creates a dictionary of columns with their dependency columns for metric use. Columns are
+    encoded with a string representing the type and index. For example:
+      {
+        "f3": ["s1", "s2"],
+        "o4": ["f3"],
+        "o5": []
+      }
+    "s1" and "s2" are SourceColumnInfos, "f3" is FeatureColumnInfo and "o4", "o5" are
+    OnDemandColumnInfos. "f3" depends on "s1" and "s2", "o5" doesn't depend on any column, etc.
+    :param column_infos: A list of ColumnInfos.
+    """
+    nodes = {info.output_name: _GraphNode(info) for info in column_infos}
+    next_node_index = 0
+    # A map from column info's output_name to its label.
+    node_label = {}
+
+    def get_node_label(node):
+        nonlocal next_node_index
+        output_name = node.output_name
+        if output_name not in node_label:
+            if node.type == COLUMN_INFO_TYPE_SOURCE:
+                type_simple_str = "s"
+            if node.type == COLUMN_INFO_TYPE_FEATURE:
+                type_simple_str = "f"
+            if node.type == COLUMN_INFO_TYPE_ON_DEMAND:
+                type_simple_str = "o"
+            new_label = type_simple_str + str(next_node_index)
+            next_node_index += 1
+            node_label[output_name] = new_label
+        return node_label[output_name]
+
+    graph_map = {}
+    for node in nodes.values():
+        label = get_node_label(node)
+        dependencies = []
+        for dep_name in sorted(node.input_names):
+            if dep_name not in nodes:
+                # skip the column if it's not in the feature spec.
+                continue
+            dep = get_node_label(nodes[dep_name])
+            dependencies.append(dep)
+        graph_map[label] = dependencies
+    return graph_map
+
+
+def assign_topological_ordering(
+    column_infos: List[ColumnInfo],
+    allow_missing_source_columns=False,
+    graph_depth_limit=DEFAULT_GRAPH_DEPTH_LIMIT,
+) -> List[ColumnInfo]:
+    """
+    Assigns the topological ordering for each ColumnInfo of the input. Returns a list of new
+    ColumnInfo objects with topological_ordering set to an integer.
+
+    :param column_infos: a list of ColumnInfos.
+    :param allow_missing_source_columns: ONLY USED BY FSE TEMPORARILY. Allow lookup key or
+    function input be missing from source columns. If true, this method will assign
+    topological_ordering to columns as if the missing sources are added in the column_infos.
+    :param graph_depth_limit raises if the given graph exceed the limit.
+    :raises ValueError if there is a cycle in the graph.
+    """
+    nodes = list(map(lambda c: _GraphNode(c), column_infos))
+    # allow_missing_source_columns is used when feature_serving_endpoint_client creates training
+    # sets. It doesn't include source columns in the dataframe.
+    # TODO[ML-33809]: clean up allow_missing_source_columns.
+    all_output_names = set([n.output_name for n in nodes])
+    all_input_names = reduce(lambda a, b: a | b, [n.input_names for n in nodes])
+    missing_inputs = all_input_names - all_output_names
+    if allow_missing_source_columns:
+        for input_name in missing_inputs:
+            if input_name not in all_output_names:
+                nodes.append(
+                    _GraphNode(ColumnInfo(SourceDataColumnInfo(input_name), False))
+                )
+    elif len(missing_inputs) > 0:
+        missing_input_names_str = ", ".join(
+            [f"'{name}'" for name in sorted(missing_inputs)]
+        )
+        raise ValueError(
+            f"Input columns {missing_input_names_str} required by FeatureLookups or "
+            "FeatureFunctions are not provided by input DataFrame or other FeatureFunctions and "
+            "FeatureLookups"
+        )
+    output_name_to_node = {node.output_name: node for node in nodes}
+    graph = {
+        node: [output_name_to_node[input_name] for input_name in node.input_names]
+        for node in nodes
+    }
+    sorted_nodes = topological_sort(graph, _column_info_sort_key, _should_be_grouped)
+    # validate depth after sorting the graph because cycle is detected during sorting.
+    _validate_graph_depth(nodes, graph_depth_limit)
+    name_to_ordering = {node.output_name: i for i, node in enumerate(sorted_nodes)}
+    return [
+        column.with_topological_ordering(name_to_ordering[column.output_name])
+        for column in column_infos
+    ]
+
+
+def get_feature_execution_groups(
+    feature_spec: FeatureSpec, df_columns: List[str] = []
+) -> List[FeatureExecutionGroup]:
+    """
+    Splits the list of column_infos in feature_spec into groups based on the topological_ordering of
+    the column_infos such that each group contains only one type of feature columns and columns
+    don't depend on other columns in the same group. The type of feature column is equivalent to the
+    class type of column_info.info field.
+    Example:
+        Given FeatureSpec with some columns, after sorting the columns by topological_ordering,
+        assuming the sorted list:
+            [source_1, feature_2, feature_3, on_demand_4, on_demand_5]
+        where feature_2 depends on feature_3. The resulting groups will be:
+            [
+                group(SOURCE, [source_1]),
+                group(FEATURE, [feature_2]),
+                group(FEATURE, [feature_3]),
+                group(ON_DEMAND, [on_demand_4, on_demand_5]),
+            ]
+
+    :param feature_spec: A FeatureSpec with topologically sorted column_infos.
+    :param df_columns: the columns from the DF used to create_training_set or score_batch.
+    """
+    # convert column infos into _GraphNode
+    nodes = list(map(lambda c: _GraphNode(c), feature_spec.column_infos))
+    if any(info.topological_ordering is None for info in feature_spec.column_infos):
+        # The old version of feature_spec may not have topological_ordering, we can safely assume
+        # they are already sorted because of validations during the feature_spec creation.
+        _logger.warning(
+            "Processing a feature spec that at least one of the column_infos has no "
+            "topological_ordering"
+        )
+    else:
+        # sort nodes by topological_ordering
+        nodes = sorted(nodes, key=lambda n: n.column_info.topological_ordering)
+    # A buffer holding the columns in a group.
+    buffer = []
+    # output names of columns in the current buffer.
+    buffered_output_names = set()
+    # Used to validate the topological sorting.
+    # df_columns is used to be backward compatible. In old FeatureSpecs, source columns might not
+    # exist. So we need to consider the df as initial resolved columns.
+    resolved_columns = set(df_columns)
+    result_list = []
+    last_type = None
+    for node in nodes:
+        if not node.input_names.issubset(resolved_columns):
+            raise ValueError(
+                "The column_infos in the FeatureSpec is not topologically sorted"
+            )
+        if node.type != last_type or buffered_output_names.intersection(
+            node.input_names
+        ):
+            # split group if the current node has a different type from the previous node OR
+            # any of the inputs are from the nodes in the current group.
+            if buffer:
+                result_list.append(FeatureExecutionGroup(last_type, buffer))
+                buffer = []
+            buffered_output_names.clear()
+            last_type = node.type
+        buffer.append(node.column_info.info)
+        resolved_columns.add(node.output_name)
+        buffered_output_names.add(node.output_name)
+    if buffer:
+        result_list.append(FeatureExecutionGroup(last_type, buffer))
+    return result_list
+
+
+def convert_to_yaml_string(feature_spec: FeatureSpec) -> str:
+    """
+    Converts the given FeatureSpec to a YAML string.
+    """
+    feature_spec_dict = feature_spec._to_dict()
+    return yaml.dump(
+        feature_spec_dict,
+        default_flow_style=False,
+        allow_unicode=True,
+        sort_keys=False,
+        Dumper=YamlSafeDumper,
+    )

wedata/feature_store/utils/feature_utils.py

@@ -0,0 +1,73 @@
+import copy
+from typing import List, Union
+
+from feature_store.entities.feature_function import FeatureFunction
+from feature_store.entities.feature_lookup import FeatureLookup
+from feature_store.spark_client.spark_client import SparkClient
+from feature_store.utils import uc_utils
+from feature_store.utils.feature_lookup_utils import get_feature_lookups_with_full_table_names
+
+
+def format_feature_lookups_and_functions(
+        _spark_client: SparkClient, features: List[Union[FeatureLookup, FeatureFunction]]
+):
+    fl_idx = []
+    ff_idx = []
+    feature_lookups = []
+    feature_functions = []
+    for idx, feature in enumerate(features):
+        if isinstance(feature, FeatureLookup):
+            fl_idx.append(idx)
+            feature_lookups.append(feature)
+        elif isinstance(feature, FeatureFunction):
+            ff_idx.append(idx)
+            feature_functions.append(feature)
+        else:
+            raise ValueError(
+                f"Expected a list of FeatureLookups for 'feature_lookups', but received type '{type(feature)}'."
+            )
+
+    # FeatureLookups and FeatureFunctions must have fully qualified table, UDF names
+    feature_lookups = get_feature_lookups_with_full_table_names(
+        feature_lookups,
+        _spark_client.get_current_catalog(),
+        _spark_client.get_current_database(),
+    )
+    feature_functions = get_feature_functions_with_full_udf_names(
+        feature_functions,
+        _spark_client.get_current_catalog(),
+        _spark_client.get_current_database(),
+    )
+
+    # Restore original order of FeatureLookups, FeatureFunctions. Copy to avoid mutating original list.
+    features = features.copy()
+    for idx, feature in zip(fl_idx + ff_idx, feature_lookups + feature_functions):
+        features[idx] = feature
+
+    return features
+
+
+def get_feature_functions_with_full_udf_names(
+        feature_functions: List[FeatureFunction], current_catalog: str, current_schema: str
+):
+    """
+    Takes in a list of FeatureFunctions, and returns copies with:
+    1. Fully qualified UDF names.
+    2. If output_name is empty, fully qualified UDF names as output_name.
+    """
+    udf_names = {ff.udf_name for ff in feature_functions}
+    uc_utils._check_qualified_udf_names(udf_names)
+    uc_utils._verify_all_udfs_in_uc(udf_names, current_catalog, current_schema)
+
+    standardized_feature_functions = []
+    for ff in feature_functions:
+        ff_copy = copy.deepcopy(ff)
+        del ff
+
+        ff_copy._udf_name = uc_utils.get_full_udf_name(
+            ff_copy.udf_name, current_catalog, current_schema
+        )
+        if not ff_copy.output_name:
+            ff_copy._output_name = ff_copy.udf_name
+        standardized_feature_functions.append(ff_copy)
+    return standardized_feature_functions

wedata/feature_store/utils/schema_utils.py

@@ -0,0 +1,117 @@
+import logging
+
+from feature_store.constants.constants import _ERROR, _WARN
+
+_logger = logging.getLogger(__name__)
+
+
+def catalog_matches_delta_schema(catalog_features, df_schema, column_filter=None):
+    """
+    Confirm that the column names and column types are the same.
+
+    Returns True if identical, False if there is a mismatch.
+
+    If column_filter is not None, only columns in column_filter must match.
+    """
+    if column_filter is not None:
+        catalog_features = [c for c in catalog_features if c.name in column_filter]
+        df_schema = [c for c in df_schema if c.name in column_filter]
+
+    catalog_schema = {
+        feature.name: feature.data_type
+        for feature in catalog_features
+    }
+    delta_schema = {
+        feature.name: feature.dataType
+        for feature in df_schema
+    }
+
+    complex_catalog_schema = get_complex_catalog_schema(
+        catalog_features, catalog_schema
+    )
+    complex_delta_schema = get_complex_delta_schema(df_schema, delta_schema)
+
+    return (
+        catalog_schema == delta_schema
+        and complex_catalog_schema == complex_delta_schema
+    )
+
+
+def get_complex_delta_schema(delta_features, delta_feature_names_to_fs_types):
+    """
+    1. Filter delta features to features that have complex datatypes.
+    2. Take the existing Spark DataType stored on the Delta features. This is later used for
+    comparison against the Catalog schema's complex Spark DataTypes.
+    3. Return a mapping of feature name to their respective complex Spark DataTypes.
+
+    :param delta_features: List[Feature]. List of features stored in Delta.
+    :param delta_feature_names_to_fs_types: Map[str, feature_store.DataType]. A mapping of feature
+    names to their respective Feature Store DataTypes.
+    :return: Map[str, spark.sql.types.DataType]. A mapping of feature names to their respective
+    Spark DataTypes.
+    """
+    complex_delta_features = [
+        feature
+        for feature in delta_features
+        if delta_feature_names_to_fs_types[feature.name] in DATA_TYPES_REQUIRES_DETAILS
+    ]
+    complex_delta_feature_names_to_spark_types = {
+        feature.name: feature.dataType for feature in complex_delta_features
+    }
+    return complex_delta_feature_names_to_spark_types
+
+
+def get_complex_catalog_schema(catalog_features, catalog_feature_names_to_fs_types):
+    """
+    1. Filter catalog features to features that have complex datatypes.
+    2. Convert the JSON string stored in each feature's data_type_details to the corresponding
+    Spark DataType. This is later used for comparison against the Delta schema's complex Spark
+    DataTypes.
+    3. Return a mapping of feature name to their respective complex Spark DataTypes.
+
+    :param catalog_features: List[Feature]. List of features stored in the Catalog.
+    :param catalog_feature_names_to_fs_types: Map[str, feature_store.DataType]. A mapping of feature
+    names to their respective Feature Store DataTypes.
+    :return: Map[str, spark.sql.types.DataType]. A mapping of feature names to their respective
+    Spark DataTypes.
+    """
+    complex_catalog_features = [
+        feature
+        for feature in catalog_features
+        if catalog_feature_names_to_fs_types[feature.name]
+        in DATA_TYPES_REQUIRES_DETAILS
+    ]
+    complex_catalog_feature_names_to_spark_types = {
+        feature.name: feature.data_type_details
+        for feature in complex_catalog_features
+    }
+    return complex_catalog_feature_names_to_spark_types
+
+
+def log_catalog_schema_not_match_delta_schema(catalog_features, df_schema, level):
+    """
+    Log the catalog schema does not match the delta table schema.
+
+    Example warning:
+    Expected recorded schema from Feature Catalog to be identical with
+    schema in delta table.Feature Catalog's schema is
+    '{'id': 'INTEGER', 'feat1': 'INTEGER'}' while delta table's
+    schema is '{'id': 'INTEGER', 'feat1': 'FLOAT'}'
+    """
+    catalog_schema = {feature.name: feature.data_type for feature in catalog_features}
+    delta_schema = {
+        feature.name: feature.dataType
+        for feature in df_schema
+    }
+    msg = (
+        f"Expected recorded schema from Feature Catalog to be identical with schema "
+        f"in Delta table. "
+        f"Feature Catalog's schema is '{catalog_schema}' while Delta table's schema "
+        f"is '{delta_schema}'"
+    )
+    if level == _WARN:
+        _logger.warning(msg)
+    elif level == _ERROR:
+        raise RuntimeError(msg)
+    else:
+        _logger.info(msg)

wedata/feature_store/utils/topological_sort.py

@@ -0,0 +1,158 @@
+from collections import defaultdict, deque
+from queue import PriorityQueue
+from typing import Callable, Dict, Hashable, List, Optional
+
+__all__ = ["find_cycle", "topological_sort"]
+
+
+class _NodeInfo:
+    def __init__(self, node):
+        self.node = node
+        # number of non-processed predecessors.
+        self.n_blockers = 0
+        # list of nodes that depend on this node.
+        self.successors = []
+
+
+def find_cycle(
+    node_dependencies: Dict[Hashable, List[Hashable]]
+) -> Optional[List[Hashable]]:
+    """
+    Finds a cycle in the node_dependencies graph. Returns a list of node(s) that forms a cycle or
+    None if no cycle can be found.
+    :param node_dependencies: A dict with hashable objects as keys and their list of dependency
+                    nodes as values.
+    """
+    # A stack used to perform DFS on the graph.
+    stack = deque()
+    # Another stack storing the path from root to current node in DFS. Used to detect cycle.
+    backtrack_stack = []
+    # A set of nodes that no cycle can be found starting from these nodes.
+    resolved = set()
+    # Create a copy of the dependency graph with defaultdict for convenience.
+    default_dependency = defaultdict(list)
+    default_dependency.update(node_dependencies)
+
+    # Perform DFS on every node in the graph.
+    for node in node_dependencies.keys():
+        if node in resolved:
+            # Skip a node if it's already resolved.
+            continue
+        # DFS from the node
+        stack.append(node)
+        while stack:
+            top = stack[-1]
+            if top not in backtrack_stack:
+                # First time visiting this node. There will be a second visit after the dependencies
+                # are resolved if it has dependencies.
+                backtrack_stack.append(top)
+            # If not expended after traversing the dependencies, meaning there is no dependency or
+            # all dependencies are resolved.
+            expanded = False
+            for depend in default_dependency[top]:
+                if depend in backtrack_stack:
+                    # found a cycle
+                    index = backtrack_stack.index(depend)
+                    return backtrack_stack[index:]
+                if depend in resolved:
+                    continue
+                # Only adding node to stack. backtrack_stack only contains nodes in the current DFS
+                # path.
+                stack.append(depend)
+                expanded = True
+            if not expanded:
+                stack.pop()
+                resolved.add(top)
+                backtrack_stack.pop()
+    return None
+
+
+def _all_items_in_queue_should_be_grouped(
+    queue: PriorityQueue, should_be_grouped: Callable
+) -> bool:
+    temp = []
+    should_group = True
+    # note: avoid using queue.qsize() because it's not guaranteed to be accurate.
+    while not queue.empty():
+        k, node = queue.get()
+        temp.append((k, node))
+        if not should_be_grouped(node):
+            should_group = False
+    for item in temp:
+        queue.put(item)
+    return should_group
+
+
+def topological_sort(
+    node_dependencies: Dict[Hashable, List[Hashable]],
+    key: Callable = None,
+    should_be_grouped: Callable = None,
+) -> List[Hashable]:
+    """
+    Topological sort the given node_dependencies graph. Returns a sorted list of nodes.
+    :param node_dependencies: A dict with hashable objects as keys and their list of dependency
+                    nodes as values.
+    :param key: a Callable that returns a sort key when called with a hashable object. The key is
+                    used to break ties in topological sorting. An object with smaller key is added
+                    to the result list first.
+    :raises ValueError if a cycle is found in the graph.
+    """
+    # Calling a dedicated find_cycle function to be able to give a detailed error message.
+    cycle = find_cycle(node_dependencies)
+    if cycle is not None:
+        raise ValueError(
+            "Following nodes form a cycle: ",
+            cycle,
+            ". Please resolve any circular dependencies before calling Feature Store.",
+        )
+
+    # A priority-queue storing the nodes whose dependency has been resolved.
+    # priority is determined by the given key function.
+    ready_queue = PriorityQueue()
+    # Map from node to _NodeInfo.
+    nodes = {}
+    if key is None:
+        key = hash  # use the built-in hash function by default
+    if should_be_grouped is None:
+        should_be_grouped = lambda _: False
+    # Perform Kahn's algorithm by traversing the graph starting from nodes without dependency.
+    # Node is removed from its successors' dependency once resolved. And node whose dependency gets
+    # all resolved is added to the priority queue.
+    for node, dependencies in node_dependencies.items():
+        # Initialize the graph to topologically sort based on the input node_dependencies.
+        # All nodes, its successors and number of predecessors should be populated.
+        if node not in nodes:
+            nodes[node] = _NodeInfo(node)
+        for dependency in dependencies:
+            if dependency not in nodes:
+                nodes[dependency] = _NodeInfo(dependency)
+            nodes[dependency].successors.append(node)
+        if len(dependencies):
+            nodes[node].n_blockers = len(dependencies)
+    # Initialize the ready_queue to start traversing the graph from nodes without any dependencies.
+    for node, node_info in nodes.items():
+        if node_info.n_blockers == 0:
+            ready_queue.put((key(node), node))
+    # At the end of the algorithm, result_list will have a topologically sorted listed of nodes.
+    result_list = []
+
+    def process_nodes(node_buffer, queue):
+        for node in node_buffer:
+            result_list.append(node)
+            for successor in nodes[node].successors:
+                s_info = nodes[successor]
+                s_info.n_blockers -= 1
+                if s_info.n_blockers == 0:
+                    queue.put((key(successor), successor))
+
+    while not ready_queue.empty():
+        if _all_items_in_queue_should_be_grouped(ready_queue, should_be_grouped):
+            batch_buffer = []
+            while not ready_queue.empty():
+                _, node = ready_queue.get()
+                batch_buffer.append(node)
+            process_nodes(batch_buffer, ready_queue)
+        else:
+            _, node = ready_queue.get()
+            process_nodes([node], ready_queue)
+    return result_list