Github user ueshin commented on a diff in the pull request:
https://github.com/apache/spark/pull/20171#discussion_r160318424
--- Diff: python/pyspark/sql/tests.py ---
@@ -3616,6 +3616,34 @@ def test_vectorized_udf_basic(self):
bool_f(col('bool')))
self.assertEquals(df.collect(), res.collect())
+ def test_register_vectorized_udf_basic(self):
+ from pyspark.sql.functions import pandas_udf
+ from pyspark.rdd import PythonEvalType
+ twoArgsPandasUDF = pandas_udf(lambda x: len(x), IntegerType())
+ self.assertEqual(twoArgsPandasUDF.deterministic, True)
+ self.assertEqual(twoArgsPandasUDF.evalType,
PythonEvalType.SQL_PANDAS_SCALAR_UDF)
+ newPandasUDF = self.spark.catalog.registerFunction(
+ "twoArgsPandasUDF", twoArgsPandasUDF, IntegerType())
+ self.assertEqual(newPandasUDF.deterministic, True)
+ self.assertEqual(newPandasUDF.evalType,
PythonEvalType.SQL_PANDAS_SCALAR_UDF)
+ [row] = self.spark.sql("SELECT twoArgsPandasUDF('test')").collect()
+ self.assertEqual(row[0], 4)
+
+ def test_register_nondeterministic_vectorized_udf_basic(self):
+ from pyspark.sql.functions import pandas_udf
+ from pyspark.rdd import PythonEvalType
+ import random
+ randomPandasUDF = pandas_udf(
+ lambda x: random.randint(6, 6) + x,
StringType()).asNondeterministic()
--- End diff --
How about the following strategy?
1. make the default value for `returnType` `None`.
2. if `returnType is None` for a Python function, use `StringType` as the
same as the current default value.
3. if `returnType is None` for UDF, use the UDF's `returnType`, otherwise
respect the user specified `returnType`.
---
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
