Skip to content

Commit 04c733e

Browse files
fang-xing-esqlfang-xing-esql
committed
more tests for lookup join and fork even they are not supported yet
1 parent bced1e6 commit 04c733e

File tree

2 files changed

+111
-0
lines changed

2 files changed

+111
-0
lines changed

‎x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/optimizer/rules/physical/local/ReplaceRoundToWithQueryAndTags.java‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -258,6 +258,7 @@ public class ReplaceRoundToWithQueryAndTags extends PhysicalOptimizerRules.Param
258258
protected PhysicalPlan rule(EvalExec evalExec, LocalPhysicalOptimizerContext ctx) {
259259
PhysicalPlan plan = evalExec;
260260
// TimeSeriesSourceOperator and LuceneTopNSourceOperator do not support QueryAndTags, skip them
261+
// Lookup join is not supported yet
261262
if (evalExec.child() instanceof EsQueryExec queryExec && queryExec.canSubstituteRoundToWithQueryBuilderAndTags()) {
262263
// Look for RoundTo and plan the push down for it.
263264
List<RoundTo> roundTos = evalExec.fields()

‎x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/optimizer/rules/physical/local/ReplaceRoundToWithQueryAndTagsTests.java‎

Lines changed: 110 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -14,12 +14,17 @@
1414
import org.elasticsearch.index.query.RangeQueryBuilder;
1515
import org.elasticsearch.xpack.esql.EsqlTestUtils;
1616
import org.elasticsearch.xpack.esql.core.expression.Alias;
17+
import org.elasticsearch.xpack.esql.core.expression.Attribute;
1718
import org.elasticsearch.xpack.esql.core.expression.Expression;
1819
import org.elasticsearch.xpack.esql.core.expression.Expressions;
1920
import org.elasticsearch.xpack.esql.core.expression.FieldAttribute;
2021
import org.elasticsearch.xpack.esql.core.expression.NamedExpression;
22+
import org.elasticsearch.xpack.esql.core.expression.ReferenceAttribute;
23+
import org.elasticsearch.xpack.esql.core.expression.function.Function;
2124
import org.elasticsearch.xpack.esql.core.tree.Source;
2225
import org.elasticsearch.xpack.esql.core.type.DataType;
26+
import org.elasticsearch.xpack.esql.expression.function.grouping.Bucket;
27+
import org.elasticsearch.xpack.esql.expression.function.scalar.date.DateTrunc;
2328
import org.elasticsearch.xpack.esql.expression.function.scalar.math.RoundTo;
2429
import org.elasticsearch.xpack.esql.optimizer.LocalPhysicalPlanOptimizerTests;
2530
import org.elasticsearch.xpack.esql.plan.physical.AggregateExec;
@@ -28,6 +33,8 @@
2833
import org.elasticsearch.xpack.esql.plan.physical.ExchangeExec;
2934
import org.elasticsearch.xpack.esql.plan.physical.FieldExtractExec;
3035
import org.elasticsearch.xpack.esql.plan.physical.LimitExec;
36+
import org.elasticsearch.xpack.esql.plan.physical.LookupJoinExec;
37+
import org.elasticsearch.xpack.esql.plan.physical.MergeExec;
3138
import org.elasticsearch.xpack.esql.plan.physical.PhysicalPlan;
3239
import org.elasticsearch.xpack.esql.plan.physical.ProjectExec;
3340
import org.elasticsearch.xpack.esql.plan.physical.TopNExec;
@@ -41,6 +48,7 @@
4148
import java.util.stream.Collectors;
4249

4350
import static org.elasticsearch.compute.aggregation.AggregatorMode.FINAL;
51+
import static org.elasticsearch.compute.aggregation.AggregatorMode.INITIAL;
4452
import static org.elasticsearch.index.query.QueryBuilders.boolQuery;
4553
import static org.elasticsearch.index.query.QueryBuilders.existsQuery;
4654
import static org.elasticsearch.index.query.QueryBuilders.matchQuery;
@@ -281,6 +289,108 @@ public void testDateTruncBucketTransformToQueryAndTagsWithOtherPushdownFunctions
281289
}
282290
}
283291

292+
// ReplaceRoundToWithQueryAndTags does not support lookup joins yet
293+
public void testDateTruncBucketNotTransformToQueryAndTagsWithLookupJoin() {
294+
for (String dateHistogram : dateHistograms) {
295+
String query = LoggerMessageFormat.format(null, """
296+
from test
297+
| rename integer as language_code
298+
| lookup join languages_lookup on language_code
299+
| stats count(*) by x = {}
300+
""", dateHistogram);
301+
PhysicalPlan plan = plannerOptimizer.plan(query, searchStats, makeAnalyzer("mapping-all-types.json"));
302+
303+
LimitExec limit = as(plan, LimitExec.class);
304+
AggregateExec agg = as(limit.child(), AggregateExec.class);
305+
assertThat(agg.getMode(), is(FINAL));
306+
List<? extends Expression> groupings = agg.groupings();
307+
NamedExpression grouping = as(groupings.get(0), NamedExpression.class);
308+
assertEquals("x", grouping.name());
309+
assertEquals(DataType.DATETIME, grouping.dataType());
310+
assertEquals(List.of("count(*)", "x"), Expressions.names(agg.aggregates()));
311+
ExchangeExec exchange = as(agg.child(), ExchangeExec.class);
312+
assertThat(exchange.inBetweenAggs(), is(true));
313+
agg = as(exchange.child(), AggregateExec.class);
314+
EvalExec eval = as(agg.child(), EvalExec.class);
315+
List<Alias> aliases = eval.fields();
316+
assertEquals(1, aliases.size());
317+
RoundTo roundTo = as(aliases.get(0).child(), RoundTo.class);
318+
assertEquals(4, roundTo.points().size());
319+
FieldExtractExec fieldExtractExec = as(eval.child(), FieldExtractExec.class);
320+
List<Attribute> attributes = fieldExtractExec.attributesToExtract();
321+
assertEquals(1, attributes.size());
322+
assertEquals("date", attributes.get(0).name());
323+
LookupJoinExec lookupJoinExec = as(fieldExtractExec.child(), LookupJoinExec.class); // this is why the rule doesn't apply
324+
// lhs of lookup join
325+
fieldExtractExec = as(lookupJoinExec.left(), FieldExtractExec.class);
326+
attributes = fieldExtractExec.attributesToExtract();
327+
assertEquals(1, attributes.size());
328+
assertEquals("integer", attributes.get(0).name());
329+
EsQueryExec esQueryExec = as(fieldExtractExec.child(), EsQueryExec.class);
330+
assertEquals("test", esQueryExec.indexPattern());
331+
List<EsQueryExec.QueryBuilderAndTags> queryBuilderAndTags = esQueryExec.queryBuilderAndTags();
332+
assertEquals(1, queryBuilderAndTags.size());
333+
EsQueryExec.QueryBuilderAndTags queryBuilder = queryBuilderAndTags.get(0);
334+
assertNull(queryBuilder.query());
335+
assertTrue(queryBuilder.tags().isEmpty());
336+
assertNull(esQueryExec.query());
337+
// rhs of lookup join
338+
esQueryExec = as(lookupJoinExec.right(), EsQueryExec.class);
339+
assertEquals("languages_lookup", esQueryExec.indexPattern());
340+
queryBuilderAndTags = esQueryExec.queryBuilderAndTags();
341+
assertEquals(1, queryBuilderAndTags.size());
342+
queryBuilder = queryBuilderAndTags.get(0);
343+
assertNull(queryBuilder.query());
344+
assertTrue(queryBuilder.tags().isEmpty());
345+
assertNull(esQueryExec.query());
346+
}
347+
}
348+
349+
// ReplaceRoundToWithQueryAndTags does not support lookup joins yet
350+
public void testDateTruncBucketNotTransformToQueryAndTagsWithFork() {
351+
for (String dateHistogram : dateHistograms) {
352+
String query = LoggerMessageFormat.format(null, """
353+
from test
354+
| fork (where integer > 100)
355+
(where keyword : "keyword")
356+
| stats count(*) by x = {}
357+
""", dateHistogram);
358+
PhysicalPlan plan = plannerOptimizer.plan(query, searchStats, makeAnalyzer("mapping-all-types.json"));
359+
360+
LimitExec limit = as(plan, LimitExec.class);
361+
AggregateExec agg = as(limit.child(), AggregateExec.class);
362+
assertThat(agg.getMode(), is(FINAL));
363+
List<? extends Expression> groupings = agg.groupings();
364+
NamedExpression grouping = as(groupings.get(0), NamedExpression.class);
365+
assertEquals("x", grouping.name());
366+
assertEquals(DataType.DATETIME, grouping.dataType());
367+
assertEquals(List.of("count(*)", "x"), Expressions.names(agg.aggregates()));
368+
agg = as(agg.child(), AggregateExec.class);
369+
assertThat(agg.getMode(), is(INITIAL));
370+
groupings = agg.groupings();
371+
grouping = as(groupings.get(0), NamedExpression.class);
372+
assertEquals("x", grouping.name());
373+
assertEquals(DataType.DATETIME, grouping.dataType());
374+
assertEquals(List.of("count(*)", "x"), Expressions.names(agg.aggregates()));
375+
EvalExec eval = as(agg.child(), EvalExec.class);
376+
List<Alias> aliases = eval.fields();
377+
assertEquals(1, aliases.size());
378+
var function = as(aliases.get(0).child(), Function.class);
379+
ReferenceAttribute fa = null; // if merge returns FieldAttribute instead of ReferenceAttribute, the rule might apply
380+
if (function instanceof DateTrunc dateTrunc) {
381+
fa = as(dateTrunc.field(), ReferenceAttribute.class);
382+
} else if (function instanceof Bucket bucket) {
383+
fa = as(bucket.field(), ReferenceAttribute.class);
384+
} else if (function instanceof RoundTo roundTo) {
385+
fa = as(roundTo.field(), ReferenceAttribute.class);
386+
}
387+
assertNotNull(fa);
388+
assertEquals("date", fa.name());
389+
assertEquals(DataType.DATETIME, fa.dataType());
390+
MergeExec mergeExec = as(eval.child(), MergeExec.class);
391+
}
392+
}
393+
284394
private static void verifyQueryAndTags(List<EsQueryExec.QueryBuilderAndTags> expected, List<EsQueryExec.QueryBuilderAndTags> actual) {
285395
assertEquals(expected.size(), actual.size());
286396
for (int i = 0; i < expected.size(); i++) {

0 commit comments

Comments
 (0)