UserDefinedFunction should validate that func is callable

apache · Feb 13, 2017 · 7d350bf · 7d350bf
1 parent ab88b24
commit 7d350bf
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 0 deletions.
diff --git a/python/pyspark/sql/functions.py b/python/pyspark/sql/functions.py
@@ -1864,6 +1864,11 @@ class UserDefinedFunction(object):
     .. versionadded:: 1.3
     """
     def __init__(self, func, returnType, name=None):
+        if not callable(func):
+            raise TypeError(
+                "Not a function or callable (__call__ is not defined): "
+                "{0}".format(type(func)))
+
         self.func = func
         self.returnType = (
             returnType if isinstance(returnType, DataType)

diff --git a/python/pyspark/sql/tests.py b/python/pyspark/sql/tests.py
@@ -504,6 +504,13 @@ def test_udf_with_string_return_type(self):
 
         self.assertTupleEqual(expected, actual)
 
+    def test_udf_shouldnt_accept_noncallable_object(self):
+        from pyspark.sql.functions import UserDefinedFunction
+        from pyspark.sql.types import StringType
+
+        non_callable = None
+        self.assertRaises(TypeError, UserDefinedFunction, non_callable, StringType())
+
     def test_basic_functions(self):
         rdd = self.sc.parallelize(['{"foo":"bar"}', '{"foo":"baz"}'])
         df = self.spark.read.json(rdd)