Skip to content

[SPARK-52877][PYTHON] Improve Python UDF Arrow Serializer Performance #51225

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
wants to merge 89 commits into from

Conversation

asl3
Copy link
Contributor

@asl3 asl3 commented Jun 19, 2025

What changes were proposed in this pull request?

This PR removes pandas <> Arrow conversion in Arrow-optimized Python UDF by directly using PyArrow, to improve arrow-optimized Python UDF performance and memory usage.

Why are the changes needed?

Python UDF arrow serializer has a lot of overhead from converting arrow batches into pandas series and converting UDF results back to a pandas dataframe.

We can instead convert Python object directly into arrow to avoid the expensive pandas conversion.

Does this PR introduce any user-facing change?

Legacy type coercion:

# +-----------------------------+--------------+------------------+--------------------+------+--------------------+-----------------------------+------------+----------------------+------------------+------------------+----------------------------+--------------------+--------------+  # noqa
    # |SQL Type \ Python Value(Type)|None(NoneType)|        True(bool)|              1(int)|a(str)|    1970-01-01(date)|1970-01-01 00:00:00(datetime)|  1.0(float)|array('i', [1])(array)|         [1](list)|       (1,)(tuple)|bytearray(b'ABC')(bytearray)|          1(Decimal)|{'a': 1}(dict)|  # noqa
    # +-----------------------------+--------------+------------------+--------------------+------+--------------------+-----------------------------+------------+----------------------+------------------+------------------+----------------------------+--------------------+--------------+  # noqa
    # |                      boolean|          None|              True|                True|     X|                   X|                            X|        True|                     X|                 X|                 X|                           X|                   X|             X|  # noqa
    # |                      tinyint|          None|                 1|                   1|     X|                   X|                            X|           1|                     X|                 X|                 X|                           X|                   1|             X|  # noqa
    # |                     smallint|          None|                 1|                   1|     X|                   X|                            X|           1|                     X|                 X|                 X|                           X|                   1|             X|  # noqa
    # |                          int|          None|                 1|                   1|     X|                   0|                            X|           1|                     X|                 X|                 X|                           X|                   1|             X|  # noqa
    # |                       bigint|          None|                 1|                   1|     X|                   X|                            0|           1|                     X|                 X|                 X|                           X|                   1|             X|  # noqa
    # |                       string|          None|            'True'|                 '1'|   'a'|        '1970-01-01'|         '1970-01-01 00:00...|       '1.0'|     "array('i', [1])"|             '[1]'|            '(1,)'|         "bytearray(b'ABC')"|                 '1'|    "{'a': 1}"|  # noqa
    # |                         date|          None|                 X|                   X|     X|datetime.date(197...|         datetime.date(197...|           X|                     X|                 X|                 X|                           X|datetime.date(197...|             X|  # noqa
    # |                    timestamp|          None|                 X|datetime.datetime...|     X|                   X|         datetime.datetime...|           X|                     X|                 X|                 X|                           X|datetime.datetime...|             X|  # noqa
    # |                        float|          None|               1.0|                 1.0|     X|                   X|                            X|         1.0|                     X|                 X|                 X|                           X|                 1.0|             X|  # noqa
    # |                       double|          None|               1.0|                 1.0|     X|                   X|                            X|         1.0|                     X|                 X|                 X|                           X|                 1.0|             X|  # noqa
    # |                       binary|          None|bytearray(b'\x00')|  bytearray(b'\x00')|     X|                   X|                            X|           X|  bytearray(b'\x01\...|bytearray(b'\x01')|bytearray(b'\x01')|           bytearray(b'ABC')|                   X|             X|  # noqa
    # |                decimal(10,0)|          None|                 X|                   X|     X|                   X|                            X|Decimal('1')|                     X|                 X|                 X|                           X|        Decimal('1')|             X|  # noqa
    # +-----------------------------+--------------+------------------+--------------------+------+--------------------+-----------------------------+------------+----------------------+------------------+------------------+----------------------------+--------------------+--------------+  # noqa

New type coercion:

# +-----------------------------+--------------+----------+--------------------+------+--------------------+-----------------------------+--------------------+----------------------+---------+-----------+----------------------------+--------------------+--------------+  # noqa
    # |SQL Type \ Python Value(Type)|None(NoneType)|True(bool)|              1(int)|a(str)|    1970-01-01(date)|1970-01-01 00:00:00(datetime)|          1.0(float)|array('i', [1])(array)|[1](list)|(1,)(tuple)|bytearray(b'ABC')(bytearray)|          1(Decimal)|{'a': 1}(dict)|  # noqa
    # +-----------------------------+--------------+----------+--------------------+------+--------------------+-----------------------------+--------------------+----------------------+---------+-----------+----------------------------+--------------------+--------------+  # noqa
    # |                      boolean|          None|      True|                True|     X|                   X|                            X|                True|                     X|        X|          X|                           X|                   X|             X|  # noqa
    # |                      tinyint|          None|         X|                   1|     X|                   X|                            X|                   1|                     X|        X|          X|                           X|                   1|             X|  # noqa
    # |                     smallint|          None|         X|                   1|     X|                   X|                            X|                   1|                     X|        X|          X|                           X|                   1|             X|  # noqa
    # |                          int|          None|         X|                   1|     X|                   0|                            X|                   1|                     X|        X|          X|                           X|                   1|             X|  # noqa
    # |                       bigint|          None|         X|                   1|     X|                   X|                            0|                   1|                     X|        X|          X|                           X|                   1|             X|  # noqa
    # |                       string|          None|    'true'|                 '1'|   'a'|        '1970-01-01'|         '1970-01-01 00:00...|               '1.0'|     "array('i', [1])"|    '[1]'|     '(1,)'|         "bytearray(b'ABC')"|                 '1'|    "{'a': 1}"|  # noqa
    # |                         date|          None|         X|datetime.date(197...|     X|datetime.date(197...|         datetime.date(197...|datetime.date(197...|                     X|        X|          X|                           X|datetime.date(197...|             X|  # noqa
    # |                    timestamp|          None|         X|                   X|     X|                   X|         datetime.datetime...|                   X|                     X|        X|          X|                           X|                   X|             X|  # noqa
    # |                        float|          None|       1.0|                 1.0|     X|                   X|                            X|                 1.0|                     X|        X|          X|                           X|                 1.0|             X|  # noqa
    # |                       double|          None|       1.0|                 1.0|     X|                   X|                            X|                 1.0|                     X|        X|          X|                           X|                 1.0|             X|  # noqa
    # |                       binary|          None|         X|                   X|     X|                   X|                            X|                   X|                     X|        X|          X|           bytearray(b'ABC')|                   X|             X|  # noqa
    # |                decimal(10,0)|          None|         X|                   X|     X|                   X|                            X|                   X|                     X|        X|          X|                           X|        Decimal('1')|             X|  # noqa
    # +-----------------------------+--------------+----------+--------------------+------+--------------------+-----------------------------+--------------------+----------------------+---------+-----------+----------------------------+--------------------+--------------+  # noqa

Table diff:
Only display differing values, <table1_value> vs. <table2_value>

# +-----------------------------+--------------+------------------------+------------------------------+------+----------------------------+-----------------------------+--------------------+----------------------+------------------------+------------------------+----------------------------+------------------------------+--------------+  # noqa
# |SQL Type \ Python Value(Type)|None(NoneType)|        True(bool)      |            1(int)            |a(str)|    1970-01-01(date)        |1970-01-01 00:00:00(datetime)|    1.0(float)      |array('i', [1])(array)|       [1](list)        |     (1,)(tuple)       |bytearray(b'ABC')(bytearray)|      1(Decimal)             |{'a': 1}(dict)|  # noqa
# +-----------------------------+--------------+------------------------+------------------------------+------+----------------------------+-----------------------------+--------------------+----------------------+------------------------+------------------------+----------------------------+------------------------------+--------------+  # noqa
# |                      tinyint|              | 1 vs. X                |                              |      |                            |                             |                    |                      |                        |                        |                            |                              |              |  # noqa
# |                     smallint|              | 1 vs. X                |                              |      |                            |                             |                    |                      |                        |                        |                            |                              |              |  # noqa
# |                          int|              | 1 vs. X                |                              |      |                            |                             |                    |                      |                        |                        |                            |                              |              |  # noqa
# |                       bigint|              | 1 vs. X                |                              |      |                            | X vs. 0                     |                    |                      |                        |                        |                            |                              |              |  # noqa
# |                       string|              | 'True' vs. 'true'      |                              |      |                            |                             |                    |                      |                        |                        |                            |                              |              |  # noqa
# |                         date|              |                        | X vs. datetime.date(...)     |      |                            |                             | X vs. datetime.date(...)|                  |                        |                        |                            |                              |              |  # noqa
# |                    timestamp|              |                        | datetime.datetime(...) vs. X |      |                            |                             |                    |                      |                        |                        |                            |                              |              |  # noqa
# |                       binary|              | bytearray(...) vs. X   | bytearray(...) vs. X         |      |                            |                             |                    |                      | bytearray(...) vs. X   | bytearray(...) vs. X   |                            |                              |              |  # noqa
# |                decimal(10,0)|              |                        |                             |      |                            |                             | Decimal('1') vs. X |                      |                        |                        |                            |                              |              |  # noqa
# +-----------------------------+--------------+------------------------+------------------------------+------+----------------------------+-----------------------------+--------------------+----------------------+------------------------+------------------------+----------------------------+------------------------------+--------------+  # noqa



How was this patch tested?

Correctness:
Added tests for both the legacy and new codepath, for arrow-batch eval.

Memory usage improvement:

According to manual benchmark, ~ 1.25x improvement in memory usage comparing the new path with the legacy pandas<>arrow conversion serialization.

Sample output from memory_profiler:

Legacy path:

Line #    Mem usage    Increment  Occurrences   Line Contents
=============================================================
    70    220.9 MiB    220.9 MiB           1       @profile
    71                                             def process_data_profiled(values, metadata):
    72    220.9 MiB      0.0 MiB           1           return complex_computation(values, metadata)


New path:

Line #    Mem usage    Increment  Occurrences   Line Contents
=============================================================
    93    175.1 MiB    175.1 MiB           1       @profile
    94                                             def process_data_profiled(values, metadata):
    95    175.1 MiB      0.0 MiB           1           return complex_computation(values, metadata)

Benchmark details here

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

No

@asl3 asl3 changed the title [DRAFT][PYTHON] Improve Python UDTF arrow serializer performance [DRAFT][PYTHON] Improve Python UDF arrow serializer performance Jun 19, 2025
@asl3 asl3 changed the title [DRAFT][PYTHON] Improve Python UDF arrow serializer performance [DRAFT][PYTHON] Improve Python UDF Arrow Serializer Performance Jun 19, 2025
@HyukjinKwon HyukjinKwon marked this pull request as draft June 23, 2025 02:52
Copy link
Member

@HyukjinKwon HyukjinKwon left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

From a cursory look, seems making sense

@asl3 asl3 changed the title [PYTHON] Improve Python UDF Arrow Serializer Performance [SPARK-52877][PYTHON] Improve Python UDF Arrow Serializer Performance Jul 19, 2025
@asl3 asl3 marked this pull request as ready for review July 19, 2025 02:09
@github-actions github-actions bot removed the DOCS label Jul 21, 2025
Copy link
Member

@viirya viirya left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think we may need to update corresponding doc especially for the type coercion difference. Maybe migration guide.

@viirya
Copy link
Member

viirya commented Jul 21, 2025

For the type coercion table, I remember we have document it in Python code doc. We may also need to update it.

@asl3 asl3 requested a review from ueshin July 21, 2025 20:32
@github-actions github-actions bot added the DOCS label Jul 21, 2025
@HyukjinKwon
Copy link
Member

Merged to master.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Projects
None yet
Development

Successfully merging this pull request may close these issues.

7 participants