Calculating SHAP values with PySpark

I'm looking for a way to reduce the computation time taken to calculate SHAP values on my large dataset (~180M rows, 6 features), and I came across this article talking about using PySpark on SHAP.

I'm new to PySpark, and I'm trying to figure out how to run my code with the snippet provided in the article.

I'm running SHAP now with the code below, where X_values was also used to fit my Isolation Forest model.

X_values = X.values
shap_values = explainer.shap_values(X_values)

Here's the snippet from the article.

def calculate_shap(iterator: Iterator[pd.DataFrame]) -> Iterator[pd.DataFrame]:
    for X in iterator:
        yield pd.DataFrame(
            explainer.shap_values(np.array(X), check_additivity=False)[0],
            columns=columns_for_shap_calculation,
        )

return_schema = StructType()
for feature in columns_for_shap_calculation:
    return_schema = return_schema.add(StructField(feature, FloatType()))

shap_values = df.mapInPandas(calculate_shap, schema=return_schema)

The article describes the snippet as follows.

The code snippet demonstrates how to parallelize applying an Explainer with a Pandas UDF in PySpark. We define a pandas UDF called calculate_shap and then pass this function to mapInPandas. This method is then used to apply the parallelized method to the PySpark dataframe. We will use this UDF to run our SHAP performance tests.

I don't quite understand how the PySpark code works.

1. In the calculate_shap UDF, what is X? Is it my X_values, because I see explainer.shap_values(np.array(X)..., but X was not passed into calculate_shap?
2. The description mentioned "apply the parallelized method to the PySpark dataframe". Which one is the "PySpark dataframe"? Where did df come from?
3. Do I need to do any pre-processing or conversion from pandas dataframe to PySpark dataframe prior to running the snippet?