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

wedata/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,
+        )

wedata/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,
+        )

wedata/feature_store/entities/function_info.py

@@ -0,0 +1,184 @@
+from collections import defaultdict
+from typing import List, Optional
+
+from pyspark.sql import Column, DataFrame
+from pyspark.sql.functions import isnull, when
+from pyspark.sql.types import StringType, StructField, StructType
+
+class FunctionParameterInfo():
+    def __init__(self, name: str, type_text: str):
+        self._name = name
+        self._type_text = type_text
+
+    @property
+    def name(self) -> str:
+        return self._name
+
+    @property
+    def type_text(self) -> str:
+        return self._type_text
+
+    @classmethod
+    def from_dict(cls, function_parameter_info_json):
+        return FunctionParameterInfo(
+            function_parameter_info_json["name"],
+            function_parameter_info_json["type_text"],
+        )
+
+
+class FunctionInfo():
+    """
+    Helper entity class that exposes properties in GetFunction's response JSON as attributes.
+    https://docs.databricks.com/api-explorer/workspace/functions/get
+
+    Note: empty fields (e.g. when 0 input parameters) are not included in the response JSON.
+    """
+
+    # Python UDFs have external_language = "Python"
+    PYTHON = "Python"
+
+    def __init__(
+        self,
+        full_name: str,
+        input_params: List[FunctionParameterInfo],
+        routine_definition: Optional[str],
+        external_language: Optional[str],
+    ):
+        self._full_name = full_name
+        self._input_params = input_params
+        self._routine_definition = routine_definition
+        self._external_language = external_language
+
+    @property
+    def full_name(self) -> str:
+        return self._full_name
+
+    @property
+    def input_params(self) -> List[FunctionParameterInfo]:
+        return self._input_params
+
+    @property
+    def routine_definition(self) -> Optional[str]:
+        return self._routine_definition
+
+    @property
+    def external_language(self) -> Optional[str]:
+        """
+        Field is None if language is SQL (not an external language).
+        """
+        return self._external_language
+
+    @classmethod
+    def from_dict(cls, function_info_json):
+        input_params = function_info_json.get("input_params", {}).get("parameters", [])
+        return FunctionInfo(
+            full_name=function_info_json["full_name"],
+            input_params=[FunctionParameterInfo.from_dict(p) for p in input_params],
+            routine_definition=function_info_json.get("routine_definition", None),
+            external_language=function_info_json.get("external_language", None),
+        )
+
+
+class InformationSchemaSparkClient:
+    """
+    Internal client to retrieve Unity Catalog metadata from system.information_schema.
+    https://docs.databricks.com/sql/language-manual/sql-ref-information-schema.html
+    """
+
+    def _get_routines_with_parameters(self, full_routine_names: List[str]) -> DataFrame:
+        """
+        Retrieve the routines with their parameters from information_schema.routines, information_schema.parameters.
+        Return DataFrame only contains routines that 1. exist and 2. the caller has GetFunction permission on.
+
+        Note: The returned DataFrame contains the cartesian product of routines and parameters.
+        For efficiency, routines table columns are only present in the first row for each routine.
+        """
+        routine_name_schema = StructType(
+            [
+                StructField("specific_catalog", StringType(), False),
+                StructField("specific_schema", StringType(), False),
+                StructField("specific_name", StringType(), False),
+            ]
+        )
+        routine_names_df = self._spark_client.createDataFrame(
+            [full_routine_name.split(".") for full_routine_name in full_routine_names],
+            routine_name_schema,
+        )
+        routines_table = self._spark_client.read_table(
+            "system.information_schema.routines"
+        )
+        parameters_table = self._spark_client.read_table(
+            "system.information_schema.parameters"
+        )
+
+        # Inner join routines table to filter out non-existent routines.
+        # Left join parameters as routines may have no parameters.
+        full_routines_with_parameters_df = routine_names_df.join(
+            routines_table, on=routine_names_df.columns, how="inner"
+        ).join(parameters_table, on=routine_names_df.columns, how="left")
+
+        # Return only relevant metadata from information_schema, sorted by routine name + parameter order.
+        # For efficiency, only preserve routine column values in the first of each routine's result rows.
+        # The first row will have parameter.ordinal_value is None (no parameters) or equals 0 (first parameter).
+        def select_if_first_row(col: Column) -> Column:
+            return when(
+                isnull(parameters_table.ordinal_position)
+                | (parameters_table.ordinal_position == 0),
+                col,
+            ).otherwise(None)
+
+        return full_routines_with_parameters_df.select(
+            routine_names_df.columns
+            + [
+                select_if_first_row(routines_table.routine_definition).alias(
+                    "routine_definition"
+                ),
+                select_if_first_row(routines_table.external_language).alias(
+                    "external_language"
+                ),
+                parameters_table.ordinal_position,
+                parameters_table.parameter_name,
+                parameters_table.full_data_type,
+            ]
+        ).sort(routine_names_df.columns + [parameters_table.ordinal_position])
+
+    def get_functions(self, full_function_names: List[str]) -> List[FunctionInfo]:
+        """
+        Retrieves and maps Unity Catalog functions' metadata as FunctionInfos.
+        """
+        # Avoid unnecessary Spark calls and return if empty.
+        if not full_function_names:
+            return []
+
+        # Collect dict of routine name -> DataFrame rows describing the routine.
+        routines_with_parameters_df = self._get_routines_with_parameters(
+            full_routine_names=full_function_names
+        )
+        routine_infos = defaultdict(list)
+        for r in routines_with_parameters_df.collect():
+            routine_name = f"{r.specific_catalog}.{r.specific_schema}.{r.specific_name}"
+            routine_infos[routine_name].append(r)
+
+        # Mock GetFunction DNE error, since information_schema does not throw.
+        for function_name in full_function_names:
+            if not function_name in routine_infos:
+                raise ValueError(f"Function '{function_name}' does not exist.")
+
+        # Map routine_infos into FunctionInfos.
+        function_infos = []
+        for function_name in full_function_names:
+            routine_info = routine_infos[function_name][0]
+            input_params = [
+                FunctionParameterInfo(name=p.parameter_name, type_text=p.full_data_type)
+                for p in routine_infos[function_name]
+                if p.ordinal_position is not None
+            ]
+            function_infos.append(
+                FunctionInfo(
+                    full_name=function_name,
+                    input_params=input_params,
+                    routine_definition=routine_info.routine_definition,
+                    external_language=routine_info.external_language,
+                )
+            )
+        return function_infos

wedata/feature_store/entities/on_demand_column_info.py

@@ -0,0 +1,44 @@
+from typing import Dict
+
+class OnDemandColumnInfo:
+    def __init__(
+        self,
+        udf_name: str,
+        input_bindings: Dict[str, str],
+        output_name: str,
+    ):
+        if not udf_name:
+            raise ValueError("udf_name must be non-empty.")
+        if not output_name:
+            raise ValueError("output_name must be non-empty.")
+
+        self._udf_name = udf_name
+        self._input_bindings = input_bindings
+        self._output_name = output_name
+
+    @property
+    def udf_name(self) -> str:
+        return self._udf_name
+
+    @property
+    def input_bindings(self) -> Dict[str, str]:
+        """
+        input_bindings is serialized as the InputBindings proto message.
+        """
+        return self._input_bindings
+
+    @property
+    def output_name(self) -> str:
+        return self._output_name
+
+    @classmethod
+    def from_proto(cls, on_demand_column_info_proto):
+        input_bindings_dict = {
+            input_binding.parameter: input_binding.bound_to
+            for input_binding in on_demand_column_info_proto.input_bindings
+        }
+        return OnDemandColumnInfo(
+            udf_name=on_demand_column_info_proto.udf_name,
+            input_bindings=input_bindings_dict,
+            output_name=on_demand_column_info_proto.output_name,
+        )

wedata/feature_store/entities/source_data_column_info.py

@@ -0,0 +1,21 @@
+
+class SourceDataColumnInfo:
+    def __init__(self, name: str):
+        if not name:
+            raise ValueError("name must be non-empty.")
+        self._name = name
+
+    @property
+    def name(self):
+        return self._name
+
+    @property
+    def output_name(self) -> str:
+        """
+        This field does not exist in the proto, and is provided for convenience.
+        """
+        return self._name
+
+    @classmethod
+    def from_proto(cls, source_data_column_info_proto):
+        return cls(name=source_data_column_info_proto.name)

wedata/feature_store/entities/training_set.py

@@ -0,0 +1,134 @@
+from typing import Dict, List, Optional
+
+from pyspark.sql import DataFrame
+
+from feature_store.entities.feature_table import FeatureTable
+from feature_store.entities.function_info import FunctionInfo
+from feature_store.utils.feature_lookup_utils import (
+    join_feature_data_if_not_overridden,
+)
+
+from feature_store.entities.feature_spec import FeatureSpec
+from feature_store.utils.feature_spec_utils import (
+    COLUMN_INFO_TYPE_FEATURE,
+    COLUMN_INFO_TYPE_ON_DEMAND,
+    COLUMN_INFO_TYPE_SOURCE,
+    get_feature_execution_groups,
+)
+
+
+class TrainingSet:
+    """
+    .. note::
+
+       Aliases: `!databricks.feature_engineering.training_set.TrainingSet`, `!databricks.feature_store.training_set.TrainingSet`
+
+    Class that defines :obj:`TrainingSet` objects.
+
+    .. note::
+
+       The :class:`TrainingSet` constructor should not be called directly. Instead,
+       call :meth:`create_training_set() <databricks.feature_engineering.client.FeatureEngineeringClient.create_training_set>`.
+    """
+
+    def __init__(
+        self,
+        feature_spec: FeatureSpec,
+        df: DataFrame,
+        labels: List[str],
+        feature_table_metadata_map: Dict[str, FeatureTable],
+        feature_table_data_map: Dict[str, DataFrame],
+        uc_function_infos: Dict[str, FunctionInfo],
+        use_spark_native_join: Optional[bool] = False,
+    ):
+        """Initialize a :obj:`TrainingSet` object."""
+        assert isinstance(
+            labels, list
+        ), f"Expected type `list` for argument `labels`. Got '{labels}' with type '{type(labels)}'."
+
+        self._feature_spec = feature_spec
+        self._df = df
+        self._labels = labels
+        self._feature_table_metadata_map = feature_table_metadata_map
+        self._feature_table_data_map = feature_table_data_map
+        self._uc_function_infos = uc_function_infos
+        self._use_spark_native_join = use_spark_native_join
+        # Perform basic validations and resolve FeatureSpec and label column data types.
+        self._validate_and_inject_dtypes()
+        self._label_data_types = {
+            name: data_type for name, data_type in df.dtypes if name in labels
+        }
+
+    @property
+    def feature_spec(self) -> FeatureSpec:
+        """Define a feature spec."""
+        return self._feature_spec
+
+    def _augment_df(self) -> DataFrame:
+        """
+        Internal helper to augment DataFrame with feature lookups and on-demand features specified in the FeatureSpec.
+        Does not drop excluded columns, and does not overwrite columns that already exist.
+        Return column order is df.columns + feature lookups + on-demand features.
+        """
+        execution_groups = get_feature_execution_groups(
+            self.feature_spec, self._df.columns
+        )
+
+        result_df = self._df
+        # Iterate over all levels and type of DAG nodes in FeatureSpec and execute them.
+        for execution_group in execution_groups:
+            if execution_group.type == COLUMN_INFO_TYPE_SOURCE:
+                continue
+            if execution_group.type == COLUMN_INFO_TYPE_FEATURE:
+                # Apply FeatureLookups
+                result_df = join_feature_data_if_not_overridden(
+                    feature_spec=self.feature_spec,
+                    df=result_df,
+                    features_to_join=execution_group.features,
+                    feature_table_metadata_map=self._feature_table_metadata_map,
+                    feature_table_data_map=self._feature_table_data_map,
+                    use_spark_native_join=self._use_spark_native_join,
+                )
+            # elif execution_group.type == COLUMN_INFO_TYPE_ON_DEMAND:
+            #     # Apply all on-demand UDFs
+            #     result_df = apply_functions_if_not_overridden(
+            #         df=result_df,
+            #         functions_to_apply=execution_group.features,
+            #         uc_function_infos=self._uc_function_infos,
+            #     )
+            else:
+                # This should never be reached.
+                raise Exception("Unknown feature execution type:", execution_group.type)
+        return result_df
+
+    def _validate_and_inject_dtypes(self):
+        """
+        Performs validations through _augment_df (e.g. Delta table exists, Delta and feature table dtypes match),
+        then inject the result DataFrame dtypes into the FeatureSpec.
+        """
+        augmented_df = self._augment_df()
+        augmented_df_dtypes = {column: dtype for column, dtype in augmented_df.dtypes}
+
+        # Inject the result DataFrame column types into the respective ColumnInfo
+        for ci in self.feature_spec.column_infos:
+            ci._data_type = augmented_df_dtypes[ci.output_name]
+
+    def load_df(self) -> DataFrame:
+        """
+        Load a :class:`DataFrame <pyspark.sql.DataFrame>`.
+
+        Return a :class:`DataFrame <pyspark.sql.DataFrame>` for training.
+
+        The returned :class:`DataFrame <pyspark.sql.DataFrame>` has columns specified
+        in the ``feature_spec`` and ``labels`` parameters provided
+        in :meth:`create_training_set() <databricks.feature_engineering.client.FeatureEngineeringClient.create_training_set>`.
+
+        :return:
+           A :class:`DataFrame <pyspark.sql.DataFrame>` for training
+        """
+        augmented_df = self._augment_df()
+        # Return only included columns in order defined by FeatureSpec + labels
+        included_columns = [
+            ci.output_name for ci in self.feature_spec.column_infos if ci.include
+        ] + self._labels
+        return augmented_df.select(included_columns)