Skip to content

Cover Spark 4.2 Arrow Python UDF columnar input [fast-ut] [reduced-it][databricks] - #15885

Merged
firestarman merged 6 commits into
NVIDIA:mainfrom
firestarman:audit/15663-arrow-udf-columnar-input
Sep 7, 2026
Merged

Cover Spark 4.2 Arrow Python UDF columnar input [fast-ut] [reduced-it][databricks]#15885
firestarman merged 6 commits into
NVIDIA:mainfrom
firestarman:audit/15663-arrow-udf-columnar-input

Conversation

@firestarman

@firestarman firestarman commented Sep 3, 2026 •
edited
Loading

Copy link
Copy Markdown
Collaborator

Fixes #15663.

Description

Spark 4.2 (SPARK-56350) lets CPU ArrowEvalPythonExec consume Arrow-backed ColumnarBatch input and skip ColumnarToRow. That does not change the GPU Python UDF path (GpuColumnVector -> cuDF Arrow IPC), but it does change how tests and CPU fallbacks must behave.

  • Keep the test Arrow DSv2 reader batches alive after reader.close(). Spark 4.2 DataSourceRDD closes the reader as soon as next() is false, while CPU ArrowEvalPythonExec can still hold pass-through ColumnVector references. Downstream consumers close those batches; the test reader closes the allocator only when allocated memory is zero.
  • Add Arrow-source pandas UDF coverage that expects GPU HostColumnarToGpu + GpuArrowEvalPythonExec.
  • Add Spark 4.2 Arrow-optimized regular UDF (evalType=101) coverage. CPU spark.sql.execution.arrow.pythonUDF.columnarInput.enabled is disabled only to avoid the Spark 4.2.0 hang from SPARK-58241; that is a CPU bug, not a GPU bug.
  • Add a fallback test with a GPU Parquet scan under CPU ArrowEvalPythonExec so GpuColumnVector is not handed to Spark's CPU Arrow UDF path.

There is no user-facing config or execution-path change.

Checklists

Documentation

  • Updated for new or modified user-facing features or behaviors
  • No user-facing change

Testing

  • Added or modified tests to cover new code paths
  • Covered by existing tests
    (Please provide the names of the existing tests in the PR description.)
  • Not required

New/updated tests:

  • test_arrow_source_pandas_udf
  • test_arrow_source_regular_udf
  • test_pandas_udf_cpu_arrow_eval_after_gpu_scan

Locally verified (datasourcev2_read_test.py and udf_test.py::test_pandas_udf_cpu_arrow_eval_after_gpu_scan):

  • 330 (Spark 3.3.0): 7 passed, 1 skipped (test_arrow_source_regular_udf + INJECT_OOM)
  • 413 (Spark 4.1.3): 7 passed, 1 skipped (same skip)
  • 420 (Spark 4.2.0): 8 passed

Performance

  • Tests ran and results are added in the PR description
  • Issue filed with a link in the PR description
  • Not required

@firestarman
Keep Arrow test batches alive for Spark 4.2 Python UDFs.
d37677a 
Spark 4.2 CPU ArrowEvalPythonExec can retain pass-through ColumnVector
references, so the test Arrow source must not free earlier batches until
the reader closes. Add GPU coverage for that Arrow input path and for a
CPU Python exec above a GPU scan.

Signed-off-by: Firestarman <firestarmanllc@gmail.com>
@firestarman
firestarman requested a review from a team September 3, 2026 06:31
@firestarman firestarman mentioned this pull request Sep 3, 2026
Closed
@greptile-apps

greptile-apps Bot commented Sep 3, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor

Greptile Summary

Adds Spark 4.2 integration coverage for Arrow-backed Python UDF inputs and CPU fallback after a GPU Parquet scan.

  • Keeps test Arrow allocators available while emitted columnar batches may still be referenced.
  • Verifies pandas and Arrow-optimized regular UDF execution through HostColumnarToGpu and GpuArrowEvalPythonExec.
  • Verifies that CPU ArrowEvalPythonExec receives a host transition after a GPU scan.

Confidence Score: 5/5

The PR appears safe to merge, with no outstanding correctness, security, or repository-rule issues identified.

The current changes provide focused integration coverage and test-resource lifetime handling, and no new changes to the reviewed PR files were introduced after the previous review SHA.

Important Files Changed

Filename Overview
integration_tests/src/main/python/datasourcev2_read_test.py Adds Spark 4.2 Arrow-source pandas and regular Python UDF parity and execution-plan coverage.
integration_tests/src/main/python/udf_test.py Adds fallback coverage ensuring CPU Arrow UDF execution does not directly consume GPU column vectors.
integration_tests/src/main/scala/com/nvidia/spark/rapids/tests/datasourcev2/parquet/TestingV2Source.scala Extends the test Arrow reader’s allocator lifetime to support retained pass-through column vectors.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart LR
  A[Arrow DSv2 batch] --> B[HostColumnarToGpu]
  B --> C[GpuArrowEvalPythonExec]
  D[GPU Parquet scan] --> E[GPU-to-host transition]
  E --> F[CPU ArrowEvalPythonExec]
Loading

Reviews (5): Last reviewed commit: "Merge branch 'main' into audit/15663-arr..." | Re-trigger Greptile

@firestarman firestarman changed the title Cover Spark 4.2 Arrow Python UDF columnar input [fast-ut] [reduced-it] Cover Spark 4.2 Arrow Python UDF columnar input [fast-ut] [reduced-it][databricks] Sep 3, 2026
@firestarman
Keep Arrow test allocators until consumers release buffers.
39a862b 
Spark 4.2 closes the partition reader before SPARK-56350 CPU ArrowEvalPythonExec
drops pass-through vectors, so closing batches in reader.close() UAFs. Close the
allocator only when allocated memory is zero, and use in-range ints for pandas
UDF Arrow casts on Spark 4.x.

Signed-off-by: Firestarman <firestarmanllc@gmail.com>
@firestarman

firestarman commented Sep 3, 2026

Copy link
Copy Markdown
Collaborator Author

build

@firestarman

firestarman commented Sep 4, 2026

Copy link
Copy Markdown
Collaborator Author

build

res-life
res-life reviewed Sep 4, 2026
View reviewed changes
Comment thread ...ests/src/main/scala/com/nvidia/spark/rapids/tests/datasourcev2/parquet/TestingV2Source.scala
private val rootAllocator = new RootAllocator(Long.MaxValue)
private val allocator: BufferAllocator =
rootAllocator.newChildAllocator(s"arrow-test-reader-$startNum", 0, Long.MaxValue)
private val allBatches = new util.ArrayList[ColumnarBatch]()

@res-life res-life Sep 4, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

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

Forget to close this columnar batch list?

@firestarman firestarman Sep 7, 2026

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

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

Not forgotten — we must not close allBatches in the reader.

Spark 4.2 DataSourceRDD calls reader.close() as soon as next() returns false, while CPU ArrowEvalPythonExec (SPARK-56350) can still hold pass-through ArrowColumnVector refs. Closing this list (and the allocator) at that point UAFs; GPU HostColumnarToGpu can also still hold ArrowBufs.

allBatches only pins the wrappers. Downstream Spark/GPU consumers close the ColumnarBatchs. The reader closes the allocator only when getAllocatedMemory == 0.

Added a comment on allBatches in 3e90357 to make that explicit.

@firestarman
Document why the test Arrow reader keeps batches unclosed.
3e90357 
The allBatches list is a lifetime pin. Closing it in reader.close() UAFs
under SPARK-56350, which can retain pass-through vectors after the reader
is closed.

Signed-off-by: Firestarman <firestarmanllc@gmail.com>
@firestarman
Merge branch 'audit/15663-arrow-udf-columnar-input' of github.com:fir…
ec5d775 
…estarman/cudf-spark into audit/15663-arrow-udf-columnar-input
@firestarman
firestarman requested review from a team and res-life September 7, 2026 01:43
@firestarman

firestarman commented Sep 7, 2026

Copy link
Copy Markdown
Collaborator Author

build

res-life
res-life approved these changes Sep 7, 2026
View reviewed changes

@res-life res-life left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

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

LGTM, but please excise at lease one Spark patch version for all the Spark minor versions: 33x, 34x,35x, 40x, 41x, 42x, 50x.

@firestarman

firestarman commented Sep 7, 2026 •
edited
Loading

Copy link
Copy Markdown
Collaborator Author

LGTM, but please excise at lease one Spark patch version for all the Spark minor versions: 33x, 34x,35x, 40x, 41x, 42x, 50x.

Locally re-ran the Arrow DSv2 / CPU Arrow UDF tests after the allocator-close change:

TESTS="datasourcev2_read_test.py udf_test.py::test_pandas_udf_cpu_arrow_eval_after_gpu_scan" TEST_PARALLEL=1 ./run_pyspark_from_build.sh

  • 330 (Spark 3.3.0): 7 passed, 1 skipped (test_arrow_source_regular_udf + INJECT_OOM)
  • 413 (Spark 4.1.3): 7 passed, 1 skipped (same skip)
  • 420 (Spark 4.2.0): 8 passed

@firestarman
firestarman merged commit dab818c into NVIDIA:main Sep 7, 2026
59 checks passed
@firestarman
firestarman deleted the audit/15663-arrow-udf-columnar-input branch September 7, 2026 07:16
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@res-life res-life res-life approved these changes

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

[AUDIT] [Investigate] Skipping ColumnarToRow for Arrow-backed input to Python UDFs

3 participants

@firestarman @res-life @nvauto