Azure Databricks not paralelizing on Spark

Question

I'm running a script using Spark which works fine both on my computer and on a Google VM using all cores available in parallel (~6000 processes). However, when I've tried to run it on Azure Databricks using a cluster of min 2 workers and max 25 with 4 cores each and running DB 9.0 it just goes one by one without parallelizing.

Is there any extra set-up needed to run Spark in Databricks other than the usual in any VM?

Here's a test script I've been using to debug the parallelization issue (the cluster load goes <10%):

import pandas as pd
import os
import numpy as np
import datetime

from pyspark.sql import SparkSession
from pyspark.sql.types import BooleanType, StructType,StructField,IntegerType,FloatType, StringType


def main(df_estructural):
    sku = df_estructural.Material.iloc[0]
    
    df = pd.read_csv(os.path.join(DATA, 'Sales', str(sku)+'.csv'), header=0, thousands=',', decimal='.', encoding = 'latin-1')
    local_unique = len(df.Local.unique())

    return pd.DataFrame(np.array([[str(FIN.value), str(sku), str(local_unique )]]), columns=('Date', 'Material', 'Local'))


if __name__ == '__main__':  
    FIN = '2021-09-27'

    DATA = '/dbfs/mnt/simulaciones/data'
    DATA_RESULTADOS = '/dbfs/mnt/simulaciones/Resultados'

    df = pd.read_csv(os.path.join(DATA, 'list.csv'), sep=',', header=0, thousands='.', decimal=',', encoding = 'latin-1')
    df = df_OTB[['Material', 'Alpha']]

    spark = SparkSession \
        .builder \
        .appName("test") \
        .config("spark.sql.execution.arrow.pyspark.enabled", "true") \
        .getOrCreate()

    FIN = spark.sparkContext.broadcast(FIN)

    schema = StructType([StructField('Date', StringType(), True),
                        StructField('Material', StringType(), True),
                        StructField('Local', StringType(), True)])

    df_spark = spark.createDataFrame(df) 

    df_estructural = df_spark \
                .groupby(['Material']) \
                .applyInPandas(main, schema=schema) \
                .collect()

    spark.createDataFrame(df_estructural).toPandas().to_csv(os.path.join(DATA_RESULTADOS,'test.csv')), index=False)

user_5070440 · Answer

Try converting Pandas to Koalas which should parallelize across your Databricks nodes much better rather than your current behavior of Pandas running on the driver node.

user_8764412 · Answer

Finally found the issue, posting it here in case someone else comes across something similar.

Given that the pandas Dataframe was only 6000 rows Databricks was not paralelizing it (maybe it thought it wasn't large enough). Therefore I forced it by adding the number of repartitions as:

   df_estructural = df_spark \
                .repartition(104) \
                .groupby(['Material']) \
                .applyInPandas(main, schema=schema) \
                .collect()

Azure Databricks not paralelizing on Spark

2 Answers