Update differential_flamegraph script to make supporting negation easy once it's possible

Signed-off-by: Dom Del Nano <ddelnano@gmail.com>

Author: ddelnano

src/pxl_scripts/px/differential_flamegraph/differential.pxl

@@ -16,35 +16,62 @@
 
 import px
 
-def export_flame_graph(start_time: str, namespace: str, pod: str, baseline_pod: str):
+negate = False
+# TODO(ddelnano): negation requires switching the type of join from right to left
+# or returning a different DataFrame. This might not be possible with pxl's current
+# functionality, but this should be implemented once it's possible.
+
+def merge_and_compute_delta(pod1, pod2, negate: bool):
+
+  diff = pod1.merge(
+    pod2,
+    how='right',
+    left_on='stack_trace',
+    right_on='stack_trace'
+    suffixes=['_1', '_2'],
+  )
+  # TODO(ddelnano): This needs to be switched with pod1 if the flamegraph should
+  # be negated.
+  percentage_agg = pod2.groupby(['pod']).agg(
+    count=('count', px.sum),
+  )
+  diff.pod = px.select(negate, diff.pod_1, diff.pod_2)
+  diff.stack_trace = px.select(negate, diff.stack_trace_1, diff.stack_trace_2)
+  diff.stack_trace = px.replace(' ', diff.stack_trace, '')
+  diff.count = px.select(negate, diff.count_1, diff.count_2)
+  diff.delta = diff.count_2 - diff.count_1
+  diff.delta = px.select(negate, px.negate(diff.delta), diff.delta)
+
+  merged = diff.merge(
+    percentage_agg,
+    how='inner',
+    left_on='pod',
+    right_on='pod',
+    suffixes=['', '_x']
+  )
+  merged.percent = 100 * merged.count / merged.count_x
+  return merged
+
+def differential_flamegraph(start_time: str, namespace: str, pod: str, baseline_pod: str):
   stack_traces = px.DataFrame(table='stack_traces.beta', start_time=start_time)
   stack_traces.namespace = stack_traces.ctx['namespace']
-  stack_traces = stack_traces[px.contains(stack_traces.namespace, namespace)]
+  stack_traces = stack_traces[stack_traces.namespace == namespace]
   stack_traces.node = px.Node(px._exec_hostname())
   stack_traces.pod = stack_traces.ctx['pod']
-  stack_traces.keep_row = px.select(px.contains(stack_traces.pod, baseline_pod), True, False)
-  stack_traces.keep_row = px.select(stack_traces.keep_row or px.contains(stack_traces.pod, pod), True, False)
+  stack_traces.keep_row = stack_traces.pod == baseline_pod
+  stack_traces.keep_row = px.select(stack_traces.keep_row or stack_traces.pod == pod, True, False)
   stack_traces = stack_traces[stack_traces.keep_row]
 
   stack_traces = stack_traces.groupby(['node', 'namespace', 'pod', 'stack_trace_id']).agg(
       stack_trace=('stack_trace', px.any),
       count=('count', px.sum)
   )
 
-  pod1 = stack_traces[px.contains(stack_traces.pod, baseline_pod)]
-  pod1 = pod1.drop(['node', 'namespace', 'pod', 'stack_trace_id'])
+  pod1 = stack_traces[stack_traces.pod == baseline_pod]
+  pod1 = pod1.drop(['node', 'namespace', 'stack_trace_id'])
 
-  pod2 = stack_traces[px.contains(stack_traces.pod, pod)]
-  pod2 = pod2.drop(['node', 'namespace', 'pod', 'stack_trace_id'])
+  pod2 = stack_traces[stack_traces.pod == pod]
+  pod2 = pod2.drop(['node', 'namespace', 'stack_trace_id'])
 
-  merged = pod1.merge(
-    pod2,
-    how='right',
-    left_on='stack_trace',
-    right_on='stack_trace'
-    suffixes=['_1', '_2'],
-  )
-  # Pixie's stack trace format permits spaces, but Brendan Gregg's scripts do not
-  merged.stack_trace_2 = px.replace(' ', merged.stack_trace_2, '')
-  merged.delta = merged.count_2 - merged.count_1
-  return merged[['stack_trace_2', 'count_1', 'count_2', 'delta']]
+  merged = merge_and_compute_delta(pod1, pod2, negate)
+  return merged[['stack_trace', 'count', 'delta', 'percent']]

src/pxl_scripts/px/differential_flamegraph/vis.json

@@ -33,7 +33,7 @@
                 "h": 6
             },
             "func": {
-                "name": "export_flame_graph",
+                "name": "differential_flamegraph",
                 "args": [
                     {
                         "name": "start_time",
@@ -55,8 +55,9 @@
             },
             "displaySpec": {
                 "@type": "types.px.dev/px.vispb.StackTraceFlameGraph",
-                "stacktraceColumn": "stack_trace_2",
-                "countColumn": "count_2",
+                "stacktraceColumn": "stack_trace",
+                "countColumn": "count",
+                "percentageColumn": "percent",
                 "differenceColumn": "delta"
             }
         }