fix: formula evaluation with non-participating queries (#5143)

pkg/query-service/postprocess/formula.go

@@ -24,10 +24,13 @@ func isSubset(super, sub map[string]string) bool {
 }
 
 // Function to find unique label sets
-func findUniqueLabelSets(results []*v3.Result) []map[string]string {
+func findUniqueLabelSets(results []*v3.Result, queriesInExpression map[string]struct{}) []map[string]string {
 	allLabelSets := make([]map[string]string, 0)
 	// The size of the `results` small, It is the number of queries in the request
 	for _, result := range results {
+		if _, ok := queriesInExpression[result.QueryName]; !ok {
+			continue
+		}
 		// The size of the `result.Series` slice is usually small, It is the number of series in the query result.
 		// We will limit the number of series in the query result to order of 100-1000.
 		for _, series := range result.Series {
@@ -120,7 +123,15 @@ func joinAndCalculate(
 			}
 		}
 
-		if len(expression.Vars()) != len(values) {
+		canEval := true
+
+		for _, v := range expression.Vars() {
+			if _, ok := values[v]; !ok {
+				canEval = false
+			}
+		}
+
+		if !canEval {
 			// not enough values for expression evaluation
 			continue
 		}
@@ -154,7 +165,12 @@ func processResults(
 	expression *govaluate.EvaluableExpression,
 	canDefaultZero map[string]bool,
 ) (*v3.Result, error) {
-	uniqueLabelSets := findUniqueLabelSets(results)
+
+	queriesInExpression := make(map[string]struct{})
+	for _, v := range expression.Vars() {
+		queriesInExpression[v] = struct{}{}
+	}
+	uniqueLabelSets := findUniqueLabelSets(results, queriesInExpression)
 	newSeries := make([]*v3.Series, 0)
 
 	for _, labelSet := range uniqueLabelSets {

pkg/query-service/postprocess/formula_test.go

@@ -11,9 +11,10 @@ import (
 
 func TestFindUniqueLabelSets(t *testing.T) {
 	tests := []struct {
-		name   string
-		result []*v3.Result
-		want   []map[string]string
+		name                string
+		result              []*v3.Result
+		want                []map[string]string
+		queriesInExpression map[string]struct{}
 	}{
 		{
 			name: "test1",
@@ -40,6 +41,10 @@ func TestFindUniqueLabelSets(t *testing.T) {
 					},
 				},
 			},
+			queriesInExpression: map[string]struct{}{
+				"A": {},
+				"B": {},
+			},
 			want: []map[string]string{
 				{
 					"service_name": "frontend",
@@ -96,6 +101,12 @@ func TestFindUniqueLabelSets(t *testing.T) {
 					},
 				},
 			},
+			queriesInExpression: map[string]struct{}{
+				"A": {},
+				"B": {},
+				"C": {},
+				"D": {},
+			},
 			want: []map[string]string{
 				{
 					"service_name": "frontend",
@@ -122,6 +133,10 @@ func TestFindUniqueLabelSets(t *testing.T) {
 					Series:    []*v3.Series{},
 				},
 			},
+			queriesInExpression: map[string]struct{}{
+				"A": {},
+				"B": {},
+			},
 			want: []map[string]string{},
 		},
 		{
@@ -160,6 +175,10 @@ func TestFindUniqueLabelSets(t *testing.T) {
 					},
 				},
 			},
+			queriesInExpression: map[string]struct{}{
+				"A": {},
+				"B": {},
+			},
 			want: []map[string]string{
 				{
 					"service_name": "frontend",
@@ -175,7 +194,7 @@ func TestFindUniqueLabelSets(t *testing.T) {
 
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
-			got := findUniqueLabelSets(tt.result)
+			got := findUniqueLabelSets(tt.result, tt.queriesInExpression)
 			if !reflect.DeepEqual(got, tt.want) {
 				t.Errorf("findUniqueLabelSets() = %v, want %v\n", got, tt.want)
 			}
@@ -1683,3 +1702,135 @@ func TestProcessResultsNoDefaultZero(t *testing.T) {
 		})
 	}
 }
+
+func TestProcessResultsMixedQueries(t *testing.T) {
+	tests := []struct {
+		name    string
+		results []*v3.Result
+		want    *v3.Result
+	}{
+		{
+			name: "test1",
+			results: []*v3.Result{
+				{
+					QueryName: "A",
+					Series: []*v3.Series{
+						{
+							Labels: map[string]string{
+								"service_name": "frontend",
+								"operation":    "GET /api",
+							},
+							Points: []v3.Point{
+								{
+									Timestamp: 1,
+									Value:     10,
+								},
+								{
+									Timestamp: 2,
+									Value:     20,
+								},
+							},
+						},
+					},
+				},
+				{
+					QueryName: "B",
+					Series: []*v3.Series{
+						{
+							Labels: map[string]string{
+								"service_name": "frontend",
+								"operation":    "GET /api",
+							},
+							Points: []v3.Point{
+								{
+									Timestamp: 1,
+									Value:     10,
+								},
+								{
+									Timestamp: 2,
+									Value:     20,
+								},
+							},
+						},
+					},
+				},
+				{
+					QueryName: "C",
+					Series: []*v3.Series{
+						{
+							Labels: map[string]string{
+								"service_name": "redis",
+							},
+							Points: []v3.Point{
+								{
+									Timestamp: 1,
+									Value:     30,
+								},
+								{
+									Timestamp: 2,
+									Value:     50,
+								},
+								{
+									Timestamp: 3,
+									Value:     45,
+								},
+							},
+						},
+					},
+				},
+			},
+			want: &v3.Result{
+				Series: []*v3.Series{
+					{
+						Labels: map[string]string{
+							"service_name": "frontend",
+							"operation":    "GET /api",
+						},
+						Points: []v3.Point{
+							{
+								Timestamp: 1,
+								Value:     1,
+							},
+							{
+								Timestamp: 2,
+								Value:     1,
+							},
+						},
+					},
+				},
+			},
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			expression, err := govaluate.NewEvaluableExpression("A / B")
+			if err != nil {
+				t.Errorf("Error parsing expression: %v", err)
+			}
+			canDefaultZero := map[string]bool{
+				"A": true,
+				"B": true,
+				"C": true,
+			}
+			got, err := processResults(tt.results, expression, canDefaultZero)
+			if err != nil {
+				t.Errorf("Error processing results: %v", err)
+			}
+			if len(got.Series) != len(tt.want.Series) {
+				t.Errorf("processResults(): number of sereis - got = %v, want %v", len(got.Series), len(tt.want.Series))
+			}
+
+			for i := range got.Series {
+				if len(got.Series[i].Points) != len(tt.want.Series[i].Points) {
+					t.Errorf("processResults(): number of points - got = %v, want %v", got, tt.want)
+				}
+				for j := range got.Series[i].Points {
+					if got.Series[i].Points[j].Value != tt.want.Series[i].Points[j].Value {
+						t.Errorf("processResults(): got = %v, want %v", got.Series[i].Points[j].Value, tt.want.Series[i].Points[j].Value)
+					}
+				}
+			}
+		})
+	}
+}