I am working on Azure Databricks. Standard_DS12_v2 (28GB memory 4 cores) compute type. I use a cluster of 1 driver and 3 worker nodes of the same type.
I have a dataset that can be grouped on country and product names and has the monthly sales for each group. I use a custom pandas_udf function to forecast the values for each group (fbprophet library for forecast).
I call the pandas udf function as follows:
result = sdf.groupBy('key').applyInPandas(forecast_groups, schema=result_schema)
rdf = result.toPandas()
Where, key is simply a column as, "Country_Name-Product_Name". I'll refer to this as group/key. I use toPandas() to initiate an action and because I use pandas on processing of the results. One important point is that this code (udf_function for forecasting) is used sequentially for each country. The dataframe sdf only has one country and multiple products at a time.
I face the following error on one country:
toPandas attempted Arrow optimization because 'spark.sql.execution.arrow.pyspark.enabled' is set to true, but has reached the error below and can not continue. Note that 'spark.sql.execution.arrow.pyspark.fallback.enabled' does not have an effect on failures in the middle of computation. An error occurred while calling o377089.getResult.
Followed by a long stack trace to
Caused by: org.apache.spark.SparkException: Checkpoint block rdd_5743_0 not found! Either the executor that originally checkpointed this partition is no longer alive, or the original RDD is unpersisted. If this problem persists, you may consider using `rdd.checkpoint()` instead, which is slower than local checkpointing but more fault-tolerant.
I face this error only in the country that I know has more data i.e. more groups or keys. All other countries work fine and yield expected results. But, for this one country the code fails. What is weird is that few months ago when I didn't have data for 2021 sales. I doubled my data i.e. till 2023 and ran the code just to kind of test the scalability. That worked, but now for some reason that country with a lot of groups/keys is not working.
I am confused whether this is a memory issue or not. I remember doubling the data and trying that and it worked in the past on the same cluster types and now only few extra months yield in to this error. I did face the RDD lost error before but I restarted the process and it worked fine. Previously this error was once in a blue moon thing now it is permanent. This area where I face the issue is a small part of a very complex pipeline.
Any help regarding possible pathways to debug this would be greatly appreciated.