pkg/query/flamegraph_arrow: Properly calculate height

Previously the height was estimated, but didn't take inlined functions and injected label nodes into account.

go.mod

@@ -1,6 +1,6 @@
 module github.com/parca-dev/parca
 
-go 1.20
+go 1.21
 
 require (
 	cloud.google.com/go/storage v1.31.0

pkg/query/flamegraph_arrow.go

@@ -100,9 +100,6 @@ func generateFlamegraphArrowRecord(ctx context.Context, mem memory.Allocator, tr
 	fb := newFlamegraphBuilder(mem, totalRows)
 	defer fb.Release()
 
-	// This keeps track of the max depth of our flame graph.
-	maxHeight := int32(0)
-
 	// these change with every iteration below
 	row := fb.builderCumulative.Len()
 
@@ -114,12 +111,6 @@ func generateFlamegraphArrowRecord(ctx context.Context, mem memory.Allocator, tr
 		for i := 0; i < int(r.Record.NumRows()); i++ {
 			beg, end := r.Locations.ValueOffsets(i)
 
-			// TODO: This height is only an estimation, inlined functions are not taken into account.
-			numLocations := int32(end - beg)
-			if numLocations > maxHeight {
-				maxHeight = numLocations
-			}
-
 			var sampleLabels map[string]string
 			for j, labelColumn := range r.LabelColumns {
 				if labelColumn.Col.IsValid(i) {
@@ -152,6 +143,9 @@ func generateFlamegraphArrowRecord(ctx context.Context, mem memory.Allocator, tr
 					}
 					fb.rootsRow[lsstring] = []int{sampleLabelRow}
 				}
+				fb.maxHeight = max(fb.maxHeight, fb.height)
+				fb.height = 1
+
 				fb.parent.Set(sampleLabelRow)
 				row = fb.builderCumulative.Len()
 			}
@@ -181,6 +175,8 @@ func generateFlamegraphArrowRecord(ctx context.Context, mem memory.Allocator, tr
 						fb.compareRows = append(fb.compareRows, fb.rootsRow[unsafeString(lsbytes)]...)
 						// append this row afterward to not compare to itself
 						fb.parent.Reset()
+						fb.maxHeight = max(fb.maxHeight, fb.height)
+						fb.height = 0
 					}
 
 					merged, err := fb.mergeUnsymbolizedRows(
@@ -193,6 +189,7 @@ func generateFlamegraphArrowRecord(ctx context.Context, mem memory.Allocator, tr
 						return nil, 0, 0, 0, err
 					}
 					if merged {
+						fb.height++
 						continue locations
 					}
 
@@ -221,13 +218,16 @@ func generateFlamegraphArrowRecord(ctx context.Context, mem memory.Allocator, tr
 						fb.compareRows = append(fb.compareRows, fb.rootsRow[unsafeString(lsbytes)]...)
 						// append this row afterward to not compare to itself
 						fb.parent.Reset()
+						fb.maxHeight = max(fb.maxHeight, fb.height)
+						fb.height = 0
 					}
 
 					merged, err := fb.mergeSymbolizedRows(r, aggregateFields, sampleLabels, i, j, k, int(end))
 					if err != nil {
 						return nil, 0, 0, 0, err
 					}
 					if merged {
+						fb.height++
 						continue stacktraces
 					}
 
@@ -248,13 +248,15 @@ func generateFlamegraphArrowRecord(ctx context.Context, mem memory.Allocator, tr
 			}
 		}
 	}
+	// the last row can also have the most height.
+	fb.maxHeight = max(fb.maxHeight, fb.height)
 
 	record, err := fb.NewRecord()
 	if err != nil {
 		return nil, 0, 0, 0, err
 	}
 
-	return record, fb.cumulative, maxHeight + 1, 0, nil
+	return record, fb.cumulative, fb.maxHeight + 1, 0, nil
 }
 
 // mergeSymbolizedRows compares the symbolized fields by function name and labels and merges them if they equal.
@@ -390,6 +392,8 @@ type flamegraphBuilder struct {
 	cumulative int64
 	// This keeps track of the total diff values so that we can set the irst row's diff value at the end.
 	diff int64
+	// This keeps track of the max height of the flame graph.
+	maxHeight int32
 	// parent keeps track of the parent of a row. This is used to build the children array.
 	parent parent
 	// This keeps track of a row's children and will be converted to an arrow array of lists at the end.
@@ -405,6 +409,8 @@ type flamegraphBuilder struct {
 	rootsRow map[string][]int
 	// compareRows are the rows that we compare to the current location against and potentially merge.
 	compareRows []int
+	// height keeps track of the current stack trace's height of the flame graph.
+	height int32
 
 	builderMappingStart       *array.Uint64Builder
 	builderMappingLimit       *array.Uint64Builder
@@ -569,6 +575,7 @@ func (fb *flamegraphBuilder) appendRow(
 	sampleRow, locationRow, lineRow int,
 	row int,
 ) error {
+	fb.height++
 	for j := range fb.rb.Fields() {
 		switch fb.schema.Field(j).Name {
 		// Mapping
@@ -711,8 +718,6 @@ func (fb *flamegraphBuilder) AppendLabelRow(r profile.RecordReader, row int, lab
 	}
 	// Add this label row to the root row's children.
 	fb.children[0] = append(fb.children[0], row)
-	//// Add the next row as child of this label row.
-	//fb.children[row] = append(fb.children[row], row+1)
 
 	fb.builderMappingStart.AppendNull()
 	fb.builderMappingLimit.AppendNull()

pkg/query/flamegraph_arrow_test.go

@@ -326,7 +326,7 @@ func TestGenerateFlamegraphArrow(t *testing.T) {
 		aggregate: []string{FlamegraphFieldLabels},
 		// expectations
 		cumulative: 10,
-		height:     5, // TODO 6
+		height:     6,
 		trimmed:    0, // TODO
 		rows: []flamegraphRow{
 			{MappingStart: 0, MappingLimit: 0, MappingOffset: 0, MappingFile: "a", MappingBuildID: "aID", LocationAddress: 0, LocationFolded: false, LocationLine: 0, FunctionStartLine: 0, FunctionName: `{"goroutine":"1"}`, FunctionSystemName: "1", FunctionFilename: "1", Cumulative: 10, Labels: nil, Children: []uint32{1, 6, 10}}, // 0
@@ -482,7 +482,7 @@ func TestGenerateFlamegraphArrowWithInlined(t *testing.T) {
 	require.NoError(t, err)
 
 	require.Equal(t, int64(1), total)
-	require.Equal(t, int32(4), height)
+	require.Equal(t, int32(5), height)
 	require.Equal(t, int64(0), trimmed)
 
 	require.Equal(t, int64(16), record.NumCols())