Yicong-Huang commented on code in PR #59211:
URL: https://github.com/apache/spark/pull/59211#discussion_r4170774294
##########
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:
Fixed in 1bc1914: the test now builds the handler with
`...mapInBatch.legacy.acceptAnyIterable.enabled=false`, so a returned DataFrame
hits the strict non-iterator rejection, and asserts
`getMessageParameters()["actual"] == "DataFrame"` to distinguish it from the
per-element path. Also added `test_legacy_accepts_any_iterable` to pin the
legacy accept-any-iterable adaptation (both modes now covered).
--
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]
