This is an automated email from the ASF dual-hosted git repository.
Yicong-Huang pushed a commit to branch branch-4.x
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/branch-4.x by this push:
new b456177af27b [SPARK-58601][PYTHON] Tighten mapInPandas return-value
contract to require a strict Iterator
b456177af27b is described below
commit b456177af27bea14f5088cd23aca3ae019451052
Author: Yicong Huang <[email protected]>
AuthorDate: Wed Aug 26 06:48:21 2026 +0000
[SPARK-58601][PYTHON] Tighten mapInPandas return-value contract to require
a strict Iterator
### 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 w [...]
### 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:
```python
df.mapInPandas(lambda it: [pdf for pdf in it], schema) # accepted
```
New behavior:
```python
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` [...]
### Was this patch authored or co-authored using generative AI tooling?
No
Closes #57800 from Yicong-Huang/SPARK-58601.
Authored-by: Yicong Huang <[email protected]>
Signed-off-by: Yicong-Huang <[email protected]>
(cherry picked from commit 69a108ba8adfef816f0652104c64062947c0df8c)
Signed-off-by: Yicong-Huang <[email protected]>
---
.../source/migration_guide/pyspark_upgrade.rst | 4 ++
python/pyspark/sql/connect/_typing.py | 6 +--
python/pyspark/sql/connect/dataframe.py | 2 +-
python/pyspark/sql/tests/arrow/test_arrow_map.py | 37 ++++++++++++++++
python/pyspark/sql/tests/pandas/test_pandas_map.py | 49 +++++++++++++++++++---
python/pyspark/worker.py | 40 ++++++++++++++++--
.../org/apache/spark/sql/internal/SQLConf.scala | 14 +++++++
.../sql/execution/python/ArrowPythonRunner.scala | 1 +
8 files changed, 140 insertions(+), 13 deletions(-)
diff --git a/python/docs/source/migration_guide/pyspark_upgrade.rst
b/python/docs/source/migration_guide/pyspark_upgrade.rst
index b257ee49b7dd..35bfe5f23c96 100644
--- a/python/docs/source/migration_guide/pyspark_upgrade.rst
+++ b/python/docs/source/migration_guide/pyspark_upgrade.rst
@@ -19,6 +19,10 @@
Upgrading PySpark
==================
+Upgrading from PySpark 4.2 to 4.3
+---------------------------------
+* In Spark 4.3, a ``mapInPandas`` UDF must return an iterator of
``pandas.DataFrame``\s; returning any other iterable such as a ``list`` now
raises ``UDF_RETURN_TYPE``, matching the existing ``mapInArrow`` behavior and
the declared ``Iterator[...]`` signature. To restore the previous behavior of
accepting any iterable for both ``mapInPandas`` and ``mapInArrow``, set
``spark.sql.execution.pythonUDF.mapInBatch.legacy.acceptAnyIterable.enabled``
to ``true``.
+
Upgrading from PySpark 4.1 to 4.2
---------------------------------
* In Spark 4.2, the minimum supported version for PyArrow has been raised from
15.0.0 to 18.0.0 in PySpark.
diff --git a/python/pyspark/sql/connect/_typing.py
b/python/pyspark/sql/connect/_typing.py
index 6c256bc65af9..707059b22501 100644
--- a/python/pyspark/sql/connect/_typing.py
+++ b/python/pyspark/sql/connect/_typing.py
@@ -17,7 +17,7 @@
import datetime
import decimal
from types import FunctionType
-from typing import Any, Callable, Iterable, NewType, Optional, Protocol,
Tuple, Union
+from typing import Any, Callable, Iterable, Iterator, NewType, Optional,
Protocol, Tuple, Union
import pyarrow
from pandas.core.frame import DataFrame as PandasDataFrame
@@ -44,9 +44,9 @@ DataTypeOrString = Union[DataType, str]
DataFrameLike = PandasDataFrame
-PandasMapIterFunction = Callable[[Iterable[DataFrameLike]],
Iterable[DataFrameLike]]
+PandasMapIterFunction = Callable[[Iterator[DataFrameLike]],
Iterator[DataFrameLike]]
-ArrowMapIterFunction = Callable[[Iterable[pyarrow.RecordBatch]],
Iterable[pyarrow.RecordBatch]]
+ArrowMapIterFunction = Callable[[Iterator[pyarrow.RecordBatch]],
Iterator[pyarrow.RecordBatch]]
PandasGroupedMapFunction = Union[
Callable[[DataFrameLike], DataFrameLike],
diff --git a/python/pyspark/sql/connect/dataframe.py
b/python/pyspark/sql/connect/dataframe.py
index b47b63dec1ac..95d6cddd50ba 100644
--- a/python/pyspark/sql/connect/dataframe.py
+++ b/python/pyspark/sql/connect/dataframe.py
@@ -2333,7 +2333,7 @@ class DataFrame(ParentDataFrame):
for f in schema.fields
]
- def foreach_partition_func(itr: Iterable[pa.RecordBatch]) ->
Iterable[pa.RecordBatch]:
+ def foreach_partition_func(itr: Iterator[pa.RecordBatch]) ->
Iterator[pa.RecordBatch]:
def flatten() -> Iterator[Row]:
for table in itr:
columnar_data = [
diff --git a/python/pyspark/sql/tests/arrow/test_arrow_map.py
b/python/pyspark/sql/tests/arrow/test_arrow_map.py
index 8f7ea52a1b07..f2c5644059b6 100644
--- a/python/pyspark/sql/tests/arrow/test_arrow_map.py
+++ b/python/pyspark/sql/tests/arrow/test_arrow_map.py
@@ -55,6 +55,43 @@ class MapInArrowTestsMixin:
expected = df.collect()
self.assertEqual(actual, expected)
+ def test_map_in_arrow_legacy_accept_any_iterable(self):
+ # With the legacy flag enabled, returning a non-Iterator iterable
(e.g. list) is accepted.
+ def list_not_iter(iterator):
+ return [batch for batch in iterator]
+
+ with self.sql_conf(
+
{"spark.sql.execution.pythonUDF.mapInBatch.legacy.acceptAnyIterable.enabled":
True}
+ ):
+ df = self.spark.range(10)
+ actual = df.mapInArrow(list_not_iter, "id long").collect()
+ expected = df.collect()
+ self.assertEqual(actual, expected)
+
+ def test_map_in_arrow_legacy_accept_sequence_protocol(self):
+ # A sequence-protocol object (implements __getitem__ but not __iter__)
is iterable via
+ # iter(...) even though it is not a collections.abc.Iterable, so the
legacy flag must
+ # accept it too.
+ class SequenceOnly:
+ def __init__(self, items):
+ self._items = items
+
+ def __getitem__(self, index):
+ return self._items[index]
+
+ self.assertFalse(hasattr(SequenceOnly([]), "__iter__"))
+
+ def returns_sequence(iterator):
+ return SequenceOnly([batch for batch in iterator])
+
+ with self.sql_conf(
+
{"spark.sql.execution.pythonUDF.mapInBatch.legacy.acceptAnyIterable.enabled":
True}
+ ):
+ df = self.spark.range(10)
+ actual = df.mapInArrow(returns_sequence, "id long").collect()
+ expected = df.collect()
+ self.assertEqual(actual, expected)
+
def test_map_in_arrow_with_limit(self):
def get_size(iterator):
for batch in iterator:
diff --git a/python/pyspark/sql/tests/pandas/test_pandas_map.py
b/python/pyspark/sql/tests/pandas/test_pandas_map.py
index 0edb929aac8f..b6cad5b7328e 100644
--- a/python/pyspark/sql/tests/pandas/test_pandas_map.py
+++ b/python/pyspark/sql/tests/pandas/test_pandas_map.py
@@ -86,11 +86,39 @@ class MapInPandasTestsMixin:
expected = df.collect()
self.assertEqual(actual, expected)
- # test returning list of DataFrames
- df = self.spark.range(10, numPartitions=3)
- actual = df.mapInPandas(lambda it: [pdf for pdf in it], "id
long").collect()
- expected = df.collect()
- self.assertEqual(actual, expected)
+ def test_map_in_pandas_legacy_accept_any_iterable(self):
+ # With the legacy flag enabled, returning a non-Iterator iterable
(e.g. list) is accepted.
+ with self.sql_conf(
+
{"spark.sql.execution.pythonUDF.mapInBatch.legacy.acceptAnyIterable.enabled":
True}
+ ):
+ df = self.spark.range(10, numPartitions=3)
+ actual = df.mapInPandas(lambda it: [pdf for pdf in it], "id
long").collect()
+ expected = df.collect()
+ self.assertEqual(actual, expected)
+
+ def test_map_in_pandas_legacy_accept_sequence_protocol(self):
+ # A sequence-protocol object (implements __getitem__ but not __iter__)
is iterable via
+ # iter(...) even though it is not a collections.abc.Iterable, so the
legacy flag must
+ # accept it too.
+ class SequenceOnly:
+ def __init__(self, items):
+ self._items = items
+
+ def __getitem__(self, index):
+ return self._items[index]
+
+ self.assertFalse(hasattr(SequenceOnly([]), "__iter__"))
+
+ def returns_sequence(iterator):
+ return SequenceOnly([pdf for pdf in iterator])
+
+ with self.sql_conf(
+
{"spark.sql.execution.pythonUDF.mapInBatch.legacy.acceptAnyIterable.enabled":
True}
+ ):
+ df = self.spark.range(10, numPartitions=3)
+ actual = df.mapInPandas(returns_sequence, "id long").collect()
+ expected = df.collect()
+ self.assertEqual(actual, expected)
def test_multiple_columns(self):
data = [(1, "foo"), (2, None), (3, "bar"), (4, "bar")]
@@ -186,6 +214,10 @@ class MapInPandasTestsMixin:
def bad_iter_elem(_):
return iter([1])
+ def list_not_iter(iterator):
+ # Iterable but not an Iterator: violates the
Iterator[pandas.DataFrame] contract.
+ return [pdf for pdf in iterator]
+
with self.assertRaisesRegex(
PythonException,
"Return type of the user-defined function should be iterator of
pandas.DataFrame, "
@@ -200,6 +232,13 @@ class MapInPandasTestsMixin:
):
(self.spark.range(10, numPartitions=3).mapInPandas(bad_iter_elem,
"a int").count())
+ with self.assertRaisesRegex(
+ PythonException,
+ "Return type of the user-defined function should be iterator of
pandas.DataFrame, "
+ "but is list",
+ ):
+ (self.spark.range(10, numPartitions=3).mapInPandas(list_not_iter,
"a int").count())
+
def test_dataframes_with_other_column_names(self):
with self.quiet():
self.check_dataframes_with_other_column_names()
diff --git a/python/pyspark/worker.py b/python/pyspark/worker.py
index e7af151cbc7b..a37c81893e03 100644
--- a/python/pyspark/worker.py
+++ b/python/pyspark/worker.py
@@ -147,6 +147,16 @@ class RunnerConf(Conf):
== "true"
)
+ @property
+ def map_in_batch_legacy_accept_any_iterable(self) -> bool:
+ return (
+ self.get(
+
"spark.sql.execution.pythonUDF.mapInBatch.legacy.acceptAnyIterable.enabled",
+ "false",
+ )
+ == "true"
+ )
+
@property
def binary_as_bytes(self) -> bool:
return self.get("spark.sql.execution.pyspark.binaryAsBytes", "true")
== "true"
@@ -2151,6 +2161,20 @@ def read_udfs(pickleSer, udf_info_list, eval_type,
runner_conf, eval_conf):
# invoke the UDF
output_batches = udf_func(input_batches)
+ # The declared signature is Iterator[...], so a strict iterator is
required by
+ # default. With the legacy flag, accept any object Python can
iterate over -- via
+ # iter(...), which honors both __iter__ and the sequence protocol
(__getitem__) --
+ # by adapting it into an iterator before the shared element-type
verification.
+ if runner_conf.map_in_batch_legacy_accept_any_iterable and not
isinstance(
+ output_batches, Iterator
+ ):
+ try:
+ output_batches = iter(output_batches)
+ except TypeError:
+ # Not iterable at all; leave it so verify_return_type
below raises the
+ # standard UDF_RETURN_TYPE error.
+ pass
+
# Post-processing
verified_iter = verify_return_type(
output_batches,
@@ -3187,11 +3211,19 @@ def read_udfs(pickleSer, udf_info_list, eval_type,
runner_conf, eval_conf):
df_for_struct=True,
)[0]
- # mapInPandas accepts any iterable (e.g. a list), not just an
- # iterator, so the standard verify_return_type (which requires an
- # Iterator) is intentionally not reused here.
result = map_udf(dataframe_iter())
- if not isinstance(result, Iterator) and not hasattr(result,
"__iter__"):
+ # The declared signature is Iterator[...], so a strict iterator is
required by
+ # default. With the legacy flag, accept any object Python can
iterate over -- via
+ # iter(...), which honors both __iter__ and the sequence protocol
(__getitem__) --
+ # by adapting it into an iterator.
+ if runner_conf.map_in_batch_legacy_accept_any_iterable and not
isinstance(
+ result, Iterator
+ ):
+ try:
+ result = iter(result)
+ except TypeError:
+ pass # Not iterable; fall through to the UDF_RETURN_TYPE
error below.
+ if not isinstance(result, Iterator):
raise PySparkTypeError(
errorClass="UDF_RETURN_TYPE",
messageParameters={
diff --git
a/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala
b/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala
index 1db7b42a01ff..fbbad38a8a5d 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala
@@ -5541,6 +5541,17 @@ object SQLConf {
.booleanConf
.createWithDefault(false)
+ val PYTHON_UDF_MAP_IN_BATCH_LEGACY_ACCEPT_ANY_ITERABLE_ENABLED =
+
buildConf("spark.sql.execution.pythonUDF.mapInBatch.legacy.acceptAnyIterable.enabled")
+ .internal()
+ .doc("When true, mapInPandas and mapInArrow UDFs may return any iterable
(e.g. a list) " +
+ "rather than a strict iterator, matching the behavior before 4.3.0.
When false, the " +
+ "returned value must be an iterator, matching the declared
Iterator[...] signatures.")
+ .version("4.3.0")
+ .withBindingPolicy(ConfigBindingPolicy.SESSION)
+ .booleanConf
+ .createWithDefault(false)
+
val PYTHON_PLANNER_EXEC_MEMORY =
buildConf("spark.sql.planner.pythonExecution.memory")
.doc("Specifies the memory allocation for executing Python code in Spark
driver, in MiB. " +
@@ -9528,6 +9539,9 @@ class SQLConf extends Serializable with Logging with
SqlApiConf {
def legacyPandasConversionUDF: Boolean =
getConf(PYTHON_UDF_LEGACY_PANDAS_CONVERSION_ENABLED)
+ def legacyMapInBatchAcceptAnyIterable: Boolean =
+ getConf(PYTHON_UDF_MAP_IN_BATCH_LEGACY_ACCEPT_ANY_ITERABLE_ENABLED)
+
def pythonPlannerExecMemory: Option[Long] =
getConf(PYTHON_PLANNER_EXEC_MEMORY)
def replaceExceptWithFilter: Boolean = getConf(REPLACE_EXCEPT_WITH_FILTER)
diff --git
a/sql/core/src/main/scala/org/apache/spark/sql/execution/python/ArrowPythonRunner.scala
b/sql/core/src/main/scala/org/apache/spark/sql/execution/python/ArrowPythonRunner.scala
index 22ff6f711819..4d380a9cf431 100644
---
a/sql/core/src/main/scala/org/apache/spark/sql/execution/python/ArrowPythonRunner.scala
+++
b/sql/core/src/main/scala/org/apache/spark/sql/execution/python/ArrowPythonRunner.scala
@@ -184,6 +184,7 @@ object ArrowPythonRunner {
SQLConf.ARROW_EXECUTION_USE_LARGE_VAR_TYPES,
SQLConf.PYTHON_TABLE_UDF_LEGACY_PANDAS_CONVERSION_ENABLED,
SQLConf.PYTHON_UDF_LEGACY_PANDAS_CONVERSION_ENABLED,
+ SQLConf.PYTHON_UDF_MAP_IN_BATCH_LEGACY_ACCEPT_ANY_ITERABLE_ENABLED,
SQLConf.PYTHON_UDF_PANDAS_INT_TO_DECIMAL_COERCION_ENABLED,
SQLConf.PYTHON_UDF_PANDAS_PREFER_INT_EXTENSION_DTYPE,
SQLConf.PYSPARK_BINARY_AS_BYTES,
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
