soaring0616 opened a new pull request, #59667:
URL: https://github.com/apache/spark/pull/59667
### What changes were proposed in this pull request?
Annotate `wrapped_train_fn` in `TorchDistributor._get_spark_task_function`
with the `mapInArrow` contract (`ArrowMapIterFunction`) and remove its `# type:
ignore[no-untyped-def]`.
Before:
def wrapped_train_fn(iterator): # type: ignore[no-untyped-def]
After:
def wrapped_train_fn(iterator: Iterator["pa.RecordBatch"]) ->
Iterator["pa.RecordBatch"]:
`pyarrow` is only imported inside the function body, so `pa.RecordBatch` is
a string annotation and `import pyarrow as pa` is added under `TYPE_CHECKING`.
Also widen `input_schema_json` of `_setup_spark_partition_data` from
`Dict[str, Any]` to `Optional[Dict[str, Any]]`.
### Why are the changes needed?
With the ignore in place, mypy skips the body of `wrapped_train_fn`.
Removing only the comment reports `Function is missing a type annotation
[no-untyped-def]` (`disallow_untyped_defs = True` in python/mypy.ini).
With the annotation added, mypy checks the body and reports [arg-type] for
the call to `_setup_spark_partition_data`: the argument is `Optional[Dict[str,
Any]]` because `schema_json` is None when no input DataFrame is given. The
helper already returns early in that case (`if input_schema_json is None:
yield; return`), so its annotation was narrower than its behavior.
The annotation has to be a string: on Python 3.12, an unquoted
`pyarrow.RecordBatch` raises NameError when the function is defined, since
`pyarrow` is only imported inside the body (on Python 3.14 it is raised when
the annotations are evaluated).
### Does this PR introduce _any_ user-facing change?
No.
### How was this patch tested?
- mypy 1.19.1 on Python 3.12
- `dev/lint-python`: compilation, custom errors check and ruff pass.
- CI on my fork: `Linters, licenses, and dependencies` (The `pyspark-ml`
job was still running when I opened this PR; I will update this section when it
finishes.)
- `pyspark.ml.torch.distributor` still imports with `pyarrow` blocked
(`sys.modules["pyarrow"] = None`).
- The new annotations survive a cloudpickle round trip (checked on a minimal
nested function with the same annotations).
- No active test covers the path with an input DataFrame: the tests that
call `_train_on_dataframe` are in `test_data_loader.py`, and
`TorchDistributorDataLoaderUnitTests` is skipped (SPARK-50864).
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: Claude Sonnet 5.5
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
