itholic commented on a change in pull request #34863:
URL: https://github.com/apache/spark/pull/34863#discussion_r776531548



##########
File path: python/pyspark/sql/tests/test_dataframe.py
##########
@@ -1138,6 +1138,39 @@ def test_create_nan_decimal_dataframe(self):
             [Row(value=None)],
         )
 
+    def test_transform_param(self):
+        def transform_fun(df, m, *, param):
+            return df.filter(col("c1") <= param).withColumn("c1", col("c1") * 
m)
+
+        def assert_equals(observed):
+            actual = observed.collect()
+            self.assertEqual(
+                [
+                    {
+                        "c1": 3,
+                    },
+                    {
+                        "c1": 6,
+                    },
+                ],
+                [row.asDict() for row in actual],
+            )
+
+        sdf = self.spark.createDataFrame([(1,), (2,), (3,)], ["c1"])
+        assert_equals(sdf.transform(transform_fun, 3, param=2))
+        assert_equals(sdf.transform(lambda df: transform_fun(df, 3, param=2)))
+        from functools import partial

Review comment:
       nit: I'm not strong feeling about this, but can we move this import to 
the top of file as PEP recommends ??




-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: [email protected]

For queries about this service, please contact Infrastructure at:
[email protected]



---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]

Reply via email to