Rolling predictions with FeatureEngineeringClient

[danielvdc]

I am performing a time series analysis, using a XGBoostRegressor with rolling predictions. I am doing so using the FeatureEngineeringClient (in combination with Unity Catalog), where I create and load in my features during training and inference, as such:

# Create seasonality features table for revenue
fe.create_table(
  name=fe_table_name_seasonality_revenue,
  primary_keys=["pm_key1","pm_key2","ts_key"],
  timestamp_keys=["ts_key"],
  df=revenue_seasonality_features_sdf,
  description="Revenue Seasonality features",
)

One of the features I am creating in the feature store is a lagged feature of my target (y). I have found that lagged information regarding my target is highly predictive for my target. This feature is created as such:

def create_lag_from_to_spark(df, var_name='y'):
    '''Create lag features for XGBoost with Spark - make sure amount of lags is not greater than prediction period'''
    
    # Create a window specification to partition by pm_key1 and order by ts_key
    window_spec = Window.partitionBy("pm_key1").orderBy("ts_key")
    
    # Loop over the lag range
    for lag in range(start, end):
        # Use the lag function to create lag columns
        df = df.withColumn(f'{var_name}_lag_{lag}', F.lag(var_name, lag).over(window_spec))
    
    return df

I am however unable to update these lagged features during inference when calling fe.score_batch(). The function called during score_batch() only works for data present in my train data, but not for the future values in my inference data. Future dates in my test set are not considered with this function. Ideally, I would edit the score_batch function below so that it re-calculates the features after every prediction.

# Predict on validation and test set
predictions_sdf = fe.score_batch(
    model_uri=f"models:/{full_model_name}/{model_reference.version}", 
    df=filtered_test_sdf)

How can I ensure rolling predictions, so that the predicted value in score_batch is recursively used to fill in my target and subsequently the lag features based on my target? I have been looking into UDFs, but based on my testing it seems like these also need to be based on data that is already present in the train set and also do not support rolling/recursive implementations based on predictions.

[stbjelcevic]

You’re running into a fundamental limitation: score_batch does point‑in‑time feature lookups and batch scoring, but it doesn’t support recursive multi‑step forecasting where predictions update features for subsequent timesteps. Feature Store looks up precomputed features “as of” your timestamp, and won’t recalculate lagged target features from predictions inside the same call.

What score_batch can and can’t do

• Automatic feature lookup: When a model is logged with Feature Engineering, score_batch retrieves the features it needs from the offline store and joins them to your input df (by primary and timestamp keys).

• Point‑in‑time correctness: If you declare a timestamp key (Workspace Feature Store) or timeseries_columns (Unity Catalog FE), the join is “as‑of” the timestamp—not an exact match—so you get the latest feature values up to that time.

• Override behavior: If you include one or more feature columns directly in df (the dataframe you pass to score_batch), those values are used instead of what’s stored in the Feature Store. This is key to enabling a rolling loop outside score_batch.

• No recursive updates: score_batch won’t iteratively feed predictions back into the feature computation to update lagged target features for future rows. You must orchestrate that loop yourself.

Two viable patterns for rolling predictions

1) Orchestrate a step‑by‑step loop around score_batch (stays inside Feature Store for lookups)

This pattern uses score_batch each step to get predictions, while you manage the lag features for the next step. It leverages the “override behavior” by passing your computed ylag* columns in df, so Feature Store uses those rather than the stored values.

High‑level approach:

1. Build a “future skeleton” df with keys (pm_key1, pm_key2) and future ts_key you want to predict.

2. Maintain a per‑entity state (e.g., deque of last y values) initialized from historical data to seed ylag* for the first future step.

3. For each future timestamp t:

   • Construct df_step containing keys and ts_key=t.
   • Add your computed lag columns ylag{k} to df_step from the current state.
   • Call fe.score_batch(model_uri, df_step). Because df_step includes ylag{k}, those values are used during prediction.
   • Update the per‑entity state by pushing the predicted ŷ(t).
   • Proceed to the next timestamp.

2) Do the full recursion inside a grouped Pandas UDF with a directly loaded model (bypasses score_batch)

If you prefer to avoid repeated Spark jobs per step, you can load the XGBoost model directly (for example via mlflow.xgboost.load_model) and run a stateful loop per entity with applyInPandas, generating predictions and updating lag features row‑by‑row. Caution: Feature Store‑packaged models aren’t meant to be loaded via mlflow.pyfunc for arbitrary predict() calls; use score_batch for FS models or log a second “native” model artifact specifically for programmatic inference.

Alternative modeling strategies (to avoid recursion)

• Direct multi‑step models: Train separate horizon‑specific models (t+1, t+2, …) so inference is non‑recursive and compatible with score_batch. This sidesteps the feedback loop into lag features.

• Exogenous‑only features: Use features that are available for future timesteps without needing the target y (e.g., calendars, promotions, covariates). Then score_batch is sufficient with time‑series tables and as‑of joins.

Key takeaways

• score_batch can’t be “edited” to recalculate features between predictions; implement a loop around it and pass your lag features in df to override the stored values for each step.

• For a fully stateful recursive approach, use applyInPandas with a natively loaded model (log an additional non‑FS‑packaged artifact if needed).

• Ensure point‑in‑time correctness by using timestamp keys/timeseries_columns in your feature tables, as you’re already doing with timestamp_keys=["ts_key"].