Github user BryanCutler commented on a diff in the pull request:
https://github.com/apache/spark/pull/18732#discussion_r142520328
--- Diff: python/pyspark/sql/tests.py ---
@@ -3376,6 +3377,132 @@ def test_vectorized_udf_empty_partition(self):
res = df.select(f(col('id')))
self.assertEquals(df.collect(), res.collect())
+ def test_vectorized_udf_varargs(self):
+ from pyspark.sql.functions import pandas_udf, col
+ df = self.spark.createDataFrame(self.sc.parallelize([Row(id=1)],
2))
+ f = pandas_udf(lambda *v: v[0], LongType())
+ res = df.select(f(col('id')))
+ self.assertEquals(df.collect(), res.collect())
+
+
[email protected](not _have_pandas or not _have_arrow, "Pandas or Arrow not
installed")
+class GroupbyApplyTests(ReusedPySparkTestCase):
+ @classmethod
+ def setUpClass(cls):
+ ReusedPySparkTestCase.setUpClass()
+ cls.spark = SparkSession(cls.sc)
+
+ @classmethod
+ def tearDownClass(cls):
+ ReusedPySparkTestCase.tearDownClass()
+ cls.spark.stop()
+
+ def assertFramesEqual(self, expected, result):
+ msg = ("DataFrames are not equal: " +
+ ("\n\nExpected:\n%s\n%s" % (expected, expected.dtypes)) +
+ ("\n\nResult:\n%s\n%s" % (result, result.dtypes)))
+ self.assertTrue(expected.equals(result), msg=msg)
+
+ @property
+ def data(self):
+ from pyspark.sql.functions import pandas_udf, array, explode, col,
lit
+ return self.spark.range(10).toDF('id') \
+ .withColumn("vs", array([lit(i) for i in range(20, 30)])) \
+ .withColumn("v", explode(col('vs'))).drop('vs')
+
+ def test_simple(self):
+ from pyspark.sql.functions import pandas_udf
+ df = self.data
+
+ def foo(df):
+ ret = df
+ ret = ret.assign(v1=df.v * df.id * 1.0)
+ ret = ret.assign(v2=df.v + df.id)
+ return ret
+
+ foo_udf = pandas_udf(
+ foo,
+ StructType(
+ [StructField('id', LongType()),
+ StructField('v', IntegerType()),
+ StructField('v1', DoubleType()),
+ StructField('v2', LongType())]))
+
+ result = df.groupby('id').apply(foo_udf).sort('id').toPandas()
+ expected =
df.toPandas().groupby('id').apply(foo).reset_index(drop=True)
+ self.assertFramesEqual(expected, result)
+
+ def test_decorator(self):
+ from pyspark.sql.functions import pandas_udf
+ df = self.data
+
+ @pandas_udf(StructType(
+ [StructField('id', LongType()),
+ StructField('v', IntegerType()),
+ StructField('v1', DoubleType()),
+ StructField('v2', LongType())]))
+ def foo(df):
+ ret = df
+ ret = ret.assign(v1=df.v * df.id * 1.0)
+ ret = ret.assign(v2=df.v + df.id)
+ return ret
+
+ result = df.groupby('id').apply(foo).sort('id').toPandas()
+ expected =
df.toPandas().groupby('id').apply(foo.func).reset_index(drop=True)
+ self.assertFramesEqual(expected, result)
+
+ def test_coerce(self):
+ from pyspark.sql.functions import pandas_udf
+ df = self.data
+
+ def foo(df):
+ ret = df
+ ret = ret.assign(v=df.v + 1)
+ return ret
+
+ @pandas_udf(StructType([StructField('id', LongType()),
StructField('v', DoubleType())]))
+ def foo(df):
--- End diff --
isn't this a redeclaration of `foo`? maybe some of these could be lamdas?
---
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
