profiles parquet encoding: fix profile column count (#3420)

Author: korniltsev

pkg/phlaredb/schemas/v1/profiles.go

@@ -503,8 +503,9 @@ func deconstructMemoryProfile(imp InMemoryProfile, row parquet.Row) parquet.Row
 			col++
 			return col
 		}
-		totalCols = 8 + (7 * len(imp.Samples.StacktraceIDs)) + len(imp.Comments)
+		totalCols = profileColumnCount(imp)
 	)
+
 	if cap(row) < totalCols {
 		row = make(parquet.Row, 0, totalCols)
 	}
@@ -614,6 +615,17 @@ func deconstructMemoryProfile(imp InMemoryProfile, row parquet.Row) parquet.Row
 	return row
 }
 
+func profileColumnCount(imp InMemoryProfile) int {
+	var totalCols = 10 + (7 * len(imp.Samples.StacktraceIDs)) + len(imp.Comments)
+	if len(imp.Comments) == 0 {
+		totalCols++
+	}
+	if len(imp.Samples.StacktraceIDs) == 0 {
+		totalCols += 7
+	}
+	return totalCols
+}
+
 func NewMergeProfilesRowReader(rowGroups []parquet.RowReader) parquet.RowReader {
 	if len(rowGroups) == 0 {
 		return phlareparquet.EmptyRowReader

pkg/phlaredb/schemas/v1/profiles_test.go

@@ -542,3 +542,41 @@ func Test_SamplesRange(t *testing.T) {
 		})
 	}
 }
+
+func TestColumnCount(t *testing.T) {
+	profiles := []InMemoryProfile{{
+		SeriesIndex: 1,
+		TimeNanos:   2,
+		Samples: Samples{
+			StacktraceIDs: []uint32{1, 2, 3},
+			Values:        []uint64{1, 2, 3},
+		},
+	},
+		{
+			SeriesIndex: 1,
+			TimeNanos:   2,
+			Samples: Samples{
+				StacktraceIDs: []uint32{1, 2, 3},
+				Values:        []uint64{1, 2, 3},
+				Spans:         []uint64{1, 2, 3},
+			},
+		},
+		{
+			SeriesIndex: 1,
+			TimeNanos:   2,
+			Samples: Samples{
+				StacktraceIDs: []uint32{1, 2, 3},
+				Values:        []uint64{1, 2, 3},
+				Spans:         []uint64{1, 2, 3},
+			},
+			Comments: []int64{1, 2, 3},
+		}}
+	for _, profile := range profiles {
+		count := profileColumnCount(profile)
+
+		row := deconstructMemoryProfile(profile, nil)
+		assert.Equal(t, len(row), count)
+		assert.Equal(t, cap(row), count)
+	}
+
+}