fixed udf and its test

python/pyspark/sql/tests.py

@@ -909,7 +909,7 @@ def foo(x):
             raise StopIteration()
 
         with self.assertRaises(Py4JJavaError):
-            self.spark.range(0, 1000).withColumn('v', udf(foo)).show()
+            self.spark.range(0, 1000).withColumn('v', udf(foo)('id')).show()
 
     def test_validate_column_types(self):
         from pyspark.sql.functions import udf, to_json

python/pyspark/sql/udf.py

@@ -25,7 +25,7 @@
 from pyspark.sql.column import Column, _to_java_column, _to_seq
 from pyspark.sql.types import StringType, DataType, StructType, _parse_datatype_string,\
     to_arrow_type, to_arrow_schema
-from pyspark.util import _get_argspec
+from pyspark.util import _get_argspec, fail_on_stopiteration
 
 __all__ = ["UDFRegistration"]
 
@@ -92,7 +92,7 @@ def __init__(self, func,
             raise TypeError(
                 "Invalid evalType: evalType should be an int but is {}".format(evalType))
 
-        self.func = func
+        self.func = fail_on_stopiteration(func)
         self._returnType = returnType
         # Stores UserDefinedPythonFunctions jobj, once initialized
         self._returnType_placeholder = None

python/pyspark/util.py

@@ -91,8 +91,8 @@ def majorMinorVersion(sparkVersion):
 
 def fail_on_stopiteration(f):
     """
-        Wraps the input function to fail on 'StopIteration' by raising a 'RuntimeError'
-        prevents silent loss of data when 'f' is used in a for loop
+    Wraps the input function to fail on 'StopIteration' by raising a 'RuntimeError'
+    prevents silent loss of data when 'f' is used in a for loop
     """
     def wrapper(*args, **kwargs):
         try: