Skip to content

ESQL: Improve local folding of aggregates #103670

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 3 commits into from
Dec 26, 2023
Merged

ESQL: Improve local folding of aggregates #103670

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
78763f8
ESQL: Improve local folding of aggregates
costin Dec 22, 2023
165a6b3
Update docs/changelog/103670.yaml
costin Dec 22, 2023
7e224c1
Use ternary operator
costin Dec 26, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 5 additions & 0 deletions docs/changelog/103670.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,5 @@
pr: 103670
summary: "ESQL: Improve local folding of aggregates"
area: ES|QL
type: bug
issues: []
55 changes: 54 additions & 1 deletion .../esql/src/main/java/org/elasticsearch/xpack/esql/optimizer/LocalLogicalPlanOptimizer.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,15 +7,24 @@

package org.elasticsearch.xpack.esql.optimizer;

import org.elasticsearch.compute.data.Block;
import org.elasticsearch.compute.data.BlockFactory;
import org.elasticsearch.compute.data.BlockUtils;
import org.elasticsearch.xpack.esql.expression.function.aggregate.Count;
import org.elasticsearch.xpack.esql.expression.function.scalar.nulls.Coalesce;
import org.elasticsearch.xpack.esql.optimizer.LogicalPlanOptimizer.PropagateEmptyRelation;
import org.elasticsearch.xpack.esql.plan.logical.Eval;
import org.elasticsearch.xpack.esql.plan.logical.TopN;
import org.elasticsearch.xpack.esql.planner.AbstractPhysicalOperationProviders;
import org.elasticsearch.xpack.esql.planner.PlannerUtils;
import org.elasticsearch.xpack.esql.stats.SearchStats;
import org.elasticsearch.xpack.ql.expression.Alias;
import org.elasticsearch.xpack.ql.expression.Attribute;
import org.elasticsearch.xpack.ql.expression.Expression;
import org.elasticsearch.xpack.ql.expression.FieldAttribute;
import org.elasticsearch.xpack.ql.expression.Literal;
import org.elasticsearch.xpack.ql.expression.NamedExpression;
import org.elasticsearch.xpack.ql.expression.function.aggregate.AggregateFunction;
import org.elasticsearch.xpack.ql.optimizer.OptimizerRules;
import org.elasticsearch.xpack.ql.plan.logical.Aggregate;
import org.elasticsearch.xpack.ql.plan.logical.EsRelation;
Expand All @@ -25,10 +34,15 @@
import org.elasticsearch.xpack.ql.plan.logical.Project;
import org.elasticsearch.xpack.ql.rule.ParameterizedRule;
import org.elasticsearch.xpack.ql.rule.ParameterizedRuleExecutor;
import org.elasticsearch.xpack.ql.type.DataType;
import org.elasticsearch.xpack.ql.type.DataTypes;

import java.util.ArrayList;
import java.util.List;

import static java.util.Arrays.asList;
import static org.elasticsearch.xpack.esql.optimizer.LogicalPlanOptimizer.cleanup;
import static org.elasticsearch.xpack.esql.optimizer.LogicalPlanOptimizer.operators;
import static org.elasticsearch.xpack.ql.optimizer.OptimizerRules.TransformDirection.UP;

public class LocalLogicalPlanOptimizer extends ParameterizedRuleExecutor<LogicalPlan, LocalLogicalOptimizerContext> {
Expand All @@ -50,10 +64,23 @@ protected List<Batch<LogicalPlan>> batches() {
var rules = new ArrayList<Batch<LogicalPlan>>();
rules.add(local);
// TODO: if the local rules haven't touched the tree, the rest of the rules can be skipped
rules.addAll(LogicalPlanOptimizer.rules());
rules.addAll(asList(operators(), cleanup()));
replaceRules(rules);
return rules;
}

private List<Batch<LogicalPlan>> replaceRules(List<Batch<LogicalPlan>> listOfRules) {
for (Batch<LogicalPlan> batch : listOfRules) {
var rules = batch.rules();
for (int i = 0; i < rules.length; i++) {
if (rules[i] instanceof PropagateEmptyRelation) {
rules[i] = new LocalPropagateEmptyRelation();
}
}
}
return listOfRules;
}

public LogicalPlan localOptimize(LogicalPlan plan) {
return execute(plan);
}
Expand Down Expand Up @@ -132,6 +159,32 @@ protected boolean skipExpression(Expression e) {
}
}

/**
* Local aggregation can only produce intermediate state that get wired into the global agg.
*/
private static class LocalPropagateEmptyRelation extends PropagateEmptyRelation {
Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The gist of this PR - override the global rule with a specific local rule that takes into account the intermediate aggregation and also takes care of handling the boolean types (seen).
This removes the weird logic downstream in the LocalExecutionPlanner as its much more contained.


/**
* Local variant of the aggregation that returns the intermediate value.
*/
@Override
protected void aggOutput(NamedExpression agg, AggregateFunction aggFunc, BlockFactory blockFactory, List<Block> blocks) {
List<Attribute> output = AbstractPhysicalOperationProviders.intermediateAttributes(List.of(agg), List.of());
for (Attribute o : output) {
DataType dataType = o.dataType();
// boolean right now is used for the internal #seen so always return true
var value = dataType == DataTypes.BOOLEAN ? true
// look for count(literal) with literal != null
: aggFunc instanceof Count count && (count.foldable() == false || count.fold() != null) ? 0L
// otherwise nullify
: null;
var wrapper = BlockUtils.wrapperFor(blockFactory, PlannerUtils.toElementType(dataType), 1);
wrapper.accept(value);
blocks.add(wrapper.builder().build());
}
}
}

abstract static class ParameterizedOptimizerRule<SubPlan extends LogicalPlan, P> extends ParameterizedRule<SubPlan, LogicalPlan, P> {

public final LogicalPlan apply(LogicalPlan plan, P context) {
Expand Down
71 changes: 34 additions & 37 deletions ...lugin/esql/src/main/java/org/elasticsearch/xpack/esql/optimizer/LogicalPlanOptimizer.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,7 @@
import org.apache.lucene.util.BytesRef;
import org.elasticsearch.common.util.Maps;
import org.elasticsearch.compute.data.Block;
import org.elasticsearch.compute.data.BlockFactory;
import org.elasticsearch.compute.data.BlockUtils;
import org.elasticsearch.xpack.esql.EsqlIllegalArgumentException;
import org.elasticsearch.xpack.esql.evaluator.predicate.operator.comparison.Equals;
Expand All @@ -24,7 +25,6 @@
import org.elasticsearch.xpack.esql.plan.logical.local.EsqlProject;
import org.elasticsearch.xpack.esql.plan.logical.local.LocalRelation;
import org.elasticsearch.xpack.esql.plan.logical.local.LocalSupplier;
import org.elasticsearch.xpack.esql.planner.AbstractPhysicalOperationProviders;
import org.elasticsearch.xpack.esql.planner.PlannerUtils;
import org.elasticsearch.xpack.esql.type.EsqlDataTypes;
import org.elasticsearch.xpack.ql.expression.Alias;
Expand Down Expand Up @@ -62,7 +62,6 @@
import org.elasticsearch.xpack.ql.rule.ParameterizedRuleExecutor;
import org.elasticsearch.xpack.ql.rule.Rule;
import org.elasticsearch.xpack.ql.tree.Source;
import org.elasticsearch.xpack.ql.type.DataType;
import org.elasticsearch.xpack.ql.type.DataTypes;
import org.elasticsearch.xpack.ql.util.CollectionUtils;
import org.elasticsearch.xpack.ql.util.Holder;
Expand Down Expand Up @@ -100,17 +99,8 @@ protected List<Batch<LogicalPlan>> batches() {
return rules();
}

protected static List<Batch<LogicalPlan>> rules() {
var substitutions = new Batch<>(
"Substitutions",
Limiter.ONCE,
new SubstituteSurrogates(),
new ReplaceRegexMatch(),
new ReplaceAliasingEvalWithProject()
// new NormalizeAggregate(), - waits on https://github.com/elastic/elasticsearch/issues/100634
);

var operators = new Batch<>(
protected static Batch<LogicalPlan> operators() {
return new Batch<>(
"Operator Optimization",
new CombineProjections(),
new CombineEvals(),
Expand Down Expand Up @@ -145,19 +135,33 @@ protected static List<Batch<LogicalPlan>> rules() {
new PruneOrderByBeforeStats(),
new PruneRedundantSortClauses()
);
}

var skip = new Batch<>("Skip Compute", new SkipQueryOnLimitZero());
var cleanup = new Batch<>(
protected static Batch<LogicalPlan> cleanup() {
return new Batch<>(
"Clean Up",
new ReplaceDuplicateAggWithEval(),
// pushing down limits again, because ReplaceDuplicateAggWithEval could create new Project nodes that can still be optimized
new PushDownAndCombineLimits(),
new ReplaceLimitAndSortAsTopN()
);
}

protected static List<Batch<LogicalPlan>> rules() {
var substitutions = new Batch<>(
"Substitutions",
Limiter.ONCE,
new SubstituteSurrogates(),
new ReplaceRegexMatch(),
new ReplaceAliasingEvalWithProject()
// new NormalizeAggregate(), - waits on https://github.com/elastic/elasticsearch/issues/100634
);

var skip = new Batch<>("Skip Compute", new SkipQueryOnLimitZero());
var defaultTopN = new Batch<>("Add default TopN", new AddDefaultTopN());
var label = new Batch<>("Set as Optimized", Limiter.ONCE, new SetAsOptimized());

return asList(substitutions, operators, skip, cleanup, defaultTopN, label);
return asList(substitutions, operators(), skip, cleanup(), defaultTopN, label);
}

// TODO: currently this rule only works for aggregate functions (AVG)
Expand Down Expand Up @@ -632,6 +636,7 @@ protected LogicalPlan rule(UnaryPlan plan) {
}
}

@SuppressWarnings("removal")
static class PropagateEmptyRelation extends OptimizerRules.OptimizerRule<UnaryPlan> {

@Override
Expand All @@ -649,39 +654,31 @@ protected LogicalPlan rule(UnaryPlan plan) {
return p;
}

private static List<Block> aggsFromEmpty(List<? extends NamedExpression> aggs) {
// TODO: Should we introduce skip operator that just never queries the source
private List<Block> aggsFromEmpty(List<? extends NamedExpression> aggs) {
List<Block> blocks = new ArrayList<>();
var blockFactory = PlannerUtils.NON_BREAKING_BLOCK_FACTORY;
int i = 0;
for (var agg : aggs) {
// there needs to be an alias
if (agg instanceof Alias a && a.child() instanceof AggregateFunction aggFunc) {
List<Attribute> output = AbstractPhysicalOperationProviders.intermediateAttributes(List.of(agg), List.of());
Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This section is now moved in the local rule.

for (Attribute o : output) {
DataType dataType = o.dataType();
// fill the boolean block later in LocalExecutionPlanner
if (dataType != DataTypes.BOOLEAN) {
// look for count(literal) with literal != null
var wrapper = BlockUtils.wrapperFor(
PlannerUtils.NON_BREAKING_BLOCK_FACTORY,
PlannerUtils.toElementType(dataType),
1
);
if (aggFunc instanceof Count count && (count.foldable() == false || count.fold() != null)) {
wrapper.accept(0L);
} else {
wrapper.accept(null);
}
blocks.add(wrapper.builder().build());
}
}
aggOutput(agg, aggFunc, blockFactory, blocks);
} else {
throw new EsqlIllegalArgumentException("Did not expect a non-aliased aggregation {}", agg);
}
}
return blocks;
}

/**
* The folded aggregation output - this variant is for the coordinator/final.
*/
protected void aggOutput(NamedExpression agg, AggregateFunction aggFunc, BlockFactory blockFactory, List<Block> blocks) {
// look for count(literal) with literal != null
Object value = aggFunc instanceof Count count && (count.foldable() == false || count.fold() != null) ? 0L : null;
var wrapper = BlockUtils.wrapperFor(blockFactory, PlannerUtils.toElementType(aggFunc.dataType()), 1);
wrapper.accept(value);
blocks.add(wrapper.builder().build());
}
}

private static LogicalPlan skipPlan(UnaryPlan plan) {
Expand Down
25 changes: 0 additions & 25 deletions ...plugin/esql/src/main/java/org/elasticsearch/xpack/esql/planner/LocalExecutionPlanner.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -56,7 +56,6 @@
import org.elasticsearch.xpack.esql.enrich.EnrichLookupService;
import org.elasticsearch.xpack.esql.evaluator.EvalMapper;
import org.elasticsearch.xpack.esql.evaluator.command.GrokEvaluatorExtracter;
import org.elasticsearch.xpack.esql.plan.logical.local.LocalSupplier;
import org.elasticsearch.xpack.esql.plan.physical.AggregateExec;
import org.elasticsearch.xpack.esql.plan.physical.DissectExec;
import org.elasticsearch.xpack.esql.plan.physical.EnrichExec;
Expand Down Expand Up @@ -88,7 +87,6 @@
import org.elasticsearch.xpack.ql.expression.NameId;
import org.elasticsearch.xpack.ql.expression.NamedExpression;
import org.elasticsearch.xpack.ql.expression.Order;
import org.elasticsearch.xpack.ql.type.DataTypes;
import org.elasticsearch.xpack.ql.util.Holder;

import java.util.ArrayList;
Expand Down Expand Up @@ -323,29 +321,6 @@ private PhysicalOperation planExchange(ExchangeExec exchangeExec, LocalExecution
private PhysicalOperation planExchangeSink(ExchangeSinkExec exchangeSink, LocalExecutionPlannerContext context) {
Objects.requireNonNull(exchangeSinkHandler, "ExchangeSinkHandler wasn't provided");
var child = exchangeSink.child();
// see https://github.com/elastic/elasticsearch/issues/100807 - handle case where the plan has been fully minimized
// to a local relation and the aggregate intermediate data erased. For this scenario, match the output the exchange output
// with that of the local relation

if (child instanceof LocalSourceExec localExec) {
var output = exchangeSink.output();
var localOutput = localExec.output();
if (output.equals(localOutput) == false) {
// the outputs are going to be similar except for the bool "seen" flags which are added in below
List<Block> blocks = new ArrayList<>(asList(localExec.supplier().get()));
if (blocks.size() > 0) {
for (int i = 0, s = output.size(); i < s; i++) {
var out = output.get(i);
if (out.dataType() == DataTypes.BOOLEAN) {
blocks.add(i, PlannerUtils.NON_BREAKING_BLOCK_FACTORY.newConstantBooleanBlockWith(true, 1));
}
}
}
var newSupplier = LocalSupplier.of(blocks.toArray(Block[]::new));

child = new LocalSourceExec(localExec.source(), output, newSupplier);
}
}

PhysicalOperation source = plan(child, context);

Expand Down
Loading