wedata-feature-engineering 0.1.0__py3-none-any.whl → 0.1.2__py3-none-any.whl

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)

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"

feature_store/entities/feature_table.py

@@ -0,0 +1,164 @@
+from typing import Dict
+
+
+
+class FeatureTable:
+    """
+    .. note::
+
+       Aliases:`!databricks.feature_engineering.entities.feature_table.FeatureTable`, `!databricks.feature_store.entities.feature_table.FeatureTable`
+
+    Value class describing one feature table.
+
+    This will typically not be instantiated directly, instead the
+    :meth:`create_table() <databricks.feature_engineering.client.FeatureEngineeringClient.create_table>`
+    will create :class:`.FeatureTable` objects.
+    """
+
+    def __init__(
+        self,
+        name,
+        table_id,
+        description,
+        primary_keys,
+        partition_columns,
+        features,
+        creation_timestamp=None,
+        online_stores=None,
+        notebook_producers=None,
+        job_producers=None,
+        table_data_sources=None,
+        path_data_sources=None,
+        custom_data_sources=None,
+        timestamp_keys=None,
+        tags=None,
+    ):
+        """Initialize a FeatureTable object."""
+        self.name = name
+        self.table_id = table_id
+        self.description = description
+        self.primary_keys = primary_keys
+        self.partition_columns = partition_columns
+        self.features = features
+        self.creation_timestamp = creation_timestamp
+        self.online_stores = online_stores if online_stores is not None else []
+        self.notebook_producers = (
+            notebook_producers if notebook_producers is not None else []
+        )
+        self.job_producers = job_producers if job_producers is not None else []
+        self.table_data_sources = (
+            table_data_sources if table_data_sources is not None else []
+        )
+        self.path_data_sources = (
+            path_data_sources if path_data_sources is not None else []
+        )
+        self.custom_data_sources = (
+            custom_data_sources if custom_data_sources is not None else []
+        )
+        self.timestamp_keys = timestamp_keys if timestamp_keys is not None else []
+        self._tags = tags
+
+    # @property
+    # @deprecated("FeatureTable.primary_keys", since="v0.3.6")
+    # def keys(self):
+    #     return self.primary_keys
+
+    @property
+    def tags(self) -> Dict[str, str]:
+        """
+        Get the tags associated with the feature table.
+
+        :return a Dictionary of all tags associated with the feature table as key/value pairs
+        """
+        if self._tags is None:
+            # If no tags are set, self._tags is expected an empty dictionary.
+            raise ValueError(
+                "Internal error: tags have not been fetched for this FeatureTable instance"
+            )
+        return self._tags
+
+
+    @classmethod
+    def from_uc_get_table_response(cls, uc_get_table_response: Dict[str, object]):
+        """Return a FeatureStore object from a UC get_table response. Note: UC does not return online_stores or tags.
+
+        :param dict uc_get_table_response: A dictionary representing a UC get_table response.
+        :return FeatureTable: a FeatureStore object from the UC response.
+        """
+        table_name = uc_get_table_response["full_name"]
+
+        if uc_get_table_response["securable_kind"] == "TABLE_ONLINE_VIEW":
+            source_table = uc_get_table_response["properties_pairs"]["properties"][
+                "source_table"
+            ]
+            raise ValueError(
+                f"Table '{table_name}' is an online view. Online Views are not feature tables. Please use the source table '{source_table}' instead."
+            )
+
+        if (
+            "table_type" in uc_get_table_response
+            and uc_get_table_response["table_type"] == "VIEW"
+        ):
+            return cls(
+                name=table_name,
+                table_id=uc_get_table_response["table_id"],
+                description=uc_get_table_response["comment"]
+                if "comment" in uc_get_table_response
+                else "",
+                primary_keys=[],
+                partition_columns=[],
+                features=[],
+                creation_timestamp=uc_get_table_response["created_at"],
+                timestamp_keys=[],
+            )
+
+        table_constraints = (
+            uc_get_table_response["table_constraints"]
+            if "table_constraints" in uc_get_table_response
+            else []
+        )
+        primary_key_constraints = [
+            c for c in table_constraints if "primary_key_constraint" in c
+        ]
+        if len(primary_key_constraints) == 0:
+            raise ValueError(
+                "Table can't be used as a feature table because it has no primary key constraint defined."
+                + " Use 'ALTER TABLE table_name ADD CONSTRAINT table_name_pk PRIMARY KEY( key_column [,...] )'"
+                + " to add a primary key constraint on the table."
+            )
+        primary_key_constraint = primary_key_constraint = primary_key_constraints[0][
+            "primary_key_constraint"
+        ]
+        timestamp_keys = (
+            primary_key_constraint["timeseries_columns"]
+            if "timeseries_columns" in primary_key_constraint
+            else []
+        )
+        primary_keys = [
+            c
+            for c in primary_key_constraint["child_columns"]
+            if c not in timestamp_keys
+        ]
+
+        columns = uc_get_table_response["columns"]
+        features = [c["name"] for c in columns]
+        partition_columns_unordered = [c for c in columns if "partition_index" in c]
+        partition_columns = [
+            c["name"]
+            for c in sorted(
+                partition_columns_unordered, key=lambda x: x["partition_index"]
+            )
+        ]
+
+        return cls(
+            name=table_name,
+            table_id=uc_get_table_response["table_id"],
+            description=uc_get_table_response["comment"]
+            if "comment" in uc_get_table_response
+            else "",
+            primary_keys=primary_keys,
+            partition_columns=partition_columns,
+            features=features,
+            creation_timestamp=uc_get_table_response["created_at"],
+            timestamp_keys=timestamp_keys,
+        )

feature_store/entities/feature_table_info.py

@@ -0,0 +1,40 @@
+from typing import Optional
+
+
+
+class FeatureTableInfo:
+    def __init__(
+        self, table_name: str, table_id: str, lookback_window: Optional[float] = None
+    ):
+        if not table_name:
+            raise ValueError("table_name must be non-empty.")
+        if not table_id:
+            raise ValueError("table_id must be non-empty.")
+        self._table_name = table_name
+        self._table_id = table_id
+        self._lookback_window = lookback_window
+
+    @property
+    def table_name(self):
+        return self._table_name
+
+    @property
+    def table_id(self):
+        return self._table_id
+
+    @property
+    def lookback_window(self):
+        return self._lookback_window
+
+    @classmethod
+    def from_proto(cls, feature_table_info_proto):
+        lookback_window = (
+            feature_table_info_proto.lookback_window
+            if feature_table_info_proto.HasField("lookback_window")
+            else None
+        )
+        return cls(
+            table_name=feature_table_info_proto.table_name,
+            table_id=feature_table_info_proto.table_id,
+            lookback_window=lookback_window,
+        )