Anna Liang created SPARK-60417:
----------------------------------
Summary: Remove type: ignore[no-untyped-def] in
pyspark.ml.torch.distributor
Key: SPARK-60417
URL: https://issues.apache.org/jira/browse/SPARK-60417
Project: Spark
Issue Type: Improvement
Components: ML, PySpark
Affects Versions: 4.4.0
Reporter: Anna Liang
`wrapped_train_fn` in `TorchDistributor._get_task_function` is passed to
`DataFrame.mapInArrow`, but it is unannotated and marked `# type:
ignore[no-untyped-def]`. Removing only the comment makes mypy report "Function
is missing a type annotation [no-untyped-def]", because python/mypy.ini sets
`disallow_untyped_defs = True`. While the function is unannotated, mypy also
skips its body.
Annotate it with the `mapInArrow` contract (`Iterator[pyarrow.RecordBatch]` in
and out, see `ArrowMapIterFunction` in `pyspark/sql/connect/_typing.py`).
`pyarrow` is only imported inside the function. So the annotation is a string
and `pyarrow` is imported under `TYPE_CHECKING`
Once the body is checked, mypy reports that the 2nd argument passed to
"_setup_spark_partition_data" has type "dict[str, Any] | None" but the parameter
(input_schema_json) expects "dict[str, Any]". The argument is None when no
input DataFrame is given, and _setup_spark_partition_data already returns early
in that case ("if input_schema_json is None: yield; return"), so its
annotation is widened to Optional.
--
This message was sent by Atlassian Jira
(v8.20.10#820010)
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
