Visitar URL original
[SPARK-60069][PYTHON] Migrate SQL_ARROW_UDTF to the EvalTypeHandler pipeline by Yicong-Huang · Pull Request #59274 · apache/spark · GitHub
Skip to content

[SPARK-60069][PYTHON] Migrate SQL_ARROW_UDTF to the EvalTypeHandler pipeline - #59274

Open
Yicong-Huang wants to merge 5 commits into
apache:masterfrom
Yicong-Huang:udtf-handler-dispatch
Open

Yicong-Huang wants to merge 5 commits into
apache:masterfrom
Yicong-Huang:udtf-handler-dispatch

Conversation

@Yicong-Huang

@Yicong-Huang Yicong-Huang commented Oct 7, 2026 •

Copy link
Copy Markdown
Contributor

What changes were proposed in this pull request?

Move SQL_ARROW_UDTF onto the EvalTypeHandler pipeline and lay out the UDTF pieces so the remaining UDTF eval types can follow.

  • Extract the shared UDTF prologue of read_udtf (AnalyzeResult constructor wrapping, UDTF instantiation incl. PARTITION BY wrapping, eval signature check) into read_single_udtf(pickleSer, udtf_info, partition_wrapper), returning (udtf, args_offsets, kwargs_offsets, return_type), the same shape as read_single_udf.
  • read_udtf now dispatches via get_eval_type_handler first, like read_udfs.
  • New eval_handlers/_udtf.py:
    • PartitionedUDTF, the base of UDTFWithPartitions and ArrowUDTFWithPartition (moved out of read_udtf), holding the shared per-partition UDTF instance, terminate/cleanup and SkipRestOfInputTableException handling.
    • UDTFEvalTypeHandler, a base handler that runs the eval -> terminate -> cleanup lifecycle and declares the partition_wrapper that read_single_udtf applies, so the prep no longer branches on eval type.
    • ArrowUDTFHandler for SQL_ARROW_UDTF (BatchEvalTypeHandler + UDTFEvalTypeHandler).
    • build_null_checker, moved out of read_udtf.
  • Move wrap_kwargs_support to eval_handlers/utils.py so handlers can use it.
PartitionedUDTF
├─ UDTFWithPartitions
└─ ArrowUDTFWithPartition

UDTFEvalTypeHandler           partition_wrapper = UDTFWithPartitions
└─ ArrowUDTFHandler           partition_wrapper = ArrowUDTFWithPartition

SQL_ARROW_TABLE_UDF and SQL_TABLE_UDF stay in read_udtf for now.

Why are the changes needed?

Step toward running UDTF and UDF eval types through one handler dispatch path and removing read_udtf (SPARK-59415).

Does this PR introduce any user-facing change?

No.

How was this patch tested?

New test_udtf_eval_type_handlers; existing test_udtf, test_arrow_udtf, and their Connect parity suites pass. ASV ArrowUDTFTimeBench, master vs this PR (median +- stdev of 15 runs, ms):

scenario identity filter count
sm_batch_few_col 0.82+-0.04 -> 0.74+-0.05 1.06+-0.03 -> 1.00+-0.03 0.59+-0.04 -> 0.55+-0.03
lg_batch_many_col 6.72+-0.26 -> 6.64+-0.27 8.78+-0.17 -> 8.94+-0.10 1.14+-0.03 -> 1.11+-0.03
pure_strings 2.23+-0.04 -> 2.38+-0.05 4.19+-0.04 -> 4.18+-0.07 1.19+-0.04 -> 1.17+-0.04

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

Generated-by: Claude Opus 5.5

This pull request and its description were written by Isaac.

Yicong-Huang and others added 2 commits October 7, 2026 18:51
Co-authored-by: Isaac <no-reply@databricks.com>
Co-authored-by: Isaac <no-reply@databricks.com>
@Yicong-Huang Yicong-Huang changed the title [PYTHON] Dispatch SQL_ARROW_UDTF through an eval type handler [SPARK-60069][PYTHON] Migrate SQL_ARROW_UDTF to the EvalTypeHandler pipeline Oct 7, 2026
Yicong-Huang and others added 3 commits October 8, 2026 00:27
Co-authored-by: Isaac <no-reply@databricks.com>
Co-authored-by: Isaac <no-reply@databricks.com>
Co-authored-by: Isaac <no-reply@databricks.com>
@Yicong-Huang
Yicong-Huang marked this pull request as ready for review October 9, 2026 08:41

@uros-b uros-b left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

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

Looks good to me, thank you @Yicong-Huang! Could you please ping @HyukjinKwon / @zhengruifeng for review too?

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants