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

feature_store/entities/function_info.py

@@ -1,184 +0,0 @@
-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

feature_store/entities/on_demand_column_info.py

@@ -1,44 +0,0 @@
-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,
-        )

feature_store/entities/source_data_column_info.py

@@ -1,21 +0,0 @@
-
-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)

feature_store/entities/training_set.py

@@ -1,134 +0,0 @@
-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)