How to use Mosaic streaming with TimeseriesDataset and PyTorch Forecasting?

Hi, I’m working on forecasting daily product demand for a retailer using the Temporal Fusion Transformer (TFT) model from the PyTorch Forecasting library. My dataset is huge, with around 10 billion rows, and is stored in a Spark DataFrame on Databricks.

I’m encountering an issue when I try to convert the Spark DataFrame to Pandas, resulting in a “driver out of memory” error. To tackle this, I attempted to use Mosaic streaming for efficient data loading, but the challenge is that the TimeseriesDataset in PyTorch Forecasting requires a Pandas DataFrame.

Could anyone guide me on how to use Mosaic streaming with the TimeseriesDataset and PyTorch Forecasting? If possible, I’d greatly appreciate an example to understand how to apply it.

Thank you in advance!