tree 69f85b7be36f5473950a2795d16fc1c4a27d4bc0
parent 027327d94b3413ffb228ac482d51e75856c88d02
author Ian Cook <ianmcook@gmail.com> 1715243134 +0900
committer Hyukjin Kwon <gurwls223@apache.org> 1715243134 +0900

[SPARK-47365][PYTHON] Add toArrow() DataFrame method to PySpark

### What changes were proposed in this pull request?
- Add a PySpark DataFrame method `toArrow()` which returns the contents of the DataFrame as a [PyArrow Table](https://arrow.apache.org/docs/python/generated/pyarrow.Table.html), for both local Spark and Spark Connect.
- Add a new entry to the **Apache Arrow in PySpark** user guide page describing usage of the `toArrow()` method.
- Add  a new option to the method `_collect_as_arrow()` to provide more useful output when there are zero records returned. (This keeps the implementation of `toArrow()` simpler.)

### Why are the changes needed?
In the Apache Arrow community, we hear from a lot of users who want to return the contents of a PySpark DataFrame as a PyArrow Table. Currently the only documented way to do this is to return the contents as a pandas DataFrame, then use PyArrow (`pa`) to convert that to a PyArrow Table.
```py
pa.Table.from_pandas(df.toPandas())
```
But going through pandas adds significant overhead which is easily avoided since internally `toPandas()` already converts the contents of Spark DataFrame to Arrow format as an intermediate step when `spark.sql.execution.arrow.pyspark.enabled` is `true`.

Currently it is also possible to use the experimental `_collect_as_arrow()` method to return the contents of a PySpark DataFrame as a list of PyArrow RecordBatches. This PR adds a new non-experimental method `toArrow()` which returns the more user-friendly PyArrow Table object.

This PR also adds a new argument `empty_list_if_zero_records` to the experimental method `_collect_as_arrow()` to control what the method returns in the case when the result data has zero rows. If set to `True` (the default), the existing behavior is preserved, and the method returns an empty Python list. If set to `False`, the method returns returns a length-one list containing an empty Arrow RecordBatch which includes the schema. This is used by `toArrow()` which requires the schema even if the data has zero rows.

For Spark Connect, there is already a `SparkSession.client.to_table()` method that returns a PyArrow table. This PR uses that to expose `toArrow()` for Spark Connect.

### Does this PR introduce _any_ user-facing change?

- It adds a DataFrame method `toArrow()` to the PySpark SQL DataFrame API.
- It adds a new argument `empty_list_if_zero_records` to the experimental DataFrame method `_collect_as_arrow()` with a default value which preserves the method's existing behavior.
- It exposes `toArrow()` for Spark Connect, via the existing `SparkSession.client.to_table()` method.
- It does not introduce any other user-facing changes.

### How was this patch tested?
This adds a new test and a new helper function for the test in `pyspark/sql/tests/test_arrow.py`.

### Was this patch authored or co-authored using generative AI tooling?
No

Closes #45481 from ianmcook/SPARK-47365.

Lead-authored-by: Ian Cook <ianmcook@gmail.com>
Co-authored-by: Hyukjin Kwon <gurwls223@gmail.com>
Signed-off-by: Hyukjin Kwon <gurwls223@apache.org>