fix vertical sharding bug in without and union operations

Signed-off-by: Ben Ye <benye@amazon.com>
thanos-io · Aug 7, 2023 · 30cf9fe · 30cf9fe
1 parent ffc42df
commit 30cf9fe
Show file tree

Hide file tree

Showing 2 changed files with 36 additions and 3 deletions.
diff --git a/pkg/querysharding/analysis.go b/pkg/querysharding/analysis.go
@@ -101,10 +101,12 @@ func without(sliceA, sliceB []string) []string {
 	if sliceA == nil {
 		return nil
 	}
-
-	if len(sliceA) == 0 || len(sliceB) == 0 {
+	if len(sliceA) == 0 {
 		return []string{}
 	}
+	if len(sliceB) == 0 {
+		return sliceA
+	}
 
 	keyMap := make(map[string]struct{}, len(sliceA))
 	for _, s := range sliceA {
@@ -123,9 +125,15 @@ func without(sliceA, sliceB []string) []string {
 }
 
 func union(sliceA, sliceB []string) []string {
-	if len(sliceA) == 0 || len(sliceB) == 0 {
+	if len(sliceA) == 0 && len(sliceB) == 0 {
 		return []string{}
 	}
+	if len(sliceA) == 0 {
+		return sliceB
+	}
+	if len(sliceB) == 0 {
+		return sliceA
+	}
 
 	keyMap := make(map[string]struct{}, len(sliceA))
 	for _, s := range sliceA {

diff --git a/pkg/querysharding/analyzer_test.go b/pkg/querysharding/analyzer_test.go
@@ -107,6 +107,11 @@ http_requests_total`,
 			expression:     `sum by (cluster, pod) (http_requests_total{code="400"}) / on (pod) sum by (cluster, pod) (http_requests_total)`,
 			shardingLabels: []string{"pod"},
 		},
+		{
+			name:           "binary expression with vector matching with outer aggregation",
+			expression:     `sum(http_requests_total{code="400"} * http_requests_total) by (pod)`,
+			shardingLabels: []string{"pod"},
+		},
 		{
 			name: "multiple binary expressions with vector matchers",
 			expression: `
@@ -180,6 +185,16 @@ sum by (container) (
 			expression:     `sum(sum_over_time(container_memory_working_set_bytes{container_name!="POD",container_name!="",namespace="kube-system"}[1d:5m])) by (instance, cluster) / avg(label_replace(sum(sum_over_time(kube_node_status_capacity_memory_bytes[1d:5m])) by (node, cluster), "instance", "$1", "node", "(.*)")) by (instance, cluster)`,
 			shardingLabels: []string{"cluster"},
 		},
+		{
+			name:           "complex query with label_replace and nested aggregations",
+			expression:     `avg(label_replace(label_replace(avg(count_over_time(kube_pod_container_resource_requests{resource="memory", unit="byte", container!="",container!="POD", node!="", }[1h] )*avg_over_time(kube_pod_container_resource_requests{resource="memory", unit="byte", container!="",container!="POD", node!="", }[1h] )) by (namespace,container,pod,node,cluster_id) , "container_name","$1","container","(.+)"), "pod_name","$1","pod","(.+)")) by (namespace,container_name,pod_name,node,cluster_id)`,
+			shardingLabels: []string{"namespace", "node", "cluster_id"},
+		},
+		{
+			name:           "complex query with label_replace, nested aggregations and binary expressions",
+			expression:     `sort_desc(avg(label_replace(label_replace(label_replace(count_over_time(container_memory_working_set_bytes{container!="", container!="POD", instance!="", }[1h] ), "node", "$1", "instance", "(.+)"), "container_name", "$1", "container", "(.+)"), "pod_name", "$1", "pod", "(.+)")*label_replace(label_replace(label_replace(avg_over_time(container_memory_working_set_bytes{container!="", container!="POD", instance!="", }[1h] ), "node", "$1", "instance", "(.+)"), "container_name", "$1", "container", "(.+)"), "pod_name", "$1", "pod", "(.+)")) by (namespace, container_name, pod_name, node, cluster_id))`,
+			shardingLabels: []string{"namespace", "cluster_id"},
+		},
 	}
 
 	shardableWithoutLabels := []testCase{
@@ -193,6 +208,16 @@ sum by (container) (
 			expression:     "max without (pod) (sum without (pod, cluster) (http_requests_total))",
 			shardingLabels: []string{"pod", "cluster"},
 		},
+		{
+			name:           "binary expression with outer without grouping",
+			expression:     `sum(http_requests_total{code="400"} * http_requests_total) without (pod)`,
+			shardingLabels: []string{"pod"},
+		},
+		{
+			name:           "binary expression with vector matching and outer without grouping",
+			expression:     `sum(http_requests_total{code="400"} * ignoring(cluster) http_requests_total) without ()`,
+			shardingLabels: []string{"__name__", "cluster"},
+		},
 		{
 			name:           "binary expression with without vector matching and grouping",
 			expression:     `sum without (cluster, pod) (http_requests_total{code="400"}) / ignoring (pod) sum without (cluster, pod) (http_requests_total)`,