Skip to content

[FLINK-40472][python] Support Arrow vectorized UDFs in DataFrame API - #29210

Open
auroflow wants to merge 12 commits into
apache:masterfrom
auroflow:auroflow/codex/flink-40472-arrow-dataframe-udf
Open

auroflow wants to merge 12 commits into
apache:masterfrom
auroflow:auroflow/codex/flink-40472-arrow-dataframe-udf

Conversation

@auroflow

@auroflow auroflow commented Sep 16, 2026

Copy link
Copy Markdown
Contributor

What is the purpose of the change

Implement FLINK-40472 by adding synchronous Arrow vectorized scalar UDFs to the PyFlink DataFrame API. UDFs receive Arrow arrays and can apply pyarrow.compute operations directly to column batches.

Arrow mode can be inferred from Arrow container annotations or selected with func_type="arrow". The Table API supports explicit Arrow mode.

Brief change log

  • Add Arrow UDF declaration and annotation inference, and document the new execution mode.
  • Distinguish general, pandas and Arrow functions in Python calculation splitting and map merging.
  • Extend the existing Arrow operator, worker and coder protocol to carry native Arrow batches while retaining pandas defaults.
  • Support nested and whole-row arguments, normalize chunked results, and validate result containers, row counts and terminal schemas.
  • Share Arrow schema mapping, preserve nested nullability and PyArrow 5 compatibility, reduce temporary validation allocations, and reset map child writers between batches.

Verifying this change

This change extends existing declaration, worker, coder, planner and Arrow reader/writer tests, plus the combined DataFrame UDF integration test. Tests cover observable values, schemas, execution plans and error behavior.

  • Focused pure-Python and compiled coder checks passed on PyArrow 5, 18 and 23.
  • Targeted Java planner, protocol, Arrow operator and reader/writer tests passed, including PythonCalcSplitFunctionKindRuleTest, ProtoUtilsTest and ArrowReaderWriterTest.
  • Java 17 reactor build passed with ./mvnw clean install -DskipTests -Dfast -Pskip-webui-build -T4. The full Maven test suite was not run.
  • Spotless, flake8 and mypy checks passed, including the pre-push checks.

Does this pull request potentially affect one of the following parts:

  • Dependencies (does it add or upgrade a dependency): no
  • The public API, i.e., is any changed class annotated with @Public(Evolving): yes — extends the Python DataFrame and Table UDF APIs.
  • The serializers: yes — Python Arrow coders, Arrow writers and additive worker-protocol metadata; managed-state serializers are unchanged.
  • The runtime per-record code paths (performance sensitive): yes — Python scalar execution and Arrow batch serialization.
  • Anything that affects deployment or recovery: JobManager (and its components), Checkpointing, Kubernetes/Yarn, ZooKeeper: no
  • The S3 file system connector: no

Documentation

  • Does this pull request introduce a new feature? yes
  • If yes, how is the feature documented? docs — DataFrame UDF reference documentation and Python API docstrings.

Was generative AI tooling used to co-author this PR?
  • Yes (please specify the tool below)

Generated-by: Codex (GPT-6)

Respect whole-row argument offsets, normalize chunked results before nested calls, skip unused pandas timezone conversion for native Arrow, and validate container nullability without copying payloads.

Generated-by: Codex (GPT-6)
Defer and memoize ancestor visibility until a NOT NULL descendant contains physical nulls. Extend the existing coder round-trip and allocation test with Boolean lists.

Generated-by: Codex (GPT-6)
Use expected Arrow fields for structure and nullability checks in both coder implementations, while retaining Flink types for pandas conversion.

Generated-by: Codex (GPT-6)
Extend existing worker tests for normalized intermediate results, selected and repeated outputs, and whole-row result references. Apply import ordering to the rebase resolutions.

Generated-by: Codex (GPT-6)
Reset key and value writers together with map vectors so consecutive batches reuse aligned offsets. Cover direct and nested maps, nulls, empty maps, and partial batches through Arrow IPC round trips.

Generated-by: Codex (GPT-6)
Use the legacy MAP constructor on PyArrow 5 while retaining logical map-value null checks. Apply validated native output schemas by rebuilding container metadata over existing buffers, preserving sliced and nested results without version-dependent casts.

Generated-by: Codex (GPT-6)
…tion

Expand parent visibility as booleans using child counts from list/map offsets. Reuse coder round-trip tests for slices, varying child counts and nullability, and remove memory-usage assertions.

Generated-by: Codex (GPT-6)
@flinkbot

flinkbot commented Sep 16, 2026

Copy link
Copy Markdown
Collaborator

CI report:

Bot commands The @flinkbot bot supports the following commands:
  • @flinkbot run azure re-run the last Azure build

Move Arrow schema validation beside existing serialization utilities and UDF result handling beside pandas execution helpers.

Generated-by: Codex (GPT-6)
@auroflow
auroflow marked this pull request as ready for review September 17, 2026 02:54
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