ESQL: Verify PushDownEval/Enrich/RegexExtract logical optimizer rules correctness

[costin]

Description

Double check PushDownEnrich correctness in being pushed down within the plan. Add unit tests both at the logical and physical plan along with more CSV tests for various scenarios: with/without filters (some that can be pushed down, some that do not), projections and evals.

Relates #104957

[elasticsearchmachine]

Pinging @elastic/es-analytical-engine (Team:Analytics)

[luigidellaquila]

I'm renaming this issue and eventually labeling it differently, because PushDownEnrich is conceptually very similar to two other rules (PushDownEval and PushDownRegexExtract) and I suspect all the three are affected by the same conceptual problem.
More context in next comments.

[luigidellaquila]

What these rules do

All these three rules swap OrderBy with Eval/Enrich/Grok/Dissect respectively, eg.

if (child instanceof OrderBy orderBy) {
    return orderBy.replaceChild(eval.replaceChild(orderBy.child()));
} 

These are not properly optimization rules, the reason for their existence is described in PushDownEval javadoc

(expand this for details)

javadoc here

the goal is to group SORT commands together
E.g. from
... | SORT a | EVAL x = b + 1 | SORT x
to
... | EVAL x = b + 1 | SORT a | SORT x

At this point, PushDownAndCombineOrderBy will merge the two SORT commands in a single one, that will eventually become a TopN.

In the example above, the first SORT has no semantic value: even if the data are sorted by a, then they will be re-sorted by x afterwards. In practice, the first SORT could be safely removed even without swapping.

The problem

The problem happens when Eval/Dissect/... overwrite existing fields, that are then used in SORT
Eg. a plan fragment like

... | SORT emp_no | EVAL emp_no = somethingElse | ...

will be converted to

... | EVAL emp_no = somethingElse | SORT emp_no |  ...

Two problems here:

• Conceptually, it changes the query semantics: the sort could be performed on different values.
• Practically, it leads to problems with the layout.

---

Actual failures

Unfortunately, reproducing the problem above is possible also with simple and reasonable queries:

EVAL

from employees 
| keep emp_no, first_name 
| sort emp_no asc 
| eval emp_no = -emp_no 
| mv_expand first_name

{
    "error": {
        "root_cause": [
            {
                "type": "null_pointer_exception",
                "reason": "Cannot invoke \"org.elasticsearch.xpack.esql.planner.Layout$ChannelSet.type()\" because the return value of \"java.util.List.get(int)\" is null"
            }
        ],
        "type": "null_pointer_exception",
        "reason": "Cannot invoke \"org.elasticsearch.xpack.esql.planner.Layout$ChannelSet.type()\" because the return value of \"java.util.List.get(int)\" is null"
    },
    "status": 500
}

[2024-02-02T15:03:41,754][TRACE][o.e.x.e.o.LogicalPlanOptimizer] [runTask-0] Rule optimizer.LogicalPlanOptimizer$PushDownEval applied
Limit[500[INTEGER]]                                               = Limit[500[INTEGER]]
\_MvExpand[emp_no{r}#36,emp_no{r}#41]                             = \_MvExpand[emp_no{r}#36,emp_no{r}#41]
  \_Limit[500[INTEGER]]                                           =   \_Limit[500[INTEGER]]
    \_Eval[[NEG(emp_no{f}#39) AS emp_no]]                         !     \_OrderBy[[Order[emp_no{f}#39,ASC,LAST]]]
      \_OrderBy[[Order[emp_no{f}#39,ASC,LAST]]]                   !       \_EsqlProject[[first_name{f}#40, emp_no{r}#36]]
        \_EsqlProject[[emp_no{f}#39, first_name{f}#40]]           !         \_Eval[[NEG(emp_no{f}#39) AS emp_no]]
          \_EsRelation[employees][emp_no{f}#39, first_name{f}#40] =           \_EsRelation[employees][emp_no{f}#39, first_name{f}#40]          

This fragment of the optimizer execution shows the problem:

• in the original plan, OrderBy uses emp_no{f}#39, that was extracted by the EsRelation.
• In the transformed plan, OrderBy still tries to use emp_no{f}#39, but it was overwritten with emp_no{r}#36 by Eval, so it's no longer in the layout, thus the NullPointerException

Why is MV_EXPAND there?

Just to delay LIMIT pushdown a bit, and prevent that OrderBy is converted into TopN before PushDownEval is even executed.
So one of the two:

• the query with MV_EXPAND fails because of the swap of OrderBy and Eval, performed by PushDownEval
• if we remove MV_EXPAND, the query succeeds because PushDownEval is never executed

As a proof that the problem is in PushDownEval, disabling that rule the query succeeds also with MV_EXPAND.
So the problem is not MV_EXPAND.

DISSECT/GROK

from employees 
| keep emp_no, first_name 
| sort first_name asc 
| dissect first_name "%{first_name}" 
| mv_expand emp_no

{
    "error": {
        "root_cause": [
            {
                "type": "null_pointer_exception",
                "reason": "Cannot invoke \"org.elasticsearch.xpack.esql.planner.Layout$ChannelSet.type()\" because the return value of \"java.util.List.get(int)\" is null"
            }
        ],
        "type": "null_pointer_exception",
        "reason": "Cannot invoke \"org.elasticsearch.xpack.esql.planner.Layout$ChannelSet.type()\" because the return value of \"java.util.List.get(int)\" is null"
    },
    "status": 500
}

[2024-02-02T15:15:22,602][TRACE][o.e.x.e.o.LogicalPlanOptimizer] [runTask-0] Rule optimizer.LogicalPlanOptimizer$PushDownRegexExtract applied
Limit[500[INTEGER]]                                                                                                           = Limit[500[INTEGER]]
\_MvExpand[emp_no{f}#125,emp_no{r}#127]                                                                                       = \_MvExpand[emp_no{f}#125,emp_no{r}#127]
  \_Limit[500[INTEGER]]                                                                                                       =   \_Limit[500[INTEGER]]
    \_Dissect[first_name{f}#126,Parser[pattern=%{first_name}, appendSeparator=, parser=org.elasticsearch.dissect.DissectParse !     \_OrderBy[[Order[first_name{f}#126,ASC,LAST]]]
r@13d9b617],[first_name{r}#121]]                                                                                              !       \_EsqlProject[[emp_no{f}#125, first_name{r}#121]]
      \_OrderBy[[Order[first_name{f}#126,ASC,LAST]]]                                                                          !         \_Dissect[first_name{f}#126,Parser[pattern=%{first_name}, appendSeparator=, parser=org.elasticsearch.dissect.DissectParse
        \_EsqlProject[[emp_no{f}#125, first_name{f}#126]]                                                                     ! r@13d9b617],[first_name{r}#121]]
          \_EsRelation[employees][emp_no{f}#125, first_name{f}#126]                                                           =           \_EsRelation[employees][emp_no{f}#125, first_name{f}#126]

Exact same problem described for Eval (see above)

ENRICH

Reproducing the problem with ENRICH seems to be more complicated with our standard test datasets.

The logic of the rule is exactly the same though, so similar scenarios seem definitely possible.

---

In conclusion, the immediate strategy to fix the problem could be to check that Eval/Enrich/Dissect... do not overwrite SORT fields before we push down.

Something like

if (child instanceof OrderBy orderBy && doNotOverlap(orderBy.order(), eval.fields())) {
    return orderBy.replaceChild(eval.replaceChild(orderBy.child()));
} 

This will still result in a query failure due to unbounded sort in some cases.

Another strategy could be, instead of pushdown, try to analyze the plan and see if some SORTs are redundant and can be safely removed.