stbjelcevic
Databricks Employee
Databricks Employee

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"].