Skip to content

Commit

Permalink
Skip CASE function from InferIsNotNull rule checks
Browse files Browse the repository at this point in the history
  • Loading branch information
@astefan
astefan committed Sep 18, 2024
1 parent 0d79a69 commit af3bfe6
Show file tree
Hide file tree
Showing 3 changed files with 85 additions and 37 deletions.
19 changes: 19 additions & 0 deletions x-pack/plugin/esql/qa/testFixtures/src/main/resources/conditional.csv-spec
View file
Original file line number Diff line number Diff line change
Expand Up @@ -261,3 +261,22 @@ warning:Line 1:38: java.lang.IllegalArgumentException: single-value function enc

a:integer | b:integer | same:boolean
;

// https://github.com/elastic/elasticsearch/issues/112704
caseOnTheValue_NotOnTheField
FROM employees
| WHERE emp_no < 10022 AND emp_no > 10012
| KEEP languages, emp_no
| EVAL eval = CASE(languages == 1, null, languages == 2, "bilingual", languages > 2, "multilingual", languages IS NULL, "languages is null")
| SORT languages, emp_no
| WHERE eval IS NOT NULL;

languages:integer| emp_no:integer|eval:keyword
2 |10016 |bilingual
2 |10017 |bilingual
2 |10018 |bilingual
5 |10014 |multilingual
5 |10015 |multilingual
null |10020 |languages is null
null |10021 |languages is null
;
3 changes: 2 additions & 1 deletion .../main/java/org/elasticsearch/xpack/esql/optimizer/rules/logical/local/InferIsNotNull.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,7 @@
import org.elasticsearch.xpack.esql.core.expression.predicate.nulls.IsNotNull;
import org.elasticsearch.xpack.esql.core.rule.Rule;
import org.elasticsearch.xpack.esql.core.util.CollectionUtils;
import org.elasticsearch.xpack.esql.expression.function.scalar.conditional.Case;
import org.elasticsearch.xpack.esql.expression.function.scalar.nulls.Coalesce;
import org.elasticsearch.xpack.esql.plan.logical.LogicalPlan;

Expand Down Expand Up @@ -106,6 +107,6 @@ private boolean doResolve(Expression exp, AttributeMap<Expression> aliases, Set<
}

private static boolean skipExpression(Expression e) {
return e instanceof Coalesce;
return e instanceof Coalesce || e instanceof Case;
}
}
100 changes: 64 additions & 36 deletions .../src/test/java/org/elasticsearch/xpack/esql/optimizer/LocalLogicalPlanOptimizerTests.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -29,6 +29,7 @@
import org.elasticsearch.xpack.esql.core.type.DataType;
import org.elasticsearch.xpack.esql.core.type.EsField;
import org.elasticsearch.xpack.esql.expression.function.EsqlFunctionRegistry;
import org.elasticsearch.xpack.esql.expression.function.scalar.conditional.Case;
import org.elasticsearch.xpack.esql.expression.function.scalar.nulls.Coalesce;
import org.elasticsearch.xpack.esql.expression.function.scalar.string.StartsWith;
import org.elasticsearch.xpack.esql.expression.predicate.operator.arithmetic.Add;
Expand Down Expand Up @@ -58,8 +59,11 @@

import static java.util.Collections.emptyMap;
import static org.elasticsearch.xpack.esql.EsqlTestUtils.L;
import static org.elasticsearch.xpack.esql.EsqlTestUtils.ONE;
import static org.elasticsearch.xpack.esql.EsqlTestUtils.TEST_SEARCH_STATS;
import static org.elasticsearch.xpack.esql.EsqlTestUtils.TEST_VERIFIER;
import static org.elasticsearch.xpack.esql.EsqlTestUtils.THREE;
import static org.elasticsearch.xpack.esql.EsqlTestUtils.TWO;
import static org.elasticsearch.xpack.esql.EsqlTestUtils.as;
import static org.elasticsearch.xpack.esql.EsqlTestUtils.getFieldAttribute;
import static org.elasticsearch.xpack.esql.EsqlTestUtils.greaterThanOf;
Expand All @@ -81,10 +85,6 @@ public class LocalLogicalPlanOptimizerTests extends ESTestCase {
private static LogicalPlanOptimizer logicalOptimizer;
private static Map<String, EsField> mapping;

private static final Literal ONE = L(1);
private static final Literal TWO = L(2);
private static final Literal THREE = L(3);

@BeforeClass
public static void init() {
parser = new EsqlParser();
Expand Down Expand Up @@ -386,38 +386,6 @@ public void testMissingFieldInFilterNoProjection() {
);
}

public void testIsNotNullOnCoalesce() {
var plan = localPlan("""
from test
| where coalesce(emp_no, salary) is not null
""");

var limit = as(plan, Limit.class);
var filter = as(limit.child(), Filter.class);
var inn = as(filter.condition(), IsNotNull.class);
var coalesce = as(inn.children().get(0), Coalesce.class);
assertThat(Expressions.names(coalesce.children()), contains("emp_no", "salary"));
var source = as(filter.child(), EsRelation.class);
}

public void testIsNotNullOnExpression() {
var plan = localPlan("""
from test
| eval x = emp_no + 1
| where x is not null
""");

var limit = as(plan, Limit.class);
var filter = as(limit.child(), Filter.class);
var inn = as(filter.condition(), IsNotNull.class);
assertThat(Expressions.names(inn.children()), contains("x"));
var eval = as(filter.child(), Eval.class);
filter = as(eval.child(), Filter.class);
inn = as(filter.condition(), IsNotNull.class);
assertThat(Expressions.names(inn.children()), contains("emp_no"));
var source = as(filter.child(), EsRelation.class);
}

public void testSparseDocument() throws Exception {
var query = """
from large
Expand Down Expand Up @@ -516,6 +484,66 @@ public void testIsNotNullOnFunctionWithTwoFields() {
assertEquals(expected, new InferIsNotNull().apply(f));
}

public void testIsNotNullOnCoalesce() {
var plan = localPlan("""
from test
| where coalesce(emp_no, salary) is not null
""");

var limit = as(plan, Limit.class);
var filter = as(limit.child(), Filter.class);
var inn = as(filter.condition(), IsNotNull.class);
var coalesce = as(inn.children().get(0), Coalesce.class);
assertThat(Expressions.names(coalesce.children()), contains("emp_no", "salary"));
var source = as(filter.child(), EsRelation.class);
}

public void testIsNotNullOnExpression() {
var plan = localPlan("""
from test
| eval x = emp_no + 1
| where x is not null
""");

var limit = as(plan, Limit.class);
var filter = as(limit.child(), Filter.class);
var inn = as(filter.condition(), IsNotNull.class);
assertThat(Expressions.names(inn.children()), contains("x"));
var eval = as(filter.child(), Eval.class);
filter = as(eval.child(), Filter.class);
inn = as(filter.condition(), IsNotNull.class);
assertThat(Expressions.names(inn.children()), contains("emp_no"));
var source = as(filter.child(), EsRelation.class);
}

public void testIsNotNullOnCase() {
var plan = localPlan("""
from test
| where case(emp_no > 10000, "1", salary < 50000, "2", first_name) is not null
""");

var limit = as(plan, Limit.class);
var filter = as(limit.child(), Filter.class);
var inn = as(filter.condition(), IsNotNull.class);
var caseF = as(inn.children().get(0), Case.class);
assertThat(Expressions.names(caseF.children()), contains("emp_no > 10000", "\"1\"", "salary < 50000", "\"2\"", "first_name"));
var source = as(filter.child(), EsRelation.class);
}

public void testIsNotNullOnCase_With_IS_NULL() {
var plan = localPlan("""
from test
| where case(emp_no IS NULL, "1", salary IS NOT NULL, "2", first_name) is not null
""");

var limit = as(plan, Limit.class);
var filter = as(limit.child(), Filter.class);
var inn = as(filter.condition(), IsNotNull.class);
var caseF = as(inn.children().get(0), Case.class);
assertThat(Expressions.names(caseF.children()), contains("emp_no IS NULL", "\"1\"", "salary IS NOT NULL", "\"2\"", "first_name"));
var source = as(filter.child(), EsRelation.class);
}

private IsNotNull isNotNull(Expression field) {
return new IsNotNull(EMPTY, field);
}
Expand Down

0 comments on commit af3bfe6

Please sign in to comment.