1.8 Training a Model with Features from a Feature Table

1.8 Training a Model with Features from a Feature Table

Building machine learning models on tabular and time-series data introduces two critical architectural hurdles:

1. Data Leakage (Lookahead Bias): Joining historical event observations with the current state of a feature table contaminates training data with future information, producing artificially optimistic validation scores that collapse in production.
2. Train/Serve Skew: If feature transformation and join logic are implemented manually in training pipelines, client applications must replicate that exact logic during real-time inference, leading to logic drift and production failures.

Databricks Feature Engineering solves both challenges through FeatureLookup specifications, Point-in-Time (As-Of) time-series joins, and automated feature packaging via fe.log_model().

The lookup selects the latest feature state recorded before each transaction, so no row can ever see a value that did not exist when the event occurred.

---

The Mechanics of Point-in-Time (As-Of) Joins

Consider an e-commerce fraud detection model. A customer completes a transaction on 2026-03-15 14:30:00. If we perform a standard SQL JOIN on customer_id against a customer feature table containing metrics updated on 2026-03-20, the model will train on information that occurred after the transaction took place. This is lookahead data leakage.

In a Point-in-Time Join, the Feature Engineering Client ensures that each observation row only joins with the latest feature record whose timestamp is less than or equal to the observation timestamp:

$\text{Feature Timestamp} \le \text{Observation Timestamp}$

Feature Table Timeline:
  [V1: $100 avg] (Jan 01) --------> [V2: $250 avg] (Mar 01) --------> [V3: $600 avg] (Apr 01)
                                            |
Observation Event:                          v (Correctly joins with V2)
  Transaction Event at Mar 15 -------------> [Point-in-Time Match!]

---

Constructing Training Sets with FeatureLookup

To construct a training set, developers define one or more FeatureLookup objects referencing Unity Catalog feature tables, and pass them to fe.create_training_set():

from databricks.feature_engineering import FeatureEngineeringClient, FeatureLookup
from datetime import timedelta

fe = FeatureEngineeringClient()

# 1. Base observation DataFrame (raw labels and timestamps)
observation_df = spark.table("prod_catalog.gold.fraud_observations")

# 2. Define static entity feature lookup
customer_static_lookup = FeatureLookup(
    table_name="prod_catalog.ml_features.customer_demographics",
    feature_names=["account_age_days", "kyc_verified", "credit_tier"],
    lookup_key="customer_id"
)

# 3. Define temporal time-series feature lookup (Point-in-Time)
transaction_temporal_lookup = FeatureLookup(
    table_name="prod_catalog.ml_features.customer_temporal_spending",
    feature_names=["rolling_avg_30d_spend", "declined_tx_count_24h"],
    lookup_key="customer_id",
    timestamp_lookup_key="transaction_timestamp",  # Matches against observation timestamp
    lookback_window=timedelta(days=30)             # Restricts search window
)

# 4. Create the unified training set
training_set = fe.create_training_set(
    df=observation_df,
    feature_lookups=[customer_static_lookup, transaction_temporal_lookup],
    label="label_is_fraud",
    exclude_columns=["raw_transaction_id", "ip_address"]
)

# 5. Materialize into PySpark DataFrame for model training
training_df = training_set.load_df()

Parameters of FeatureLookup

• table_name: Fully qualified 3-level name of the feature table in Unity Catalog.
• feature_names: List of specific feature columns to extract. If omitted, all non-key feature columns are retrieved.
• lookup_key: Join key column name(s) in the observation DataFrame matching the feature table's primary key(s).
• timestamp_lookup_key: (Optional) Column name in the observation DataFrame containing the event timestamp used for point-in-time time-series matching.
• lookback_window: (Optional) timedelta object defining the maximum age of valid feature records. Prevents joining with ancient, stale records.

---

Packaging Models with fe.log_model()

Standard MLflow logging (mlflow.sklearn.log_model) records only the trained estimator weights. When using Databricks Feature Engineering, data scientists log models using fe.log_model().

fe.log_model() packages the trained model artifact alongside the complete FeatureLookup lineage metadata:

import mlflow
import mlflow.sklearn
from sklearn.ensemble import GradientBoostingClassifier

# Train model on materialized features
X = training_df.drop("customer_id", "transaction_timestamp", "label_is_fraud").toPandas()
y = training_df.select("label_is_fraud").toPandas()["label_is_fraud"]

model = GradientBoostingClassifier(n_estimators=100)
model.fit(X, y)

# Log model with packaged Feature Engineering metadata
with mlflow.start_run(run_name="fraud_gbm_fe"):
    fe.log_model(
        model=model,
        artifact_path="model",
        flavor=mlflow.sklearn,
        training_set=training_set,
        registered_model_name="prod_catalog.ml_models.fraud_detection_gbm"
    )

The Superpower of fe.log_model: Zero-Code Feature Retrieval at Scoring Time

Because the model artifact contains the feature table references and join keys, inference pipelines do not need to manually query or join feature tables. The scoring engine automatically resolves lookups:

# Scoring DataFrame only needs the primary keys and event timestamps!
scoring_requests_df = spark.createDataFrame([
    ("cust_9921", "2026-03-19 10:00:00"),
    ("cust_4412", "2026-03-19 10:05:00")
], ["customer_id", "transaction_timestamp"])

# Automatically fetches features and scores predictions
predictions_df = fe.score_batch(
    model_uri="models:/prod_catalog.ml_models.fraud_detection_gbm@champion",
    df=scoring_requests_df
)

---

The Rest of the create_training_set Signature

Beyond the four arguments used above, three more show up in scenario questions:

exclude_columns deserves attention. Lookup keys and event timestamps must be present in the observation DataFrame for the join to happen, yet they are almost never legitimate model inputs: a raw customer_id invites the model to memorise individuals, and a raw timestamp lets it learn the calendar instead of the behaviour. Excluding them at training-set construction is cleaner than dropping them later, because the packaged model then records that they were never features.

Name collisions and rename_outputs

When two feature tables both expose a column called score, or a feature name collides with a column already present in the observation DataFrame, rename_outputs maps the feature to a new name in the training set:

FeatureLookup(
    table_name="prod_catalog.ml_features.risk_scores",
    feature_names=["score"],
    lookup_key="customer_id",
    rename_outputs={"score": "risk_score"},
)

What lookback_window returns when nothing qualifies

lookback_window takes a datetime.timedelta and defaults to None, meaning any feature value is eligible regardless of age. When a window is set and no feature row falls inside it, the join produces null — it does not reach further back and it does not raise. That is deliberate, because a 30-day spend average computed 400 days ago is not a usable feature, but it means the pipeline has to handle nulls for new entities. A sudden rise in null features usually points at a stalled feature job rather than at a modelling problem.

On-Demand Features with FeatureFunction

Some features cannot be precomputed at all, because they depend on values that only exist at request time: the distance between a user's current location and a restaurant, or the ratio of this transaction's amount to the stored 30-day average. Those are declared as FeatureFunction entries alongside the lookups, pointing at a Python UDF registered in Unity Catalog:

from databricks.feature_engineering import FeatureFunction

amount_ratio = FeatureFunction(
    udf_name="prod_catalog.ml_features.amount_to_avg_ratio",
    input_bindings={"amount": "transaction_amount", "avg_spend": "rolling_avg_30d_spend"},
    output_name="amount_ratio",
)

training_set = fe.create_training_set(
    df=observation_df,
    feature_lookups=[customer_static_lookup, transaction_temporal_lookup, amount_ratio],
    label="label_is_fraud",
)

input_bindings maps each UDF parameter to either a column of the observation DataFrame or a feature produced by an earlier lookup, and output_name is the column the computed value lands in. The property that matters is that the same UDF runs at training time and at scoring time, so an on-demand feature cannot drift the way a transformation duplicated inside an application inevitably does.