(spark) branch master updated: [SPARK-46252][PYTHON][TESTS] Improve test coverage of memory_profiler.py

[gurwls223]

This is an automated email from the ASF dual-hosted git repository.

gurwls223 pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/spark.git


The following commit(s) were added to refs/heads/master by this push:
     new 44036f2be3d [SPARK-46252][PYTHON][TESTS] Improve test coverage of 
memory_profiler.py
44036f2be3d is described below

commit 44036f2be3d5da7d8e20ee5082868dd70eb35884
Author: Xinrong Meng <xinr...@apache.org>
AuthorDate: Wed Dec 6 17:15:07 2023 +0900

    [SPARK-46252][PYTHON][TESTS] Improve test coverage of memory_profiler.py
    
    ### What changes were proposed in this pull request?
    Improve test coverage of memory_profiler.py
    
    ### Why are the changes needed?
    Subtasks of 
[SPARK-46041](https://issues.apache.org/jira/browse/SPARK-46041) to improve 
test coverage
    
    ### Does this PR introduce _any_ user-facing change?
    No.
    
    ### How was this patch tested?
    Test changes only.
    
    ### Was this patch authored or co-authored using generative AI tooling?
    No
    
    Closes #44167 from xinrong-meng/test_profiler.
    
    Authored-by: Xinrong Meng <xinr...@apache.org>
    Signed-off-by: Hyukjin Kwon <gurwls...@apache.org>
---
 python/pyspark/tests/test_memory_profiler.py | 26 ++++++++++++++++++++++++++
 1 file changed, 26 insertions(+)

diff --git a/python/pyspark/tests/test_memory_profiler.py 
b/python/pyspark/tests/test_memory_profiler.py
index 31ad62e11d9..c2c185b9863 100644
--- a/python/pyspark/tests/test_memory_profiler.py
+++ b/python/pyspark/tests/test_memory_profiler.py
@@ -45,6 +45,32 @@ class MemoryProfilerTests(PySparkTestCase):
         self.sc = SparkContext("local[4]", class_name, conf=conf)
         self.spark = SparkSession(sparkContext=self.sc)
 
+    def test_code_map(self):
+        from pyspark.profiler import CodeMapForUDF
+
+        code_map = CodeMapForUDF(include_children=False, backend="psutil")
+
+        def f(x):
+            return x + 1
+
+        code = f.__code__
+        code_map.add(code)
+        code_map.add(code)  # no-op, will return directly
+
+        self.assertIn(code, code_map)
+        self.assertEqual(len(code_map._toplevel), 1)
+
+    def test_udf_line_profiler(self):
+        from pyspark.profiler import UDFLineProfiler
+
+        profiler = UDFLineProfiler()
+
+        def f(x):
+            return x + 1
+
+        profiler.add_function(f)
+        self.assertTrue(profiler.code_map)
+
     def test_memory_profiler(self):
         self.exec_python_udf()
 


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscr...@spark.apache.org
For additional commands, e-mail: commits-h...@spark.apache.org