Demand forecasting for a retailer with hundreds of thousands to millions of SKUs is a distributed systems problem as much as a machine learning problem. Training a single model per SKU is not feasible at scale — training 1M models sequentially would take weeks. This post covers the practical architecture for building a production forecasting system that handles millions of SKUs: how to structure the data, when to use Dask vs. PySpark, how to run distributed training on Databricks, and what the operational pitfalls look like.
Consider a mid-size retail operation: 800K active SKUs, 3 years of daily sales history, and a weekly forecast cycle. The data volume is:
A pandas-based approach fails immediately at this scale. You need distributed computation for data processing and either vectorized batch prediction or parallelized model training.
| Use Dask when... | Use PySpark / Databricks when... |
|---|---|
| Your team is Python-first and knows pandas | Your org already has Databricks or EMR infrastructure |
| Data fits on a cluster of moderate size (1-100GB) | Data is in the hundreds of GB to TB range |
| You need fast iteration during development | You need enterprise governance, Delta Lake, or Unity Catalog |
| You want to run on-prem or on generic cloud VMs | You need tight integration with cloud storage and MLflow |
For a 30GB dataset, Dask is sufficient and faster to develop with. For anything above 200GB or with a Databricks contract already in place, PySpark is the better choice.
import dask.dataframe as dd
import pandas as pd
import numpy as np
from dask.distributed import Client
# Start a local Dask cluster (or connect to an existing one)
client = Client(n_workers=8, threads_per_worker=2, memory_limit="8GB")
print(client.dashboard_link)
# Load sales data — Dask reads in parallel across partitions
sales = dd.read_parquet(
"s3://my-bucket/sales/daily/*.parquet",
columns=["sku_id", "date", "units_sold", "store_id", "price", "on_promotion"],
)
# Feature engineering at scale
def build_lag_features(df: pd.DataFrame, lags: list[int] = [7, 14, 28]) -> pd.DataFrame:
"""Applied per-partition — each partition should contain complete history for its SKUs."""
df = df.sort_values("date")
for lag in lags:
df[f"units_lag_{lag}"] = df.groupby("sku_id")["units_sold"].shift(lag)
df["units_roll_7d"] = (
df.groupby("sku_id")["units_sold"]
.transform(lambda x: x.shift(1).rolling(7).mean())
)
df["units_roll_28d"] = (
df.groupby("sku_id")["units_sold"]
.transform(lambda x: x.shift(1).rolling(28).mean())
)
return df
# Repartition by sku_id so each partition contains all history for a set of SKUs
n_partitions = 200
sales = sales.assign(partition_key=sales["sku_id"].map_partitions(lambda s: s.astype("category").cat.codes % n_partitions))
sales = sales.repartition(partition_col="partition_key", npartitions=n_partitions)
meta = sales._meta.copy()
for lag in [7, 14, 28]:
meta[f"units_lag_{lag}"] = 0.0
meta["units_roll_7d"] = 0.0
meta["units_roll_28d"] = 0.0
features = sales.map_partitions(build_lag_features, meta=meta)
features = features.persist() # trigger computation and keep in memory
print(f"Feature matrix: ~{features.shape[0].compute():,} rows")
from dask.distributed import get_client
import lightgbm as lgb
from lightgbm.dask import DaskLGBMRegressor
# Option A: LightGBM native Dask support (best for a single global model)
X = features.drop(["sku_id", "date", "units_sold"], axis=1)
y = features["units_sold"]
model = DaskLGBMRegressor(
n_estimators=500,
learning_rate=0.05,
num_leaves=63,
n_jobs=2,
client=client,
)
model.fit(X, y)
# Option B: Per-SKU models via dask.delayed (for many independent models)
import dask
from sklearn.linear_model import Ridge
from sklearn.preprocessing import StandardScaler
from sklearn.pipeline import Pipeline
import pickle
@dask.delayed
def train_sku_model(sku_df: pd.DataFrame) -> bytes:
"""Train a model for a single SKU. Returns serialized model bytes."""
sku_df = sku_df.dropna()
if len(sku_df) < 30:
return None # not enough history
feature_cols = [c for c in sku_df.columns if c.startswith("units_") and c != "units_sold"]
X = sku_df[feature_cols].values
y = sku_df["units_sold"].values
pipe = Pipeline([
("scaler", StandardScaler()),
("model", Ridge(alpha=1.0)),
])
pipe.fit(X, y)
return pickle.dumps(pipe)
# Group by SKU and train — this runs in parallel across workers
sku_groups = features.groupby("sku_id").apply(lambda df: df, meta=features._meta)
# Collect a sample for demonstration — in production you'd iterate over partitions
sample_skus = features.compute().groupby("sku_id")
delayed_models = [train_sku_model(group) for _, group in list(sample_skus)[:1000]]
results = dask.compute(*delayed_models)
print(f"Trained {sum(1 for r in results if r is not None)} models")
from pyspark.sql import SparkSession
from pyspark.sql import functions as F
from pyspark.sql.window import Window
import pandas as pd
spark = SparkSession.builder.appName("SalesForecasting").getOrCreate()
# Read from Delta Lake (Databricks)
sales = spark.read.format("delta").load("dbfs:/data/sales/daily")
# Window functions for lag features — native Spark SQL, highly optimized
sku_window = Window.partitionBy("sku_id").orderBy("date")
sku_rolling = Window.partitionBy("sku_id").orderBy("date").rowsBetween(-27, -1)
sales_features = (
sales
.withColumn("units_lag_7", F.lag("units_sold", 7).over(sku_window))
.withColumn("units_lag_14", F.lag("units_sold", 14).over(sku_window))
.withColumn("units_lag_28", F.lag("units_sold", 28).over(sku_window))
.withColumn("units_roll_28d", F.avg("units_sold").over(sku_rolling))
.withColumn("week_of_year", F.weekofyear("date"))
.withColumn("month", F.month("date"))
.filter(F.col("date") >= "2022-01-01")
.dropna()
)
# Cache the feature table to Delta for downstream use
sales_features.write.format("delta").mode("overwrite").saveAsTable("ml.sales_features")
from pyspark.sql.functions import pandas_udf, PandasUDFType
from pyspark.sql.types import StructType, StructField, StringType, FloatType, IntegerType
import lightgbm as lgb
# Define output schema
forecast_schema = StructType([
StructField("sku_id", StringType()),
StructField("forecast_date", StringType()),
StructField("forecast_units", FloatType()),
StructField("model_type", StringType()),
])
@pandas_udf(forecast_schema, PandasUDFType.GROUPED_MAP)
def train_and_forecast(sku_df: pd.DataFrame) -> pd.DataFrame:
"""
For each SKU group: train a LightGBM model, generate 12-week forecast.
Runs in parallel across Spark workers — each executor handles a partition of SKUs.
"""
sku_id = sku_df["sku_id"].iloc[0]
sku_df = sku_df.sort_values("date")
feature_cols = [c for c in sku_df.columns if c not in ["sku_id", "date", "units_sold"]]
X = sku_df[feature_cols].values
y = sku_df["units_sold"].values
if len(sku_df) < 52: # less than a year of weekly data
# Fall back to simple seasonal average
results = []
for week in range(1, 13):
results.append({
"sku_id": sku_id,
"forecast_date": f"W+{week}",
"forecast_units": float(y[-52:].mean()),
"model_type": "seasonal_avg",
})
return pd.DataFrame(results)
# Train LightGBM on all available history
split = int(len(X) * 0.85)
params = {"objective": "regression", "metric": "rmse", "verbosity": -1, "n_estimators": 200}
model = lgb.LGBMRegressor(**params)
model.fit(X[:split], y[:split], eval_set=[(X[split:], y[split:])], callbacks=[lgb.early_stopping(20)])
# Iterative forecast: generate week t+1, use it to generate t+2, etc.
last_row = sku_df[feature_cols].iloc[-1].copy()
results = []
for week in range(1, 13):
pred = float(model.predict(last_row.values.reshape(1, -1))[0])
results.append({
"sku_id": sku_id,
"forecast_date": f"W+{week}",
"forecast_units": max(0, pred), # no negative demand
"model_type": "lightgbm",
})
# Shift lag features forward
if "units_lag_28" in last_row.index:
last_row["units_lag_28"] = last_row.get("units_lag_14", last_row["units_lag_28"])
if "units_lag_14" in last_row.index:
last_row["units_lag_14"] = last_row.get("units_lag_7", last_row["units_lag_14"])
if "units_lag_7" in last_row.index:
last_row["units_lag_7"] = pred
return pd.DataFrame(results)
# Apply the UDF — runs in parallel across the cluster
forecasts = sales_features.groupby("sku_id").apply(train_and_forecast)
forecasts.write.format("delta").mode("overwrite").saveAsTable("ml.sales_forecasts")
print(f"Forecasts written: {forecasts.count():,} rows")
Many SKUs have intermittent demand — they sell zero units on most days, with occasional spikes. Per-SKU models perform poorly on intermittent series. The solution is hierarchical forecasting: forecast at the category, subcategory, and SKU level simultaneously, then reconcile the forecasts so they add up consistently.
from statsforecast import StatsForecast
from statsforecast.models import AutoARIMA, AutoETS, CrostonOptimized
import pandas as pd
# Load a sample of SKUs for hierarchical forecasting
df = pd.read_parquet("s3://bucket/sales/weekly_aggregated.parquet")
# statsforecast expects: unique_id, ds, y
df_sf = df.rename(columns={"sku_id": "unique_id", "week": "ds", "units_sold": "y"})
# CrostonOptimized for intermittent demand (SKUs that sell rarely)
# AutoETS for regular seasonal patterns
# The library automatically selects the best model per series
sf = StatsForecast(
models=[
CrostonOptimized(), # for intermittent (many zeros)
AutoETS(season_length=52), # for seasonal annual patterns
],
freq="W",
n_jobs=-1, # parallelize across available cores
)
forecasts = sf.forecast(df=df_sf, h=12, level=[80, 95]) # forecast 12 weeks ahead
print(forecasts.head(20))
# Outputs per-SKU forecast with prediction intervals
At 800K SKUs, even a 10ms per-SKU operation adds up to 2.2 hours of sequential processing. Distributed execution is not optional. Key cost optimization patterns for Databricks:
OPTIMIZE on large tables after bulk writes.-- After writing forecasts, optimize the Delta table
OPTIMIZE ml.sales_forecasts ZORDER BY (sku_id, forecast_date);
-- Check file statistics
DESCRIBE DETAIL ml.sales_forecasts;
Sales forecasting at millions of SKUs requires distributed processing (Dask for moderate scale, PySpark/Databricks for large scale), a model strategy that handles intermittent demand (Croston or related methods for sparse SKUs, LightGBM for active ones), and careful attention to partition design to avoid hot spots or idle workers. The algorithmic choice — LightGBM vs. AutoETS vs. Croston — is secondary to the data engineering. A well-partitioned feature pipeline with proper lag feature construction consistently outperforms a sophisticated model running on a poorly structured dataset.