Fix distributed queries with explicit timestamp

The distributed optimizer does not consider queries with a specific
timestamp which causes them to be incorrect for some cases.

Signed-off-by: Filip Petkovski <[email protected]>

Author: fpetkovski

engine/distributed_test.go

@@ -245,6 +245,10 @@ func TestDistributedAggregations(t *testing.T) {
 		{name: "subquery over distributed binary expression", query: `max_over_time((bar / bar)[30s:15s])`},
 		{name: "timestamp", query: `timestamp(bar)`},
 		{name: "timestamp - step invariant", query: `timestamp(bar @ 6000.000)`},
+		{name: "query with @start() absolute timestamp", query: `sum(bar @ start())`},
+		{name: "query with @end() timestamp", query: `sum(bar @ end())`},
+		{name: "query with numeric timestamp", query: `sum(bar @ 140)`},
+		{name: "query with range and @end() timestamp", query: `sum(count_over_time(bar[1h] @ end()))`},
 	}
 
 	lookbackDeltas := []time.Duration{0, 30 * time.Second, 5 * time.Minute}

logicalplan/distribute.go

@@ -95,7 +95,7 @@ func (r RemoteExecution) String() string {
 	if r.QueryRangeStart.UnixMilli() == 0 {
 		return fmt.Sprintf("remote(%s)", r.Query)
 	}
-	return fmt.Sprintf("remote(%s) [%s]", r.Query, r.QueryRangeStart.UTC().String())
+	return fmt.Sprintf("remote(%s) [%s, %s]", r.Query, r.QueryRangeStart.UTC().String())
 }
 
 func (r RemoteExecution) Type() NodeType { return RemoteExecutionNode }
@@ -308,14 +308,16 @@ func newRemoteAggregation(rootAggregation *Aggregation, engines []api.RemoteEngi
 // For each engine which matches the time range of the query, it creates a RemoteExecution scoped to the range of the engine.
 // All remote executions are wrapped in a Deduplicate logical node to make sure that results from overlapping engines are deduplicated.
 func (m DistributedExecutionOptimizer) distributeQuery(expr *Node, engines []api.RemoteEngine, opts *query.Options, allowedStartOffset time.Duration) Node {
+	if ts := getQueryTimestamp(expr); ts != nil {
+		return *expr
+	}
 	startOffset := calculateStartOffset(expr, opts.LookbackDelta)
 	if allowedStartOffset < startOffset {
 		return *expr
 	}
 	if IsConstantExpr(*expr) {
 		return *expr
 	}
-
 	var globalMinT int64 = math.MaxInt64
 	for _, e := range engines {
 		if e.MinT() < globalMinT {
@@ -464,8 +466,10 @@ func calculateStartOffset(expr *Node, lookbackDelta time.Duration) time.Duration
 		return lookbackDelta
 	}
 
-	var selectRange time.Duration
-	var offset time.Duration
+	var (
+		selectRange time.Duration
+		offset      time.Duration
+	)
 	Traverse(expr, func(node *Node) {
 		switch n := (*node).(type) {
 		case *Subquery:
@@ -479,6 +483,20 @@ func calculateStartOffset(expr *Node, lookbackDelta time.Duration) time.Duration
 	return maxDuration(offset+selectRange, lookbackDelta)
 }
 
+func getQueryTimestamp(expr *Node) *int64 {
+	var timestamp *int64
+	Traverse(expr, func(node *Node) {
+		switch n := (*node).(type) {
+		case *VectorSelector:
+			if n.Timestamp != nil {
+				timestamp = n.Timestamp
+				return
+			}
+		}
+	})
+	return timestamp
+}
+
 func numSteps(start, end time.Time, step time.Duration) int64 {
 	return (end.UnixMilli()-start.UnixMilli())/step.Milliseconds() + 1
 }