[SPARK-58601][PYTHON] Tighten mapInPandas return-value contract to require a strict Iterator

[Yicong-Huang]

What changes were proposed in this pull request?

Tighten the mapInPandas (SQL_MAP_PANDAS_ITER_UDF) return-value contract in worker.py to require a strict Iterator, matching mapInArrow (SQL_MAP_ARROW_ITER_UDF). Previously the runtime check was isinstance(result, Iterator) or hasattr(result, "__iter__"), which accepted any iterable (e.g. a returned list); it is now isinstance(result, Iterator).

A legacy escape hatch is added to restore the old behavior for both mapInPandas and mapInArrow, since they share the MapInBatchExec evaluator:

• New config spark.sql.execution.pythonUDF.mapInBatch.legacy.acceptAnyIterable.enabled (internal, default false). When true, both UDFs may return any iterable rather than a strict iterator, matching the pre-4.3.0 behavior.
• The flag is threaded to the Python workers via ArrowPythonRunner.getPythonRunnerConfMap.
• The change is documented in the PySpark 4.2-to-4.3 migration guide.

Why are the changes needed?

The declared signatures have always been PandasMapIterFunction = Callable[[Iterator[DataFrameLike]], Iterator[DataFrameLike]] and ArrowMapIterFunction = Callable[[Iterator[pyarrow.RecordBatch]], Iterator[pyarrow.RecordBatch]], and the DataFrame.mapInPandas docstring states the function "outputs an iterator of pandas.DataFrames". The runtime, however, leniently accepted any iterable for mapInPandas, diverging from the documented contract. This aligns the mapInPandas runtime with the declared Iterator[...] signature (the same tightening mapInArrow already received in SPARK-56612), while the legacy flag gives users a migration path off the old iterable-accepting behavior for both APIs.

Does this PR introduce any user-facing change?

Yes. A mapInPandas UDF that returns a non-Iterator iterable (e.g. a list) is now rejected with UDF_RETURN_TYPE ("iterator of pandas.DataFrame"). Previously such a return value was accepted. Users who relied on the old behavior can set spark.sql.execution.pythonUDF.mapInBatch.legacy.acceptAnyIterable.enabled=true to restore it (this flag also covers mapInArrow).

Previous behavior:

df.mapInPandas(lambda it: [pdf for pdf in it], schema)  # accepted

New behavior:

df.mapInPandas(lambda it: [pdf for pdf in it], schema)  # PySparkTypeError: UDF_RETURN_TYPE
# to keep the old behavior:
spark.conf.set("spark.sql.execution.pythonUDF.mapInBatch.legacy.acceptAnyIterable.enabled", True)

How was this patch tested?

Updated test_pandas_map.py: removed the "returning list of DataFrames" positive case from test_map_in_pandas; added a list_not_iter negative case to check_other_than_dataframe_iter asserting a returned list is rejected (mirroring test_arrow_map.py::test_other_than_recordbatch_iter); and added test_map_in_pandas_legacy_accept_any_iterable asserting the legacy flag restores acceptance. Added the matching test_map_in_arrow_legacy_accept_any_iterable to test_arrow_map.py. Ran both suites plus the mapInPandas and mapInArrow Connect parity suites. The behavior change is documented in the PySpark 4.2-to-4.3 migration guide.

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

No

[uros-b]

Should we update the migration guide? To explain mapInPandas UDF returning a plain list of DataFrames (e.g. lambda it: [pdf for pdf in it]) previously succeeded and now raises UDF_RETURN_TYPE with "but is list"

Otherwise looks good, thank you @Yicong-Huang!

[Yicong-Huang]

Should we update the migration guide? To explain mapInPandas UDF returning a plain list of DataFrames (e.g. lambda it: [pdf for pdf in it]) previously succeeded and now raises UDF_RETURN_TYPE with "but is list"

Thanks @uros-b! I am on the fence here. The public doc declares Callable[[Iterator[DataFrameLike]], and accepting list which is an iterable not iterator has always been a silently allowed side effect. This PR tightens the allowed return type to its declared type, so not sure how to mention it in migration guide. Maybe @HyukjinKwon @cloud-fan can suggest action here?

Actually I think we should just keep the change for now, and maybe issue a warning or sth

[cloud-fan]

0 blocking, 1 non-blocking, 0 nits.
The implementation and focused regression test are sound, but the observable compatibility change should be called out in the PySpark migration guide.

Suggestions (1)

• General: Add a PySpark 4.3 migration note explaining that mapInPandas now rejects lists and other non-Iterator iterables, and recommend returning iter(...).

Verification

Verified that the public mapInPandas documentation requires an iterator return, the worker now performs a strict collections.abc.Iterator check, and the closest mapInArrow analogue enforces the same outer-container contract through verify_return_type. The new test covers the formerly accepted list case and preserves existing coverage for invalid iterator elements.

PR metadata suggestions

• Correct the user-facing-change section: returning a list or another non-Iterator iterable now changes from success to UDF_RETURN_TYPE.

[Yicong-Huang]

I changed it to describe as behavior change. Also added a flag to revert back to the old behavior. the same flag will guard mapInArrow as well. A migration guide from 4.2 to 4.3 is added.

Note that I kept the flag spark.sql.execution.pythonUDF.mapInBatch.legacy.acceptAnyIterable.enabled to be false by default. meaning that the change will take effect and choose to use the old behavior by explcitily setting the flag to be true.

@uros-b @HyukjinKwon @cloud-fan could you please check again?

[cloud-fan]

1 addressed, 0 remaining, 1 new. (1 newly introduced, 0 late catches, 0 previously raised.)
1 blocking, 0 non-blocking, 0 nits.
The runtime and compatibility paths are coherent, but Spark Connect's public callback annotations need to be aligned with the newly enforced contract.

Correctness (1)

• python/docs/source/migration_guide/pyspark_upgrade.rst:24: Spark Connect's callback aliases still advertise Iterable returns even though the shared worker now requires Iterator. -- see inline

Verification

Traced both MapInBatch evaluation types from MapInBatchExec through the Arrow runner configuration into worker.read_udfs. The strict and legacy branches match for pandas and Arrow, and focused tests cover both modes. Spark Connect dispatches those same evaluation types while its callback aliases still declare Iterable returns.

PR metadata suggestions

• Correct the claim that the declared signatures have always used Iterator: Spark Connect's current callback aliases use Iterable, or describe that discrepancy until those aliases are updated.

[cloud-fan]

Please update the Spark Connect callback aliases to require Iterator as well. connect/_typing.py still declares both PandasMapIterFunction and ArrowMapIterFunction with Iterable returns, but Connect dispatches these same strict evaluation types, so its public annotations currently tell users that returning a list is valid when the worker rejects it.

[HyukjinKwon]

4.4 since branch-4.3 is already cut out