xinrong-meng commented on code in PR #44775:
URL: https://github.com/apache/spark/pull/44775#discussion_r1468269380
##########
python/pyspark/profiler.py:
##########
@@ -196,16 +197,41 @@ def add(
for subcode in filter(inspect.iscode, code.co_consts):
self.add(subcode, toplevel_code=toplevel_code)
+ class CodeMapForUDFV2(CodeMap):
+ def add(
+ self,
+ code: Any,
+ toplevel_code: Optional[Any] = None,
+ ) -> None:
+ if code in self:
+ return
+
+ if toplevel_code is None:
+ toplevel_code = code
+ filename = code.co_filename
+ self._toplevel.append((filename, code))
+ self[code] = {}
+ else:
+ self[code] = self[toplevel_code]
+ for subcode in filter(inspect.iscode, code.co_consts):
+ self.add(subcode, toplevel_code=toplevel_code)
+
+ def items(self) -> Iterator[Tuple[str, Iterator[Tuple[int, Any]]]]:
+ """Iterate on the toplevel code blocks."""
+ for filename, code in self._toplevel:
+ measures = self[code]
+ if not measures:
+ continue # skip if no measurement
+ linenos = range(min(measures), max(measures) + 1)
Review Comment:
```
============================================================
Profile of UDF<id=2>
============================================================
Filename:
/var/folders/h_/60n1p_5s7751jx1st4_sk0780000gp/T/ipykernel_69451/109011680.py
Line # Mem usage Increment Occurrences Line Contents
=============================================================
8 147.7 MiB 147.7 MiB 20 @udf("string")
9 def a(x):
10 149.6 MiB 1.8 MiB 20 if
TaskContext.get().partitionId() % 2 == 0:
11 59.9 MiB 0.1 MiB 8 return str(x)
12 else:
13 89.9 MiB 0.1 MiB 12 return None
```
tested on Jupyter.
--
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]
