[GitHub] [spark] lu-wang-dl commented on a diff in pull request #41946: [SPARK-44264] FunctionPickler Class

["lu-wang-dl (via GitHub)" <gi...@apache.org>]

lu-wang-dl commented on code in PR #41946:
URL: https://github.com/apache/spark/pull/41946#discussion_r1262766128


##########
python/pyspark/ml/tests/test_util.py:
##########
@@ -73,6 +80,136 @@ def _check_uid_set_equal(stages, expected_stages):
         )
 
 
+class TestFunctionPickler(unittest.TestCase):
+
+
+    # Function that will be used to test pickling.
+    @staticmethod
+    def _test_function(x: float, y: float) -> float:
+        return x**2 + y**2
+
+    def _check_if_test_function_pickled(self, f, og_fn, output_value, *arguments, **key_word_args):
+        fn, args, kwargs = cloudpickle.load(f)
+        self.assertEqual(fn, og_fn)
+        self.assertEqual(args, arguments)
+        self.assertEqual(kwargs, key_word_args)
+        fn_output = fn(*args, **kwargs)
+        self.assertEqual(fn_output, output_value)
+
+    def test_pickle_fn_and_save(self):
+        x, y = 1, 3 # args of test_function
+        tmp_dir = "silly_goose"
+        os.mkdir(tmp_dir)
+        file_path_to_save = "silly_bear"
+
+        with self.subTest(msg="See if it pickles correctly if no file_path or save_dir are specified"):

Review Comment:
   Use `parameterized ` to avoid duplicates.



-- 
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: reviews-unsubscribe@spark.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org