[SPARK-39301][SQL][PYTHON] Leverage LocalRelation and respect Arrow batch size in createDataFrame with Arrow optimization

[HyukjinKwon]

What changes were proposed in this pull request?

This PR proposes to use LocalRelation instead of LogicalRDD when creating a (small) DataFrame with Arrow optimization, which passes the data as a local data in the driver side (which is consistent with Scala code path).

Namely:

import pandas as pd
spark.conf.set("spark.sql.execution.arrow.pyspark.enabled", True)
spark.createDataFrame(pd.DataFrame({'a': [1, 2, 3, 4]})).explain(True)

Before

== Parsed Logical Plan ==
LogicalRDD [a#0L], false

== Analyzed Logical Plan ==
a: bigint
LogicalRDD [a#0L], false

== Optimized Logical Plan ==
LogicalRDD [a#0L], false

== Physical Plan ==
*(1) Scan ExistingRDD arrow[a#0L]

After

== Parsed Logical Plan ==
LocalRelation [a#0L]

== Analyzed Logical Plan ==
a: bigint
LocalRelation [a#0L]

== Optimized Logical Plan ==
LocalRelation [a#0L]

== Physical Plan ==
LocalTableScan [a#0L]

This is controlled by a new configuration spark.sql.execution.arrow.localRelationThreshold defaulting to 48MB. This default was picked by benchmark I ran below.

In addition, this PR also fixes createDataFrame to respect spark.sql.execution.arrow.maxRecordsPerBatch configuration when creating Arrow bathes. Previously, we divided the input pandas DataFrame by the default partition number which forced users to set spark.rpc.message.maxSize when the input pandas DataFrame is too large. See the benchmark performed below.

Why are the changes needed?

We have some nice optimization for LocalRelation (e.g., ConvertToLocalRelation). For example, the stats are fully known when you use LocalRelation. With LogicalRDD, many optimizations cannot be applied. Even in some cases (e.g., executeCollect), we can avoid creating RDDs too.

For respecting spark.sql.execution.arrow.maxRecordsPerBatch, 1. we can avoid forcing users to set spark.rpc.message.maxSize, and 2. I believe the configuration is supposed to be respected for all code path that creates Arrow batches if possible.

Does this PR introduce any user-facing change?

No, it is an optimization. The number of partitions can be different, but that should be internal.

How was this patch tested?

• Manually tested.
• Added a unittest.
• I did two benchmark tests with 1 Driver & 4 Workers (i3.xlarge), see below.

Benchmark 1 (best cases)

import time
import random
import string

import pandas as pd

names = [random.choice(list(string.ascii_lowercase)) for i in range(1000)]
ages = [random.randint(0, 100) for i in range(1000)]
l = list(zip(names, ages))
d = [{'name': a_name, 'age': an_age} for a_name, an_age in l]
pdf = pd.DataFrame({'name': names, 'age': ages})
spark.range(1).count()  # heat up

start = time.time()
for _ in range(100):
    _ = spark.createDataFrame(pdf)

end = time.time()
print(end - start)

Before

10.250491698582968

After

6.004616181055705

Benchmark 2 (worst cases)

curl -O https://eforexcel.com/wp/wp-content/uploads/2020/09/HR2m.zip
unzip HR2m.zip

import pandas as pd
pdf = pd.read_csv("HR2m.csv")
pdf23 = pdf.iloc[:int(len(pdf)/32)]
pdf45 = pdf.iloc[:int(len(pdf)/16)]
pdf90 = pdf.iloc[:int(len(pdf)/8)]
pdf175 = pdf.iloc[:int(len(pdf)/4)]
pdf350 = pdf.iloc[:int(len(pdf)/2)]
pdf700 = pdf.iloc[:int(len(pdf))]
pdf2gb = pd.concat([pdf, pdf, pdf])
pdf5gb = pd.concat([pdf2gb, pdf2gb])

spark.createDataFrame(pdf23)._jdf.rdd().count()  # explicitly create RDD.
...

Before

23MB: 1.02 seconds
45MB: 1.69 seconds
90MB: 2.38 seconds
175MB: 3.19 seconds
350MB: 6.10 seconds
2GB: 43.21 seconds
5GB: X (threw an exception that says to set 'spark.rpc.message.size' higher)

After

23MB: 1.31 seconds (local collection is used)
45MB: 2.47 seconds (local collection is used)
90MB: 1.79 seconds
175MB: 3.22 seconds
350MB: 6.41 seconds
2GB: 47.12 seconds
5GB: 1.29 minutes

NOTE that the performance varies depending on network stability, and the numbers above are from second run (it's not the average).

[HyukjinKwon]

cc @ueshin @viirya @BryanCutler FYI

[HyukjinKwon]

Intentionally I used Iterator to avoid Py4J copies Array into Python driver side.

[HyukjinKwon]

Wow, thanks for reviews guys. Let me think a bit more and push some changes soon.

[HyukjinKwon]

The reason of doing this is to avoid reconfiguring spark.rpc.message.maxSize. When the batch is too large, it throws an exception with complaining spark.rpc.message.maxSize is too small.

[BryanCutler]

I thought this was to control how many partitions were in the rdd? Each partition could have multiple batches, and probably should be capped at arrowMaxRecordsPerBatch, but since it was coming from a local Pandas DataFrame already in memory, that didn't seem to be a big deal.

[HyukjinKwon]

Yeah, that's true... perf diff seems trivial in any event and seems it works around the spark.rpc.message.maxSize issue.

[BryanCutler]

I believe it was like this to create the same number of partitions as when arrow is disabled, although that might have changed since. If the DataFrame is split with arrowMaxRecordsPerBatch and a user wanted to create a certain number of partitions, then would they have to look at the size of the input and then adjust arrowMaxRecordsPerBatch accordingly?

[HyukjinKwon]

Yeah, that's true .. but I wonder if the default number of partitions is something we should consider given that it wasn't already configurable before, and SparkSession.createDataFrame does not expose the number of partitions too.

If they really need, users might want to create an RDD with an explicit parallelism .. we don't support this now though (see also #29719).

[HyukjinKwon]

BTW, just to extra clarify, when the pandas DataFrame is small (lower than the threshold), the number of partitions remains same (configured by spark.sql.leafNodeDefaultParallelism that falls back to sparkContext.defaultParallelism if not set).

The number of partitions is only different when the input DataFrame is large, which I think makes more sense in general ..

[HyukjinKwon]

This PR is ready for a look now. TL;DR:

• I added one new configuration spark.sql.execution.arrow.localRelationThreshold as a threshold. Default value was picked given the benchmark.
• I fixed the code path to respect spark.sql.execution.arrow.maxRecordsPerBatch when creating Arrow batches to avoid forcing users to set a larger spark.rpc.message.maxSize.

[HyukjinKwon]

Maybe 32 MB? Don't have a strong preference.

[BryanCutler]

Is this the max size of each batch or all batches together?

[HyukjinKwon]

It's all together ..so pretty small

[HyukjinKwon]

Gentle ping for a review :-). I know it has some trade-off but I believe this addresses more common cases and benefit more users.

[BryanCutler]

Looks like it's a good optimization for smaller files, thanks @HyukjinKwon !

[BryanCutler]

Is this the max size of each batch or all batches together?

[BryanCutler]

I thought this was to control how many partitions were in the rdd? Each partition could have multiple batches, and probably should be capped at arrowMaxRecordsPerBatch, but since it was coming from a local Pandas DataFrame already in memory, that didn't seem to be a big deal.

[HyukjinKwon]

cc @mengxr and @WeichenXu123 in case you guys have some comments.

[HyukjinKwon]

Let me merge this in few days ... assuming that we're all good. Hopefully my benchmark is good enough.

[HyukjinKwon]

Rebased

[HyukjinKwon]

Let me get this in. It's the early stage of Spark 3.4 so should be good timing to merge such stuff.

Merged to master.