ueshin commented on code in PR #45378:
URL: https://github.com/apache/spark/pull/45378#discussion_r1515240180


##########
python/pyspark/sql/profiler.py:
##########
@@ -236,18 +236,22 @@ def clear_perf_profiles(self, id: Optional[int] = None) 
-> None:
             The UDF ID whose profiling results should be cleared.
             If not specified, all the results will be cleared.
         """
-        ids_to_remove = [
-            result_id
-            for result_id, (perf, _, *_) in self._profile_results.items()
-            if perf is not None
-        ]
         with self._lock:
             if id is not None:
-                if id in ids_to_remove:
-                    self._profile_results.pop(id, None)
+                if id in self._profile_results:
+                    perf, mem, *rest = self._profile_results[id]
+                    self._profile_results[id] = (None, mem, *rest)
+                    if mem is None:
+                        self._profile_results.pop(id, None)
             else:
-                for id_to_remove in ids_to_remove:
-                    self._profile_results.pop(id_to_remove, None)
+                ids_to_remove = []
+                for id, (perf, mem, *rest) in 
list(self._profile_results.items()):
+                    self._profile_results[id] = (None, mem, *rest)
+                    if mem is None:
+                        ids_to_remove.append(id)

Review Comment:
   nit: Can't we pop it here?



##########
python/pyspark/sql/profiler.py:
##########
@@ -262,15 +266,21 @@ def clear_memory_profiles(self, id: Optional[int] = None) 
-> None:
             If not specified, all the results will be cleared.
         """
         with self._lock:
-            ids_to_remove = [
-                id for id, (_, mem, *_) in self._profile_results.items() if 
mem is not None
-            ]
             if id is not None:
-                if id in ids_to_remove:
-                    self._profile_results.pop(id, None)
+                if id in self._profile_results:
+                    perf, mem, *rest = self._profile_results[id]
+                    self._profile_results[id] = (perf, None, *rest)
+                    if perf is None:
+                        self._profile_results.pop(id, None)
             else:
-                for id_to_remove in ids_to_remove:
-                    self._profile_results.pop(id_to_remove, None)
+                ids_to_remove = []
+                for id, (perf, mem, *rest) in 
list(self._profile_results.items()):
+                    self._profile_results[id] = (perf, None, *rest)
+                    if perf is None:
+                        ids_to_remove.append(id)

Review Comment:
   ditto.



-- 
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-unsubscr...@spark.apache.org

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


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

Reply via email to