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

wedata/feature_store/entities/feature_lookup.py

@@ -0,0 +1,179 @@
+import copy
+import datetime
+import logging
+from typing import Dict, List, Optional, Union
+
+from feature_store.utils import common_utils
+
+_logger = logging.getLogger(__name__)
+
+
+class FeatureLookup:
+
+    """
+    特征查找类
+
+    特征查找类用于指定特征表中的特征，并将其与训练集中的特征进行关联。
+
+    特征查找类有以下属性：
+
+    - table_name：特征表的名称。
+    - lookup_key：用于在特征表和训练集之间进行联接的键。lookup_key必须是训练集中的列。lookup_key的类型和顺序必须与特征表的主键匹配。
+    - feature_names：要从特征表中查找的特征的名称。如果您的模型需要主键作为特征，则可以将它们声明为独立的FeatureLookups。
+    - rename_outputs：如果提供，则会将特征重命名为 :meth:`create_training_set() <databricks.feature_engineering.client.FeatureEngineeringClient.create_training_set>`返回的 :class:`TrainingSet <databricks.ml_features.training_set.TrainingSet>` 中的特征。
+    - timestamp_lookup_key：用于在特征表和训练集之间进行联接的时间戳键。timestamp_lookup_key必须是训练集中的列。timestamp_lookup_key的类型必须与特征表的时间戳键的类型匹配。
+    - lookback_window: 当对特征表执行时间点查找时使用的回溯窗口，该查找针对传递给 :meth:`create_training_set() <databricks.feature_engineering.client.FeatureEngineeringClient.create_training_set>` 方法的数据帧。特征存储将检索在数据帧的``timestamp_lookup_key``指定时间戳之前且在``lookback_window``时间范围内的最新特征值，如果不存在这样的特征值则返回null。当设置为0时，仅返回特征表中的精确匹配项。
+    - feature_name：特征名称。**已弃用**。使用 `feature_names`。
+    - output_name：如果提供，则会将此特征重命名为 :meth:`create_training_set() <databricks.feature_engineering.client.FeatureEngineeringClient.create_training_set>` 返回的 :class:`TrainingSet <databricks.ml_features.training_set.TrainingSet>` 中的特征。**已弃用**。使用 `rename_outputs`。
+
+    示例：
+
+    from databricks.feature_store import FeatureLookup
+
+    lookup = FeatureLookup(
+        table_name="my_feature_table",
+        lookup_key="my_lookup_key",
+        feature_names=["my_feature_1", "my_feature_2"],
+        rename_outputs={"my_feature_1": "my_feature_1_renamed"},
+        timestamp_lookup_key="my_timestamp_lookup_key",
+        lookback_window=datetime.timedelta(days=1)
+    )
+
+    """
+
+    def __init__(
+        self,
+        table_name: str,
+        lookup_key: Union[str, List[str]],
+        *,
+        feature_names: Union[str, List[str], None] = None,
+        rename_outputs: Optional[Dict[str, str]] = None,
+        timestamp_lookup_key: Optional[str] = None,
+        lookback_window: Optional[datetime.timedelta] = None,
+        **kwargs,
+    ):
+        """Initialize a FeatureLookup object. See class documentation."""
+
+        self._feature_name_deprecated = kwargs.pop("feature_name", None)
+        self._output_name_deprecated = kwargs.pop("output_name", None)
+
+        if kwargs:
+            raise TypeError(
+                f"FeatureLookup got unexpected keyword argument(s): {list(kwargs.keys())}"
+            )
+
+        self._table_name = table_name
+
+        if type(timestamp_lookup_key) is list:
+            if len(timestamp_lookup_key) == 0:
+                timestamp_lookup_key = None
+            elif len(timestamp_lookup_key) == 1:
+                timestamp_lookup_key = timestamp_lookup_key[0]
+            else:
+                raise ValueError(
+                    f"Setting multiple timestamp lookup keys is not supported."
+                )
+
+        if rename_outputs is not None and not isinstance(rename_outputs, dict):
+            raise ValueError(
+                f"Unexpected type for rename_outputs: {type(rename_outputs)}"
+            )
+
+        self._feature_names = common_utils.as_list(feature_names, default=[])
+
+        # Make sure the user didn't accidentally pass in any nested lists/dicts in feature_names
+        for fn in self._feature_names:
+            if not isinstance(fn, str):
+                raise ValueError(
+                    f"Unexpected type for element in feature_names: {type(self._feature_names)}, only strings allowed in list"
+                )
+
+        if lookback_window is not None:
+            if not timestamp_lookup_key:
+                raise ValueError(
+                    f"Unexpected lookback_window value: {lookback_window}, lookback windows can only be applied on time series "
+                    f"feature tables. Use timestamp_lookup_key to perform point-in-time lookups with lookback window."
+                )
+            if not isinstance(
+                lookback_window, datetime.timedelta
+            ) or lookback_window < datetime.timedelta(0):
+                raise ValueError(
+                    f"Unexpected value for lookback_window: {lookback_window}, only non-negative datetime.timedelta allowed."
+                )
+
+        self._lookup_key = copy.copy(lookup_key)
+        self._timestamp_lookup_key = copy.copy(timestamp_lookup_key)
+        self._lookback_window = copy.copy(lookback_window)
+
+        self._rename_outputs = {}
+        if rename_outputs is not None:
+            self._rename_outputs = rename_outputs.copy()
+
+        self._inject_deprecated_feature_name()
+        self._inject_deprecated_output_name()
+
+    @property
+    def table_name(self):
+        """The table name to use in this FeatureLookup."""
+        return self._table_name
+
+    @property
+    def lookup_key(self):
+        """The lookup key(s) to use in this FeatureLookup."""
+        return self._lookup_key
+
+    @property
+    def feature_name(self):
+        """The feature name to use in this FeatureLookup. **Deprecated**. Use `feature_names`."""
+        return self._feature_name_deprecated
+
+    @property
+    def feature_names(self):
+        """The feature names to use in this FeatureLookup."""
+        return self._feature_names
+
+    @property
+    def output_name(self):
+        """The output name to use in this FeatureLookup. **Deprecated**. Use `feature_names`."""
+        if self._output_name_deprecated:
+            return self._output_name_deprecated
+        else:
+            return self._feature_name_deprecated
+
+    @property
+    def timestamp_lookup_key(self):
+        return self._timestamp_lookup_key
+
+    @property
+    def lookback_window(self):
+        """A lookback window applied only for point-in-time lookups."""
+        return self._lookback_window
+
+    def _get_feature_names(self):
+        return self._feature_names
+
+    def _get_output_name(self, feature_name):
+        """Lookup the renamed output, or fallback to the feature name itself if no mapping is present"""
+        return self._rename_outputs.get(feature_name, feature_name)
+
+    def _inject_deprecated_feature_name(self):
+        if self._feature_name_deprecated:
+            if len(self._feature_names) > 0:
+                raise ValueError(
+                    "Use either feature_names or feature_name parameter, but not both."
+                )
+            _logger.warning(
+                f'The feature_name parameter is deprecated. Use "feature_names".'
+            )
+            self._feature_names = [self._feature_name_deprecated]
+
+    def _inject_deprecated_output_name(self):
+        if len(self._feature_names) == 1 and self._output_name_deprecated:
+            if len(self._rename_outputs) > 0:
+                raise ValueError(
+                    "Use either output_name or rename_outputs parameter, but not both."
+                )
+            _logger.warning(
+                f'The output_name parameter is deprecated.  Use "rename_outputs".'
+            )
+            self._rename_outputs[self._feature_names[0]] = self._output_name_deprecated

wedata/feature_store/entities/feature_spec.py

@@ -0,0 +1,454 @@
+import os
+from typing import Any, Dict, List, Type, Union
+
+import mlflow
+from databricks.sdk.service.catalog import FunctionInfo
+from google.protobuf.json_format import MessageToDict, ParseDict
+from mlflow.utils.file_utils import TempDir, read_yaml, write_yaml
+
+from feature_store.entities.column_info import ColumnInfo
+from feature_store.entities.feature_column_info import FeatureColumnInfo
+from feature_store.entities.feature_spec_constants import (
+    BOUND_TO,
+    DATA_TYPE,
+    FEATURE_COLUMN_INFO,
+    FEATURE_STORE,
+    INCLUDE,
+    INPUT_BINDINGS,
+    INPUT_COLUMNS,
+    INPUT_FUNCTIONS,
+    INPUT_TABLES,
+    NAME,
+    ON_DEMAND_COLUMN_INFO,
+    ON_DEMAND_FEATURE,
+    OUTPUT_NAME,
+    PARAMETER,
+    SERIALIZATION_VERSION,
+    SOURCE,
+    SOURCE_DATA_COLUMN_INFO,
+    TABLE_NAME,
+    TOPOLOGICAL_ORDERING,
+    TRAINING_DATA,
+    UDF_NAME,
+)
+from feature_store.entities.feature_table_info import FeatureTableInfo
+from feature_store.entities.on_demand_column_info import OnDemandColumnInfo
+from feature_store.entities.source_data_column_info import SourceDataColumnInfo
+from feature_store.utils import common_utils
+
+# Change log for serialization version. Please update for each serialization version.
+# 1. Initial.
+# 2. (2021/06/16): Record feature_store_client_version to help us make backward compatible changes in the future.
+# 3. (2021/08/25): Record table_id to handle feature table lineage stability if tables are deleted.
+# 4. (2021/09/25): Record timestamp_lookup_key to handle point-in-time lookups.
+# 5. (2021/02/15): Record include flag for column info if False.
+#                  Record input functions as FunctionInfo and function computation as OnDemandColumnInfo.
+#                  Remove redundant fields: table_name from table_infos, output_name from column_infos.
+# 6. (2023/04/21): Record lookback_window in table info for point-in-time lookups.
+# 7. (2023/05/05): Record the Spark data type for all columns to track model signatures.
+# 8. (2023/08/14): Record the topological_ordering for all columns to support chained transform and lookup.
+# 9. (2023/09/11): Change the type of lookback_window from int to double for sub-second values
+
+
+class FeatureSpec:
+
+    FEATURE_ARTIFACT_FILE = "feature_spec.yaml"
+    SERIALIZATION_VERSION_NUMBER = 9
+
+    def __init__(
+        self,
+        column_infos: List[ColumnInfo],
+        table_infos: List[FeatureTableInfo],
+        function_infos: List[FunctionInfo],
+        workspace_id: int = None,
+        feature_store_client_version: str = None,
+        serialization_version: int = None,
+    ):
+        self._column_infos = column_infos
+        self._table_infos = table_infos
+        self._function_infos = function_infos
+        self._workspace_id = workspace_id
+        # The Feature Store Python client version which wrote this FeatureSpec.
+        # If empty, the client version is <=0.3.1.
+        self._feature_store_client_version = feature_store_client_version
+        self._serialization_version = serialization_version
+
+        # Perform validations
+        self._validate_column_infos()
+        self._validate_table_infos()
+        self._validate_function_infos()
+
+    def _validate_column_infos(self):
+        if not self.column_infos:
+            raise ValueError("column_infos must be non-empty.")
+
+        for column_info in self.column_infos:
+            if not isinstance(column_info, ColumnInfo):
+                raise ValueError(
+                    f"Expected all elements of column_infos to be instances of ColumnInfo. "
+                    f"'{column_info}' is of the wrong type."
+                )
+            if (
+                self._serialization_version >= 8
+                and column_info.topological_ordering is not None
+            ):
+                ordering = column_info.topological_ordering
+                if not isinstance(ordering, int) or ordering < 0:
+                    raise ValueError(
+                        "The topological_ordering of column_info must be non non-negative integers."
+                    )
+
+    def _validate_table_infos(self):
+        if self.table_infos is None:
+            raise ValueError("Internal Error: table_infos must be provided.")
+
+        # table_infos should not be duplicated
+        common_utils.validate_strings_unique(
+            [table_info.table_name for table_info in self.table_infos],
+            "Internal Error: Expect all table_names in table_infos to be unique. Found duplicates {}",
+        )
+
+        # Starting FeatureSpec v3, unique table names in table_infos must match those in column_infos.
+        if self.serialization_version >= 3:
+            unique_table_names = set(
+                [table_info.table_name for table_info in self.table_infos]
+            )
+            unique_column_table_names = set(
+                [fci.table_name for fci in self.feature_column_infos]
+            )
+            if unique_table_names != unique_column_table_names:
+                raise Exception(
+                    f"Internal Error: table_names from table_infos {sorted(unique_table_names)} "
+                    f"must match those from column_infos {sorted(unique_column_table_names)}"
+                )
+
+    def _validate_function_infos(self):
+        if self.function_infos is None:
+            raise ValueError("Internal Error: function_infos must be provided.")
+
+        # function_infos should not be duplicated
+        common_utils.validate_strings_unique(
+            [function_info.udf_name for function_info in self.function_infos],
+            "Internal Error: Expect all udf_names in function_infos to be unique. Found duplicates {}",
+        )
+
+        # Unique UDF names in function_infos must match those in column_infos.
+        # No version check is required as both fields were added simultaneously in FeatureSpec v5.
+        unique_udf_names = set(
+            [function_info.udf_name for function_info in self.function_infos]
+        )
+        unique_column_udf_names = set(
+            [odci.udf_name for odci in self.on_demand_column_infos]
+        )
+        if unique_udf_names != unique_column_udf_names:
+            raise Exception(
+                f"Internal Error: udf_names from function_infos {sorted(unique_udf_names)} "
+                f"must match those from column_infos {sorted(unique_column_udf_names)}"
+            )
+
+    @property
+    def column_infos(self):
+        return self._column_infos
+
+    @property
+    def table_infos(self):
+        return self._table_infos
+
+    @property
+    def function_infos(self):
+        return self._function_infos
+
+    @property
+    def workspace_id(self):
+        return self._workspace_id
+
+    @property
+    def feature_column_infos(self) -> List[FeatureColumnInfo]:
+        return self._get_infos_of_type(FeatureColumnInfo)
+
+    @property
+    def on_demand_column_infos(self) -> List[OnDemandColumnInfo]:
+        return self._get_infos_of_type(OnDemandColumnInfo)
+
+    @property
+    def serialization_version(self) -> int:
+        return self._serialization_version
+
+    def _get_infos_of_type(
+            self,
+            info_type: Union[
+                Type[SourceDataColumnInfo],
+                Type[FeatureColumnInfo],
+                Type[OnDemandColumnInfo],
+            ],
+    ):
+        """
+        Helper method to return the ColumnInfo.info subinfo field based on its type.
+        """
+        return [
+            column_info.info
+            for column_info in self.column_infos
+            if isinstance(column_info.info, info_type)
+        ]
+
+    @classmethod
+    def from_proto(cls, feature_spec_proto):
+        # Serialization version is not deserialized from the proto as there is currently only one
+        # possible version.
+        column_infos = [
+            ColumnInfo.from_proto(column_info_proto)
+            for column_info_proto in feature_spec_proto.input_columns
+        ]
+        table_infos = [
+            FeatureTableInfo.from_proto(table_info_proto)
+            for table_info_proto in feature_spec_proto.input_tables
+        ]
+        function_infos = [
+            FunctionInfo.from_proto(function_info_proto)
+            for function_info_proto in feature_spec_proto.input_functions
+        ]
+        return cls(
+            column_infos=column_infos,
+            table_infos=table_infos,
+            function_infos=function_infos,
+            workspace_id=feature_spec_proto.workspace_id,
+            feature_store_client_version=feature_spec_proto.feature_store_client_version,
+            serialization_version=feature_spec_proto.serialization_version,
+        )
+
+
+    @staticmethod
+    def _input_columns_proto_to_yaml_dict(column_info: Dict[str, Any]):
+        """
+        Converts a single ColumnInfo's proto dict to the expected element in FeatureSpec YAML's input_columns.
+        To keep the YAML clean, unnecessary fields are removed (e.g. SourceDataColumnInfo.name field, ColumnInfo.include when True).
+
+        Example of a column_info transformation. Note that "name" and "include" attributes were excluded.
+        {"source_data_column_info": {"name": "source_column"}, "include": True} -> {"source_column": {"source": "training_data"}}
+
+        Order of elements in the YAML dict should be:
+        1. Attributes present in ColumnInfo.info, using the proto field order
+        2. Remaining attributes of ColumnInfo, using the proto field order
+        3. Feature Store source type
+        """
+        # Parse oneof field ColumnInfo.info level attributes as column_info_attributes; record column_name, source
+        if SOURCE_DATA_COLUMN_INFO in column_info:
+            column_info_attributes = column_info[SOURCE_DATA_COLUMN_INFO]
+            # pop NAME attribute and use as the YAML key for this column_info to avoid redundancy in YAML
+            column_name, source = column_info_attributes.pop(NAME), TRAINING_DATA
+        elif FEATURE_COLUMN_INFO in column_info:
+            column_info_attributes = column_info[FEATURE_COLUMN_INFO]
+            # pop OUTPUT_NAME attribute and use as the YAML key for this column_info to avoid redundancy in YAML
+            column_name, source = column_info_attributes.pop(OUTPUT_NAME), FEATURE_STORE
+        elif ON_DEMAND_COLUMN_INFO in column_info:
+            column_info_attributes = column_info[ON_DEMAND_COLUMN_INFO]
+            # Map InputBindings message dictionary to {parameter: bound_to} KV dictionary if defined
+            if INPUT_BINDINGS in column_info_attributes:
+                column_info_attributes[INPUT_BINDINGS] = {
+                    ib[PARAMETER]: ib[BOUND_TO]
+                    for ib in column_info_attributes[INPUT_BINDINGS]
+                }
+            # pop OUTPUT_NAME attribute and use as the YAML key for this column_info to avoid redundancy in YAML
+            column_name, source = (
+                column_info_attributes.pop(OUTPUT_NAME),
+                ON_DEMAND_FEATURE,
+            )
+        else:
+            raise ValueError(
+                f"Expected column_info to be keyed by a valid ColumnInfo.info type. "
+                f"'{column_info}' has key '{list(column_info)[0]}'."
+            )
+
+        # Parse and insert ColumnInfo level attributes
+        # Note: the ordering of fields in the result yaml file is undefined but in reality, they are
+        # in the same order as they are added in the column_info_attributes dict.
+
+        # DATA_TYPE is supported starting FeatureSpec v7 and is not guaranteed to exist.
+        if DATA_TYPE in column_info:
+            column_info_attributes[DATA_TYPE] = column_info[DATA_TYPE]
+        if not column_info[INCLUDE]:
+            column_info_attributes[INCLUDE] = False
+        # TOPOLOGICAL_ORDERING is supported starting FeatureSpec v8.
+        if TOPOLOGICAL_ORDERING in column_info:
+            column_info_attributes[TOPOLOGICAL_ORDERING] = column_info[
+                TOPOLOGICAL_ORDERING
+            ]
+
+        # Insert source; return YAML keyed by column_name
+        column_info_attributes[SOURCE] = source
+        return {column_name: column_info_attributes}
+
+    def _to_dict(self):
+        """
+        Convert FeatureSpec to a writeable YAML artifact. Uses MessageToDict to convert FeatureSpec proto to dict.
+        Sanitizes and modifies the dict as follows:
+        1. Remove redundant or unnecessary information for cleanliness in the YAML
+        2. Modifies the dict to be of the format {column_name: column_attributes_dict}
+
+        :return: Sanitized FeatureSpec dictionary of {column_name: column_attributes}
+        """
+        yaml_dict = MessageToDict(self.to_proto(), preserving_proto_field_name=True)
+        yaml_dict[INPUT_COLUMNS] = [
+            self._input_columns_proto_to_yaml_dict(column_info)
+            for column_info in yaml_dict[INPUT_COLUMNS]
+        ]
+
+        if INPUT_TABLES in yaml_dict:
+            # pop TABLE_NAME attribute and use as the YAML key for each table_info to avoid redundancy in YAML
+            yaml_dict[INPUT_TABLES] = [
+                {table_info.pop(TABLE_NAME): table_info}
+                for table_info in yaml_dict[INPUT_TABLES]
+            ]
+        if INPUT_FUNCTIONS in yaml_dict:
+            # pop UDF_NAME attribute and use as the YAML key for each table_info to avoid redundancy in YAML
+            yaml_dict[INPUT_FUNCTIONS] = [
+                {function_info.pop(UDF_NAME): function_info}
+                for function_info in yaml_dict[INPUT_FUNCTIONS]
+            ]
+
+        # For readability, place SERIALIZATION_VERSION last in the dictionary.
+        yaml_dict[SERIALIZATION_VERSION] = yaml_dict.pop(SERIALIZATION_VERSION)
+        return yaml_dict
+
+    def save(self, path: str):
+        """
+        Convert spec to a YAML artifact and store at given `path` location.
+        :param path: Root path to where YAML artifact is expected to be stored.
+        :return: None
+        """
+        write_yaml(
+            root=path,
+            file_name=self.FEATURE_ARTIFACT_FILE,
+            data=self._to_dict(),
+            sort_keys=False,
+        )
+
+    @staticmethod
+    def _input_columns_yaml_to_proto_dict(column_info: Dict[str, Any]):
+        """
+        Convert the FeatureSpec YAML dictionary to the expected ColumnInfo proto dictionary.
+
+        Example of a column_info transformation.
+        {"source_column": {"source": "training_data"}} -> {"source_data_column_info": {"name": "source_column"}}
+        """
+        if len(column_info) != 1:
+            raise ValueError(
+                f"Expected column_info dictionary to only have one key, value pair. "
+                f"'{column_info}' has length {len(column_info)}."
+            )
+        column_name, column_data = list(column_info.items())[0]
+        if not column_data:
+            raise ValueError(
+                f"Expected values of '{column_name}' dictionary to be non-empty."
+            )
+        if SOURCE not in column_data:
+            raise ValueError(
+                f"Expected values of column_info dictionary to include the source. No source found "
+                f"for '{column_name}'."
+            )
+
+        # Parse oneof field ColumnInfo.info level attributes
+        source = column_data.pop(SOURCE)
+        if source == TRAINING_DATA:
+            column_data[NAME] = column_name
+            column_info_dict = {SOURCE_DATA_COLUMN_INFO: column_data}
+        elif source == FEATURE_STORE:
+            column_data[OUTPUT_NAME] = column_name
+            column_info_dict = {FEATURE_COLUMN_INFO: column_data}
+        elif source == ON_DEMAND_FEATURE:
+            column_data[OUTPUT_NAME] = column_name
+            # Map {parameter_val: bound_to_val} dictionary to InputBindings(parameter, bound_to) message dictionary.
+            column_data[INPUT_BINDINGS] = [
+                {PARAMETER: parameter, BOUND_TO: bound_to}
+                for parameter, bound_to in column_data.get(INPUT_BINDINGS, {}).items()
+            ]
+            column_info_dict = {ON_DEMAND_COLUMN_INFO: column_data}
+        else:
+            raise ValueError(
+                f"Internal Error: Expected column_info to have source matching oneof ColumnInfo.info. "
+                f"'{column_info}' has source of '{source}'."
+            )
+
+        # Parse ColumnInfo level attributes
+        # TOPOLOGICAL_ORDERING is supported starting FeatureSpec v8.
+        if TOPOLOGICAL_ORDERING in column_data:
+            column_info_dict[TOPOLOGICAL_ORDERING] = column_data.pop(
+                TOPOLOGICAL_ORDERING
+            )
+        # DATA_TYPE is supported starting FeatureSpec v7 and is not guaranteed to exist.
+        if DATA_TYPE in column_data:
+            column_info_dict[DATA_TYPE] = column_data.pop(DATA_TYPE)
+        # INCLUDE is supported starting FeatureSpec v5 and only present in the YAML when INCLUDE = False
+        if INCLUDE in column_data:
+            column_info_dict[INCLUDE] = column_data.pop(INCLUDE)
+        return column_info_dict
+
+    # @classmethod
+    # def _from_dict(cls, spec_dict):
+    #     """
+    #     Convert YAML artifact to FeatureSpec. Transforms YAML artifact to dict keyed by
+    #     source_data_column_info or feature_column_info, such that ParseDict can convert the dict to
+    #     a proto message, and from_proto can convert the proto message to a FeatureSpec object
+    #     :return: :py:class:`~databricks.ml_features_common.entities.feature_spec.FeatureSpec`
+    #     """
+    #     if INPUT_COLUMNS not in spec_dict:
+    #         raise ValueError(
+    #             f"{INPUT_COLUMNS} must be a key in {cls.FEATURE_ARTIFACT_FILE}."
+    #         )
+    #     if not spec_dict[INPUT_COLUMNS]:
+    #         raise ValueError(
+    #             f"{INPUT_COLUMNS} in {cls.FEATURE_ARTIFACT_FILE} must be non-empty."
+    #         )
+    #     spec_dict[INPUT_COLUMNS] = [
+    #         cls._input_columns_yaml_to_proto_dict(column_info)
+    #         for column_info in spec_dict[INPUT_COLUMNS]
+    #     ]
+    #
+    #     # feature_spec.yaml doesn't include input_tables, input_functions if any are true:
+    #     # 1. The YAML is written by an older client that does not support the functionality.
+    #     # 2. The FeatureSpec does not contain FeatureLookups (input_tables), FeatureFunctions (input_functions).
+    #     input_tables = []
+    #     for input_table in spec_dict.get(INPUT_TABLES, []):
+    #         table_name, attributes = list(input_table.items())[0]
+    #         input_tables.append({TABLE_NAME: table_name, **attributes})
+    #     spec_dict[INPUT_TABLES] = input_tables
+    #
+    #     input_functions = []
+    #     for input_function in spec_dict.get(INPUT_FUNCTIONS, []):
+    #         udf_name, attributes = list(input_function.items())[0]
+    #         input_functions.append({UDF_NAME: udf_name, **attributes})
+    #     spec_dict[INPUT_FUNCTIONS] = input_functions
+    #
+    #     return cls.from_proto(
+    #         ParseDict(spec_dict, ProtoFeatureSpec(), ignore_unknown_fields=True)
+    #     )
+
+    @classmethod
+    def _read_file(cls, path: str):
+        """
+        Read the YAML artifact from a file path.
+        """
+        parent_dir, file = os.path.split(path)
+        spec_dict = read_yaml(parent_dir, file)
+        return cls._from_dict(spec_dict)
+
+    @classmethod
+    def load(cls, path: str):
+        """
+        Load the FeatureSpec YAML artifact in the provided root directory (at path/feature_spec.yaml).
+
+        :param path: Root path to the YAML artifact. This can be a MLflow artifact path or file path.
+        :return: :py:class:`~databricks.ml_features_common.entities.feature_spec.FeatureSpec`
+        """
+        # Create the full file path to the FeatureSpec.
+        path = os.path.join(path, cls.FEATURE_ARTIFACT_FILE)
+
+        if common_utils.is_artifact_uri(path):
+            with TempDir() as tmp_location:
+                # Returns a file and not directory since the artifact_uri is a single file.
+                local_path = mlflow.artifacts.download_artifacts(
+                    artifact_uri=path, dst_path=tmp_location.path()
+                )
+                return FeatureSpec._read_file(local_path)
+        else:
+            return FeatureSpec._read_file(path)

wedata/feature_store/entities/feature_spec_constants.py

@@ -0,0 +1,25 @@
+# Field names from feature_spec.proto.
+SOURCE_DATA_COLUMN_INFO = "source_data_column_info"
+FEATURE_COLUMN_INFO = "feature_column_info"
+ON_DEMAND_COLUMN_INFO = "on_demand_column_info"
+INPUT_COLUMNS = "input_columns"
+NAME = "name"
+OUTPUT_NAME = "output_name"
+INPUT_TABLES = "input_tables"
+TABLE_NAME = "table_name"
+TABLE_ID = "table_id"
+SERIALIZATION_VERSION = "serialization_version"
+INPUT_FUNCTIONS = "input_functions"
+INCLUDE = "include"
+DATA_TYPE = "data_type"
+TOPOLOGICAL_ORDERING = "topological_ordering"
+UDF_NAME = "udf_name"
+INPUT_BINDINGS = "input_bindings"
+PARAMETER = "parameter"
+BOUND_TO = "bound_to"
+
+# FeatureSpec YAML source field and allowed values
+SOURCE = "source"
+TRAINING_DATA = "training_data"
+FEATURE_STORE = "feature_store"
+ON_DEMAND_FEATURE = "on_demand_feature"