There is a ready feature engineering function for that:
# on non ML runtime please install databricks-feature-engineering>=0.13.0a3"
from databricks.feature_engineering import FeatureEngineeringClient
fe = FeatureEngineeringClient()
from databricks.feature_engineering import FeatureLookup
# The `FeatureSpec` can be accessed in Unity Catalog as a function.
# `FeatureSpec`s can be used to create training sets or feature serving endpoints.
fe.create_feature_spec(
name = f"{CATALOG}.{SCHEMA}.feature_spec",
features=[
FeatureLookup(
table_name=f"{CATALOG}.{SCHEMA}.offline_feature_table",
lookup_key="id",
),
],
)
## add serving endpoint (can be done through UI too)
rom databricks.feature_engineering.entities.feature_serving_endpoint import (
ServedEntity,
EndpointCoreConfig,
)
fe.create_feature_serving_endpoint(
name="my-feature-serving-endpoint",
config=EndpointCoreConfig(
served_entities=ServedEntity(
feature_spec_name=f"{CATALOG}.{SCHEMA}.feature_spec",
workload_size="Small",
scale_to_zero_enabled=True,
instance_profile_arn=None,
)
)
)
## inference
import mlflow.deployments
client = mlflow.deployments.get_deploy_client("databricks")
response = client.predict(
endpoint="my-feature-serving-endpoint",
inputs={
"dataframe_records": [
{"id": 1},
{"id": 7},
{"id": 12345},
]
},
)
print(response)