From 127fb5bc26e0e9f64b0b67f19519bd9b7e3bd3e1 Mon Sep 17 00:00:00 2001
From: Thomas Jackson <thomas@conduit.xyz>
Date: Fri, 1 Nov 2024 11:34:16 -0700
Subject: [PATCH] Correctly override node in aggregateBinaryExpr

In #676 we added support for query push down for BinaryExp where it is a
literal and either a VectorSelector or a subset of AggregateExprs.
However this initial implementation introduced a bug for AggregateExprs.
Specifically it was doing the query push down but not doing the node
replace properly -- such that the literal operation was done twice (e.g.
add twice, multiply twice, etc.).

This patch will have the aggregateBinaryExpr replace the Binary Expr
(i.e. `(max(foo) * 100)` ) with the "filled in" AggregateExpr -- such
that the re-execution of the aggregation still occurs without re-doing
the literal calculation.

Fixes #688
---
 pkg/proxystorage/proxy.go      | 11 ++++------
 pkg/proxystorage/proxy_test.go | 38 ++++++++++++++++++++++++++++++----
 2 files changed, 38 insertions(+), 11 deletions(-)

diff --git a/pkg/proxystorage/proxy.go b/pkg/proxystorage/proxy.go
index 818032f9..b3655db9 100644
--- a/pkg/proxystorage/proxy.go
+++ b/pkg/proxystorage/proxy.go
@@ -795,7 +795,7 @@ func (p *ProxyStorage) NodeReplacer(ctx context.Context, s *parser.EvalStmt, nod
 		// aggregateBinaryExpr will send the node as a query to the downstream and
 		// replace the aggregate expr with the resulting data. This will cause the aggregation
 		// (min, max, topk, bottomk) to be re-run against the expression.
-		aggregateBinaryExpr := func(agg *parser.AggregateExpr) error {
+		aggregateBinaryExpr := func(agg *parser.AggregateExpr) (parser.Node, error) {
 			logrus.Debugf("BinaryExpr (AggregateExpr + Literal): %v", n)
 
 			removeOffsetFn()
@@ -816,7 +816,7 @@ func (p *ProxyStorage) NodeReplacer(ctx context.Context, s *parser.EvalStmt, nod
 				result, warnings, err = state.client.Query(ctx, n.String(), s.Start.Add(-offset))
 			}
 			if err != nil {
-				return err
+				return nil, err
 			}
 
 			iterators := promclient.IteratorsForValue(result)
@@ -833,7 +833,7 @@ func (p *ProxyStorage) NodeReplacer(ctx context.Context, s *parser.EvalStmt, nod
 			ret.UnexpandedSeriesSet = proxyquerier.NewSeriesSet(series, promhttputil.WarningsConvert(warnings), err)
 
 			agg.Expr = ret
-			return nil
+			return agg, nil
 		}
 
 		// Only valid if the other side is either `NumberLiteral` or `StringLiteral`
@@ -853,10 +853,7 @@ func (p *ProxyStorage) NodeReplacer(ctx context.Context, s *parser.EvalStmt, nod
 			case *parser.AggregateExpr:
 				switch otherTyped.Op {
 				case parser.MIN, parser.MAX, parser.TOPK, parser.BOTTOMK:
-					if err := aggregateBinaryExpr(otherTyped); err != nil {
-						return nil, err
-					}
-					return n, nil
+					return aggregateBinaryExpr(otherTyped)
 				}
 			}
 		}
diff --git a/pkg/proxystorage/proxy_test.go b/pkg/proxystorage/proxy_test.go
index 301f0fe0..862ab9db 100644
--- a/pkg/proxystorage/proxy_test.go
+++ b/pkg/proxystorage/proxy_test.go
@@ -192,32 +192,62 @@ func TestNodeReplacer(t *testing.T) {
 		// we expect to re-do the AggregateExpr but have replaced the VectorSelector with the query
 		{
 			in:  "min(foo{}) > 1",
-			out: "min() > 1",
+			out: "min()",
 			queries: []string{
 				"min(foo) > 1 @ 10000",
 			},
 		},
 		{
 			in:  "max(foo{}) > 1",
-			out: "max() > 1",
+			out: "max()",
 			queries: []string{
 				"max(foo) > 1 @ 10000",
 			},
 		},
 		{
 			in:  "topk(5, foo{}) > 1",
-			out: "topk(5, ) > 1",
+			out: "topk(5, )",
 			queries: []string{
 				"topk(5, foo) > 1 @ 10000",
 			},
 		},
 		{
 			in:  "bottomk(5, foo{}) > 1",
-			out: "bottomk(5, ) > 1",
+			out: "bottomk(5, )",
 			queries: []string{
 				"bottomk(5, foo) > 1 @ 10000",
 			},
 		},
+
+		{
+			in:  "min(foo{}) * 1000",
+			out: "min()",
+			queries: []string{
+				"min(foo) * 1000 @ 10000",
+			},
+		},
+		{
+			in:  "max(foo{}) * 1000",
+			out: "max()",
+			queries: []string{
+				"max(foo) * 1000 @ 10000",
+			},
+		},
+		{
+			in:  "topk(5, foo{}) * 1000",
+			out: "topk(5, )",
+			queries: []string{
+				"topk(5, foo) * 1000 @ 10000",
+			},
+		},
+		{
+			in:  "bottomk(5, foo{}) * 1000",
+			out: "bottomk(5, )",
+			queries: []string{
+				"bottomk(5, foo) * 1000 @ 10000",
+			},
+		},
+
 		// Check that some others do nothing
 		{
 			in: "avg(foo) > 1",