Greetings @javeed ,   You’re right to call out the friction between a PySpark DataFrame and many Python ML libraries like statsmodels; most Python ML stacks expect pandas, while Spark is distributed-first. Here’s how to bridge that gap efficiently for 50GB time series workloads on Databricks.

What works at scale on Databricks

• Pandas API on Spark lets you write pandas-style code that executes on Spark, avoiding single-node pandas OOM and scaling out across a cluster while preserving familiar idioms. Import as pyspark.pandas as ps and use it when you can keep your logic “pandas-like.”
• Pandas UDFs and pandas function APIs (applyInPandas, mapInPandas, cogroup.applyInPandas) let you run arbitrary Python/pandas code (including statsmodels) against Spark partitions or groups with Arrow-optimized transfer. This is the main interoperability pattern: you group by your series key(s), materialize each group as a pandas DataFrame inside the UDF, and run statsmodels there. It’s proven in production and documented with examples and constraints.
• Statsmodels with Spark via grouped Pandas UDFs: Databricks’ examples show OLS per group with statsmodels in a grouped Pandas UDF, and SARIMAX time-series forecasting scaled across thousands of SKUs using groupBy + applyInPandas. The same pattern applies to decomposition (trend/seasonal/residual) or ARIMA/ETS per series.

Recommended design pattern for time series at scale

• Partition by series key and parallelize per series: df.repartition(n_tasks, key...).groupBy(key...).applyInPandas(fn, schema) so each series (or series group) is processed independently with your statsmodels code. Temporarily disabling AQE helps ensure one Spark task per group when maximizing parallelism in training steps.
• Control per-group memory: All rows in a group are loaded to a worker before your function runs; keep groups bounded and watch skew. If a few series are huge, sub-sample, window, or split them to avoid OOM within the worker.
• Tune Arrow batch size for pandas UDFs: Adjust spark.sql.execution.arrow.maxRecordsPerBatch when many columns (wide features) drive high JVM memory during Arrow conversion.
• Mind timestamps: Spark stores timestamps as UTC; prefer pandas-native time series operations inside the UDF for best performance and correctness with time zones.
• Log your runs: Use MLflow autologging to capture parameters/metrics/artifacts automatically; it supports statsmodels, scikit-learn, Spark MLlib, and others in Databricks Runtime ML.

When PySpark ML lacks the model you need

• Keep the data distributed, but run Python ML per group via pandas UDFs (statsmodels, sktime components that are series-local, etc.). This avoids a bulk conversion of the whole table to pandas, and lets you scale horizontally by adding workers.
• Use distributed-capable alternatives where available: For tree/GBM, prefer XGBoost’s Spark estimators (xgboost.spark) to train models in parallel across the cluster instead of single-node xgboost Python APIs.
• If single-node training is unavoidable, constrain to sampled subsets or per-series batch sizes that fit in memory, and parallelize across many series with pandas UDFs; cluster scale helps for large row counts, while adding columns often hurts more—plan sizing accordingly.
•  

Minimal pattern: statsmodels per series with applyInPandas

Practical tips

• Prefer Pandas API on Spark for exploratory and feature engineering when it fits your workflow; it’s pandas-equivalent on Spark and avoids full pandas conversion.
• Use pandas UDFs for the “bridge” to libraries that don’t support Spark DataFrames (statsmodels, SHAP, etc.), and tune partitions to use all cores across nodes efficiently.
• Don’t collect 50GB to the driver; only move per-group batches as pandas via Arrow inside UDFs, and keep outputs as Spark DataFrames/Delta tables.
• Track experiments with MLflow and, where possible, use Databricks Runtime ML for pre-installed ML libraries and optimized MLlib algorithms.