Skip to content

Commit cd844d0

Browse files
kanoshioukanoshiou
committed
Add more tests for UnionAll cases
1 parent b5ede88 commit cd844d0

File tree

2 files changed

+232
-0
lines changed

2 files changed

+232
-0
lines changed

‎x-pack/plugin/esql/qa/testFixtures/src/main/resources/mv_expand.csv-spec‎

Lines changed: 68 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -567,3 +567,71 @@ b:integer | a:integer
567567
2 | 12
568568
3 | 12
569569
;
570+
571+
572+
testPushDownMvExpandPastProject5
573+
required_capability: fork_v9
574+
required_capability: subquery_in_from_command
575+
required_capability: push_down_mv_expand_past_project
576+
from employees, (from employees | keep salary)
577+
| eval salary = salary::keyword
578+
| keep salary
579+
| mv_expand salary
580+
| sort salary
581+
| limit 6
582+
;
583+
584+
salary:keyword
585+
25324
586+
25324
587+
25945
588+
25945
589+
25976
590+
25976
591+
;
592+
593+
594+
testPushDownMvExpandPastProject6
595+
required_capability: fork_v9
596+
required_capability: subquery_in_from_command
597+
required_capability: push_down_mv_expand_past_project
598+
from employees, (from employees | keep salary)
599+
| eval salary = salary::keyword
600+
| eval tmp = salary
601+
| keep salary, tmp
602+
| mv_expand salary
603+
| sort salary
604+
| limit 6
605+
;
606+
607+
salary:keyword | tmp:keyword
608+
25324 | 25324
609+
25324 | 25324
610+
25945 | 25945
611+
25945 | 25945
612+
25976 | 25976
613+
25976 | 25976
614+
;
615+
616+
617+
testPushDownMvExpandPastProject7
618+
required_capability: fork_v9
619+
required_capability: subquery_in_from_command
620+
required_capability: push_down_mv_expand_past_project
621+
from employees, (from employees | keep salary)
622+
| eval salary = salary::keyword
623+
| eval tmp = salary
624+
| keep salary, tmp
625+
| mv_expand tmp
626+
| sort salary
627+
| limit 6
628+
;
629+
630+
salary:keyword | tmp:keyword
631+
25324 | 25324
632+
25324 | 25324
633+
25945 | 25945
634+
25945 | 25945
635+
25976 | 25976
636+
25976 | 25976
637+
;

‎x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/optimizer/LogicalPlanOptimizerTests.java‎

Lines changed: 164 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -132,9 +132,11 @@
132132
import org.elasticsearch.xpack.esql.plan.logical.Project;
133133
import org.elasticsearch.xpack.esql.plan.logical.Row;
134134
import org.elasticsearch.xpack.esql.plan.logical.Sample;
135+
import org.elasticsearch.xpack.esql.plan.logical.Subquery;
135136
import org.elasticsearch.xpack.esql.plan.logical.TimeSeriesAggregate;
136137
import org.elasticsearch.xpack.esql.plan.logical.TopN;
137138
import org.elasticsearch.xpack.esql.plan.logical.UnaryPlan;
139+
import org.elasticsearch.xpack.esql.plan.logical.UnionAll;
138140
import org.elasticsearch.xpack.esql.plan.logical.inference.Completion;
139141
import org.elasticsearch.xpack.esql.plan.logical.inference.Rerank;
140142
import org.elasticsearch.xpack.esql.plan.logical.join.InlineJoin;
@@ -3209,6 +3211,168 @@ public void testPushDownMvExpandPastProject4() {
32093211
as(limit.child(), LocalRelation.class);
32103212
}
32113213

3214+
/**
3215+
* <pre>{@code
3216+
* Project[[$$salary$temp_name$57{r$}#58 AS salary#53]]
3217+
* \_Limit[1000[INTEGER],true,false]
3218+
* \_MvExpand[$$salary$converted_to$keyword{r$}#56,$$salary$temp_name$57{r$}#58]
3219+
* \_Limit[1000[INTEGER],false,false]
3220+
* \_UnionAll[[_meta_field{r}#42, emp_no{r}#43, first_name{r}#44, gender{r}#45, hire_date{r}#46, job{r}#47, job.raw{r}#48, l
3221+
* anguages{r}#49, last_name{r}#50, long_noidx{r}#51, salary{r}#52, $$salary$converted_to$keyword{r$}#56]]
3222+
* |_EsqlProject[[_meta_field{f}#16, emp_no{f}#10, first_name{f}#11, gender{f}#12, hire_date{f}#17, job{f}#18, job.raw{f}#19, l
3223+
* anguages{f}#13, last_name{f}#14, long_noidx{f}#20, salary{f}#15, $$salary$converted_to$keyword{r}#54]]
3224+
* | \_Eval[[TOSTRING(salary{f}#15) AS $$salary$converted_to$keyword#54]]
3225+
* | \_Limit[1000[INTEGER],false,false]
3226+
* | \_EsRelation[employees][_meta_field{f}#16, emp_no{f}#10, first_name{f}#11, ..]
3227+
* \_EsqlProject[[_meta_field{r}#32, emp_no{r}#33, first_name{r}#34, gender{r}#35, hire_date{r}#36, job{r}#37, job.raw{r}#38, l
3228+
* anguages{r}#39, last_name{r}#40, long_noidx{r}#41, salary{f}#26, $$salary$converted_to$keyword{r}#55]]
3229+
* \_Eval[[null[KEYWORD] AS _meta_field#32, null[INTEGER] AS emp_no#33, null[KEYWORD] AS first_name#34, null[TEXT] AS ge
3230+
* nder#35, null[DATETIME] AS hire_date#36, null[TEXT] AS job#37, null[KEYWORD] AS job.raw#38, null[INTEGER] AS languages#39,
3231+
* null[KEYWORD] AS last_name#40, null[LONG] AS long_noidx#41, TOSTRING(salary{f}#26) AS $$salary$converted_to$keyword#55]]
3232+
* \_Subquery[]
3233+
* \_EsqlProject[[salary{f}#26]]
3234+
* \_Limit[1000[INTEGER],false,false]
3235+
* \_EsRelation[employees][_meta_field{f}#27, emp_no{f}#21, first_name{f}#22, ..]
3236+
* }</pre>
3237+
*/
3238+
public void testPushDownMvExpandPastProject5() {
3239+
assumeTrue("Requires subquery in FROM command support", EsqlCapabilities.Cap.SUBQUERY_IN_FROM_COMMAND.isEnabled());
3240+
LogicalPlan plan = optimizedPlan("""
3241+
from employees, (from employees | keep salary)
3242+
| eval salary = salary::keyword
3243+
| keep salary
3244+
| mv_expand salary
3245+
""");
3246+
var project = as(plan, Project.class);
3247+
var limit = asLimit(project.child(), 1000, true);
3248+
var mvExpand = as(limit.child(), MvExpand.class);
3249+
limit = asLimit(mvExpand.child(), 1000, false);
3250+
var union = as(limit.child(), UnionAll.class);
3251+
assertThat(union.children(), hasSize(2));
3252+
3253+
var branch1 = as(union.children().getFirst(), EsqlProject.class);
3254+
var eval1 = as(branch1.child(), Eval.class);
3255+
var limit1 = asLimit(eval1.child(), 1000, false);
3256+
as(limit1.child(), EsRelation.class);
3257+
3258+
var branch2 = as(union.children().get(1), EsqlProject.class);
3259+
var eval2 = as(branch2.child(), Eval.class);
3260+
var subquery = as(eval2.child(), Subquery.class);
3261+
var subProject = as(subquery.child(), EsqlProject.class);
3262+
var subLimit = asLimit(subProject.child(), 1000, false);
3263+
as(subLimit.child(), EsRelation.class);
3264+
}
3265+
3266+
/**
3267+
* <pre>{@code
3268+
* Project[[$$salary$temp_name$61{r$}#62 AS salary#57, $$salary$converted_to$keyword{r$}#60 AS tmp#9]]
3269+
* \_Limit[1000[INTEGER],true,false]
3270+
* \_MvExpand[$$salary$converted_to$keyword$salary$0{r}#63,$$salary$temp_name$61{r$}#62]
3271+
* \_Eval[[$$salary$converted_to$keyword{r$}#60 AS $$salary$converted_to$keyword$salary$0#63]]
3272+
* \_Limit[1000[INTEGER],false,false]
3273+
* \_UnionAll[[_meta_field{r}#46, emp_no{r}#47, first_name{r}#48, gender{r}#49, hire_date{r}#50, job{r}#51, job.raw{r}#52, l
3274+
* anguages{r}#53, last_name{r}#54, long_noidx{r}#55, salary{r}#56, $$salary$converted_to$keyword{r$}#60]]
3275+
* |_EsqlProject[[_meta_field{f}#20, emp_no{f}#14, first_name{f}#15, gender{f}#16, hire_date{f}#21, job{f}#22, job.raw{f}#23,l
3276+
* anguages{f}#17, last_name{f}#18, long_noidx{f}#24, salary{f}#19, $$salary$converted_to$keyword{r}#58]]
3277+
* | \_Eval[[TOSTRING(salary{f}#19) AS $$salary$converted_to$keyword#58]]
3278+
* | \_Limit[1000[INTEGER],false,false]
3279+
* | \_EsRelation[employees][_meta_field{f}#20, emp_no{f}#14, first_name{f}#15, ..]
3280+
* \_EsqlProject[[_meta_field{r}#36, emp_no{r}#37, first_name{r}#38, gender{r}#39, hire_date{r}#40, job{r}#41, job.raw{r}#42,l
3281+
* anguages{r}#43, last_name{r}#44, long_noidx{r}#45, salary{f}#30, $$salary$converted_to$keyword{r}#59]]
3282+
* \_Eval[[null[KEYWORD] AS _meta_field#36, null[INTEGER] AS emp_no#37, null[KEYWORD] AS first_name#38, null[TEXT] AS ge
3283+
* nder#39, null[DATETIME] AS hire_date#40, null[TEXT] AS job#41, null[KEYWORD] AS job.raw#42, null[INTEGER] AS languages#43,
3284+
* null[KEYWORD] AS last_name#44, null[LONG] AS long_noidx#45, TOSTRING(salary{f}#30) AS $$salary$converted_to$keyword#59]]
3285+
* \_Subquery[]
3286+
* \_EsqlProject[[salary{f}#30]]
3287+
* \_Limit[1000[INTEGER],false,false]
3288+
* \_EsRelation[employees][_meta_field{f}#31, emp_no{f}#25, first_name{f}#26, ..]
3289+
* }</pre>
3290+
*/
3291+
public void testPushDownMvExpandPastProject6() {
3292+
assumeTrue("Requires subquery in FROM command support", EsqlCapabilities.Cap.SUBQUERY_IN_FROM_COMMAND.isEnabled());
3293+
LogicalPlan plan = optimizedPlan("""
3294+
from employees, (from employees | keep salary)
3295+
| eval salary = salary::keyword
3296+
| eval tmp = salary
3297+
| keep salary, tmp
3298+
| mv_expand salary
3299+
""");
3300+
var project = as(plan, Project.class);
3301+
var limit = asLimit(project.child(), 1000, true);
3302+
var mvExpand = as(limit.child(), MvExpand.class);
3303+
var eval = as(mvExpand.child(), Eval.class);
3304+
limit = asLimit(eval.child(), 1000, false);
3305+
var union = as(limit.child(), UnionAll.class);
3306+
assertThat(union.children(), hasSize(2));
3307+
3308+
var branch1 = as(union.children().getFirst(), EsqlProject.class);
3309+
var eval1 = as(branch1.child(), Eval.class);
3310+
var limit1 = asLimit(eval1.child(), 1000, false);
3311+
as(limit1.child(), EsRelation.class);
3312+
3313+
var branch2 = as(union.children().get(1), EsqlProject.class);
3314+
var eval2 = as(branch2.child(), Eval.class);
3315+
var subquery = as(eval2.child(), Subquery.class);
3316+
var subProject = as(subquery.child(), EsqlProject.class);
3317+
var subLimit = asLimit(subProject.child(), 1000, false);
3318+
as(subLimit.child(), EsRelation.class);
3319+
}
3320+
3321+
/**
3322+
* <pre>{@code
3323+
* Project[[$$salary$converted_to$keyword{r$}#60 AS salary#6, tmp{r}#57]]
3324+
* \_Limit[1000[INTEGER],true,false]
3325+
* \_MvExpand[$$salary$converted_to$keyword$tmp$0{r}#61,tmp{r}#57]
3326+
* \_Eval[[$$salary$converted_to$keyword{r$}#60 AS $$salary$converted_to$keyword$tmp$0#61]]
3327+
* \_Limit[1000[INTEGER],false,false]
3328+
* \_UnionAll[[_meta_field{r}#46, emp_no{r}#47, first_name{r}#48, gender{r}#49, hire_date{r}#50, job{r}#51, job.raw{r}#52, l
3329+
* anguages{r}#53, last_name{r}#54, long_noidx{r}#55, salary{r}#56, $$salary$converted_to$keyword{r$}#60]]
3330+
* |_EsqlProject[[_meta_field{f}#20, emp_no{f}#14, first_name{f}#15, gender{f}#16, hire_date{f}#21, job{f}#22, job.raw{f}#23,l
3331+
* anguages{f}#17, last_name{f}#18, long_noidx{f}#24, salary{f}#19, $$salary$converted_to$keyword{r}#58]]
3332+
* | \_Eval[[TOSTRING(salary{f}#19) AS $$salary$converted_to$keyword#58]]
3333+
* | \_Limit[1000[INTEGER],false,false]
3334+
* | \_EsRelation[employees][_meta_field{f}#20, emp_no{f}#14, first_name{f}#15, ..]
3335+
* \_EsqlProject[[_meta_field{r}#36, emp_no{r}#37, first_name{r}#38, gender{r}#39, hire_date{r}#40, job{r}#41, job.raw{r}#42,l
3336+
* anguages{r}#43, last_name{r}#44, long_noidx{r}#45, salary{f}#30, $$salary$converted_to$keyword{r}#59]]
3337+
* \_Eval[[null[KEYWORD] AS _meta_field#36, null[INTEGER] AS emp_no#37, null[KEYWORD] AS first_name#38, null[TEXT] AS ge
3338+
* nder#39, null[DATETIME] AS hire_date#40, null[TEXT] AS job#41, null[KEYWORD] AS job.raw#42, null[INTEGER] AS languages#43,
3339+
* null[KEYWORD] AS last_name#44, null[LONG] AS long_noidx#45, TOSTRING(salary{f}#30) AS $$salary$converted_to$keyword#59]]
3340+
* \_Subquery[]
3341+
* \_EsqlProject[[salary{f}#30]]
3342+
* \_Limit[1000[INTEGER],false,false]
3343+
* \_EsRelation[employees][_meta_field{f}#31, emp_no{f}#25, first_name{f}#26, ..]
3344+
* }</pre>
3345+
*/
3346+
public void testPushDownMvExpandPastProject7() {
3347+
assumeTrue("Requires subquery in FROM command support", EsqlCapabilities.Cap.SUBQUERY_IN_FROM_COMMAND.isEnabled());
3348+
LogicalPlan plan = optimizedPlan("""
3349+
from employees, (from employees | keep salary)
3350+
| eval salary = salary::keyword
3351+
| eval tmp = salary
3352+
| keep salary, tmp
3353+
| mv_expand tmp
3354+
""");
3355+
var project = as(plan, Project.class);
3356+
var limit = asLimit(project.child(), 1000, true);
3357+
var mvExpand = as(limit.child(), MvExpand.class);
3358+
var eval = as(mvExpand.child(), Eval.class);
3359+
limit = asLimit(eval.child(), 1000, false);
3360+
var union = as(limit.child(), UnionAll.class);
3361+
assertThat(union.children(), hasSize(2));
3362+
3363+
var branch1 = as(union.children().getFirst(), EsqlProject.class);
3364+
var eval1 = as(branch1.child(), Eval.class);
3365+
var limit1 = asLimit(eval1.child(), 1000, false);
3366+
as(limit1.child(), EsRelation.class);
3367+
3368+
var branch2 = as(union.children().get(1), EsqlProject.class);
3369+
var eval2 = as(branch2.child(), Eval.class);
3370+
var subquery = as(eval2.child(), Subquery.class);
3371+
var subProject = as(subquery.child(), EsqlProject.class);
3372+
var subLimit = asLimit(subProject.child(), 1000, false);
3373+
as(subLimit.child(), EsRelation.class);
3374+
}
3375+
32123376
public void testTopNEnrich() {
32133377
LogicalPlan plan = optimizedPlan("""
32143378
from test

0 commit comments

Comments
 (0)