Github user holdenk commented on a diff in the pull request:
https://github.com/apache/spark/pull/20908#discussion_r178606547
--- Diff: python/pyspark/sql/tests.py ---
@@ -3966,6 +3967,23 @@ def random_udf(v):
random_udf = random_udf.asNondeterministic()
return random_udf
+ def test_pandas_udf_tokenize(self):
+ from pyspark.sql.functions import pandas_udf
+ tokenize = pandas_udf(lambda s: s.apply(lambda str: str.split('
')),
+ ArrayType(StringType()))
+ self.assertEqual(tokenize.returnType, ArrayType(StringType()))
+ df = self.spark.createDataFrame([("hi boo",), ("bye boo",)],
["vals"])
+ result = df.select(tokenize("vals").alias("hi"))
+ self.assertEqual([Row(hi=[u'hi', u'boo']), Row(hi=[u'bye',
u'boo'])], result.collect())
+
+ def test_pandas_udf_nested_arrays_does_not_work(self):
+ from pyspark.sql.functions import pandas_udf
+ tokenize = pandas_udf(lambda s: s.apply(lambda str: [str.split('
')]),
+ ArrayType(ArrayType(StringType())))
+ result = df.select(tokenize("vals").alias("hi"))
+ # If we start supporting nested arrays we should update the
documentation in functions.py
+ self.assertRaises(ArrowTypeError, result.collect())
--- End diff --
Sure, sounds good.
---
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscr...@spark.apache.org
For additional commands, e-mail: reviews-h...@spark.apache.org
