databricks-feature-engineering 0.8.0__tar.gz → 0.12.1__tar.gz

{databricks_feature_engineering-0.8.0 → databricks_feature_engineering-0.12.1}/PKG-INFO

@@ -1,6 +1,6 @@
-Metadata-Version: 2.1
+Metadata-Version: 2.4
 Name: databricks-feature-engineering
-Version: 0.8.0
+Version: 0.12.1
 Summary: Databricks Feature Engineering Client
 Author-email: Databricks <feedback@databricks.com>
 License: Databricks Proprietary License
@@ -10,15 +10,16 @@ Requires-Python: >=3.8.10
 Description-Content-Type: text/markdown
 License-File: LICENSE.md
 License-File: NOTICE.md
-Requires-Dist: mlflow-skinny[databricks]<3,>=2.11.0
+Requires-Dist: mlflow-skinny[databricks]<4,>=2.16.0
 Requires-Dist: pyyaml<7,>=6
 Requires-Dist: boto3<2,>=1.16.7
 Requires-Dist: dbl-tempo<1,>=0.1.26
 Requires-Dist: azure-cosmos==4.3.1
-Requires-Dist: numpy<2,>=1.19.2
-Requires-Dist: protobuf<5,>=3.12.0
+Requires-Dist: numpy<3,>=1.19.2
+Requires-Dist: protobuf<6,>=3.12.0
 Requires-Dist: flask<3,>=1.1.2
 Requires-Dist: sqlparse<1,>=0.4.0
+Dynamic: license-file
 
 The Databricks Feature Engineering client is used to:
 * Create, read, and write feature tables

{databricks_feature_engineering-0.8.0 → databricks_feature_engineering-0.12.1}/databricks/feature_engineering/client.py

@@ -1,4 +1,5 @@
 import logging
+import os
 from types import ModuleType
 from typing import Any, Dict, List, Optional, Union
 
@@ -58,6 +59,7 @@ from databricks.ml_features.constants import (
     _DEFAULT_PUBLISH_STREAM_TRIGGER,
     _DEFAULT_WRITE_STREAM_TRIGGER,
     MERGE,
+    MODEL_DATA_PATH_ROOT,
 )
 from databricks.ml_features.entities.feature_serving_endpoint import (
     EndpointCoreConfig,
@@ -74,6 +76,7 @@ from databricks.ml_features.utils.feature_utils import (
 )
 from databricks.ml_features.utils.request_context import RequestContext
 from databricks.ml_features.utils.utils import as_list, enable_if
+from databricks.ml_features_common import mlflow_model_constants
 from databricks.ml_features_common.utils import uc_utils
 
 _logger = logging.getLogger(__name__)
@@ -372,6 +375,7 @@ class FeatureEngineeringClient:
         feature_spec: Optional[str] = None,
         label: Union[str, List[str], None],
         exclude_columns: Optional[List[str]] = None,
+        use_spark_native_join: bool = False,
         **kwargs,
     ) -> TrainingSet:
         """
@@ -382,10 +386,11 @@ class FeatureEngineeringClient:
           :class:`FeatureLookups <databricks.ml_features.entities.feature_lookup.FeatureLookup>` are
           joined into the :class:`DataFrame <pyspark.sql.DataFrame>`, and
           :class:`FeatureFunctions <databricks.ml_features.entities.feature_function.FeatureFunction>` are
-          computed on-demand.
-        :param feature_spec: Name of the feature spec for training set creation.
+          computed on-demand. `feature_lookups` cannot be specified if `feature_spec` is provided.
+        :param feature_spec: Full name of the FeatureSpec in Unity Catalog. `feature_spec` cannot be specified if `feature_lookups` is provided.
         :param label: Names of column(s) in :class:`DataFrame <pyspark.sql.DataFrame>` that contain training set labels. To create a training set without a label field, i.e. for unsupervised training set, specify label = None.
         :param exclude_columns: Names of the columns to drop from the :class:`TrainingSet <databricks.ml_features.training_set.TrainingSet>` :class:`DataFrame <pyspark.sql.DataFrame>`.
+        :param use_spark_native_join: Use spark to optimize table joining performance. The optimization is only applicable when `Photon <https://docs.databricks.com/en/compute/photon.html>` is enabled.
         :return: A :class:`TrainingSet <databricks.ml_features.training_set.TrainingSet>` object.
         """
 
@@ -411,6 +416,7 @@ class FeatureEngineeringClient:
                 label=label,
                 exclude_columns=exclude_columns,
                 client_name=request_context.FEATURE_ENGINEERING_CLIENT,
+                use_spark_native_join=use_spark_native_join,
                 **kwargs,
             )
 
@@ -432,6 +438,7 @@ class FeatureEngineeringClient:
             label=label,
             exclude_columns=exclude_columns,
             client_name=request_context.FEATURE_ENGINEERING_CLIENT,
+            use_spark_native_join=use_spark_native_join,
             **kwargs,
         )
 
@@ -602,6 +609,19 @@ class FeatureEngineeringClient:
 
           .. note::
             - ``signature`` is not recommended and it's preferred to use ``infer_input_example``.
+            - ``output_schema``: When logging a model without labels in the training set, you must pass ``output_schema`` to ``log_model`` to suggest the output schema explicitly. For example:
+
+                .. code-block:: python
+
+                    from mlflow.types import ColSpec, DataType, Schema
+
+                    output_schema = Schema([ColSpec(DataType.???)]) # Refer to mlflow signature types for the right choice of type here
+                    ...
+                    fe.log_model(
+                        ...
+                        output_schema=output_schema
+                    )
+
 
         :return: `None`
         """
@@ -728,6 +748,7 @@ class FeatureEngineeringClient:
         result_type: str = "double",
         env_manager: str = LOCAL,
         params: Optional[dict[str, Any]] = None,
+        use_spark_native_join: bool = False,
         **kwargs,
     ) -> DataFrame:
         """
@@ -827,6 +848,7 @@ class FeatureEngineeringClient:
         :param env_manager: The environment manager to use in order to create the python environment for model inference.
            See :func:`mlflow.pyfunc.spark_udf` env_manager.
         :param params: Additional parameters to pass to the model for inference.
+        :param use_spark_native_join: Use spark to optimize table joining performance. The optimization is only applicable when `Photon <https://docs.databricks.com/en/compute/photon.html>` is enabled.
         :return: A :class:`DataFrame <pyspark.sql.DataFrame>`
            containing:
 
@@ -844,6 +866,7 @@ class FeatureEngineeringClient:
             env_manager=env_manager,
             client_name=request_context.FEATURE_ENGINEERING_CLIENT,
             params=params,
+            use_spark_native_join=use_spark_native_join,
             **kwargs,
         )
 
@@ -1023,3 +1046,22 @@ class FeatureEngineeringClient:
         )
 
         return materialized_view_info
+
+    def build_model_env(self, model_uri: str, save_path: str) -> str:
+        """
+        Prebuild the model Python environment required by the given model and generate an archive file saved to the specified ``save_path``.
+        The resulting environment can then be used in :meth:`FeatureEngineeringClient.score_batch` as the ``prebuilt_env_uri`` parameter.
+
+        :param model_uri: URI of the model used to build the Python environment.
+        :param save_path: Directory path to save the prebuilt model environment archive file.
+          This can be a local directory path, a mounted DBFS path (e.g., '/dbfs/...'),
+          or a mounted UC volume path (e.g., '/Volumes/...').
+
+        :return: The path of the archive file containing the Python environment data.
+        """
+        artifact_path = os.path.join(mlflow.pyfunc.DATA, MODEL_DATA_PATH_ROOT)
+        model_data_path = os.path.join(model_uri, artifact_path)
+        raw_model_path = os.path.join(
+            model_data_path, mlflow_model_constants.RAW_MODEL_FOLDER
+        )
+        return mlflow.pyfunc.build_model_env(raw_model_path, save_path)

{databricks_feature_engineering-0.8.0 → databricks_feature_engineering-0.12.1}/databricks/feature_store/mlflow_model.py

@@ -7,11 +7,9 @@ from pyspark.sql import DataFrame, SparkSession
 
 from databricks.feature_engineering.client import FeatureEngineeringClient
 from databricks.ml_features._spark_client._spark_client import SparkClient
-from databricks.ml_features.constants import (
-    _NO_RESULT_TYPE_PASSED,
-    MODEL_DATA_PATH_ROOT,
-)
+from databricks.ml_features.constants import MODEL_DATA_PATH_ROOT
 from databricks.ml_features.utils import request_context
+from databricks.ml_features_common.mlflow_model_constants import _NO_RESULT_TYPE_PASSED
 
 
 class _FeatureStoreModelWrapper:

{databricks_feature_engineering-0.8.0 → databricks_feature_engineering-0.12.1}/databricks/ml_features/_catalog_client/_catalog_client.py

@@ -14,6 +14,7 @@ from databricks.ml_features.api.proto.feature_catalog_pb2 import (
     CreateFeatures,
     CreateFeatureSpec,
     CreateFeatureTable,
+    DefaultValue,
     DeleteDataSources,
     DeleteFeatureSpec,
     DeleteFeatureTable,
@@ -77,6 +78,7 @@ from databricks.ml_features_common.entities.online_feature_table import (
     OnlineFeatureTable,
 )
 from databricks.ml_features_common.entities.store_type import StoreType
+from databricks.ml_features_common.utils.data_type_utils import serialize_default_value
 from databricks.ml_features_common.utils.uc_utils import reformat_full_table_name
 
 _METHOD_TO_INFO = extract_api_info_for_service(
@@ -629,6 +631,10 @@ class CatalogClient:
             if feature.lookback_window is None
             else feature.lookback_window.total_seconds()
         )
+        default_values = [
+            DefaultValue(feature_name=k, default_value=serialize_default_value(v))
+            for k, v in feature.default_values.items()
+        ]
         return FeatureLookupProto(
             table_name=feature.table_name,
             lookup_keys=self._resolve_string_or_list(feature.lookup_key),
@@ -638,6 +644,7 @@ class CatalogClient:
                 feature.timestamp_lookup_key
             ),
             lookback_window_seconds=lookback_window_seconds,
+            default_values=default_values,
         )
 
     @staticmethod
@@ -705,14 +712,12 @@ class CatalogClient:
         features: List[Union[FeatureLookup, FeatureFunction]],
         exclude_columns: List[str],
         input_columns: List[str],
-        feature_spec_yaml: Optional[str],
         req_context: RequestContext,
     ):
         req_body = GenerateFeatureSpecYaml(
             features=self._convert_features_to_proto(features),
             exclude_columns=self._resolve_string_or_list(exclude_columns),
             input_columns=input_columns,
-            feature_spec_yaml=feature_spec_yaml,
         )
         response_proto = self._call_endpoint(
             GenerateFeatureSpecYaml, req_body, req_context

{databricks_feature_engineering-0.8.0 → databricks_feature_engineering-0.12.1}/databricks/ml_features/_publish_client/_publish_client.py

@@ -249,7 +249,10 @@ class PublishClient:
         if filter_condition:
             # force filtering before proceeding with other operations.
             # e.g. computing latest snapshot
-            publish_df = publish_df.filter(filter_condition).cache()
+            publish_df = publish_df.filter(filter_condition)
+            # [ML-49081] caching the filtered dataframe does not work if it is streaming
+            if not streaming:
+                publish_df = publish_df.cache()
 
         # Only publish features that exist in the catalog.
         if len(features_not_in_catalog) > 0:

{databricks_feature_engineering-0.8.0 → databricks_feature_engineering-0.12.1}/databricks/ml_features/_spark_client/_spark_client.py

@@ -13,6 +13,7 @@ from pyspark.sql.functions import col, count, lit
 from pyspark.sql.types import StructType
 from pyspark.sql.utils import AnalysisException
 
+from databricks.ml_features.constants import _PREBUILT_ENV_URI
 from databricks.ml_features.utils import utils
 from databricks.ml_features.utils.utils import (
     sanitize_identifier,
@@ -557,6 +558,7 @@ class SparkClient:
         result_type=None,
         env_manager=None,
         params: Optional[dict[str, Any]] = None,
+        prebuilt_env_uri: Optional[str] = None,
     ):
         kwargs = {}
         if result_type:
@@ -565,6 +567,8 @@ class SparkClient:
             kwargs["env_manager"] = env_manager
         if params:
             kwargs["params"] = params
+        if prebuilt_env_uri:
+            kwargs[_PREBUILT_ENV_URI] = prebuilt_env_uri
 
         return spark_udf(self._spark, model_uri, **kwargs)
 

{databricks_feature_engineering-0.8.0 → databricks_feature_engineering-0.12.1}/databricks/ml_features/_training_scoring_client/_training_scoring_client.py

@@ -8,7 +8,7 @@ from typing import Any, Dict, List, Optional, Set, Union
 import mlflow
 import yaml
 from mlflow.models import Model, ModelSignature
-from mlflow.utils.file_utils import TempDir, YamlSafeDumper, read_yaml
+from mlflow.utils.file_utils import TempDir
 from pyspark.sql import DataFrame
 from pyspark.sql.functions import struct
 
@@ -25,7 +25,7 @@ from databricks.ml_features._information_schema_spark_client._information_schema
 from databricks.ml_features._spark_client._spark_client import SparkClient
 from databricks.ml_features._spark_client._spark_client_helper import SparkClientHelper
 from databricks.ml_features.constants import (
-    _NO_RESULT_TYPE_PASSED,
+    _PREBUILT_ENV_URI,
     _USE_SPARK_NATIVE_JOIN,
     _WARN,
     MODEL_DATA_PATH_ROOT,
@@ -36,6 +36,7 @@ from databricks.ml_features.entities.feature_lookup import FeatureLookup
 from databricks.ml_features.entities.feature_table import FeatureTable
 from databricks.ml_features.training_set import TrainingSet
 from databricks.ml_features.utils import (
+    input_example_utils,
     request_context,
     training_scoring_utils,
     utils,
@@ -46,6 +47,10 @@ from databricks.ml_features.utils.signature_utils import (
     drop_signature_inputs_and_invalid_params,
     get_mlflow_signature_from_feature_spec,
 )
+from databricks.ml_features.utils.version_compatibility_utils import (
+    is_log_model_artifact_path_deprecated,
+    mlflow_log_model_starts_run,
+)
 from databricks.ml_features_common import mlflow_model_constants
 from databricks.ml_features_common.entities.column_info import ColumnInfo
 from databricks.ml_features_common.entities.feature_column_info import FeatureColumnInfo
@@ -53,12 +58,14 @@ from databricks.ml_features_common.entities.feature_spec import FeatureSpec
 from databricks.ml_features_common.entities.source_data_column_info import (
     SourceDataColumnInfo,
 )
+from databricks.ml_features_common.mlflow_model_constants import _NO_RESULT_TYPE_PASSED
 from databricks.ml_features_common.utils import uc_utils, utils_common
 from databricks.ml_features_common.utils.feature_spec_utils import (
     assign_topological_ordering,
     convert_to_yaml_string,
     get_encoded_graph_map,
 )
+from databricks.ml_features_common.utils.yaml_utils import read_yaml
 
 _logger = logging.getLogger(__name__)
 
@@ -140,6 +147,9 @@ class TrainingScoringClient:
         self._catalog_client_helper.add_consumer(
             consumer_feature_table_map, add_consumer_req_context
         )
+        # Spark query planning is known to cause spark driver to crash if there are many feature tables to PiT join.
+        # See https://docs.google.com/document/d/1EyA4vvlWikTJMeinsLkxmRAVNlXoF1eqoZElOdqlWyY/edit
+        # So we disable native join by default.
         training_scoring_utils.warn_if_non_photon_for_native_spark(
             kwargs.get(_USE_SPARK_NATIVE_JOIN, False), self._spark_client
         )
@@ -150,10 +160,7 @@ class TrainingScoringClient:
             ft_metadata.feature_table_metadata_map,
             ft_metadata.feature_table_data_map,
             uc_function_infos,
-            kwargs.get(
-                _USE_SPARK_NATIVE_JOIN,
-                training_scoring_utils.get_default_use_spark_native_join(feature_spec),
-            ),
+            kwargs.get(_USE_SPARK_NATIVE_JOIN, False),
         )
 
     def create_training_set_from_feature_lookups(
@@ -210,29 +217,37 @@ class TrainingScoringClient:
             label_names,
         )
 
-        # Build feature_spec locally for comparison with the feature spec yaml generated by the
-        # FeatureStore backend. This will be removed once the migration is validated.
-        local_feature_spec = training_scoring_utils.build_feature_spec(
-            feature_lookups,
-            ft_metadata,
-            column_infos,
-            exclude_columns,
-            self._catalog_client.feature_store_workspace_id,
-        )
-        local_generated_yaml_string = convert_to_yaml_string(local_feature_spec)
-        try:
+        if client_name == request_context.FEATURE_ENGINEERING_CLIENT:
+            # Feature Engineering client moved the FeatureSpec generation logic to the server.
             input_columns = [col for col in df.columns if col not in label_names]
-            server_generated_yaml = self._catalog_client.generate_feature_spec_yaml(
+            feature_spec_yaml = self._catalog_client.generate_feature_spec_yaml(
                 features,
                 exclude_columns,
                 input_columns,
-                local_generated_yaml_string,
                 req_context,
             )
-            feature_spec = FeatureSpec._from_dict(yaml.safe_load(server_generated_yaml))
-        except Exception as e:
-            # Fall back to local generated yaml string if server generated yaml is not available
-            feature_spec = local_feature_spec
+            feature_spec = FeatureSpec._from_dict(yaml.safe_load(feature_spec_yaml))
+        else:
+            if any(
+                [
+                    hasattr(f, "default_values") and f.default_values
+                    for f in feature_lookups
+                ]
+            ):
+                _logger.warning(
+                    "FeatureStore Client does not support default values for FeatureLookups, Any specified default values will be ignored. Please use Feature Engineering Client"
+                )
+                # clear default values from feature lookups
+                for feature_lookup in feature_lookups:
+                    feature_lookup.default_values = None
+            # As server only supports FeatureSpec generation on UC tables. Feature Store client will continue to generate feature spec locally.
+            feature_spec = training_scoring_utils.build_feature_spec(
+                feature_lookups,
+                ft_metadata,
+                column_infos,
+                exclude_columns,
+                self._catalog_client.feature_store_workspace_id,
+            )
 
         return self.create_training_set(
             feature_spec,
@@ -388,6 +403,7 @@ class TrainingScoringClient:
                 result_type=result_type,
                 env_manager=env_manager,
                 params=params,
+                prebuilt_env_uri=kwargs.get(_PREBUILT_ENV_URI, None),
             )
             # TODO (ML-17260) Consider reading the timestamp from the backend instead of feature store artifacts
             ml_model = Model.load(
@@ -490,10 +506,7 @@ class TrainingScoringClient:
             feature_table_metadata_map=feature_table_metadata_map,
             feature_table_data_map=feature_table_data_map,
             uc_function_infos=uc_function_infos,
-            use_spark_native_join=kwargs.get(
-                _USE_SPARK_NATIVE_JOIN,
-                training_scoring_utils.get_default_use_spark_native_join(feature_spec),
-            ),
+            use_spark_native_join=kwargs.get(_USE_SPARK_NATIVE_JOIN, False),
         )._augment_df()
 
         # Only included FeatureSpec columns should be part of UDF inputs for scoring.
@@ -600,7 +613,9 @@ class TrainingScoringClient:
                 training_set.feature_spec
             )
             label_type_map = training_set._label_data_types
-            df_head = training_set._df.head()
+
+            labels = training_set._labels
+            df_head = training_set._df.drop(*labels).head()
         else:
             # FeatureSpec.load expects the root directory of feature_spec.yaml
             root_dir, file_name = os.path.split(feature_spec_path)
@@ -714,8 +729,14 @@ class TrainingScoringClient:
                 )
 
             try:
-                if df_head is not None and infer_input_example:
-                    input_example = df_head.asDict()
+                if (
+                    df_head is not None
+                    and signature is not None
+                    and infer_input_example
+                ):
+                    input_example = input_example_utils.infer_input_example(
+                        df_head, signature
+                    )
                 else:
                     input_example = None
             except Exception:
@@ -723,15 +744,23 @@ class TrainingScoringClient:
 
             feature_spec.save(data_path)
 
-            # Log the packaged model. If no run is active, this call will create an active run.
+            if (
+                mlflow.tracking.fluent.active_run() is None
+                and not mlflow_log_model_starts_run()
+            ):
+                mlflow.start_run()
+            compatible_parameters = {}
+            if is_log_model_artifact_path_deprecated():
+                compatible_parameters["name"] = artifact_path
+            else:
+                compatible_parameters["artifact_path"] = artifact_path
             mlflow.pyfunc.log_model(
-                artifact_path=artifact_path,
                 loader_module=mlflow_model_constants.MLFLOW_MODEL_NAME,
                 data_path=data_path,
-                code_path=None,
                 conda_env=conda_env,
                 signature=signature,
                 input_example=input_example,
+                **compatible_parameters,
             )
         if registered_model_name is not None:
             # The call to mlflow.pyfunc.log_model will create an active run, so it is safe to