[SPARK-44856][PYTHON] Improve Python UDTF arrow serializer performance

[HyukjinKwon]

What changes were proposed in this pull request?

This PR removes pandas <> Arrow <> pandas conversion in Arrow-optimized Python UDTF by directly using PyArrow.

Why are the changes needed?

Currently, there is a lot of overhead in the arrow serializer for Python UDTFs. The overhead is largely from converting arrow batches into pandas series and converting UDTF's results back to a pandas dataframe.

We should try directly converting Python object into arrow and vice versa to avoid the expensive pandas conversion.

Does this PR introduce any user-facing change?

Yes. Previously the conversion was

How was this patch tested?

Existing tests.

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

No.

[allisonwang-db]

Why do we need to check both error classes?

[HyukjinKwon]

The problem is that, previously, we were able to create a pandas DataFrame with different schema but Arrow does not seem allowing it.

More specifically previous code path:

                        yield verify_result(
                            pd.DataFrame(check_return_value(res))
                        ), arrow_return_type

did not throw an error at pd.DataFrame(...).

However, new code path

convert_to_arrow(func())

                    ret = LocalDataToArrowConversion.convert(
                       data, return_type, prefers_large_var_types
                   ).to_batches()

this throw an error at LocalDataToArrowConversion.convert(...).

We could further improve this but I would prefer to get out of this scope in this PR considering that there are already behaviour differences with/without Arrow.

[HyukjinKwon]

Let me add a legacy conf ... to be safer ..

[dongjoon-hyun]

Could you fix the remaining failures?

pyspark.errors.exceptions.base.PySparkRuntimeError: [UDTF_ARROW_TYPE_CAST_ERROR] Cannot convert the output value of the column 'x' with type 'object' to the specified return type of the column: 'list<element: int32>'. Please check if the data types match and try again.

Stale

[dongjoon-hyun]

+1, LGTM. Thank you again for making it work, @HyukjinKwon .

[ueshin]

Can we move this to before the above line?

[HyukjinKwon]

Actually no .. otherwise, the tests fail. There is a test that throw an exception when the batch is empty. LocalDataToArrowConversion.convert checks if data is empty, and the exception will be wrapped by PySparkRuntimeError.

[ueshin]

Can this return multiple batches? What happens in that case?

[HyukjinKwon]

Yeah, if it grows over the default size (https://arrow.apache.org/docs/python/generated/pyarrow.dataset.Dataset.html#pyarrow.dataset.Dataset.to_batches) it can be multiple batches. It should work though - I wrote the codes that it should work via ArrowStreamUDFSerializer.dump_stream.

[ueshin]

IIRC, it must return exactly one batch per input row. convert_to_arrow should always return one batch. cc @allisonwang-db

[HyukjinKwon]

From what I read about the codes, it seems fine... but would be great if we can confirm this.

[allisonwang-db]

For directly UDTF usage like MyUDTF(lit(1), lit(2)) it's fine to return multiple batches, but for lateral joins like SELECT * FROM t, LATERAL MyUDTF(a, b)), we must match each input row with all output of the UDTF for that row. If we return multiple batches, then we can't distinguish which batch to join with which input row.

[HyukjinKwon]

Just checked. https://github.com/apache/arrow/blob/d2ddee62329eb711572b4d71d6380673d7f7edd1/cpp/src/arrow/table.cc#L612-L638

The batch size will be long max by default, which I believe it's pretty safe. Arrow batch cannot contain # of rows larger than long in any way.

Stale.

[HyukjinKwon]

Merged to master.