Sharding optimizations I: AST mapping

[owen-d]

What

This PR starts implementing shard-based query optimizations. In essence, this takes advantage of the v9 schema onwards which splits storage across n shards. For example,

sum(rate({foo="bar"} |="id=123" [1m])) maps into

sum(
  downstream<sum(rate(({foo="bar"}|="id=123")[1m])), shard=0_of_2> ++ 
  downstream<sum(rate(({foo="bar"}|="id=123")[1m])), shard=1_of_2>
)

where downstream<x> indicates that it should be executed independently on a downstream querier and x ++ y indicates that the resulting vectors should be concatenated.

This PR hinges on the new interface

// ASTMapper is the exported interface for mapping between multiple AST representations
type ASTMapper interface {
	Map(Expr) (Expr, error)
}

which maps one AST into a functionally equivalent AST. We use this to map incoming queries into versions which are more parallelizable.

Remaining Work

I'm splitting my work into a few different PRs to lighten the cognitive load while reviewing. These will be

1. AST mapping PR (current pr). This won’t expose options or anything, but will just add astmapping & related testing for sharding.
2. some internal pkg refactoring to make way for (3). This also won’t change anything but internal organization.
3. code which wires together & downstreams sharded queries + options to enable sharding. This one will activate the code paths.

Inspiration

cortexproject/cortex#1878

Misc

There is a bit of Evaluator refactoring in this PR which isn't technically part of the AST mapping, but it was burdensome to revert the changes here and they aren't too obtrusive.

[cyriltovena]

I love those changes.

[cyriltovena]

I wonder if evaluator is big enough to become a package. Just a thought, not action required.

[owen-d]

yeah, it's a good idea, let's see where everything settles and then do some refactoring

[codecov-io]

Codecov Report

Merging #1846 into master will increase coverage by 0.06%.
The diff coverage is 76.37%.

@@            Coverage Diff             @@
##           master    #1846      +/-   ##
==========================================
+ Coverage   64.86%   64.93%   +0.06%     
==========================================
  Files         122      125       +3     
  Lines        9239     9391     +152     
==========================================
+ Hits         5993     6098     +105     
- Misses       2833     2867      +34     
- Partials      413      426      +13     

Impacted Files	Coverage Δ
pkg/logql/astmapper.go	0.00% <0.00%> (ø)
pkg/logql/ast.go	88.40% <55.55%> (-0.54%)	⬇️
pkg/logql/shardmapper.go	77.41% <77.41%> (ø)
pkg/logql/evaluator.go	90.69% <83.33%> (-1.12%)	⬇️
pkg/logql/engine.go	90.06% <100.00%> (ø)
pkg/logql/sharding.go	100.00% <100.00%> (ø)
pkg/promtail/targets/tailer.go	73.86% <0.00%> (-4.55%)	⬇️
pkg/promtail/targets/filetarget.go	68.71% <0.00%> (-1.85%)	⬇️
... and 3 more

[cyriltovena]

I'd argue marchersExpr requires shard mapping. Those will just increase the amount of data we process with their overhead, this is because the result is limited and it will always be super fast to answer unless there is a lot streams in the query.

However here it seems that the function is recursive so in case of metrics rate({app="foo"}[5m]) we probably want sharding mapping.

It's also difficult to imagine because I have yet to discover where will the parallelization take place. So I'll need more to judge on that.

[owen-d]

Good point. I'll think about how to more effectively short circuit this rather than relying on the frontend.

[cyriltovena]

LGTM