Skip to content

Commit

Permalink
ESQL: Refactor local logical optimizer rules (elastic#108015)
Browse files Browse the repository at this point in the history
Do not rely on any optimization rules from the ql project in LocalLogicalPlanOptimizer; for this:
* Copy InferIsNotNull to esql project
* Remove obsolete inheritance for InferIsNotNull, inheriting directly from Rule<P, Q>
  • Loading branch information
alex-spies authored Apr 30, 2024
1 parent d85fa15 commit bd75c0e
Showing 1 changed file with 86 additions and 2 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,7 @@
import org.elasticsearch.xpack.esql.type.EsqlDataTypes;
import org.elasticsearch.xpack.ql.expression.Alias;
import org.elasticsearch.xpack.ql.expression.Attribute;
import org.elasticsearch.xpack.ql.expression.AttributeMap;
import org.elasticsearch.xpack.ql.expression.Expression;
import org.elasticsearch.xpack.ql.expression.FieldAttribute;
import org.elasticsearch.xpack.ql.expression.Literal;
Expand All @@ -40,15 +41,19 @@
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.rule.Rule;
import org.elasticsearch.xpack.ql.type.DataType;
import org.elasticsearch.xpack.ql.type.DataTypes;
import org.elasticsearch.xpack.ql.util.CollectionUtils;

import java.util.ArrayList;
import java.util.LinkedHashSet;
import java.util.List;
import java.util.Map;
import java.util.Set;

import static java.util.Arrays.asList;
import static java.util.Collections.emptySet;
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;
Expand Down Expand Up @@ -171,10 +176,89 @@ else if (plan instanceof Project project) {
}
}

static class InferIsNotNull extends OptimizerRules.InferIsNotNull {
/**
* Simplify IsNotNull targets by resolving the underlying expression to its root fields with unknown
* nullability.
* e.g.
* (x + 1) / 2 IS NOT NULL --> x IS NOT NULL AND (x+1) / 2 IS NOT NULL
* SUBSTRING(x, 3) > 4 IS NOT NULL --> x IS NOT NULL AND SUBSTRING(x, 3) > 4 IS NOT NULL
* When dealing with multiple fields, a conjunction/disjunction based on the predicate:
* (x + y) / 4 IS NOT NULL --> x IS NOT NULL AND y IS NOT NULL AND (x + y) / 4 IS NOT NULL
* This handles the case of fields nested inside functions or expressions in order to avoid:
* - having to evaluate the whole expression
* - not pushing down the filter due to expression evaluation
* IS NULL cannot be simplified since it leads to a disjunction which prevents the filter to be
* pushed down:
* (x + 1) IS NULL --> x IS NULL OR x + 1 IS NULL
* and x IS NULL cannot be pushed down
* <br/>
* Implementation-wise this rule goes bottom-up, keeping an alias up to date to the current plan
* and then looks for replacing the target.
*/
static class InferIsNotNull extends Rule<LogicalPlan, LogicalPlan> {

@Override
protected boolean skipExpression(Expression e) {
public LogicalPlan apply(LogicalPlan plan) {
// the alias map is shared across the whole plan
AttributeMap<Expression> aliases = new AttributeMap<>();
// traverse bottom-up to pick up the aliases as we go
plan = plan.transformUp(p -> inspectPlan(p, aliases));
return plan;
}

private LogicalPlan inspectPlan(LogicalPlan plan, AttributeMap<Expression> aliases) {
// inspect just this plan properties
plan.forEachExpression(Alias.class, a -> aliases.put(a.toAttribute(), a.child()));
// now go about finding isNull/isNotNull
LogicalPlan newPlan = plan.transformExpressionsOnlyUp(IsNotNull.class, inn -> inferNotNullable(inn, aliases));
return newPlan;
}

private Expression inferNotNullable(IsNotNull inn, AttributeMap<Expression> aliases) {
Expression result = inn;
Set<Expression> refs = resolveExpressionAsRootAttributes(inn.field(), aliases);
// no refs found or could not detect - return the original function
if (refs.size() > 0) {
// add IsNull for the filters along with the initial inn
var innList = CollectionUtils.combine(refs.stream().map(r -> (Expression) new IsNotNull(inn.source(), r)).toList(), inn);
result = Predicates.combineAnd(innList);
}
return result;
}

/**
* Unroll the expression to its references to get to the root fields
* that really matter for filtering.
*/
protected Set<Expression> resolveExpressionAsRootAttributes(Expression exp, AttributeMap<Expression> aliases) {
Set<Expression> resolvedExpressions = new LinkedHashSet<>();
boolean changed = doResolve(exp, aliases, resolvedExpressions);
return changed ? resolvedExpressions : emptySet();
}

private boolean doResolve(Expression exp, AttributeMap<Expression> aliases, Set<Expression> resolvedExpressions) {
boolean changed = false;
// check if the expression can be skipped or is not nullabe
if (skipExpression(exp)) {
resolvedExpressions.add(exp);
} else {
for (Expression e : exp.references()) {
Expression resolved = aliases.resolve(e, e);
// found a root attribute, bail out
if (resolved instanceof Attribute a && resolved == e) {
resolvedExpressions.add(a);
// don't mark things as change if the original expression hasn't been broken down
changed |= resolved != exp;
} else {
// go further
changed |= doResolve(resolved, aliases, resolvedExpressions);
}
}
}
return changed;
}

private static boolean skipExpression(Expression e) {
return e instanceof Coalesce;
}
}
Expand Down

0 comments on commit bd75c0e

Please sign in to comment.