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)