What's the best workflow for training scikit-learn models on Spark DataFrames?

Summary

• Convert Spark DataFrames to pandas on a single machine for scikit-learn training if data fits in memory (< 10 GB).
• Use distributed alternatives like Spark ML or parallel scikit-learn (joblib) if your data is too large for a single machine.
• Combine scikit-learn with PySpark UDFs for distributed inference and feature engineering across large datasets.

What's the best workflow for training scikit-learn models on Spark DataFrames?

Scikit-learn trains on a single machine, so the workflow depends on your data size. For data that fits in memory, convert your Spark DataFrame to pandas and train normally. For larger data, use distributed alternatives or split training across Spark workers. The key is matching scikit-learn's memory model to your data size.

Option 1: Convert to pandas and train (best for data < 10 GB)

This is the simplest and often fastest approach for moderate-size datasets.

import pandas as pd
from sklearn.ensemble import RandomForestClassifier
from sklearn.model_selection import train_test_split
import mlflow

    # Read Spark DataFrame as pandas
spark_df = spark.table("my_feature_table")
df_pandas = spark_df.toPandas()  # Collects to driver

    # Train scikit-learn model
X = df_pandas.drop("target", axis=1)
y = df_pandas["target"]
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2)

model = RandomForestClassifier(n_estimators=100)
model.fit(X_train, y_train)

    # Log to MLflow
with mlflow.start_run():
    mlflow.sklearn.log_model(model, "model")
    mlflow.log_metric("accuracy", model.score(X_test, y_test))

Advantages:

• Simple, familiar scikit-learn API
• Fast for moderate data sizes
• Full access to scikit-learn ecosystem

Limitations:

• Data must fit in driver memory
• No parallelization of training

Option 2: Distributed training with Spark ML (best for data > 10 GB)

If your data is too large for a single machine, use Spark ML (native Spark implementation of ML algorithms) or distribute training across workers.

from pyspark.ml.feature import VectorAssembler
from pyspark.ml.classification import RandomForestClassifier as SparkRF
from pyspark.ml import Pipeline

spark_df = spark.table("my_feature_table")

    # Spark ML pipeline
assembler = VectorAssembler(inputCols=feature_cols, outputCol="features")
rf_model = SparkRF(labelCol="target", featuresCol="features", numTrees=100)

pipeline = Pipeline(stages=[assembler, rf_model])
model = pipeline.fit(spark_df)

    # Make predictions
predictions = model.transform(spark_df)

Advantages:

• Scales to terabytes
• Distributed training across workers
• Integrates with Spark ecosystem

Limitations:

• Different API from scikit-learn
• Fewer algorithm options than scikit-learn

Option 3: Parallel scikit-learn with joblib (hybrid approach)

For medium-sized data that doesn't fit in memory, use Spark to parallelize scikit-learn training across workers via joblib.

from sklearn.ensemble import RandomForestClassifier
from sklearn.model_selection import train_test_split
import joblib

    # Define training function
def train_sklearn_model(partition):
    import pandas as pd
    from sklearn.ensemble import RandomForestClassifier
    
    pdf = partition.toPandas()
    X = pdf.drop("target", axis=1)
    y = pdf["target"]
    
    model = RandomForestClassifier(n_estimators=50)
    model.fit(X, y)
    return model

    # Partition Spark DataFrame and train in parallel
spark_df = spark.table("my_feature_table")
partitions = spark_df.repartition(4)  # 4 partitions across 4 workers

models = (partitions
    .mapPartitions(train_sklearn_model)
    .collect())

    # Ensemble the models
from sklearn.ensemble import VotingClassifier
ensemble = VotingClassifier(estimators=[(f"model_{i}", m) for i, m in enumerate(models)])

Advantages:

• Uses scikit-learn directly
• Parallelizes across Spark workers
• Middle ground between simplicity and scale

Limitations:

• Requires careful partitioning
• Ensemble of small models, not one large model

Recommended workflow

Data size	Best approach	Why
< 1 GB	Pandas + scikit-learn	Fast, simple, full scikit-learn API
1-10 GB	Pandas + scikit-learn	Still fits in driver memory
10-100 GB	Spark ML (distributed training)	Data too large for single machine; need distribution
100 GB - 1 TB	Spark ML or distributed ensemble	Full distribution needed
> 1 TB	Spark ML + feature sampling	Train on sampled features or subsets

Best practices

1. Cache your Spark DataFrame before training

spark_df.cache()  # Avoids recomputing from source
spark_df.count()  # Trigger cache

2. Use MLflow with either approach

with mlflow.start_run():
    mlflow.log_param("data_size_gb", spark_df.memory_estimate_gb())
    mlflow.sklearn.log_model(model, "model")
    mlflow.log_metric("accuracy", accuracy)

3. Sample large datasets for rapid iteration

sample_df = spark_df.sample(fraction=0.1)  # Smaller sample for iteration
pandas_df = sample_df.toPandas()
    # Iterate on sample, then train on full data

4. Use Feature Store for consistent features

from databricks.feature_store import FeatureStoreClient
fs = FeatureStoreClient()
df = fs.score_batch("my_features", df_with_keys)
df_pandas = df.toPandas()
    # Train scikit-learn model on consistent features

Example: end-to-end scikit-learn with Spark

from databricks.feature_store import FeatureStoreClient
import mlflow
from sklearn.ensemble import RandomForestClassifier
from sklearn.model_selection import train_test_split

    # 1. Get features from Feature Store
fs = FeatureStoreClient()
train_df = fs.score_batch("fraud_features", training_keys_df)

    # 2. Convert to pandas (only if < 10 GB)
df_pandas = train_df.toPandas()
X = df_pandas.drop("fraud_label", axis=1)
y = df_pandas["fraud_label"]

    # 3. Train-test split
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2)

    # 4. Train scikit-learn model
model = RandomForestClassifier(n_estimators=100, max_depth=10)
model.fit(X_train, y_train)

    # 5. Log to MLflow
with mlflow.start_run():
    mlflow.sklearn.log_model(model, "model")
    mlflow.log_metric("test_accuracy", model.score(X_test, y_test))
    mlflow.log_param("n_estimators", 100)

FAQs

Should I use scikit-learn or Spark ML for large datasets?

For data > 10 GB that doesn't fit in driver memory, Spark ML is the better choice because it distributes training. Scikit-learn is better if you prefer its API or need specific algorithms not in Spark ML.

Can I serve scikit-learn models trained on Spark data?

Yes. Log your scikit-learn model to MLflow and deploy via Model Serving. The model size doesn't change; serving complexity depends on prediction latency and throughput requirements.

Can I fine-tune scikit-learn on Spark Data?

Yes. Use GridSearchCV or RandomizedSearchCV on pandas data after converting from Spark. Parallelize hyperparameter search using Spark's distributed training if needed.