pkg/query: Fix panic when diffing two profiles

pkg/parcacol/querier.go

@@ -1491,7 +1491,7 @@ func (q *Querier) GetProfileMetadataMappings(
 
 			values := locations.ListValues().(*array.Dictionary)
 
-			compactedDict, err := compactDictionary.CompactDictionary(memory.DefaultAllocator, values)
+			compactedDict, err := compactDictionary.CompactDictionary(q.pool, values)
 			if err != nil {
 				fmt.Println("failed to compact dictionary", err)
 				return err
@@ -1523,7 +1523,6 @@ func (q *Querier) GetProfileMetadataMappings(
 
 func (q *Querier) GetProfileMetadataLabels(
 	ctx context.Context,
-	match []string,
 	start, end time.Time,
 ) ([]string, error) {
 	ctx, span := q.tracer.Start(ctx, "Querier/Labels")

pkg/query/columnquery.go

@@ -49,7 +49,7 @@ type Querier interface {
 	QuerySingle(ctx context.Context, query string, time time.Time, invertCallStacks bool) (profile.Profile, error)
 	QueryMerge(ctx context.Context, query string, start, end time.Time, aggregateByLabels, invertCallStacks bool) (profile.Profile, error)
 	GetProfileMetadataMappings(ctx context.Context, query string, start, end time.Time) ([]string, error)
-	GetProfileMetadataLabels(ctx context.Context, match []string, start, end time.Time) ([]string, error)
+	GetProfileMetadataLabels(ctx context.Context, start, end time.Time) ([]string, error)
 }
 
 var (
@@ -274,29 +274,46 @@ func (q *ColumnQueryAPI) Query(ctx context.Context, req *pb.QueryRequest) (*pb.Q
 	}()
 
 	if req.GetReportType() == pb.QueryRequest_REPORT_TYPE_PROFILE_METADATA {
-		mappingFiles, err := q.getMappingFiles(ctx, req.GetMerge())
+		var profileMetadata *pb.ProfileMetadata
 
-		labels, labels_err := q.getLabels(ctx, &pb.LabelsRequest{
-			Match: []string{},
-			Start: req.GetMerge().Start,
-			End:   req.GetMerge().End,
-		})
+		switch req.Mode {
+		case pb.QueryRequest_MODE_MERGE:
+			mappingFiles, labels, err := getMappingFilesAndLabels(ctx, q.querier, req.GetMerge().Query, req.GetMerge().Start.AsTime(), req.GetMerge().End.AsTime())
+			if err != nil {
+				return nil, err
+			}
 
-		if err != nil {
-			return nil, fmt.Errorf("getting mapping files: %w", err)
-		}
+			profileMetadata = &pb.ProfileMetadata{
+				MappingFiles: mappingFiles,
+				Labels:       labels,
+			}
+
+		case pb.QueryRequest_MODE_DIFF:
+			mappingFiles_a, labels_a, err := getMappingFilesAndLabels(ctx, q.querier, req.GetDiff().A.GetMerge().GetQuery(), req.GetDiff().A.GetMerge().Start.AsTime(), req.GetDiff().A.GetMerge().End.AsTime())
+			if err != nil {
+				return nil, err
+			}
 
-		if labels_err != nil {
-			return nil, fmt.Errorf("getting labels: %w", labels_err)
+			mappingFiles_b, labels_b, err := getMappingFilesAndLabels(ctx, q.querier, req.GetDiff().B.GetMerge().GetQuery(), req.GetDiff().B.GetMerge().Start.AsTime(), req.GetDiff().B.GetMerge().End.AsTime())
+			if err != nil {
+				return nil, err
+			}
+
+			mergedMappingFiles := MergeTwoSortedSlices(mappingFiles_a, mappingFiles_b)
+			mergedLabels := MergeTwoSortedSlices(labels_a, labels_b)
+
+			profileMetadata = &pb.ProfileMetadata{
+				MappingFiles: mergedMappingFiles,
+				Labels:       mergedLabels,
+			}
+		default:
+			return nil, status.Error(codes.InvalidArgument, "unknown query mode")
 		}
 
 		return &pb.QueryResponse{
 			Total:    0,
 			Filtered: 0,
-			Report: &pb.QueryResponse_ProfileMetadata{ProfileMetadata: &pb.ProfileMetadata{
-				MappingFiles: mappingFiles,
-				Labels:       labels,
-			}},
+			Report:   &pb.QueryResponse_ProfileMetadata{ProfileMetadata: profileMetadata},
 		}, nil
 	}
 
@@ -918,31 +935,58 @@ func sliceRecord(r arrow.Record, indices []int64) []arrow.Record {
 	return slices
 }
 
-func (q *ColumnQueryAPI) getMappingFiles(
+func getMappingFilesAndLabels(
 	ctx context.Context,
-	m *pb.MergeProfile,
-) ([]string, error) {
-	p, err := q.querier.GetProfileMetadataMappings(
-		ctx,
-		m.Query,
-		m.Start.AsTime(),
-		m.End.AsTime(),
-	)
+	q Querier,
+	query string,
+	startTime, endTime time.Time,
+) ([]string, []string, error) {
+	mappingFiles, err := q.GetProfileMetadataMappings(ctx, query, startTime, endTime)
 	if err != nil {
-		return nil, err
+		return nil, nil, err
 	}
 
-	return p, nil
+	labels, err := q.GetProfileMetadataLabels(ctx, startTime, endTime)
+	if err != nil {
+		return nil, nil, err
+	}
+
+	return mappingFiles, labels, nil
 }
 
-func (q *ColumnQueryAPI) getLabels(
-	ctx context.Context,
-	req *pb.LabelsRequest,
-) ([]string, error) {
-	l, err := q.querier.GetProfileMetadataLabels(ctx, req.Match, req.Start.AsTime(), req.End.AsTime())
-	if err != nil {
-		return nil, err
+// This is a deduplicating k-way merge.
+// The two slices that are passed in are assumed to be sorted.
+func MergeTwoSortedSlices(arr1, arr2 []string) []string {
+	merged := make([]string, 0, len(arr1)+len(arr2))
+	i, j := 0, 0
+
+	for i < len(arr1) && j < len(arr2) {
+		if arr1[i] < arr2[j] {
+			if len(merged) == 0 || merged[len(merged)-1] != arr1[i] {
+				merged = append(merged, arr1[i])
+			}
+			i++
+		} else {
+			if len(merged) == 0 || merged[len(merged)-1] != arr2[j] {
+				merged = append(merged, arr2[j])
+			}
+			j++
+		}
+	}
+
+	for i < len(arr1) {
+		if len(merged) == 0 || merged[len(merged)-1] != arr1[i] {
+			merged = append(merged, arr1[i])
+		}
+		i++
+	}
+
+	for j < len(arr2) {
+		if len(merged) == 0 || merged[len(merged)-1] != arr2[j] {
+			merged = append(merged, arr2[j])
+		}
+		j++
 	}
 
-	return l, nil
+	return merged
 }

pkg/query/columnquery_test.go

@@ -1713,3 +1713,12 @@ func BenchmarkFilterData(t *testing.B) {
 		}
 	}
 }
+
+func TestKwayMerge(t *testing.T) {
+	arr1 := []string{"a", "c", "e"}
+	arr2 := []string{"f", "i", "m", "o", "r"}
+
+	merged := MergeTwoSortedSlices(arr1, arr2)
+
+	require.Equal(t, []string{"a", "c", "e", "f", "i", "m", "o", "r"}, merged)
+}