Fix root cumulative value in flamegraph trimming (#6028)

* Fix root cumulative value in flamegraph trimming

Updates the flamegraph trimming logic to correctly set the root row's cumulative value as the sum of its direct children's cumulative values.

* fix failing test

* use dynamic sizing for value type

* implement feedback on using largest type value

Author: yomete

pkg/query/flamegraph_arrow.go

@@ -357,7 +357,7 @@ func generateFlamegraphArrowRecord(ctx context.Context, mem memory.Allocator, tr
 		fb.trimmedLocationLine.AppendNull()
 		fb.trimmedFunctionStartLine = array.NewUint8Builder(fb.pool)
 		fb.trimmedFunctionStartLine.AppendNull()
-		fb.trimmedCumulative = array.NewUint8Builder(fb.pool)
+		fb.trimmedCumulative = array.NewUint64Builder(fb.pool)
 		fb.trimmedCumulative.AppendNull()
 		fb.trimmedFlat = array.NewUint8Builder(fb.pool)
 		fb.trimmedFlat.AppendNull()
@@ -1646,8 +1646,7 @@ func (fb *flamegraphBuilder) trim(ctx context.Context, tracer trace.Tracer, thre
 	trimmedFunctionNameIndices := array.NewInt32Builder(fb.pool)
 	trimmedFunctionSystemNameIndices := array.NewInt32Builder(fb.pool)
 	trimmedFunctionFilenameIndices := array.NewInt32Builder(fb.pool)
-	trimmedCumulativeType := smallestUnsignedTypeFor(largestCumulativeValue)
-	trimmedCumulative := array.NewBuilder(fb.pool, trimmedCumulativeType)
+	trimmedCumulative := builder.NewOptInt64Builder(arrow.PrimitiveTypes.Int64)
 	trimmedFlatType := smallestUnsignedTypeFor(largestFlatValue)
 	trimmedFlat := array.NewBuilder(fb.pool, trimmedFlatType)
 	trimmedDiffType := smallestSignedTypeFor(smallestDiffValue, largestDiffValue)
@@ -1660,7 +1659,8 @@ func (fb *flamegraphBuilder) trim(ctx context.Context, tracer trace.Tracer, thre
 	// Parent indices use int32 to allow for -1 (null parent)
 	trimmedParent := array.NewInt32Builder(fb.pool)
 
-	valueOffset := array.NewBuilder(fb.pool, trimmedCumulativeType)
+	valueOffsetType := smallestUnsignedTypeFor(largestCumulativeValue)
+	valueOffset := array.NewBuilder(fb.pool, valueOffsetType)
 
 	releasers = append(releasers,
 		trimmedMappingFileIndices,
@@ -1705,6 +1705,7 @@ func (fb *flamegraphBuilder) trim(ctx context.Context, tracer trace.Tracer, thre
 	trimmingQueue.elements = trimmingQueue.elements[:0]
 	trimmingQueue.push(trimmingElement{row: 0})
 
+	rootCumulativeTrimmedValue := int64(0)
 	// keep processing new elements until the queue is empty
 	for trimmingQueue.len() > 0 {
 		// pop the first item from the queue
@@ -1728,18 +1729,7 @@ func (fb *flamegraphBuilder) trim(ctx context.Context, tracer trace.Tracer, thre
 
 		// The following two will never be null.
 		cum := fb.builderCumulative.Value(te.row)
-		switch b := trimmedCumulative.(type) {
-		case *array.Uint64Builder:
-			b.Append(uint64(cum))
-		case *array.Uint32Builder:
-			b.Append(uint32(cum))
-		case *array.Uint16Builder:
-			b.Append(uint16(cum))
-		case *array.Uint8Builder:
-			b.Append(uint8(cum))
-		default:
-			panic(fmt.Errorf("unsupported type %T", b))
-		}
+		trimmedCumulative.Append(cum)
 
 		switch b := trimmedFlat.(type) {
 		case *array.Uint64Builder:
@@ -1812,6 +1802,10 @@ func (fb *flamegraphBuilder) trim(ctx context.Context, tracer trace.Tracer, thre
 			} else {
 				trimmedChildren[te.parent] = append(trimmedChildren[te.parent], row)
 			}
+			// Only accumulate direct children of root (parent row 0) for root's cumulative value
+			if te.parent == 0 {
+				rootCumulativeTrimmedValue += cum
+			}
 		}
 
 		cumThreshold := float32(cum) * threshold
@@ -1921,6 +1915,14 @@ func (fb *flamegraphBuilder) trim(ctx context.Context, tracer trace.Tracer, thre
 	}
 	fb.labels = trimmedLabels
 
+	trimmedCumulative.Set(0, rootCumulativeTrimmedValue)
+	tempTrimmedCumulative := array.NewBuilder(fb.pool, smallestUnsignedTypeFor(uint64(rootCumulativeTrimmedValue)))
+	for i := 0; i < trimmedCumulative.Len(); i++ {
+		copyInt64BuilderValueToUnknownUnsigned(trimmedCumulative, tempTrimmedCumulative, i)
+	}
+
+	trimmedCumulative.Release()
+
 	release(
 		fb.builderLabelsOnly,
 		fb.builderLabelsExist,
@@ -1939,7 +1941,7 @@ func (fb *flamegraphBuilder) trim(ctx context.Context, tracer trace.Tracer, thre
 	fb.builderInlined = trimmedInlined
 	fb.trimmedLocationLine = trimmedLocationLine
 	fb.trimmedFunctionStartLine = trimmedFunctionStartLine
-	fb.trimmedCumulative = trimmedCumulative
+	fb.trimmedCumulative = tempTrimmedCumulative
 	fb.trimmedFlat = trimmedFlat
 	fb.trimmedDiff = trimmedDiff
 	fb.trimmedChildren = trimmedChildren

pkg/query/flamegraph_arrow_test.go

@@ -931,6 +931,86 @@ func TestGenerateFlamegraphArrowTrimming(t *testing.T) {
 	fc.compare(expectedColumns)
 }
 
+func TestGenerateFlamegraphArrowTrimmingRootCumulative(t *testing.T) {
+	ctx := context.Background()
+	mem := memory.NewGoAllocator()
+	var err error
+
+	mappings := []*pprofprofile.Mapping{{
+		ID:   1,
+		File: "a",
+	}}
+
+	functions := []*pprofprofile.Function{{
+		ID:   1,
+		Name: "func1",
+	}, {
+		ID:   2,
+		Name: "func2",
+	}, {
+		ID:   3,
+		Name: "func3",
+	}, {
+		ID:   4,
+		Name: "func4",
+	}}
+
+	locations := []*pprofprofile.Location{{
+		ID:      1,
+		Mapping: mappings[0],
+		Line:    []pprofprofile.Line{{Function: functions[0]}},
+	}, {
+		ID:      2,
+		Mapping: mappings[0],
+		Line:    []pprofprofile.Line{{Function: functions[1]}},
+	}, {
+		ID:      3,
+		Mapping: mappings[0],
+		Line:    []pprofprofile.Line{{Function: functions[2]}},
+	}, {
+		ID:      4,
+		Mapping: mappings[0],
+		Line:    []pprofprofile.Line{{Function: functions[3]}},
+	}}
+
+	// Create a profile with:
+	// root -> func1 (100) -> func2 (80) -> func3 (60)
+	//                     -> func4 (10) [trimmed]
+	p, err := PprofToSymbolizedProfile(
+		profile.Meta{Duration: (10 * time.Second).Nanoseconds()},
+		&pprofprofile.Profile{
+			Sample: []*pprofprofile.Sample{{
+				Location: []*pprofprofile.Location{locations[2], locations[1], locations[0]},
+				Value:    []int64{60},
+			}, {
+				Location: []*pprofprofile.Location{locations[3], locations[1], locations[0]},
+				Value:    []int64{10},
+			}, {
+				Location: []*pprofprofile.Location{locations[1], locations[0]},
+				Value:    []int64{20},
+			}, {
+				Location: []*pprofprofile.Location{locations[0]},
+				Value:    []int64{10},
+			}},
+		},
+		0, []string{},
+	)
+	require.NoError(t, err)
+
+	tracer := noop.NewTracerProvider().Tracer("")
+
+	fa, cumulative, _, _, err := generateFlamegraphArrowRecord(ctx, mem, tracer, p, []string{FlamegraphFieldFunctionName}, float32(0.5))
+	require.NoError(t, err)
+	defer fa.Release()
+
+	cumulativeCol := fa.Column(fa.Schema().FieldIndices("cumulative")[0]).(*array.Uint8)
+	rootCumulative := int64(cumulativeCol.Value(0))
+
+	require.Equal(t, cumulative, rootCumulative, "Returned cumulative should match root row cumulative value")
+
+	require.Equal(t, int64(100), rootCumulative, "Root cumulative should be 100 (sum of direct child func1)")
+}
+
 func TestParents(t *testing.T) {
 	p := parent(-1)
 	require.Equal(t, -1, p.Get())