HyukjinKwon commented on code in PR #59211:
URL: https://github.com/apache/spark/pull/59211#discussion_r4169274191
##########
python/pyspark/eval_handlers/tests/test_pandas_eval_type_handlers.py:
##########
@@ -156,6 +186,85 @@ def
test_forwards_use_large_var_types_to_output_conversion(self):
self.assertTrue(pa.types.is_large_string(out[0].schema.field(0).type))
[email protected](not (have_pandas and have_pyarrow), _missing_message)
+class PandasScalarIterUDFHandlerTests(unittest.TestCase):
+ def test_invokes_udf_over_batch_stream(self):
+ def add_one(series_iter):
+ for s in series_iter:
+ yield s + 1
+
+ handler = _scalar_iter_handler(add_one)
+ out = list(handler.run(0, iter([_batch(a=[1, 2]), _batch(a=[3])])))
+ self.assertEqual([b.column(0).to_pylist() for b in out], [[2, 3], [4]])
+
+ def test_coerces_output_to_return_type(self):
+ # The UDF yields int32, but the declared return type is LongType
(int64).
+ def add_one(series_iter):
+ for s in series_iter:
+ yield (s + 1).astype("int32")
+
+ handler = _scalar_iter_handler(add_one)
+ out = list(handler.run(0, iter([_batch(a=[10, 20])])))
+ self.assertEqual(out[0].schema.field(0).type, pa.int64())
+ self.assertEqual(out[0].column(0).to_pylist(), [11, 21])
+
+ def test_rejects_row_count_mismatch(self):
+ # Emitting more rows than were consumed must fail (fail-fast row
limit).
+ def too_many(series_iter):
+ for s in series_iter:
+ yield pd.Series(list(range(len(s) + 1)))
+
+ handler = _scalar_iter_handler(too_many)
+ with self.assertRaises(PySparkRuntimeError):
+ list(handler.run(0, iter([_batch(a=[1, 2])])))
+
+ def test_struct_return_takes_dataframe(self):
+ # struct_in_pandas="dict" + df_for_struct=True: a struct return is a
DataFrame.
+ struct_type = StructType([StructField("x", LongType()),
StructField("y", LongType())])
+
+ def to_struct(series_iter):
+ for s in series_iter:
+ yield pd.DataFrame({"x": s, "y": s * 10})
+
+ handler = _scalar_iter_handler(to_struct, return_type=struct_type)
+ out = list(handler.run(0, iter([_batch(a=[1, 2])])))
+ self.assertEqual(
+ out[0].column(0).to_pylist(),
+ [{"x": 1, "y": 10}, {"x": 2, "y": 20}],
+ )
+
+
[email protected](not (have_pandas and have_pyarrow), _missing_message)
+class PandasMapUDFHandlerTests(unittest.TestCase):
+ _return_type = StructType([StructField("v", LongType())])
+
+ def test_maps_batch_stream(self):
+ # mapInPandas expands the wire struct into a DataFrame and re-wraps
the output.
+ def double_v(df_iter):
+ for df in df_iter:
+ yield pd.DataFrame({"v": df["v"] * 2})
+
+ handler = _map_handler(double_v, self._return_type)
+ out = list(handler.run(0, iter([_struct_batch(v=[1, 2, 3])])))
+ self.assertEqual(out[0].column(0).field("v").to_pylist(), [2, 4, 6])
+
+ def test_rejects_non_iterator_result(self):
+ # A UDF returning a DataFrame (not an iterator of them) is rejected.
+ handler = _map_handler(lambda df_iter: pd.DataFrame({"v": [1]}),
self._return_type)
Review Comment:
**Nit (P3):** Under the default RunnerConf({}) built here, the legacy
accept-any-iterable flag is on (it defaults to "true" in
RunnerConf.map_in_batch_legacy_accept_any_iterable). So the handler calls
iter() on the returned DataFrame, which iterates its column labels, and the
PySparkTypeError comes from the per-element check ('expected iterator of
pandas.DataFrame, actual iterator of str'). The non-Iterator branch this test
is named for never runs. That makes it a duplicate of
test_rejects_wrong_element_type, and it would keep passing if strict
non-iterator rejection regressed.
Could you pass a RunnerConf with
spark.sql.execution.pythonUDF.mapInBatch.legacy.acceptAnyIterable.enabled set
to "false" here? Also assert the error parameters (for example, check_error
with messageParameters actual='DataFrame', or assertRaisesRegex on 'but is
DataFrame') so the test distinguishes the two branches. If you want both modes
covered at this level, a second case with the flag left at its default can pin
the current legacy behavior.
**Verification:**
- **Regression:** A map-handler unit test with the legacy flag disabled
returns a pandas DataFrame from the UDF and asserts UDF_RETURN_TYPE with actual
type DataFrame. The test would fail if strict non-iterator rejection were
bypassed.
- **Behavior:** The full pandas eval-type handler unit suite passes,
including the existing wrong-element-type test, with no production code change.
--
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]
