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())
    +
    +
    +@unittest.skipIf(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: reviews-unsubscr...@spark.apache.org
For additional commands, e-mail: reviews-h...@spark.apache.org

Reply via email to