How to speed up downloading a CSV file locally with PySpark (databricks)? - pyspark

Question

We created an ImageClassifier to predict wether certain instagram images are of a certain class. Running this model works fine.
#creating deep image feauturizer using the InceptionV3 lib
featurizer = DeepImageFeaturizer(inputCol="image",
outputCol="features",
modelName="InceptionV3")
#using lr for speed and reliability
lr = LogisticRegression(maxIter=5, regParam=0.03,
elasticNetParam=0.5, labelCol="label")
#define Pipeline
sparkdn = Pipeline(stages=[featurizer, lr])
spark_model = sparkdn.fit(df)
We made this seperately from our basetable (which runs on a higher cluster). We need to extract the spark_model predictions as a csv to import it back in the other notebook and merge it with our basetable.
To do this we have tried the following
image_final_estimation = spark_model.transform(image_final)
display(image_final_estimation) #since this gives an option in databricks to
download the csv
AND
image_final_estimation.coalesce(1).write.csv(path = 'imagesPred2.csv') #and then we would be able to read it back in with spark.read.csv
The thing is these operations take very long (probably due to the nature of the task) and they crash our cluster. We are able to show our outcome, but not only with '.show()', not with the display() method.
Is there any other way to save this csv locally? Or how can we improve the speed of these tasks?
Please note that we use the community edition of Databricks.

Answer 1

When storing DataFrames on files, a good way to parallelize the writing is defining an appropiate number of Partitions related to that DataFrame/RDD.
In the code shown by you, you are using the coalesce function (which it basically reduces the number of partitions to 1, which reduces the effects of parallelism).
On Databricks Community Edition, I tried the following test with a CSV Dataset provided by Databricks( https://docs.databricks.com/getting-started/databricks-datasets.html). The idea is to measure the time elapsed by writing the data into a csv using one partition vs using many partitions.
carDF = spark.read.option("header", True).csv("dbfs:/databricks-datasets/Rdatasets/data-001/csv/car/*")
print("Total count of Rows {0}".format(carDF.count()))
print("Original Partitions Number: {0}".format(carDF.rdd.getNumPartitions()))
>>Total count of Rows 39005
>>Original Partitions Number: 7
%timeit carDF.write.format("csv").mode("overwrite").save("/tmp/caroriginal")
>>2.79 s ± 180 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
So, up to now, writing the dataset on a local file using 7 partitions took 2.79 seconds
newCarDF = carDF.coalesce(1)
print("Total count of Rows {0}".format(newCarDF.count()))
print("New Partitions Number: {0}".format(newCarDF.rdd.getNumPartitions()))
>>Total count of Rows 39005
>>New Partitions Number: 1
%timeit newCarDF.write.format("csv").mode("overwrite").save("/tmp/carmodified")
>>4.13 s ± 172 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
So, for the same DataFrame, writing to a csv with one partition took 4,13 seconds.
In conclusion, in this case the "coalesce(1)" part is affecting on the writing performance.
Hope this helps