Github user BryanCutler commented on a diff in the pull request:
https://github.com/apache/spark/pull/18732#discussion_r142519186
--- Diff: python/pyspark/sql/tests.py ---
@@ -3376,6 +3377,133 @@ 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
--- End diff --
could you simplify to
```
return df.assign(v1=df.v * df.id * 1.0, v2=df.v + df.id)
```
---
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
