From 223380d8cf1e349f9ed88718a5a5c5018e693f30 Mon Sep 17 00:00:00 2001 From: Yicong-Huang <17627829+Yicong-Huang@users.noreply.github.com> Date: Wed, 5 Aug 2026 22:17:43 +0000 Subject: [PATCH 1/3] refactor: require strict Iterator return from mapInPandas --- .../pyspark/sql/tests/pandas/test_pandas_map.py | 17 +++++++++++------ python/pyspark/worker.py | 5 +---- 2 files changed, 12 insertions(+), 10 deletions(-) diff --git a/python/pyspark/sql/tests/pandas/test_pandas_map.py b/python/pyspark/sql/tests/pandas/test_pandas_map.py index bfcedc6c8899f..7f313b78b94b3 100644 --- a/python/pyspark/sql/tests/pandas/test_pandas_map.py +++ b/python/pyspark/sql/tests/pandas/test_pandas_map.py @@ -86,12 +86,6 @@ def test_map_in_pandas(self): 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_multiple_columns(self): data = [(1, "foo"), (2, None), (3, "bar"), (4, "bar")] df = self.spark.createDataFrame(data, "a int, b string") @@ -186,6 +180,10 @@ def no_iter(_): 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 +198,13 @@ def bad_iter_elem(_): ): (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 9ac8abbfa4856..a1f9d5f9ccbdf 100644 --- a/python/pyspark/worker.py +++ b/python/pyspark/worker.py @@ -2698,11 +2698,8 @@ def dataframe_iter(): 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__"): + if not isinstance(result, Iterator): raise PySparkTypeError( errorClass="UDF_RETURN_TYPE", messageParameters={ From 17785f805c553e9c93c87b776869e181d5c363eb Mon Sep 17 00:00:00 2001 From: Yicong-Huang <17627829+Yicong-Huang@users.noreply.github.com> Date: Thu, 6 Aug 2026 18:35:22 +0000 Subject: [PATCH 2/3] feat: add legacy flag to accept any iterable from mapInPandas and mapInArrow --- .../pyspark/sql/tests/arrow/test_arrow_map.py | 13 +++++++++ .../sql/tests/pandas/test_pandas_map.py | 10 +++++++ python/pyspark/worker.py | 27 ++++++++++++++++++- .../apache/spark/sql/internal/SQLConf.scala | 13 +++++++++ .../execution/python/ArrowPythonRunner.scala | 1 + 5 files changed, 63 insertions(+), 1 deletion(-) diff --git a/python/pyspark/sql/tests/arrow/test_arrow_map.py b/python/pyspark/sql/tests/arrow/test_arrow_map.py index 5119e0e827f6d..311deca9007e3 100644 --- a/python/pyspark/sql/tests/arrow/test_arrow_map.py +++ b/python/pyspark/sql/tests/arrow/test_arrow_map.py @@ -55,6 +55,19 @@ def func(iterator): 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_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 7f313b78b94b3..011b126e67744 100644 --- a/python/pyspark/sql/tests/pandas/test_pandas_map.py +++ b/python/pyspark/sql/tests/pandas/test_pandas_map.py @@ -86,6 +86,16 @@ def test_map_in_pandas(self): 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_multiple_columns(self): data = [(1, "foo"), (2, None), (3, "bar"), (4, "bar")] df = self.spark.createDataFrame(data, "a int, b string") diff --git a/python/pyspark/worker.py b/python/pyspark/worker.py index a1f9d5f9ccbdf..6c93f0e00324e 100644 --- a/python/pyspark/worker.py +++ b/python/pyspark/worker.py @@ -147,6 +147,16 @@ def use_legacy_pandas_udtf_conversion(self) -> bool: == "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" @@ -1923,6 +1933,16 @@ def func(split_index: int, data: Iterator[pa.RecordBatch]) -> Iterator[pa.Record # invoke the UDF output_batches = udf_func(input_batches) + # The declared signature is Iterator[...], so a strict iterator is required. The + # legacy flag restores the pre-4.3.0 behavior of accepting any iterable (e.g. list) + # 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) + and hasattr(output_batches, "__iter__") + ): + output_batches = iter(output_batches) + # Post-processing verified_iter = verify_return_type( output_batches, @@ -2699,7 +2719,12 @@ def dataframe_iter(): )[0] result = map_udf(dataframe_iter()) - if not isinstance(result, Iterator): + # The declared signature is Iterator[...], so a strict iterator is required. The + # legacy flag restores the pre-4.3.0 behavior of accepting any iterable (e.g. list). + is_iterator = isinstance(result, Iterator) + if runner_conf.map_in_batch_legacy_accept_any_iterable: + is_iterator = is_iterator or hasattr(result, "__iter__") + if not is_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 8a0d88346b63c..a172c9603f8a7 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 @@ -5390,6 +5390,16 @@ 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") + .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. " + @@ -9273,6 +9283,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 75b8465e2607a..a1588959a0611 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 @@ -163,6 +163,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, From a0368d75c6611b57031d476d93b8eec36aafcda3 Mon Sep 17 00:00:00 2001 From: Yicong-Huang <17627829+Yicong-Huang@users.noreply.github.com> Date: Thu, 6 Aug 2026 18:38:35 +0000 Subject: [PATCH 3/3] docs: note mapInPandas strict-iterator change in PySpark 4.3 migration guide --- python/docs/source/migration_guide/pyspark_upgrade.rst | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/python/docs/source/migration_guide/pyspark_upgrade.rst b/python/docs/source/migration_guide/pyspark_upgrade.rst index b257ee49b7ddc..35bfe5f23c961 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.