[SPARK-59622][PYTHON] Migrate Arrow map/iter UDF eval types to the EvalTypeHandler pipeline - #58900
Yicong-Huang wants to merge 20 commits into
Conversation
Co-authored-by: Isaac <no-reply@databricks.com>
…nnotations Import pyarrow at the top of _arrow.py and guard its import in the package __init__ behind pyarrow availability. Move extract_key_value_indexes to the driver-safe _base.py (worker_util is worker-only under SPARK_TESTING). Use from __future__ import annotations so annotations are real names, not strings, and drop the type: ignore in verification.py via cast. Co-authored-by: Isaac <no-reply@databricks.com>
…port guard worker_util backs both the python worker and the eval handlers, so relax the SPARK_TESTING import guard to allow the handler import path (imported on the driver during test collection) instead of relocating the helper to _base. Co-authored-by: Isaac <no-reply@databricks.com>
Co-authored-by: Isaac <no-reply@databricks.com>
…d-map output The declared struct return type maps to the same pa.Schema via the existing to_arrow_schema util, so the local helper is unnecessary. Co-authored-by: Isaac <no-reply@databricks.com>
Move the private _top_level_package to the top of verification.py and order the verify_* functions alphabetically; order the handler classes in _arrow.py the same way. Co-authored-by: Isaac <no-reply@databricks.com>
…rt worker_util guard _util is a leaf module (no pyspark imports), safe to import from both the worker and the handlers on driver and executor, so the handlers no longer import the worker-only worker_util and its SPARK_TESTING guard returns to worker-only. Co-authored-by: Isaac <no-reply@databricks.com>
…andlers Mirror the source layout: framework (_base) tests in test_base_eval_type_handlers and Arrow-flavor tests in test_arrow_eval_type_handlers, leaving room for future flavors (e.g. pandas). The framework tests no longer depend on the Arrow handlers. Co-authored-by: Isaac <no-reply@databricks.com>
…lers The whole module is Arrow-only and gated on have_pyarrow, so import pyarrow once under that guard instead of in every test method. Co-authored-by: Isaac <no-reply@databricks.com>
|
@Spenserrrr @zhengruifeng @gaogaotiantian please help review! this PR moves more arrow eval types. |
…ot the test helper have_pyarrow lives in pyspark.testing.utils; use the production require_minimum_pyarrow_version so the package init does not pull the test module into worker startup. Tests keep using have_pyarrow. Co-authored-by: Isaac <no-reply@databricks.com>
Co-authored-by: Isaac <no-reply@databricks.com>
…nderscore) Co-authored-by: Isaac <no-reply@databricks.com>
Build handler input through shared helpers (_batch, _struct_batch, _one_group, _one_cogroup, _grouped_arg_offsets) and construct handlers via _scalar_handler / _grouped_handler, so every test follows the same shape and drops the ad hoc arg_offsets literals. Cover both the values-only and with-key branches uniformly for the grouped, grouped-iter, and co-grouped handlers. Co-authored-by: Isaac <no-reply@databricks.com>
They exercise SQL Arrow-UDF machinery, so they belong in the pyspark-sql module rather than the top-level pyspark-core tests: the framework tests move to pyspark.sql.tests and the Arrow handler tests to pyspark.sql.tests.arrow. Co-authored-by: Isaac <no-reply@databricks.com>
Put the tests in a tests subpackage next to the code they exercise, mirroring the per-component tests convention (pyspark.sql.tests, pyspark.ml.tests, ...), instead of the top-level pyspark tests. Co-authored-by: Isaac <no-reply@databricks.com>
Keep the pyspark.sql.tests.* entries contiguous by listing the pyspark.eval_handlers.tests goals as their own block at the end rather than interleaved. Co-authored-by: Isaac <no-reply@databricks.com>
Co-authored-by: Isaac <no-reply@databricks.com>
Co-authored-by: Isaac <no-reply@databricks.com>
…ckages Co-authored-by: Isaac <no-reply@databricks.com>
There was a problem hiding this comment.
Thanks for working on this refactor. I compared the migrated handlers with the previous read_udfs branches and verified that the serializer selection and execution behavior are preserved. I left one inline question about PyArrow-dependent handler registration; everything else looks good to me.
| from pyspark.sql.pandas.utils import require_minimum_pyarrow_version | ||
|
|
||
| require_minimum_pyarrow_version() | ||
| except Exception: |
There was a problem hiding this comment.
Could we decouple handler registration from PyArrow availability? This catch silently leaves the built-in Arrow handlers unregistered, so a valid Arrow eval type can later fall through to ValueError("Unknown eval type") instead of reporting the missing or incompatible PyArrow dependency. Could _arrow remain importable for registration and defer the PyArrow check until a handler is instantiated or run?
There was a problem hiding this comment.
Thanks @Spenserrrr I had some debates in my mind and I think you are right. I was hoping to get rid of some type ignores and string based type annotations, and wished to have arrow imported once for all arrow eval types so that it can be bundled with registration: if we have arrow, arrow eval type handlers will be registered, if not, arrow eval types handlers will not exist. I do think that might be better in the long run.
However, as you mentioned, currently when arrow is missing, we have a regression on its error message and "unknown eval type" would be more confusing than something like "arrow does not exist". As a result, I decided to revert the above changes and defer to later.
if you have better ways to support the long term goal without compromising on the error cases, please feel free to improve with new PR!
There was a problem hiding this comment.
Thanks, the current change looks good to me. For the longer-term idea, I think your intuition makes sense. We may need a richer registry entry that stores both the handler (or a lazy loader) and its dependency check, instead of dropping the entry when PyArrow is unavailable. That would let lookup distinguish an unknown eval type from a known Arrow handler that cannot run, and report the PyArrow error.
What do you think? If this direction makes sense, I’m happy to try it in a separate PR.
There was a problem hiding this comment.
I currently don't have a good idea about how we can achieve lazy evaluation of handlers and user friendly error message at the same time, in an elegant way. I think registering a handler successfully and failed later on pyarrow is not available is not ideal either: it's better to fail early, at registration. If you have a good solution feel free to raise a PR!
Decouple handler registration from pyarrow availability: _arrow imports pyarrow lazily (inside run and TYPE_CHECKING) so its handlers always register, and each handler calls require_minimum_pyarrow_version in __init__. A missing or too-old pyarrow now raises a clear error when the Arrow eval type runs, instead of leaving it unregistered and failing later with "Unknown eval type". The package __init__ imports _arrow unconditionally again. Co-authored-by: Isaac <no-reply@databricks.com>
…alTypeHandler pipeline ### What changes were proposed in this pull request? Continues the umbrella refactor [SPARK-59415](https://issues.apache.org/jira/browse/SPARK-59415), which replaces the `if/elif` dispatcher in `read_udfs` with the `pyspark.eval_handlers` framework from #58729. This PR migrates the five Arrow-native map/iter eval types to `EvalTypeHandler` subclasses in `python/pyspark/eval_handlers/_arrow.py`, each declaring its `eval_type` and self-registering: `ArrowScalarIterUDFHandler` (`SQL_SCALAR_ARROW_ITER_UDF`), `ArrowMapUDFHandler` (`SQL_MAP_ARROW_ITER_UDF`), `ArrowGroupedMapUDFHandler` (`SQL_GROUPED_MAP_ARROW_UDF`), `ArrowGroupedMapIterUDFHandler` (`SQL_GROUPED_MAP_ARROW_ITER_UDF`), and `ArrowCoGroupedMapUDFHandler` (`SQL_COGROUPED_MAP_ARROW_UDF`). Their `read_udfs` branches and the now-dead serializer-selection entries are removed. Supporting changes: the shared result-verification helpers move from `worker.py` to the leaf `eval_handlers/verification.py`; `extract_key_value_indexes` moves to a new leaf `eval_handlers/_util.py`; `_arrow.py` imports pyarrow at module top and the package `__init__` only imports it when `have_pyarrow`; handler logic is a direct move of the existing per-branch code. ### Why are the changes needed? To make each eval type's execution self-contained and testable and to remove the central `if/elif` dispatch, so eval types can be migrated one at a time. See the parent JIRA. ### Does this PR introduce _any_ user-facing change? No. Internal worker refactor; the migrated paths are behavior-identical, with no change to the UDF API or the on-the-wire format. ### How was this patch tested? New `pyspark.tests.test_base_eval_type_handlers` (framework) and `pyspark.tests.test_arrow_eval_type_handlers` (per-handler run tests, including the scalar-iter row-count failure path) suites. The existing per-eval-type Arrow suites (`test_arrow_map`, `test_arrow_grouped_map`, `test_arrow_cogrouped_map`, `test_arrow_udf_scalar`) cover the migrated paths and pass locally. ### Was this patch authored or co-authored using generative AI tooling? Generated-by: Opus 4.8 Closes #58900 from Yicong-Huang/SPARK-59415-arrow-map-iter. Authored-by: Yicong Huang <17627829+Yicong-Huang@users.noreply.github.com> Signed-off-by: Yicong-Huang <17627829+Yicong-Huang@users.noreply.github.com> (cherry picked from commit 203017c) Signed-off-by: Yicong-Huang <17627829+Yicong-Huang@users.noreply.github.com>
4 participants
What changes were proposed in this pull request?
Continues the umbrella refactor SPARK-59415, which replaces the
if/elifdispatcher inread_udfswith thepyspark.eval_handlersframework from #58729. This PR migrates the five Arrow-native map/iter eval types toEvalTypeHandlersubclasses inpython/pyspark/eval_handlers/_arrow.py, each declaring itseval_typeand self-registering:ArrowScalarIterUDFHandler(SQL_SCALAR_ARROW_ITER_UDF),ArrowMapUDFHandler(SQL_MAP_ARROW_ITER_UDF),ArrowGroupedMapUDFHandler(SQL_GROUPED_MAP_ARROW_UDF),ArrowGroupedMapIterUDFHandler(SQL_GROUPED_MAP_ARROW_ITER_UDF), andArrowCoGroupedMapUDFHandler(SQL_COGROUPED_MAP_ARROW_UDF). Theirread_udfsbranches and the now-dead serializer-selection entries are removed.Supporting changes: the shared result-verification helpers move from
worker.pyto the leafeval_handlers/verification.py;extract_key_value_indexesmoves to a new leafeval_handlers/_util.py;_arrow.pyimports pyarrow at module top and the package__init__only imports it whenhave_pyarrow; handler logic is a direct move of the existing per-branch code.Why are the changes needed?
To make each eval type's execution self-contained and testable and to remove the central
if/elifdispatch, so eval types can be migrated one at a time. See the parent JIRA.Does this PR introduce any user-facing change?
No. Internal worker refactor; the migrated paths are behavior-identical, with no change to the UDF API or the on-the-wire format.
How was this patch tested?
New
pyspark.tests.test_base_eval_type_handlers(framework) andpyspark.tests.test_arrow_eval_type_handlers(per-handler run tests, including the scalar-iter row-count failure path) suites. The existing per-eval-type Arrow suites (test_arrow_map,test_arrow_grouped_map,test_arrow_cogrouped_map,test_arrow_udf_scalar) cover the migrated paths and pass locally.Was this patch authored or co-authored using generative AI tooling?
Generated-by: Opus 4.8