wip

Author: gatorsmile

python/pyspark/sql/catalog.py

@@ -255,11 +255,25 @@ def registerFunction(self, name, f, returnType=StringType()):
         >>> _ = spark.udf.register("stringLengthInt", len, IntegerType())
         >>> spark.sql("SELECT stringLengthInt('test')").collect()
         [Row(stringLengthInt(test)=4)]
+
+        >>> import random
+        >>> from pyspark.sql.functions import udf
+        >>> from pyspark.sql.types import IntegerType, StringType
+        >>> random_udf = udf(lambda: int(random.random() * 100), IntegerType()).asNondeterministic()
+        >>> spark.catalog.registerFunction("random_udf", random_udf, StringType())
+        >>> spark.sql("SELECT random_udf()").collect()  # doctest: +SKIP
+        [Row(random_udf()=u'82')]
         """
-        udf = UserDefinedFunction(f, returnType=returnType, name=name,
-                                  evalType=PythonEvalType.SQL_BATCHED_UDF)
-        self._jsparkSession.udf().registerPython(name, udf._judf)
-        return udf._wrapped()
+
+        if hasattr(f, 'asNondeterministic'):
+            udf = f._set_name(name, returnType)
+            self._jsparkSession.udf().registerPython(name, udf._judf)
+            return udf._wrapped()
+        else:
+            udf = UserDefinedFunction(f, returnType=returnType, name=name,
+                                      evalType=PythonEvalType.SQL_BATCHED_UDF)
+            self._jsparkSession.udf().registerPython(name, udf._judf)
+            return udf._wrapped()
 
     @since(2.0)
     def isCached(self, tableName):

python/pyspark/sql/udf.py

@@ -138,6 +138,19 @@ def __call__(self, *cols):
         sc = SparkContext._active_spark_context
         return Column(judf.apply(_to_seq(sc, cols, _to_java_column)))
 
+    def _set_name(self, name, returnType=StringType()):
+        """
+        Updates the name of UserDefinedFunction.
+        """
+        # reset _judf
+        self._judf_placeholder = None
+        self._returnType_placeholder = None
+        self._name = name or (
+            func.__name__ if hasattr(func, '__name__')
+            else func.__class__.__name__)
+        self._returnType = returnType
+        return self
+
     def _wrapped(self):
         """
         Wrap this udf with a function and attach docstring from func
@@ -163,6 +176,10 @@ def wrapper(*args):
         wrapper.returnType = self.returnType
         wrapper.evalType = self.evalType
         wrapper.asNondeterministic = self.asNondeterministic
+        wrapper._judf = self._judf
+        wrapper._create_judf = self._create_judf
+        wrapper._wrapped = self._wrapped
+        wrapper._set_name = self._set_name
 
         return wrapper