Databricks ML Professional Practice Question 1: Per Group Model Training
Databricks Certified Machine Learning Professional · Test 1 · question 1 of 10
A retailer needs a separate demand-forecasting model for each of its 2,000 stores. The training data is a 300-million-row Delta table with a store_id column, and any single store's history fits comfortably in the memory of one worker. The team wants to keep using scikit-learn, train the 2,000 models in parallel across the cluster's workers, and log each store's model and metrics to MLflow.
Which approach meets all of these requirements?
Show answer and explanation
Answer:
groupBy().applyInPandas() implements split-apply-combine: Spark shuffles the rows so that all rows for one store_id arrive together as a single pandas DataFrame, and it runs the function for each group as tasks spread across the workers. Inside the function you can use any single-node library (scikit-learn here), log the model to MLflow, and return a small DataFrame that matches the declared output schema. The docs warn that the whole group is loaded into memory. That is not a problem here because the stem says each store's history fits on one worker. Skewed groups are the usual failure mode for this pattern.
Option A fails the per-store requirement. mapInPandas hands the function an iterator of Arrow batches cut by partition and batch size, not by key, so a batch can hold several stores or only part of one.
Option B collects 300 million rows onto the driver. That risks running out of driver memory, and the threads only use the driver's cores, so the workers are not used at all.
Option D produces one global model, not 2,000 per-store models. That is a different design, and it does not meet the stated requirement.
Reference: pandas function APIs