diff --git a/docs/changelog/132512.yaml b/docs/changelog/132512.yaml new file mode 100644 index 0000000000000..66b55d45e36ec --- /dev/null +++ b/docs/changelog/132512.yaml @@ -0,0 +1,5 @@ +pr: 132512 +summary: Rewrite `RoundTo` to `QueryAndTags` +area: ES|QL +type: enhancement +issues: [] diff --git a/docs/changelog/132781.yaml b/docs/changelog/132781.yaml new file mode 100644 index 0000000000000..32bc98aff0863 --- /dev/null +++ b/docs/changelog/132781.yaml @@ -0,0 +1,6 @@ +pr: 132781 +summary: Rewrite `RoundTo` to `QueryAndTags` and merge range queries on `RoundTo` + field +area: ES|QL +type: enhancement +issues: [] diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/optimizer/LocalPhysicalPlanOptimizer.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/optimizer/LocalPhysicalPlanOptimizer.java index af36963ac54a3..9176a18a1a092 100644 --- a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/optimizer/LocalPhysicalPlanOptimizer.java +++ b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/optimizer/LocalPhysicalPlanOptimizer.java @@ -18,6 +18,7 @@ import org.elasticsearch.xpack.esql.optimizer.rules.physical.local.PushSampleToSource; import org.elasticsearch.xpack.esql.optimizer.rules.physical.local.PushStatsToSource; import org.elasticsearch.xpack.esql.optimizer.rules.physical.local.PushTopNToSource; +import org.elasticsearch.xpack.esql.optimizer.rules.physical.local.ReplaceRoundToWithQueryAndTags; import org.elasticsearch.xpack.esql.optimizer.rules.physical.local.ReplaceSourceAttributes; import org.elasticsearch.xpack.esql.optimizer.rules.physical.local.SpatialDocValuesExtraction; import org.elasticsearch.xpack.esql.optimizer.rules.physical.local.SpatialShapeBoundsExtraction; @@ -60,7 +61,7 @@ protected List> batches() { } protected static List> rules(boolean optimizeForEsSource) { - List> esSourceRules = new ArrayList<>(6); + List> esSourceRules = new ArrayList<>(7); esSourceRules.add(new ReplaceSourceAttributes()); if (optimizeForEsSource) { esSourceRules.add(new PushTopNToSource()); @@ -74,6 +75,20 @@ protected static List> rules(boolean optimizeForEsSource) { // execute the rules multiple times to improve the chances of things being pushed down @SuppressWarnings("unchecked") var pushdown = new Batch("Push to ES", esSourceRules.toArray(Rule[]::new)); + List> substitutionRules = new ArrayList<>(1); + if (optimizeForEsSource) { + substitutionRules.add(new ReplaceRoundToWithQueryAndTags()); + } + // execute the SubstituteRoundToWithQueryAndTags rule once after all the other pushdown rules are applied, as this rule generate + // multiple QueryBuilders according the number of RoundTo points, it should be applied after all the other eligible pushdowns are + // done. + @SuppressWarnings("unchecked") + var substituteRoundToWithQueryAndTags = new Batch( + "Substitute RoundTo with QueryAndTags", + Limiter.ONCE, + substitutionRules.toArray(Rule[]::new) + ); + // add the field extraction in just one pass // add it at the end after all the other rules have ran var fieldExtraction = new Batch<>( @@ -84,6 +99,6 @@ protected static List> rules(boolean optimizeForEsSource) { new SpatialShapeBoundsExtraction(), new ParallelizeTimeSeriesSource() ); - return List.of(pushdown, fieldExtraction); + return List.of(pushdown, substituteRoundToWithQueryAndTags, fieldExtraction); } } diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/optimizer/rules/physical/local/PushFiltersToSource.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/optimizer/rules/physical/local/PushFiltersToSource.java index ba382b9800ece..9ee93fa6c8223 100644 --- a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/optimizer/rules/physical/local/PushFiltersToSource.java +++ b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/optimizer/rules/physical/local/PushFiltersToSource.java @@ -126,10 +126,10 @@ private static PhysicalPlan rewrite( queryExec.indexMode(), queryExec.indexNameWithModes(), queryExec.output(), - query, queryExec.limit(), queryExec.sorts(), - queryExec.estimatedRowSize() + queryExec.estimatedRowSize(), + List.of(new EsQueryExec.QueryBuilderAndTags(query, List.of())) ); // If the eval contains other aliases, not just field attributes, we need to keep them in the plan PhysicalPlan plan = evalFields.isEmpty() ? queryExec : new EvalExec(filterExec.source(), queryExec, evalFields); diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/optimizer/rules/physical/local/ReplaceRoundToWithQueryAndTags.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/optimizer/rules/physical/local/ReplaceRoundToWithQueryAndTags.java new file mode 100644 index 0000000000000..3fbd9fc161ad6 --- /dev/null +++ b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/optimizer/rules/physical/local/ReplaceRoundToWithQueryAndTags.java @@ -0,0 +1,852 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the Elastic License + * 2.0; you may not use this file except in compliance with the Elastic License + * 2.0. + */ + +package org.elasticsearch.xpack.esql.optimizer.rules.physical.local; + +import org.elasticsearch.index.query.BoolQueryBuilder; +import org.elasticsearch.index.query.ExistsQueryBuilder; +import org.elasticsearch.index.query.QueryBuilder; +import org.elasticsearch.index.query.RangeQueryBuilder; +import org.elasticsearch.xpack.esql.core.expression.Alias; +import org.elasticsearch.xpack.esql.core.expression.Attribute; +import org.elasticsearch.xpack.esql.core.expression.Expression; +import org.elasticsearch.xpack.esql.core.expression.Expressions; +import org.elasticsearch.xpack.esql.core.expression.FieldAttribute; +import org.elasticsearch.xpack.esql.core.expression.Literal; +import org.elasticsearch.xpack.esql.core.querydsl.query.Query; +import org.elasticsearch.xpack.esql.core.tree.Source; +import org.elasticsearch.xpack.esql.core.type.DataType; +import org.elasticsearch.xpack.esql.core.type.EsField; +import org.elasticsearch.xpack.esql.core.util.Queries; +import org.elasticsearch.xpack.esql.expression.function.scalar.math.RoundTo; +import org.elasticsearch.xpack.esql.expression.predicate.Range; +import org.elasticsearch.xpack.esql.expression.predicate.nulls.IsNull; +import org.elasticsearch.xpack.esql.expression.predicate.operator.comparison.GreaterThanOrEqual; +import org.elasticsearch.xpack.esql.expression.predicate.operator.comparison.LessThan; +import org.elasticsearch.xpack.esql.optimizer.LocalPhysicalOptimizerContext; +import org.elasticsearch.xpack.esql.optimizer.PhysicalOptimizerRules; +import org.elasticsearch.xpack.esql.plan.physical.EsQueryExec; +import org.elasticsearch.xpack.esql.plan.physical.EvalExec; +import org.elasticsearch.xpack.esql.plan.physical.PhysicalPlan; +import org.elasticsearch.xpack.esql.querydsl.query.SingleValueQuery; + +import java.time.ZoneId; +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; +import java.util.Map; + +import static org.elasticsearch.xpack.esql.core.type.DataType.DATETIME; +import static org.elasticsearch.xpack.esql.core.type.DataType.DATE_NANOS; +import static org.elasticsearch.xpack.esql.core.type.DataTypeConverter.safeToLong; +import static org.elasticsearch.xpack.esql.planner.TranslatorHandler.TRANSLATOR_HANDLER; +import static org.elasticsearch.xpack.esql.type.EsqlDataTypeConverter.DEFAULT_DATE_NANOS_FORMATTER; +import static org.elasticsearch.xpack.esql.type.EsqlDataTypeConverter.DEFAULT_DATE_TIME_FORMATTER; +import static org.elasticsearch.xpack.esql.type.EsqlDataTypeConverter.dateNanosToLong; +import static org.elasticsearch.xpack.esql.type.EsqlDataTypeConverter.dateTimeToLong; + +/** + * {@code ReplaceRoundToWithQueryAndTags} builds a list of ranges and associated tags base on the rounding points defined in a + * {@code RoundTo} function. It then rewrites the {@code EsQueryExec.query()} into a corresponding list of {@code QueryBuilder}s and tags, + * each mapped to its respective range. + * + * Here are some examples: + * + * 1. Aggregation with {@code DateTrunc}. + * The {@code DateTrunc} function in the query below can be rewritten to {@code RoundTo} by {@code ReplaceDateTruncBucketWithRoundTo}. + * This rule pushes down the {@code RoundTo} function by creating a list of {@code QueryBuilderAndTags}, so that + * {@code EsPhysicalOperationProviders} can build {@code LuceneSliceQueue} with the corresponding list of {@code QueryAndTags} to process + * further. + * | STATS COUNT(*) BY d = DATE_TRUNC(1 day, date) + * becomes, the rounding points are calculated according to SearchStats and predicates from the query. + * | EVAL d = ROUND_TO(hire_date, 1697760000000, 1697846400000, 1697932800000) + * | STATS COUNT(*) BY d + * becomes + * [QueryBuilderAndTags[query={ + * "esql_single_value" : { + * "field" : "date", + * "next" : { + * "range" : { + * "date" : { + * "lt" : "2023-10-21T00:00:00.000Z", + * "time_zone" : "Z", + * "format" : "strict_date_optional_time", + * "boost" : 0.0 + * } + * } + * }, + * "source" : "date_trunc(1 day, date)@2:25" + * } + * }, tags=[1697760000000]], QueryBuilderAndTags[query={ + * "esql_single_value" : { + * "field" : "date", + * "next" : { + * "range" : { + * "date" : { + * "gte" : "2023-10-21T00:00:00.000Z", + * "lt" : "2023-10-22T00:00:00.000Z", + * "time_zone" : "Z", + * "format" : "strict_date_optional_time", + * "boost" : 0.0 + * } + * } + * }, + * "source" : "date_trunc(1 day, date)@2:25" + * } + * }, tags=[1697846400000]], QueryBuilderAndTags[query={ + * "esql_single_value" : { + * "field" : "date", + * "next" : { + * "range" : { + * "date" : { + * "gte" : "2023-10-22T00:00:00.000Z", + * "time_zone" : "Z", + * "format" : "strict_date_optional_time", + * "boost" : 0.0 + * } + * } + * }, + * "source" : "date_trunc(1 day, date)@2:25" + * } + * }, tags=[1697932800000]], QueryBuilderAndTags[query={ + * "bool" : { + * "must_not" : [ + * { + * "exists" : { + * "field" : "date", + * "boost" : 0.0 + * } + * } + * ], + * "boost" : 1.0 + * } + * }, tags=[null]]] + * + * 2. Aggregation with {@code DateTrunc} and the other pushdown functions + * When there are other functions that can also be pushed down to Lucene, this rule combines the main query with the {@code RoundTo} + * ranges to create a list of {@code QueryBuilderAndTags}. The main query is then applied to each query leg. + * | WHERE keyword : "keyword" + * | STATS COUNT(*) BY d = DATE_TRUNC(1 day, date) + * becomes + * | EVAL d = ROUND_TO(hire_date, 1697760000000, 1697846400000, 1697932800000) + * | STATS COUNT(*) BY d + * becomes + * [QueryBuilderAndTags[query={ + * "bool" : { + * "filter" : [ + * { + * "match" : { + * "keyword" : { + * "query" : "keyword", + * "lenient" : true + * } + * } + * }, + * { + * "esql_single_value" : { + * "field" : "date", + * "next" : { + * "range" : { + * "date" : { + * "lt" : "2023-10-21T00:00:00.000Z", + * "time_zone" : "Z", + * "format" : "strict_date_optional_time", + * "boost" : 0.0 + * } + * } + * }, + * "source" : "date_trunc(1 day, date)@3:25" + * } + * } + * ], + * "boost" : 1.0 + * } + * }, tags=[1697760000000]], QueryBuilderAndTags[query={ + * "bool" : { + * "filter" : [ + * { + * "match" : { + * "keyword" : { + * "query" : "keyword", + * "lenient" : true + * } + * } + * }, + * { + * "esql_single_value" : { + * "field" : "date", + * "next" : { + * "range" : { + * "date" : { + * "gte" : "2023-10-21T00:00:00.000Z", + * "lt" : "2023-10-22T00:00:00.000Z", + * "time_zone" : "Z", + * "format" : "strict_date_optional_time", + * "boost" : 0.0 + * } + * } + * }, + * "source" : "date_trunc(1 day, date)@3:25" + * } + * } + * ], + * "boost" : 1.0 + * } + * }, tags=[1697846400000]], QueryBuilderAndTags[query={ + * "bool" : { + * "filter" : [ + * { + * "match" : { + * "keyword" : { + * "query" : "keyword", + * "lenient" : true + * } + * } + * }, + * { + * "esql_single_value" : { + * "field" : "date", + * "next" : { + * "range" : { + * "date" : { + * "gte" : "2023-10-22T00:00:00.000Z", + * "time_zone" : "Z", + * "format" : "strict_date_optional_time", + * "boost" : 0.0 + * } + * } + * }, + * "source" : "date_trunc(1 day, date)@3:25" + * } + * } + * ], + * "boost" : 1.0 + * } + * }, tags=[1697932800000]], QueryBuilderAndTags[query={ + * "bool" : { + * "filter" : [ + * { + * "match" : { + * "keyword" : { + * "query" : "keyword", + * "lenient" : true + * } + * } + * }, + * { + * "bool" : { + * "must_not" : [ + * { + * "exists" : { + * "field" : "date", + * "boost" : 0.0 + * } + * } + * ], + * "boost" : 0.0 + * } + * } + * ], + * "boost" : 1.0 + * } + * }, tags=[null]]] + * + * 3. Aggregation with {@code DateTrunc} and predicates on {@code DateTrunc} field + * When there are range predicates on the {@code RoundTo} field that can also be pushed down to Lucene, this rule combines the range + * query in the main query with the {@code RoundTo} ranges to create a list of {@code QueryBuilderAndTags}. The main query is then + * applied to each query leg. + * | WHERE date >= "2023-10-20" and date <= "2023-10-24" + * | STATS COUNT(*) BY d = DATE_TRUNC(1 day, date) + * becomes + * | WHERE date >= "2023-10-20" and date <= "2023-10-24" + * | EVAL d = ROUND_TO(hire_date, 1697760000000, 1697846400000, 1697932800000) + * | STATS COUNT(*) BY d + * becomes + * [QueryBuilderAndTags[query={ + * "bool" : { + * "filter" : [ + * { + * "esql_single_value" : { + * "field" : "date", + * "next" : { + * "range" : { + * "date" : { + * "gt" : "2023-10-19T00:00:00.000Z", + * "lt" : "2023-10-24T00:00:00.000Z", + * "time_zone" : "Z", + * "format" : "strict_date_optional_time", + * "boost" : 0.0 + * } + * } + * }, + * "source" : "date > \"2023-10-19T00:00:00.000Z\"@2:9" + * } + * }, + * { + * "esql_single_value" : { + * "field" : "date", + * "next" : { + * "range" : { + * "date" : { + * "lt" : "2023-10-21T00:00:00.000Z", + * "time_zone" : "Z", + * "format" : "strict_date_optional_time", + * "boost" : 0.0 + * } + * } + * }, + * "source" : "date_trunc(1 day, date)@3:25" + * } + * } + * ], + * "boost" : 1.0 + * } + * }, tags=[1697760000000]], QueryBuilderAndTags[query={ + * "esql_single_value" : { + * "field" : "date", + * "next" : { + * "range" : { + * "date" : { + * "gte" : "2023-10-21T00:00:00.000Z", + * "lt" : "2023-10-22T00:00:00.000Z", + * "time_zone" : "Z", + * "format" : "strict_date_optional_time", + * "boost" : 0.0 + * } + * } + * }, + * "source" : "date_trunc(1 day, date)@3:25" + * } + * }, tags=[1697846400000]], QueryBuilderAndTags[query={ + * "bool" : { + * "filter" : [ + * { + * "esql_single_value" : { + * "field" : "date", + * "next" : { + * "range" : { + * "date" : { + * "gt" : "2023-10-19T00:00:00.000Z", + * "lt" : "2023-10-24T00:00:00.000Z", + * "time_zone" : "Z", + * "format" : "strict_date_optional_time", + * "boost" : 0.0 + * } + * } + * }, + * "source" : "date > \"2023-10-19T00:00:00.000Z\"@2:9" + * } + * }, + * { + * "esql_single_value" : { + * "field" : "date", + * "next" : { + * "range" : { + * "date" : { + * "gte" : "2023-10-22T00:00:00.000Z", + * "time_zone" : "Z", + * "format" : "strict_date_optional_time", + * "boost" : 0.0 + * } + * } + * }, + * "source" : "date_trunc(1 day, date)@3:25" + * } + * } + * ], + * "boost" : 1.0 + * } + * }, tags=[1697932800000]]] + * There are some restrictions: + * 1. Tags are not supported by {@code LuceneTopNSourceOperator}, if the sort is pushed down to Lucene, this rewrite does not apply. + * 2. Tags are not supported by {@code TimeSeriesSourceOperator}, this rewrite does not apply to timeseries indices. + * 3. Tags are not supported by {@code LuceneCountOperator}, this rewrite does not apply to {@code EsStatsQueryExec}, count with grouping + * is not supported by {@code EsStatsQueryExec} today. + */ +public class ReplaceRoundToWithQueryAndTags extends PhysicalOptimizerRules.ParameterizedOptimizerRule< + EvalExec, + LocalPhysicalOptimizerContext> { + + @Override + protected PhysicalPlan rule(EvalExec evalExec, LocalPhysicalOptimizerContext ctx) { + PhysicalPlan plan = evalExec; + // TimeSeriesSourceOperator and LuceneTopNSourceOperator do not support QueryAndTags, skip them + if (evalExec.child() instanceof EsQueryExec queryExec && queryExec.canSubstituteRoundToWithQueryBuilderAndTags()) { + // Look for RoundTo and plan the push down for it. + List roundTos = evalExec.fields() + .stream() + .map(Alias::child) + .filter(RoundTo.class::isInstance) + .map(RoundTo.class::cast) + .toList(); + // It is not clear how to push down multiple RoundTos, dealing with multiple RoundTos is out of the scope of this PR. + if (roundTos.size() == 1) { + plan = planRoundTo(roundTos.get(0), evalExec, queryExec, ctx); + } + } + return plan; + } + + /** + * Rewrite the {@code RoundTo} to a list of {@code QueryBuilderAndTags} as input to {@code EsPhysicalOperationProviders}. + */ + private static PhysicalPlan planRoundTo(RoundTo roundTo, EvalExec evalExec, EsQueryExec queryExec, LocalPhysicalOptimizerContext ctx) { + // Usually EsQueryExec has only one QueryBuilder, one Lucene query, without RoundTo push down. + // If the RoundTo can be pushed down, a list of QueryBuilders with tags will be added into EsQueryExec, and it will be sent to + // EsPhysicalOperationProviders.sourcePhysicalOperation to create a list of LuceneSliceQueue.QueryAndTags + List queryBuilderAndTags = queryBuilderAndTags(roundTo, queryExec, ctx); + if (queryBuilderAndTags == null || queryBuilderAndTags.isEmpty()) { + return evalExec; + } + + FieldAttribute fieldAttribute = (FieldAttribute) roundTo.field(); + String tagFieldName = Attribute.rawTemporaryName( + // $$fieldName$round_to$dateType + fieldAttribute.fieldName().string(), + "round_to", + roundTo.field().dataType().typeName() + ); + FieldAttribute tagField = new FieldAttribute( + roundTo.source(), + tagFieldName, + new EsField(tagFieldName, roundTo.dataType(), Map.of(), false) + ); + // Add new tag field to attributes/output + List newAttributes = new ArrayList<>(queryExec.attrs()); + newAttributes.add(tagField); + + // create a new EsQueryExec with newAttributes/output and queryBuilderAndTags + EsQueryExec queryExecWithTags = new EsQueryExec( + queryExec.source(), + queryExec.indexPattern(), + queryExec.indexMode(), + queryExec.indexNameWithModes(), + newAttributes, + queryExec.limit(), + queryExec.sorts(), + queryExec.estimatedRowSize(), + queryBuilderAndTags + ); + + // Replace RoundTo with new tag field in EvalExec + List updatedFields = evalExec.fields() + .stream() + .map(alias -> alias.child() instanceof RoundTo ? alias.replaceChild(tagField) : alias) + .toList(); + + return new EvalExec(evalExec.source(), queryExecWithTags, updatedFields); + } + + private static List queryBuilderAndTags( + RoundTo roundTo, + EsQueryExec queryExec, + LocalPhysicalOptimizerContext ctx + ) { + LucenePushdownPredicates pushdownPredicates = LucenePushdownPredicates.from(ctx.searchStats(), ctx.flags()); + Expression field = roundTo.field(); + if (pushdownPredicates.isPushableFieldAttribute(field) == false) { + return null; + } + List roundingPoints = roundTo.points(); + int count = roundingPoints.size(); + DataType dataType = roundTo.dataType(); + // sort rounding points before building range queries on them + List points = resolveRoundingPoints(dataType, roundingPoints); + if (points.size() != count || points.isEmpty()) { + return null; + } + List queries = new ArrayList<>(count); + + Source source = roundTo.source(); + Queries.Clause clause = queryExec.hasScoring() ? Queries.Clause.MUST : Queries.Clause.FILTER; + Number tag = points.get(0); + if (points.size() == 1) { // if there is only one rounding point, just tag the main query + EsQueryExec.QueryBuilderAndTags queryBuilderAndTags = tagOnlyBucket(queryExec, tag); + queries.add(queryBuilderAndTags); + } else { + Number lower = null; + Number upper = null; + ZoneId zoneId = ctx.configuration().zoneId(); + for (int i = 1; i < count; i++) { + upper = points.get(i); + // build predicates and range queries for RoundTo ranges + queries.add(rangeBucket(source, field, dataType, lower, upper, tag, zoneId, queryExec, pushdownPredicates, clause)); + lower = upper; + tag = upper; + } + // build the last/gte bucket + queries.add(rangeBucket(source, field, dataType, lower, null, lower, zoneId, queryExec, pushdownPredicates, clause)); + } + // if the field is nullable, build bucket with null tag + if (isRoundToFieldNullable(queryExec.query(), field)) { + queries.add(nullBucket(source, field, queryExec, pushdownPredicates, clause)); + } + return queries; + } + + /** + * Sort the rounding points in ascending order. + */ + private static List resolveRoundingPoints(DataType dataType, List roundingPoints) { + return switch (dataType) { + case LONG, DATETIME, DATE_NANOS -> sortedLongRoundingPoints(roundingPoints); + case INTEGER -> sortedIntRoundingPoints(roundingPoints); + case DOUBLE -> sortedDoubleRoundingPoints(roundingPoints); + default -> List.of(); + }; + } + + private static List sortedLongRoundingPoints(List roundingPoints) { + List points = new ArrayList<>(roundingPoints.size()); + for (Expression e : roundingPoints) { + if (e instanceof Literal l && l.value() instanceof Number n) { + points.add(safeToLong(n)); + } + } + Collections.sort(points); + return points; + } + + private static List sortedIntRoundingPoints(List roundingPoints) { + List points = new ArrayList<>(roundingPoints.size()); + for (Expression e : roundingPoints) { + if (e instanceof Literal l) { + points.add((Integer) l.value()); + } + } + Collections.sort(points); + return points; + } + + private static List sortedDoubleRoundingPoints(List roundingPoints) { + List points = new ArrayList<>(roundingPoints.size()); + for (Expression e : roundingPoints) { + if (e instanceof Literal l) { + points.add((Double) l.value()); + } + } + Collections.sort(points); + return points; + } + + /** + * Create a GTE, LT or range expression, and build the corresponding range queries. + */ + private static Expression createRangeExpression( + Source source, + Expression field, + DataType dataType, + Object lower, + Object upper, + ZoneId zoneId + ) { + Literal lowerValue = new Literal(source, lower, dataType); + Literal upperValue = new Literal(source, upper, dataType); + if (lower == null) { + return new LessThan(source, field, upperValue, zoneId); + } else if (upper == null) { + return new GreaterThanOrEqual(source, field, lowerValue, zoneId); + } else { + // lower and upper should not be both null, should an assert be added here + return new Range(source, field, lowerValue, true, upperValue, false, dataType.isDate() ? zoneId : null); + } + } + + /** + * There is only one rounding point, just attach the tag/rounding point to the main query + */ + private static EsQueryExec.QueryBuilderAndTags tagOnlyBucket(EsQueryExec queryExec, Object tag) { + return new EsQueryExec.QueryBuilderAndTags(queryExec.query(), List.of(tag)); + } + + private static EsQueryExec.QueryBuilderAndTags nullBucket( + Source source, + Expression field, + EsQueryExec queryExec, + LucenePushdownPredicates pushdownPredicates, + Queries.Clause clause + ) { + IsNull isNull = new IsNull(source, field); + List nullTags = new ArrayList<>(1); + nullTags.add(null); + return buildCombinedQueryAndTags(queryExec, pushdownPredicates, isNull, clause, nullTags); + } + + private static EsQueryExec.QueryBuilderAndTags rangeBucket( + Source source, + Expression field, + DataType dataType, + Object lower, + Object upper, + Object tag, + ZoneId zoneId, + EsQueryExec queryExec, + LucenePushdownPredicates pushdownPredicates, + Queries.Clause clause + ) { + Expression range = createRangeExpression(source, field, dataType, lower, upper, zoneId); + return buildCombinedQueryAndTags(queryExec, pushdownPredicates, range, clause, List.of(tag)); + } + + private static EsQueryExec.QueryBuilderAndTags buildCombinedQueryAndTags( + EsQueryExec queryExec, + LucenePushdownPredicates pushdownPredicates, + Expression expression, + Queries.Clause clause, + List tags + ) { + Query queryDSL = TRANSLATOR_HANDLER.asQuery(pushdownPredicates, expression); + QueryBuilder mainQuery = queryExec.query(); + QueryBuilder newQuery = queryDSL.toQueryBuilder(); + QueryBuilder combinedQuery; + if (mainQuery == null) { + combinedQuery = Queries.combine(clause, List.of(newQuery)); + } else { + combinedQuery = Queries.combine(clause, List.of(mainQuery, newQuery)); + if (newQuery instanceof SingleValueQuery.Builder newSingleValueQueryBuilder + && newSingleValueQueryBuilder.next() instanceof RangeQueryBuilder roundToRangeQueryBuilder) { + if (mainQuery instanceof SingleValueQuery.Builder mainSingleValueQueryBuilder) { + // mainQuery has only one subQuery, if it is a superset of roundToRangeQueryBuilder, replace it + QueryBuilder main = mainSingleValueQueryBuilder.next(); + if (main instanceof RangeQueryBuilder mainRange) { + if (replaceRangeQueryInMainQuery( + mainRange.from(), + mainRange.to(), + mainRange.includeLower(), + mainRange.format(), + roundToRangeQueryBuilder.from(), + roundToRangeQueryBuilder.to(), + roundToRangeQueryBuilder.format() + )) { + combinedQuery = Queries.combine(clause, List.of(newQuery)); + } + } + } else if (mainQuery instanceof BoolQueryBuilder boolQueryBuilder) { + // try merging the subQueries of mainQuery with roundToRangeQueryBuilder + combinedQuery = mergeRangeQueriesOnRoundToField(clause, boolQueryBuilder, newSingleValueQueryBuilder); + } + } + } + return new EsQueryExec.QueryBuilderAndTags(combinedQuery, tags); + } + + /** + * Look for range queries on the RoundTo field in the main query, if it already covers the range of the roundTo range query, remove it + * from the main query. + */ + private static QueryBuilder mergeRangeQueriesOnRoundToField( + Queries.Clause clause, + BoolQueryBuilder originalBoolQueryBuilder, + SingleValueQuery.Builder roundToRangeQueryBuilder + ) { + BoolQueryBuilder newBoolQueryBuilder = new BoolQueryBuilder(); + for (QueryBuilder q : originalBoolQueryBuilder.mustNot()) { + newBoolQueryBuilder.mustNot(q); + } + for (QueryBuilder q : originalBoolQueryBuilder.should()) { + newBoolQueryBuilder.should(q); + } + // Update the range query in the bool query if it is a superset of the range query in the roundTo + boolean mainQueryModified = false; + for (QueryBuilder queryBuilder : originalBoolQueryBuilder.must()) { + if (queryBuilder instanceof SingleValueQuery.Builder mainSingleValueQueryBuilder + && mainSingleValueQueryBuilder.next() instanceof RangeQueryBuilder main + && roundToRangeQueryBuilder instanceof SingleValueQuery.Builder roundToSingleValueQueryBuilder + && roundToSingleValueQueryBuilder.next() instanceof RangeQueryBuilder roundTo + && main.fieldName().equalsIgnoreCase(roundTo.fieldName())) { + if (replaceRangeQueryInMainQuery( + main.from(), + main.to(), + main.includeLower(), + main.format(), + roundTo.from(), + roundTo.to(), + roundTo.format() + )) { + // The range query from main queries covers the range of the roundTo range, replace it with roundTo range query + newBoolQueryBuilder.must(roundToSingleValueQueryBuilder); + mainQueryModified = true; + } else { + newBoolQueryBuilder.must(queryBuilder); + } + } else { + newBoolQueryBuilder.must(queryBuilder); + } + } + + // TODO filter should be processed similarly as must, how to validate this code path? + for (QueryBuilder queryBuilder : originalBoolQueryBuilder.filter()) { + if (queryBuilder instanceof SingleValueQuery.Builder mainSingleValueQueryBuilder + && mainSingleValueQueryBuilder.next() instanceof RangeQueryBuilder main + && roundToRangeQueryBuilder instanceof SingleValueQuery.Builder roundToSingleValueQueryBuilder + && roundToSingleValueQueryBuilder.next() instanceof RangeQueryBuilder roundTo + && main.fieldName().equalsIgnoreCase(roundTo.fieldName())) { + // If the range query from main queries covers the range of the roundTo range, remove it from the main queries + if (replaceRangeQueryInMainQuery( + main.from(), + main.to(), + main.includeLower(), + main.format(), + roundTo.from(), + roundTo.to(), + roundTo.format() + )) { // update the range query in the main query with the lower/upper/includeLower/includeUpper from the roundToQuery + newBoolQueryBuilder.filter(roundToSingleValueQueryBuilder); + mainQueryModified = true; + } else { + newBoolQueryBuilder.filter(queryBuilder); + } + } else { + newBoolQueryBuilder.filter(queryBuilder); + } + } + return mainQueryModified + ? newBoolQueryBuilder + : Queries.combine(clause, List.of(originalBoolQueryBuilder, roundToRangeQueryBuilder)); + } + + private static boolean rangeQueriesOnTheSameField(QueryBuilder mainQueryBuilder, QueryBuilder roundToQueryBuilder) { + return mainQueryBuilder instanceof SingleValueQuery.Builder mainSingleValueQueryBuilder + && mainSingleValueQueryBuilder.next() instanceof RangeQueryBuilder main + && roundToQueryBuilder instanceof SingleValueQuery.Builder roundToSingleValueQueryBuilder + && roundToSingleValueQueryBuilder.next() instanceof RangeQueryBuilder roundTo + && main.fieldName().equalsIgnoreCase(roundTo.fieldName()); + } + + /** + * Update the range query in the main query if its range is a superset of the range in the roundTo range query + */ + private static boolean replaceRangeQueryInMainQuery( + Object mainLower, + Object mainUpper, + boolean mainIncludeLower, + String mainFormat, + Object roundToLower, + Object roundToUpper, + String roundToFormat + ) { + boolean isNumeric = mainFormat == null && roundToFormat == null; + boolean isDate = false; + boolean isDateNanos = false; + + if (isNumeric == false) { // check if this is a date or date_nanos + isDate = mainFormat != null + && mainFormat.equalsIgnoreCase(DEFAULT_DATE_TIME_FORMATTER.pattern()) + && roundToFormat != null + && roundToFormat.equalsIgnoreCase(DEFAULT_DATE_TIME_FORMATTER.pattern()); + isDateNanos = mainFormat != null + && mainFormat.equalsIgnoreCase(DEFAULT_DATE_NANOS_FORMATTER.pattern()) + && roundToFormat != null + && roundToFormat.equalsIgnoreCase(DEFAULT_DATE_NANOS_FORMATTER.pattern()); + } + + if ((isNumeric || isDate || isDateNanos) == false) { + return false; + } + + // check lower bound + if (mainLower != null) { // main has lower bound > or >= + if (roundToLower == null) { // roundTo is < + return false; + } else { // roundTo has >=, it may also have upper bound + int cmp; + if (isNumeric && roundToLower instanceof Number roundTo && mainLower instanceof Number main) { + cmp = compare(roundTo, main); + } else if (roundToLower instanceof String roundTo && mainLower instanceof String main) { + cmp = compare(roundTo, main, isDate ? DATETIME : DATE_NANOS); + } else { + return false; + } + if (cmp < 0) { + return false; + } + if (cmp == 0 && mainIncludeLower == false) { // lower bound is the same, roundTo has >=, main has > + return false; + } + } + } + + // check upper bound + if (mainUpper != null) { // main has < or <= + if (roundToUpper == null) { // roundTo is >= + return false; + } else { // roundTo has <, it may also have lower bound + int cmp; + if (isNumeric && roundToUpper instanceof Number roundTo && mainUpper instanceof Number main) { + cmp = compare(roundTo, main); + } else if (roundToUpper instanceof String roundTo && mainUpper instanceof String main) { + cmp = compare(roundTo, main, isDate ? DATETIME : DATE_NANOS); + } else { + return false; + } + return cmp <= 0; + } + } + + return true; + } + + private static int compare(Number left, Number right) { + double diff = left.doubleValue() - right.doubleValue(); + if (diff < 0) return -1; + if (diff > 0) return 1; + return 0; + } + + private static int compare(String left, String right, DataType dataType) { + Long leftValue = dataType == DATETIME ? dateTimeToLong(left) : dateNanosToLong(left); + Long rightValue = dataType == DATETIME ? dateTimeToLong(right) : dateNanosToLong(right); + long diff = leftValue - rightValue; + if (diff < 0) return -1; + if (diff > 0) return 1; + return 0; + } + + private static boolean isRoundToFieldNullable(QueryBuilder mainQuery, Expression roundToField) { + String roundToFieldName = Expressions.name(roundToField); + if (roundToFieldName == null) { // cannot find the field name ? treat it as nullable + return true; + } + // if there is range or isNotNull predicate on the field, this field is not nullable + if (hasRangeOrIsNotNull(mainQuery, roundToFieldName)) { + return false; + } + // check the subQueries of bool query + if (mainQuery instanceof BoolQueryBuilder boolQueryBuilder) { + List mainQueries = mainQueries(boolQueryBuilder); + if (mainQueries.isEmpty() == false) { + for (QueryBuilder queryBuilder : mainQueries) { + // if there is range or isNotNull predicate on the field, this field is not nullable + if (hasRangeOrIsNotNull(queryBuilder, roundToFieldName)) { + return false; + } + } + } + } + return true; + } + + private static boolean hasRangeOrIsNotNull(QueryBuilder queryBuilder, String fieldName) { + // Check if there is range or isNotNull predicate on the field + return ((queryBuilder instanceof SingleValueQuery.Builder singleValueQueryBuilder + && singleValueQueryBuilder.next() instanceof RangeQueryBuilder rangeQueryBuilder + && rangeQueryBuilder.fieldName().equals(fieldName)) + || (queryBuilder instanceof ExistsQueryBuilder existsQueryBuilder && existsQueryBuilder.fieldName().equals(fieldName))); + } + + /** + * Extract subqueries from the main query according to the clause types. + * SHOULD/OR is not supported, as only boolean term is considered. + * MUST_NOT is not supported as it does not represent a range or not null predicate on a field. + */ + private static List mainQueries(BoolQueryBuilder boolQueryBuilder) { + List mainQueries = new ArrayList<>(); + mainQueries.addAll(boolQueryBuilder.must()); + mainQueries.addAll(boolQueryBuilder.filter()); + return mainQueries; + } +} diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plan/PlanWritables.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plan/PlanWritables.java index 9312f2abdf509..03c9957754b89 100644 --- a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plan/PlanWritables.java +++ b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plan/PlanWritables.java @@ -36,7 +36,6 @@ import org.elasticsearch.xpack.esql.plan.physical.AggregateExec; import org.elasticsearch.xpack.esql.plan.physical.DissectExec; import org.elasticsearch.xpack.esql.plan.physical.EnrichExec; -import org.elasticsearch.xpack.esql.plan.physical.EsQueryExec; import org.elasticsearch.xpack.esql.plan.physical.EsSourceExec; import org.elasticsearch.xpack.esql.plan.physical.EvalExec; import org.elasticsearch.xpack.esql.plan.physical.ExchangeExec; @@ -105,7 +104,6 @@ public static List physical() { CompletionExec.ENTRY, DissectExec.ENTRY, EnrichExec.ENTRY, - EsQueryExec.ENTRY, EsSourceExec.ENTRY, EvalExec.ENTRY, ExchangeExec.ENTRY, diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plan/physical/EsQueryExec.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plan/physical/EsQueryExec.java index 2e74c7153f77e..5800fee09eccc 100644 --- a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plan/physical/EsQueryExec.java +++ b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plan/physical/EsQueryExec.java @@ -7,8 +7,6 @@ package org.elasticsearch.xpack.esql.plan.physical; -import org.elasticsearch.common.Strings; -import org.elasticsearch.common.io.stream.NamedWriteableRegistry; import org.elasticsearch.common.io.stream.StreamInput; import org.elasticsearch.common.io.stream.StreamOutput; import org.elasticsearch.index.IndexMode; @@ -28,32 +26,19 @@ import org.elasticsearch.xpack.esql.core.type.DataType; import org.elasticsearch.xpack.esql.core.type.EsField; import org.elasticsearch.xpack.esql.expression.Order; -import org.elasticsearch.xpack.esql.index.EsIndex; -import org.elasticsearch.xpack.esql.io.stream.PlanStreamInput; -import org.elasticsearch.xpack.esql.plan.logical.EsRelation; import java.io.IOException; import java.util.List; import java.util.Map; import java.util.Objects; -import static org.elasticsearch.TransportVersions.ESQL_SKIP_ES_INDEX_SERIALIZATION; - public class EsQueryExec extends LeafExec implements EstimatesRowSize { - public static final NamedWriteableRegistry.Entry ENTRY = new NamedWriteableRegistry.Entry( - PhysicalPlan.class, - "EsQueryExec", - EsQueryExec::readFrom - ); - public static final EsField DOC_ID_FIELD = new EsField("_doc", DataType.DOC_DATA_TYPE, Map.of(), false); - public static final List NO_SORTS = List.of(); // only exists to mimic older serialization, but we no longer serialize sorts private final String indexPattern; private final IndexMode indexMode; private final Map indexNameWithModes; private final List attrs; - private final QueryBuilder query; private final Expression limit; private final List sorts; @@ -63,6 +48,12 @@ public class EsQueryExec extends LeafExec implements EstimatesRowSize { */ private final Integer estimatedRowSize; + /** + * queryBuilderAndTags is build on data node from the {@code RoundTo} function by {@code PushRoundToToSource} rule. + * It will be used by {@code sourcePhysicalOperation} to create {@code LuceneSliceQueue.QueryAndTags} + */ + private final List queryBuilderAndTags; + public interface Sort { SortBuilder sortBuilder(); @@ -112,6 +103,8 @@ public FieldAttribute field() { } } + public record QueryBuilderAndTags(QueryBuilder query, List tags) {}; + public EsQueryExec( Source source, String indexPattern, @@ -120,7 +113,17 @@ public EsQueryExec( List attributes, QueryBuilder query ) { - this(source, indexPattern, indexMode, indexNameWithModes, attributes, query, null, null, null); + this( + source, + indexPattern, + indexMode, + indexNameWithModes, + attributes, + null, + null, + null, + List.of(new QueryBuilderAndTags(query, List.of())) + ); } public EsQueryExec( @@ -129,76 +132,31 @@ public EsQueryExec( IndexMode indexMode, Map indexNameWithModes, List attrs, - QueryBuilder query, Expression limit, List sorts, - Integer estimatedRowSize + Integer estimatedRowSize, + List queryBuilderAndTags ) { super(source); this.indexPattern = indexPattern; this.indexMode = indexMode; this.indexNameWithModes = indexNameWithModes; this.attrs = attrs; - this.query = query; this.limit = limit; this.sorts = sorts; this.estimatedRowSize = estimatedRowSize; - } - - /** - * The matching constructor is used during physical plan optimization and needs valid sorts. But we no longer serialize sorts. - * If this cluster node is talking to an older instance it might receive a plan with sorts, but it will ignore them. - */ - private static EsQueryExec readFrom(StreamInput in) throws IOException { - var source = Source.readFrom((PlanStreamInput) in); - String indexPattern; - Map indexNameWithModes; - if (in.getTransportVersion().onOrAfter(ESQL_SKIP_ES_INDEX_SERIALIZATION)) { - indexPattern = in.readString(); - indexNameWithModes = in.readMap(IndexMode::readFrom); - } else { - var index = EsIndex.readFrom(in); - indexPattern = index.name(); - indexNameWithModes = index.indexNameWithModes(); - } - var indexMode = EsRelation.readIndexMode(in); - var attrs = in.readNamedWriteableCollectionAsList(Attribute.class); - var query = in.readOptionalNamedWriteable(QueryBuilder.class); - var limit = in.readOptionalNamedWriteable(Expression.class); - in.readOptionalCollectionAsList(EsQueryExec::readSort); - var rowSize = in.readOptionalVInt(); - // Ignore sorts from the old serialization format - return new EsQueryExec(source, indexPattern, indexMode, indexNameWithModes, attrs, query, limit, NO_SORTS, rowSize); - } - - private static Sort readSort(StreamInput in) throws IOException { - return FieldSort.readFrom(in); - } - - private static void writeSort(StreamOutput out, Sort sort) { - throw new IllegalStateException("sorts are no longer serialized"); + // cannot keep the ctor with QueryBuilder as it has the same number of arguments as this ctor, EsqlNodeSubclassTests will fail + this.queryBuilderAndTags = queryBuilderAndTags; } @Override public void writeTo(StreamOutput out) throws IOException { - Source.EMPTY.writeTo(out); - if (out.getTransportVersion().onOrAfter(ESQL_SKIP_ES_INDEX_SERIALIZATION)) { - out.writeString(indexPattern); - out.writeMap(indexNameWithModes, (o, v) -> IndexMode.writeTo(v, out)); - } else { - new EsIndex(indexPattern, Map.of(), indexNameWithModes).writeTo(out); - } - EsRelation.writeIndexMode(out, indexMode()); - out.writeNamedWriteableCollection(output()); - out.writeOptionalNamedWriteable(query()); - out.writeOptionalNamedWriteable(limit()); - out.writeOptionalCollection(NO_SORTS, EsQueryExec::writeSort); - out.writeOptionalVInt(estimatedRowSize()); + throw new UnsupportedOperationException("not serialized"); } @Override public String getWriteableName() { - return ENTRY.name; + throw new UnsupportedOperationException("not serialized"); } public static boolean isSourceAttribute(Attribute attr) { @@ -223,10 +181,10 @@ protected NodeInfo info() { indexMode, indexNameWithModes, attrs, - query, limit, sorts, - estimatedRowSize + estimatedRowSize, + queryBuilderAndTags ); } @@ -242,8 +200,13 @@ public Map indexNameWithModes() { return indexNameWithModes; } + /** + * query is merged into queryBuilderAndTags, keep this method as it is called by too many places. + * If this method is called, the caller looks for the original queryBuilder, before {@code ReplaceRoundToWithQueryAndTags} converts it + * to multiple queries with tags. + */ public QueryBuilder query() { - return query; + return queryWithoutTag(); } @Override @@ -285,13 +248,23 @@ public PhysicalPlan estimateRowSize(State state) { } return Objects.equals(this.estimatedRowSize, size) ? this - : new EsQueryExec(source(), indexPattern, indexMode, indexNameWithModes, attrs, query, limit, sorts, size); + : new EsQueryExec(source(), indexPattern, indexMode, indexNameWithModes, attrs, limit, sorts, size, queryBuilderAndTags); } public EsQueryExec withLimit(Expression limit) { return Objects.equals(this.limit, limit) ? this - : new EsQueryExec(source(), indexPattern, indexMode, indexNameWithModes, attrs, query, limit, sorts, estimatedRowSize); + : new EsQueryExec( + source(), + indexPattern, + indexMode, + indexNameWithModes, + attrs, + limit, + sorts, + estimatedRowSize, + queryBuilderAndTags + ); } public boolean canPushSorts() { @@ -305,18 +278,79 @@ public EsQueryExec withSorts(List sorts) { } return Objects.equals(this.sorts, sorts) ? this - : new EsQueryExec(source(), indexPattern, indexMode, indexNameWithModes, attrs, query, limit, sorts, estimatedRowSize); + : new EsQueryExec( + source(), + indexPattern, + indexMode, + indexNameWithModes, + attrs, + limit, + sorts, + estimatedRowSize, + queryBuilderAndTags + ); } + /** + * query is merged into queryBuilderAndTags, keep this method as it is called by too many places. + * If this method is called, the caller looks for the original queryBuilder, before {@code ReplaceRoundToWithQueryAndTags} converts it + * to multiple queries with tags. + */ public EsQueryExec withQuery(QueryBuilder query) { - return Objects.equals(this.query, query) + QueryBuilder thisQuery = queryWithoutTag(); + return Objects.equals(thisQuery, query) ? this - : new EsQueryExec(source(), indexPattern, indexMode, indexNameWithModes, attrs, query, limit, sorts, estimatedRowSize); + : new EsQueryExec( + source(), + indexPattern, + indexMode, + indexNameWithModes, + attrs, + limit, + sorts, + estimatedRowSize, + List.of(new QueryBuilderAndTags(query, List.of())) + ); + } + + public List queryBuilderAndTags() { + return queryBuilderAndTags; + } + + public boolean canSubstituteRoundToWithQueryBuilderAndTags() { + // TimeSeriesSourceOperator and LuceneTopNSourceOperator do not support QueryAndTags + return indexMode != IndexMode.TIME_SERIES && (sorts == null || sorts.isEmpty()); + } + + /** + * Returns the original queryBuilder before {@code ReplaceRoundToWithQueryAndTags} converts it to multiple queryBuilder with tags. + * If we reach here, the caller is looking for the original query before the rule converts it. If there are multiple queries in + * queryBuilderAndTags or if the single query in queryBuilderAndTags already has a tag, that means + * {@code ReplaceRoundToWithQueryAndTags} already applied to the original query, the original query cannot be retrieved any more, + * exception will be thrown. + */ + private QueryBuilder queryWithoutTag() { + QueryBuilder queryWithoutTag; + if (queryBuilderAndTags == null || queryBuilderAndTags.isEmpty()) { + return null; + } else if (queryBuilderAndTags.size() == 1) { + QueryBuilderAndTags firstQuery = this.queryBuilderAndTags.get(0); + if (firstQuery.tags().isEmpty()) { + queryWithoutTag = firstQuery.query(); + } else { + throw new UnsupportedOperationException("query is converted to query with tags: " + "[" + firstQuery + "]"); + } + } else { + throw new UnsupportedOperationException( + "query is converted to multiple queries and tags: " + "[" + this.queryBuilderAndTags + "]" + ); + } + return queryWithoutTag; } @Override public int hashCode() { - return Objects.hash(indexPattern, indexMode, indexNameWithModes, attrs, query, limit, sorts); + return Objects.hash(indexPattern, indexMode, indexNameWithModes, attrs, limit, sorts, queryBuilderAndTags); } @Override @@ -334,10 +368,10 @@ public boolean equals(Object obj) { && Objects.equals(indexMode, other.indexMode) && Objects.equals(indexNameWithModes, other.indexNameWithModes) && Objects.equals(attrs, other.attrs) - && Objects.equals(query, other.query) && Objects.equals(limit, other.limit) && Objects.equals(sorts, other.sorts) - && Objects.equals(estimatedRowSize, other.estimatedRowSize); + && Objects.equals(estimatedRowSize, other.estimatedRowSize) + && Objects.equals(queryBuilderAndTags, other.queryBuilderAndTags); } @Override @@ -349,9 +383,6 @@ public String nodeString() { + "indexMode[" + indexMode + "], " - + "query[" - + (query != null ? Strings.toString(query, false, true) : "") - + "]" + NodeUtils.limitedToString(attrs) + ", limit[" + (limit != null ? limit.toString() : "") @@ -359,6 +390,8 @@ public String nodeString() { + (sorts != null ? sorts.toString() : "") + "] estimatedRowSize[" + estimatedRowSize + + "] queryBuilderAndTags [" + + (queryBuilderAndTags != null ? queryBuilderAndTags.toString() : "") + "]"; } diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/planner/EsPhysicalOperationProviders.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/planner/EsPhysicalOperationProviders.java index e0b570267899b..117e65f19a92b 100644 --- a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/planner/EsPhysicalOperationProviders.java +++ b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/planner/EsPhysicalOperationProviders.java @@ -258,6 +258,18 @@ public Function List.of(new LuceneSliceQueue.QueryAndTags(shardContexts.get(ctx.index()).toQuery(qb), List.of())); } + public Function> querySupplier( + List queryAndTagsFromEsQueryExec + ) { + return ctx -> queryAndTagsFromEsQueryExec.stream().map(queryBuilderAndTags -> { + QueryBuilder qb = queryBuilderAndTags.query(); + return new LuceneSliceQueue.QueryAndTags( + shardContexts.get(ctx.index()).toQuery(qb == null ? QueryBuilders.matchAllQuery().boost(0.0f) : qb), + queryBuilderAndTags.tags() + ); + }).toList(); + } + @Override public final PhysicalOperation sourcePhysicalOperation(EsQueryExec esQueryExec, LocalExecutionPlannerContext context) { if (esQueryExec.indexMode() == IndexMode.TIME_SERIES) { @@ -277,6 +289,8 @@ public final PhysicalOperation sourcePhysicalOperation(EsQueryExec esQueryExec, for (Sort sort : sorts) { sortBuilders.add(sort.sortBuilder()); } + // LuceneTopNSourceOperator does not support QueryAndTags, if there are multiple queries or if the single query has tags, + // UnsupportedOperationException will be thrown by esQueryExec.query() luceneFactory = new LuceneTopNSourceOperator.Factory( shardContexts, querySupplier(esQueryExec.query()), @@ -290,7 +304,7 @@ public final PhysicalOperation sourcePhysicalOperation(EsQueryExec esQueryExec, } else { luceneFactory = new LuceneSourceOperator.Factory( shardContexts, - querySupplier(esQueryExec.query()), + querySupplier(esQueryExec.queryBuilderAndTags()), context.queryPragmas().dataPartitioning(physicalSettings.defaultDataPartitioning()), context.queryPragmas().taskConcurrency(), context.pageSize(rowEstimatedSize), diff --git a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/optimizer/LocalPhysicalPlanOptimizerTests.java b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/optimizer/LocalPhysicalPlanOptimizerTests.java index cd6371e4d4d5e..6a9d3a2822729 100644 --- a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/optimizer/LocalPhysicalPlanOptimizerTests.java +++ b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/optimizer/LocalPhysicalPlanOptimizerTests.java @@ -173,7 +173,7 @@ public class LocalPhysicalPlanOptimizerTests extends MapperServiceTestCase { public static final String MATCH_OPERATOR_QUERY = "from test | where %s:%s"; public static final String MATCH_FUNCTION_QUERY = "from test | where match(%s, %s)"; - private TestPlannerOptimizer plannerOptimizer; + protected TestPlannerOptimizer plannerOptimizer; private TestPlannerOptimizer plannerOptimizerDateDateNanosUnionTypes; private Analyzer timeSeriesAnalyzer; private final Configuration config; @@ -262,7 +262,7 @@ private Analyzer makeAnalyzer(String mappingFileName, EnrichResolution enrichRes ); } - private Analyzer makeAnalyzer(String mappingFileName) { + protected Analyzer makeAnalyzer(String mappingFileName) { return makeAnalyzer(mappingFileName, new EnrichResolution()); } @@ -2511,7 +2511,7 @@ private boolean isMultiTypeEsField(Expression e) { return e instanceof FieldAttribute fa && fa.field() instanceof MultiTypeEsField; } - private QueryBuilder wrapWithSingleQuery(String query, QueryBuilder inner, String fieldName, Source source) { + protected static QueryBuilder wrapWithSingleQuery(String query, QueryBuilder inner, String fieldName, Source source) { return FilterTests.singleValueQuery(query, inner, fieldName, source); } diff --git a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/optimizer/rules/logical/local/ReplaceDateTruncBucketWithRoundToTests.java b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/optimizer/rules/logical/local/ReplaceDateTruncBucketWithRoundToTests.java index 225cab8330b84..aef6ba3c2abab 100644 --- a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/optimizer/rules/logical/local/ReplaceDateTruncBucketWithRoundToTests.java +++ b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/optimizer/rules/logical/local/ReplaceDateTruncBucketWithRoundToTests.java @@ -8,7 +8,6 @@ package org.elasticsearch.xpack.esql.optimizer.rules.logical.local; import org.elasticsearch.common.logging.LoggerMessageFormat; -import org.elasticsearch.test.junit.annotations.TestLogging; import org.elasticsearch.xpack.esql.EsqlTestUtils; import org.elasticsearch.xpack.esql.core.expression.Alias; import org.elasticsearch.xpack.esql.core.expression.Expression; @@ -34,7 +33,7 @@ import static org.elasticsearch.xpack.esql.EsqlTestUtils.as; import static org.elasticsearch.xpack.esql.core.type.DataType.DATETIME; -@TestLogging(value = "org.elasticsearch.xpack.esql:TRACE", reason = "debug") +//@TestLogging(value = "org.elasticsearch.xpack.esql:TRACE", reason = "debug") public class ReplaceDateTruncBucketWithRoundToTests extends LocalLogicalPlanOptimizerTests { // Key is the predicate, diff --git a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/optimizer/rules/physical/local/PushTopNToSourceTests.java b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/optimizer/rules/physical/local/PushTopNToSourceTests.java index 00f8dfb9aaacc..e91273520eae1 100644 --- a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/optimizer/rules/physical/local/PushTopNToSourceTests.java +++ b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/optimizer/rules/physical/local/PushTopNToSourceTests.java @@ -582,7 +582,17 @@ public TestPhysicalPlanBuilder limit(int limit) { public TopNExec build() { List attributes = new ArrayList<>(fields.values()); - PhysicalPlan child = new EsQueryExec(Source.EMPTY, this.index, indexMode, Map.of(), attributes, null, null, List.of(), 0); + PhysicalPlan child = new EsQueryExec( + Source.EMPTY, + this.index, + indexMode, + Map.of(), + attributes, + null, + List.of(), + 0, + List.of(new EsQueryExec.QueryBuilderAndTags(null, List.of())) + ); if (aliases.isEmpty() == false) { child = new EvalExec(Source.EMPTY, child, aliases); } diff --git a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/optimizer/rules/physical/local/ReplaceRoundToWithQueryAndTagsTests.java b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/optimizer/rules/physical/local/ReplaceRoundToWithQueryAndTagsTests.java new file mode 100644 index 0000000000000..10f52b23e694f --- /dev/null +++ b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/optimizer/rules/physical/local/ReplaceRoundToWithQueryAndTagsTests.java @@ -0,0 +1,746 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the Elastic License + * 2.0; you may not use this file except in compliance with the Elastic License + * 2.0. + */ + +package org.elasticsearch.xpack.esql.optimizer.rules.physical.local; + +import org.elasticsearch.common.logging.LoggerMessageFormat; +import org.elasticsearch.index.query.BoolQueryBuilder; +import org.elasticsearch.index.query.ExistsQueryBuilder; +import org.elasticsearch.index.query.MatchQueryBuilder; +import org.elasticsearch.index.query.QueryBuilder; +import org.elasticsearch.index.query.RangeQueryBuilder; +import org.elasticsearch.xpack.esql.EsqlTestUtils; +import org.elasticsearch.xpack.esql.core.expression.Alias; +import org.elasticsearch.xpack.esql.core.expression.Expression; +import org.elasticsearch.xpack.esql.core.expression.Expressions; +import org.elasticsearch.xpack.esql.core.expression.FieldAttribute; +import org.elasticsearch.xpack.esql.core.expression.NamedExpression; +import org.elasticsearch.xpack.esql.core.tree.Source; +import org.elasticsearch.xpack.esql.core.type.DataType; +import org.elasticsearch.xpack.esql.core.util.Queries; +import org.elasticsearch.xpack.esql.expression.function.scalar.math.RoundTo; +import org.elasticsearch.xpack.esql.optimizer.LocalPhysicalPlanOptimizerTests; +import org.elasticsearch.xpack.esql.plan.physical.AggregateExec; +import org.elasticsearch.xpack.esql.plan.physical.EsQueryExec; +import org.elasticsearch.xpack.esql.plan.physical.EvalExec; +import org.elasticsearch.xpack.esql.plan.physical.ExchangeExec; +import org.elasticsearch.xpack.esql.plan.physical.FieldExtractExec; +import org.elasticsearch.xpack.esql.plan.physical.LimitExec; +import org.elasticsearch.xpack.esql.plan.physical.PhysicalPlan; +import org.elasticsearch.xpack.esql.plan.physical.ProjectExec; +import org.elasticsearch.xpack.esql.plan.physical.TopNExec; +import org.elasticsearch.xpack.esql.querydsl.query.SingleValueQuery; +import org.elasticsearch.xpack.esql.session.Configuration; +import org.elasticsearch.xpack.esql.stats.SearchStats; + +import java.util.ArrayList; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.stream.Collectors; + +import static org.elasticsearch.compute.aggregation.AggregatorMode.FINAL; +import static org.elasticsearch.index.query.QueryBuilders.boolQuery; +import static org.elasticsearch.index.query.QueryBuilders.existsQuery; +import static org.elasticsearch.index.query.QueryBuilders.matchQuery; +import static org.elasticsearch.index.query.QueryBuilders.rangeQuery; +import static org.elasticsearch.index.query.QueryBuilders.termQuery; +import static org.elasticsearch.xpack.esql.EsqlTestUtils.as; +import static org.elasticsearch.xpack.esql.core.util.Queries.Clause.MUST; +import static org.elasticsearch.xpack.esql.type.EsqlDataTypeConverter.DEFAULT_DATE_NANOS_FORMATTER; +import static org.elasticsearch.xpack.esql.type.EsqlDataTypeConverter.DEFAULT_DATE_TIME_FORMATTER; +import static org.elasticsearch.xpack.esql.type.EsqlDataTypeConverter.dateNanosToLong; +import static org.elasticsearch.xpack.esql.type.EsqlDataTypeConverter.dateTimeToLong; +import static org.hamcrest.Matchers.is; + +//@TestLogging(value = "org.elasticsearch.xpack.esql:TRACE", reason = "debug") +public class ReplaceRoundToWithQueryAndTagsTests extends LocalPhysicalPlanOptimizerTests { + + public ReplaceRoundToWithQueryAndTagsTests(String name, Configuration config) { + super(name, config); + } + + private static final List dateHistograms = List.of( + "date_trunc(1 day, date)", + "bucket(date, 1 day)", + "round_to(date, \"2023-10-20\", \"2023-10-21\", \"2023-10-22\", \"2023-10-23\")" + ); + + private static final Map> roundToAllTypes = new HashMap<>( + Map.ofEntries( + Map.entry("byte", List.of(1, 2, 3, 4)), + Map.entry("short", List.of(1, 2, 3, 4)), + Map.entry("integer", List.of(1, 2, 3, 4)), + Map.entry("long", List.of(1697760000000L, 1697846400000L, 1697932800000L, 1698019200000L)), + Map.entry("float", List.of(1.0, 2.0, 3.0, 4.0)), + Map.entry("half_float", List.of(1.0, 2.0, 3.0, 4.0)), + Map.entry("scaled_float", List.of(1.0, 2.0, 3.0, 4.0)), + Map.entry("double", List.of(1.0, 2.0, 3.0, 4.0)), + Map.entry("date", List.of("\"2023-10-20\"::date", "\"2023-10-21\"::date", "\"2023-10-22\"::date", "\"2023-10-23\"::date")), + Map.entry( + "date_nanos", + List.of( + "\"2023-10-20\"::date_nanos", + "\"2023-10-21\"::date_nanos", + "\"2023-10-22\"::date_nanos", + "\"2023-10-23\"::date_nanos" + ) + ) + ) + ); + + private static final Map otherPushDownFunctions = new HashMap<>( + Map.ofEntries( + Map.entry("where keyword == \"keyword\"", termQuery("keyword", "keyword").boost(0)), + Map.entry("where keyword : \"keyword\"", matchQuery("keyword", "keyword").lenient(true)) + ) + ); + + // The date range of SearchStats is from 2023-10-20 to 2023-10-23. + private static final SearchStats searchStats = searchStats(); + + // DateTrunc/Bucket is transformed to RoundTo first and then to QueryAndTags + public void testDateTruncBucketTransformToQueryAndTags() { + for (String dateHistogram : dateHistograms) { + String query = LoggerMessageFormat.format(null, """ + from test + | stats count(*) by x = {} + """, dateHistogram); + PhysicalPlan plan = plannerOptimizer.plan(query, searchStats, makeAnalyzer("mapping-all-types.json")); + + LimitExec limit = as(plan, LimitExec.class); + AggregateExec agg = as(limit.child(), AggregateExec.class); + assertThat(agg.getMode(), is(FINAL)); + List groupings = agg.groupings(); + NamedExpression grouping = as(groupings.get(0), NamedExpression.class); + assertEquals("x", grouping.name()); + assertEquals(DataType.DATETIME, grouping.dataType()); + assertEquals(List.of("count(*)", "x"), Expressions.names(agg.aggregates())); + ExchangeExec exchange = as(agg.child(), ExchangeExec.class); + assertThat(exchange.inBetweenAggs(), is(true)); + agg = as(exchange.child(), AggregateExec.class); + EvalExec eval = as(agg.child(), EvalExec.class); + List aliases = eval.fields(); + assertEquals(1, aliases.size()); + FieldAttribute roundToTag = as(aliases.get(0).child(), FieldAttribute.class); + assertEquals("$$date$round_to$datetime", roundToTag.name()); + EsQueryExec esQueryExec = as(eval.child(), EsQueryExec.class); + List queryBuilderAndTags = esQueryExec.queryBuilderAndTags(); + List expectedQueryBuilderAndTags = expectedQueryBuilderAndTags( + query, + "date", + List.of(), + new Source(2, 24, dateHistogram), + null, + true + ); + verifyQueryAndTags(expectedQueryBuilderAndTags, queryBuilderAndTags); + assertThrows(UnsupportedOperationException.class, esQueryExec::query); + } + } + + // DateTrunc is transformed to RoundTo first but cannot be transformed to QueryAndTags, when the TopN is pushed down to EsQueryExec + public void testDateTruncNotTransformToQueryAndTags() { + for (String dateHistogram : dateHistograms) { + if (dateHistogram.contains("bucket")) { // bucket cannot be used out side of stats + continue; + } + String query = LoggerMessageFormat.format(null, """ + from test + | sort date + | eval x = {} + | keep alias_integer, date, x + | limit 5 + """, dateHistogram); + + PhysicalPlan plan = plannerOptimizer.plan(query, searchStats, makeAnalyzer("mapping-all-types.json")); + + ProjectExec projectExec = as(plan, ProjectExec.class); + TopNExec topNExec = as(projectExec.child(), TopNExec.class); + ExchangeExec exchangeExec = as(topNExec.child(), ExchangeExec.class); + projectExec = as(exchangeExec.child(), ProjectExec.class); + FieldExtractExec fieldExtractExec = as(projectExec.child(), FieldExtractExec.class); + EvalExec evalExec = as(fieldExtractExec.child(), EvalExec.class); + List aliases = evalExec.fields(); + assertEquals(1, aliases.size()); + RoundTo roundTo = as(aliases.get(0).child(), RoundTo.class); + assertEquals(4, roundTo.points().size()); + fieldExtractExec = as(evalExec.child(), FieldExtractExec.class); + EsQueryExec esQueryExec = as(fieldExtractExec.child(), EsQueryExec.class); + List queryBuilderAndTags = esQueryExec.queryBuilderAndTags(); + assertEquals(1, queryBuilderAndTags.size()); + EsQueryExec.QueryBuilderAndTags queryBuilder = queryBuilderAndTags.get(0); + assertNull(queryBuilder.query()); + assertTrue(queryBuilder.tags().isEmpty()); + assertNull(esQueryExec.query()); + } + } + + // RoundTo(all numeric data types) is transformed to QueryAndTags + public void testRoundToTransformToQueryAndTags() { + for (Map.Entry> roundTo : roundToAllTypes.entrySet()) { + String fieldName = roundTo.getKey(); + List roundingPoints = roundTo.getValue(); + String expression = "round_to(" + + fieldName + + ", " + + roundingPoints.stream().map(Object::toString).collect(Collectors.joining(",")) + + ")"; + String query = LoggerMessageFormat.format(null, """ + from test + | stats count(*) by x = {} + """, expression); + PhysicalPlan plan = plannerOptimizer.plan(query, searchStats, makeAnalyzer("mapping-all-types.json")); + + LimitExec limit = as(plan, LimitExec.class); + AggregateExec agg = as(limit.child(), AggregateExec.class); + assertThat(agg.getMode(), is(FINAL)); + List groupings = agg.groupings(); + NamedExpression grouping = as(groupings.get(0), NamedExpression.class); + assertEquals("x", grouping.name()); + assertEquals(List.of("count(*)", "x"), Expressions.names(agg.aggregates())); + ExchangeExec exchange = as(agg.child(), ExchangeExec.class); + assertThat(exchange.inBetweenAggs(), is(true)); + agg = as(exchange.child(), AggregateExec.class); + EvalExec eval = as(agg.child(), EvalExec.class); + List aliases = eval.fields(); + assertEquals(1, aliases.size()); + FieldAttribute roundToTag = as(aliases.get(0).child(), FieldAttribute.class); + assertTrue(roundToTag.name().startsWith("$$" + fieldName + "$round_to$")); + EsQueryExec esQueryExec = as(eval.child(), EsQueryExec.class); + List queryBuilderAndTags = esQueryExec.queryBuilderAndTags(); + List expectedQueryBuilderAndTags = expectedQueryBuilderAndTags( + query, + fieldName, + roundingPoints, + new Source(2, 24, expression), + null, + true + ); + verifyQueryAndTags(expectedQueryBuilderAndTags, queryBuilderAndTags); + assertThrows(UnsupportedOperationException.class, esQueryExec::query); + } + } + + // test if the combine query is generated correctly when there are other functions not on the same field, can be pushed down + public void testDateTruncBucketTransformToQueryAndTagsWithOtherPushdownFunctions() { + for (String dateHistogram : dateHistograms) { + for (Map.Entry otherPushDownFunction : otherPushDownFunctions.entrySet()) { + String predicate = otherPushDownFunction.getKey(); + QueryBuilder qb = otherPushDownFunction.getValue(); + String query = LoggerMessageFormat.format(null, """ + from test + | {} + | stats count(*) by x = {} + """, predicate, dateHistogram); + QueryBuilder mainQueryBuilder = qb instanceof MatchQueryBuilder + ? qb + : wrapWithSingleQuery(query, qb, "keyword", new Source(2, 8, predicate.substring(6))); + + PhysicalPlan plan = plannerOptimizer.plan(query, searchStats, makeAnalyzer("mapping-all-types.json")); + + LimitExec limit = as(plan, LimitExec.class); + AggregateExec agg = as(limit.child(), AggregateExec.class); + assertThat(agg.getMode(), is(FINAL)); + List groupings = agg.groupings(); + NamedExpression grouping = as(groupings.get(0), NamedExpression.class); + assertEquals("x", grouping.name()); + assertEquals(DataType.DATETIME, grouping.dataType()); + assertEquals(List.of("count(*)", "x"), Expressions.names(agg.aggregates())); + ExchangeExec exchange = as(agg.child(), ExchangeExec.class); + assertThat(exchange.inBetweenAggs(), is(true)); + agg = as(exchange.child(), AggregateExec.class); + EvalExec eval = as(agg.child(), EvalExec.class); + List aliases = eval.fields(); + assertEquals(1, aliases.size()); + FieldAttribute roundToTag = as(aliases.get(0).child(), FieldAttribute.class); + assertEquals("$$date$round_to$datetime", roundToTag.name()); + EsQueryExec esQueryExec = as(eval.child(), EsQueryExec.class); + List queryBuilderAndTags = esQueryExec.queryBuilderAndTags(); + List expectedQueryBuilderAndTags = expectedQueryBuilderAndTags( + query, + "date", + List.of(), + new Source(3, 24, dateHistogram), + mainQueryBuilder, + true + ); + verifyQueryAndTags(expectedQueryBuilderAndTags, queryBuilderAndTags); + assertThrows(UnsupportedOperationException.class, esQueryExec::query); + } + } + } + + // test the range queries on the roundTo field are merged into one when possible + public void testMergeRangePredicatesOnRoundToField() { + for (Map.Entry> roundTo : roundToAllTypes.entrySet()) { + String fieldName = roundTo.getKey(); + List roundingPoints = roundTo.getValue(); + String roundToExpression = "round_to(" + + fieldName + + ", " + + roundingPoints.stream().map(Object::toString).collect(Collectors.joining(",")) + + ")"; + + Object lowerInPredicate = null; + Object upperInPredicate = null; + var lower = roundingPoints.get(0); + var upper = roundingPoints.get(3); + if (lower instanceof String && upper instanceof String) { // date or date_nanos + lowerInPredicate = "\"2023-10-19T00:00:00.000Z\""; + upperInPredicate = "\"2023-10-24T00:00:00.000Z\""; + } else if (lower instanceof Integer lowerBound && upper instanceof Integer upperBound) { + lowerInPredicate = lowerBound - 1; + upperInPredicate = upperBound + 1; + } else if (lower instanceof Long lowerBound && upper instanceof Long upperBound) { + lowerInPredicate = lowerBound - 1L; + upperInPredicate = upperBound + 1L; + } else if (lower instanceof Double lowerBound && upper instanceof Double upperBound) { + lowerInPredicate = lowerBound - 1.0; + upperInPredicate = upperBound + 1.0; + } + + assertNotNull(lowerInPredicate); + assertNotNull(upperInPredicate); + + String rangePredicate = fieldName + " > " + lowerInPredicate + " and " + fieldName + " < " + upperInPredicate; + + String query = LoggerMessageFormat.format(null, """ + from test + | where {} + | stats count(*) by x = {} + """, rangePredicate, roundToExpression); + + RangeQueryBuilder rangeQueryBuilder = rangeQuery(fieldName).from( + lowerInPredicate instanceof String s ? s.replaceAll("\"", "") : lowerInPredicate + ) + .to(upperInPredicate instanceof String s ? s.replaceAll("\"", "") : upperInPredicate) + .includeLower(false) + .includeUpper(false) + .boost(0) + .timeZone("Z"); + + if (fieldName.equalsIgnoreCase("date")) { + rangeQueryBuilder.format(DEFAULT_DATE_TIME_FORMATTER.pattern()); + } + + if (fieldName.equalsIgnoreCase("date_nanos")) { + rangeQueryBuilder.format(DEFAULT_DATE_NANOS_FORMATTER.pattern()); + } + + QueryBuilder mainQueryBuilder = wrapWithSingleQuery( + query, + rangeQueryBuilder, + fieldName, + new Source(2, 8, fieldName + " > " + lowerInPredicate) + ); + + PhysicalPlan plan = plannerOptimizer.plan(query, searchStats, makeAnalyzer("mapping-all-types.json")); + + LimitExec limit = as(plan, LimitExec.class); + AggregateExec agg = as(limit.child(), AggregateExec.class); + assertThat(agg.getMode(), is(FINAL)); + List groupings = agg.groupings(); + NamedExpression grouping = as(groupings.get(0), NamedExpression.class); + assertEquals("x", grouping.name()); + assertEquals(List.of("count(*)", "x"), Expressions.names(agg.aggregates())); + ExchangeExec exchange = as(agg.child(), ExchangeExec.class); + assertThat(exchange.inBetweenAggs(), is(true)); + agg = as(exchange.child(), AggregateExec.class); + EvalExec eval = as(agg.child(), EvalExec.class); + List aliases = eval.fields(); + assertEquals(1, aliases.size()); + FieldAttribute roundToTag = as(aliases.get(0).child(), FieldAttribute.class); + assertTrue(roundToTag.name().startsWith("$$" + fieldName + "$round_to$")); + EsQueryExec esQueryExec = as(eval.child(), EsQueryExec.class); + List queryBuilderAndTags = esQueryExec.queryBuilderAndTags(); + List expectedQueryBuilderAndTags = expectedQueryBuilderAndTags( + query, + fieldName, + roundingPoints, + new Source(3, 24, roundToExpression), + mainQueryBuilder, + false + ); + verifyQueryAndTags(expectedQueryBuilderAndTags, queryBuilderAndTags); + assertThrows(UnsupportedOperationException.class, esQueryExec::query); + } + } + + // test the null bucket does not need to be added if the roundTo field is not null + public void testRoundToFieldsIsNotNull() { + for (Map.Entry> roundTo : roundToAllTypes.entrySet()) { + String fieldName = roundTo.getKey(); + List roundingPoints = roundTo.getValue(); + String roundToExpression = "round_to(" + + fieldName + + ", " + + roundingPoints.stream().map(Object::toString).collect(Collectors.joining(",")) + + ")"; + + String query = LoggerMessageFormat.format(null, """ + from test + | where {} is not null + | stats count(*) by x = {} + """, fieldName, roundToExpression); + + ExistsQueryBuilder mainQueryBuilder = existsQuery(fieldName).boost(0); + + PhysicalPlan plan = plannerOptimizer.plan(query, searchStats, makeAnalyzer("mapping-all-types.json")); + + LimitExec limit = as(plan, LimitExec.class); + AggregateExec agg = as(limit.child(), AggregateExec.class); + assertThat(agg.getMode(), is(FINAL)); + List groupings = agg.groupings(); + NamedExpression grouping = as(groupings.get(0), NamedExpression.class); + assertEquals("x", grouping.name()); + assertEquals(List.of("count(*)", "x"), Expressions.names(agg.aggregates())); + ExchangeExec exchange = as(agg.child(), ExchangeExec.class); + assertThat(exchange.inBetweenAggs(), is(true)); + agg = as(exchange.child(), AggregateExec.class); + EvalExec eval = as(agg.child(), EvalExec.class); + List aliases = eval.fields(); + assertEquals(1, aliases.size()); + FieldAttribute roundToTag = as(aliases.get(0).child(), FieldAttribute.class); + assertTrue(roundToTag.name().startsWith("$$" + fieldName + "$round_to$")); + EsQueryExec esQueryExec = as(eval.child(), EsQueryExec.class); + List queryBuilderAndTags = esQueryExec.queryBuilderAndTags(); + List expectedQueryBuilderAndTags = expectedQueryBuilderAndTags( + query, + fieldName, + roundingPoints, + new Source(3, 24, roundToExpression), + mainQueryBuilder, + false + ); + verifyQueryAndTags(expectedQueryBuilderAndTags, queryBuilderAndTags); + assertThrows(UnsupportedOperationException.class, esQueryExec::query); + } + } + + public void testMergeRangePredicatesOnRoundToFieldWithOtherPushdownFunctions() { + for (Map.Entry> roundTo : roundToAllTypes.entrySet()) { + for (Map.Entry otherPushDownFunction : otherPushDownFunctions.entrySet()) { + String otherPredicate = otherPushDownFunction.getKey(); + QueryBuilder otherQueryBuilder = otherPushDownFunction.getValue(); + String fieldName = roundTo.getKey(); + List roundingPoints = roundTo.getValue(); + String roundToExpression = "round_to(" + + fieldName + + ", " + + roundingPoints.stream().map(Object::toString).collect(Collectors.joining(",")) + + ")"; + Object lowerInPredicate = null; + Object upperInPredicate = null; + var lower = roundingPoints.get(0); + var upper = roundingPoints.get(3); + if (lower instanceof String && upper instanceof String) { // date or date_nanos + lowerInPredicate = "\"2023-10-19T00:00:00.000Z\""; + upperInPredicate = "\"2023-10-24T00:00:00.000Z\""; + } else if (lower instanceof Integer lowerBound && upper instanceof Integer upperBound) { + lowerInPredicate = lowerBound - 1; + upperInPredicate = upperBound + 1; + } else if (lower instanceof Long lowerBound && upper instanceof Long upperBound) { + lowerInPredicate = lowerBound - 1L; + upperInPredicate = upperBound + 1L; + } else if (lower instanceof Double lowerBound && upper instanceof Double upperBound) { + lowerInPredicate = lowerBound - 1.0; + upperInPredicate = upperBound + 1.0; + } + assertNotNull(lowerInPredicate); + assertNotNull(upperInPredicate); + String rangePredicate = fieldName + " > " + lowerInPredicate + " and " + fieldName + " < " + upperInPredicate; + + String query = LoggerMessageFormat.format(null, """ + from test + | {} and {} + | stats count(*) by x = {} + """, otherPredicate, rangePredicate, roundToExpression); + + otherQueryBuilder = otherQueryBuilder instanceof MatchQueryBuilder + ? otherQueryBuilder + : wrapWithSingleQuery( + query, + otherQueryBuilder, + "keyword", + new Source(2, 8, otherPredicate.contains("and") ? otherPredicate.substring(6, 29) : otherPredicate.substring(6)) + ); + + RangeQueryBuilder rangeQueryBuilder = rangeQuery(fieldName).from( + lowerInPredicate instanceof String s ? s.replaceAll("\"", "") : lowerInPredicate + ) + .to(upperInPredicate instanceof String s ? s.replaceAll("\"", "") : upperInPredicate) + .includeLower(false) + .includeUpper(false) + .boost(0) + .timeZone("Z"); + + if (fieldName.equalsIgnoreCase("date")) { + rangeQueryBuilder.format(DEFAULT_DATE_TIME_FORMATTER.pattern()); + } + + if (fieldName.equalsIgnoreCase("date_nanos")) { + rangeQueryBuilder.format(DEFAULT_DATE_NANOS_FORMATTER.pattern()); + } + + QueryBuilder mainQueryBuilder = wrapWithSingleQuery( + query, + rangeQueryBuilder, + fieldName, + new Source(2, 7 + otherPredicate.length(), fieldName + " > " + lowerInPredicate) + ); + + mainQueryBuilder = Queries.combine(MUST, List.of(otherQueryBuilder, mainQueryBuilder)); + + PhysicalPlan plan = plannerOptimizer.plan(query, searchStats, makeAnalyzer("mapping-all-types.json")); + + LimitExec limit = as(plan, LimitExec.class); + AggregateExec agg = as(limit.child(), AggregateExec.class); + assertThat(agg.getMode(), is(FINAL)); + List groupings = agg.groupings(); + NamedExpression grouping = as(groupings.get(0), NamedExpression.class); + assertEquals("x", grouping.name()); + assertEquals(List.of("count(*)", "x"), Expressions.names(agg.aggregates())); + ExchangeExec exchange = as(agg.child(), ExchangeExec.class); + assertThat(exchange.inBetweenAggs(), is(true)); + agg = as(exchange.child(), AggregateExec.class); + EvalExec eval = as(agg.child(), EvalExec.class); + List aliases = eval.fields(); + assertEquals(1, aliases.size()); + FieldAttribute roundToTag = as(aliases.get(0).child(), FieldAttribute.class); + assertTrue(roundToTag.name().startsWith("$$" + fieldName + "$round_to$")); + EsQueryExec esQueryExec = as(eval.child(), EsQueryExec.class); + List queryBuilderAndTags = esQueryExec.queryBuilderAndTags(); + List expectedQueryBuilderAndTags = expectedQueryBuilderAndTags( + query, + fieldName, + roundingPoints, + new Source(3, 24, roundToExpression), + mainQueryBuilder, + false + ); + verifyQueryAndTags(expectedQueryBuilderAndTags, queryBuilderAndTags); + assertThrows(UnsupportedOperationException.class, esQueryExec::query); + } + } + } + + public void testRoundToFieldsIsNotNullWithOtherPushDownFunctions() { + for (Map.Entry> roundTo : roundToAllTypes.entrySet()) { + for (Map.Entry otherPushDownFunction : otherPushDownFunctions.entrySet()) { + String otherPredicate = otherPushDownFunction.getKey(); + QueryBuilder otherQueryBuilder = otherPushDownFunction.getValue(); + String fieldName = roundTo.getKey(); + List roundingPoints = roundTo.getValue(); + String roundToExpression = "round_to(" + + fieldName + + ", " + + roundingPoints.stream().map(Object::toString).collect(Collectors.joining(",")) + + ")"; + + String query = LoggerMessageFormat.format(null, """ + from test + | {} and {} is not null + | stats count(*) by x = {} + """, otherPredicate, fieldName, roundToExpression); + + otherQueryBuilder = otherQueryBuilder instanceof MatchQueryBuilder + ? otherQueryBuilder + : wrapWithSingleQuery( + query, + otherQueryBuilder, + "keyword", + new Source(2, 8, otherPredicate.contains("and") ? otherPredicate.substring(6, 29) : otherPredicate.substring(6)) + ); + + ExistsQueryBuilder existsQueryBuilder = existsQuery(fieldName).boost(0); + + QueryBuilder mainQueryBuilder = Queries.combine( + MUST, + otherPredicate.contains("and") + ? List.of(existsQueryBuilder, otherQueryBuilder) + : List.of(otherQueryBuilder, existsQueryBuilder) + ); + + PhysicalPlan plan = plannerOptimizer.plan(query, searchStats, makeAnalyzer("mapping-all-types.json")); + + LimitExec limit = as(plan, LimitExec.class); + AggregateExec agg = as(limit.child(), AggregateExec.class); + assertThat(agg.getMode(), is(FINAL)); + List groupings = agg.groupings(); + NamedExpression grouping = as(groupings.get(0), NamedExpression.class); + assertEquals("x", grouping.name()); + assertEquals(List.of("count(*)", "x"), Expressions.names(agg.aggregates())); + ExchangeExec exchange = as(agg.child(), ExchangeExec.class); + assertThat(exchange.inBetweenAggs(), is(true)); + agg = as(exchange.child(), AggregateExec.class); + EvalExec eval = as(agg.child(), EvalExec.class); + List aliases = eval.fields(); + assertEquals(1, aliases.size()); + FieldAttribute roundToTag = as(aliases.get(0).child(), FieldAttribute.class); + assertTrue(roundToTag.name().startsWith("$$" + fieldName + "$round_to$")); + EsQueryExec esQueryExec = as(eval.child(), EsQueryExec.class); + List queryBuilderAndTags = esQueryExec.queryBuilderAndTags(); + List expectedQueryBuilderAndTags = expectedQueryBuilderAndTags( + query, + fieldName, + roundingPoints, + new Source(3, 24, roundToExpression), + mainQueryBuilder, + false + ); + verifyQueryAndTags(expectedQueryBuilderAndTags, queryBuilderAndTags); + assertThrows(UnsupportedOperationException.class, esQueryExec::query); + } + } + } + + private static void verifyQueryAndTags(List expected, List actual) { + assertEquals(expected.size(), actual.size()); + for (int i = 0; i < expected.size(); i++) { + EsQueryExec.QueryBuilderAndTags expectedItem = expected.get(i); + EsQueryExec.QueryBuilderAndTags actualItem = actual.get(i); + assertEquals(expectedItem.query().toString(), actualItem.query().toString()); + assertEquals(expectedItem.tags().get(0), actualItem.tags().get(0)); + } + } + + private static List expectedQueryBuilderAndTags( + String query, + String fieldName, + List roundingPoints, + Source source, + QueryBuilder mainQueryBuilder, + boolean fieldIsNullable + ) { + List expected = new ArrayList<>(5); + boolean isDateField = fieldName.equals("date"); + boolean isDateNanosField = fieldName.equals("date_nanos"); + boolean isNumericField = (isDateField || isDateNanosField) == false; + List> rangeAndTags = isNumericField ? numericBuckets(roundingPoints) : dateBuckets(isDateField); + boolean mergeRangeQuery = hasRangeQuery(mainQueryBuilder, fieldName); + for (int i = 0; i < rangeAndTags.size(); i++) { + List rangeAndTag = rangeAndTags.get(i); + Object lower = rangeAndTag.get(0); + Object upper = rangeAndTag.get(1); + Object tag = rangeAndTag.get(2); + RangeQueryBuilder rangeQueryBuilder; + if (isNumericField) { + rangeQueryBuilder = rangeQuery(fieldName).boost(0); + } else if (isDateField) { // date + rangeQueryBuilder = rangeQuery(fieldName).boost(0).timeZone("Z").format(DEFAULT_DATE_TIME_FORMATTER.pattern()); + } else { // date_nanos + rangeQueryBuilder = rangeQuery(fieldName).boost(0).timeZone("Z").format(DEFAULT_DATE_NANOS_FORMATTER.pattern()); + } + if (upper != null) { + rangeQueryBuilder = rangeQueryBuilder.lt(upper); + } + if (lower != null) { + rangeQueryBuilder = rangeQueryBuilder.gte(lower); + } + + QueryBuilder qb = wrapWithSingleQuery(query, rangeQueryBuilder, fieldName, source); + if (mainQueryBuilder != null) { + if (mergeRangeQuery == false || (i == 0 || i == (rangeAndTags.size() - 1))) { + qb = Queries.combine(Queries.Clause.FILTER, List.of(mainQueryBuilder, qb)); + } else { + if (mainQueryBuilder instanceof BoolQueryBuilder boolQueryBuilder) { + BoolQueryBuilder newBool = new BoolQueryBuilder(); + for (QueryBuilder queryBuilder : boolQueryBuilder.must()) { + if (queryBuilder instanceof SingleValueQuery.Builder singleValueQueryBuilder + && singleValueQueryBuilder.next() instanceof RangeQueryBuilder rangeQueryBuilder1 + && rangeQueryBuilder1.fieldName().equalsIgnoreCase(fieldName)) { + newBool.must(qb); + } else { + newBool.must(queryBuilder); + } + } + qb = newBool; + } + } + } + expected.add(new EsQueryExec.QueryBuilderAndTags(qb, List.of(tag))); + } + if (fieldIsNullable) { + // add null bucket + BoolQueryBuilder isNullQueryBuilder = boolQuery().mustNot(existsQuery(fieldName).boost(0)); + List nullTags = new ArrayList<>(1); + nullTags.add(null); + if (mainQueryBuilder != null) { + isNullQueryBuilder = boolQuery().filter(mainQueryBuilder).filter(isNullQueryBuilder.boost(0)); + } + expected.add(new EsQueryExec.QueryBuilderAndTags(isNullQueryBuilder, nullTags)); + } + return expected; + } + + private static boolean hasRangeQuery(QueryBuilder queryBuilder, String fieldName) { + if (queryBuilder instanceof SingleValueQuery.Builder singleValueQueryBuilder + && singleValueQueryBuilder.next() instanceof RangeQueryBuilder rangeQueryBuilder) { + return rangeQueryBuilder.fieldName().equalsIgnoreCase(fieldName); + } + if (queryBuilder instanceof BoolQueryBuilder boolQueryBuilder) { + for (QueryBuilder qb : boolQueryBuilder.must()) { + if (qb instanceof SingleValueQuery.Builder singleValueQueryBuilder + && singleValueQueryBuilder.next() instanceof RangeQueryBuilder rangeQueryBuilder + && rangeQueryBuilder.fieldName().equalsIgnoreCase(fieldName)) { + return true; + } + } + } + return false; + } + + private static List> dateBuckets(boolean isDate) { + // Date rounding points + String[] dates = { "2023-10-20T00:00:00.000Z", "2023-10-21T00:00:00.000Z", "2023-10-22T00:00:00.000Z", "2023-10-23T00:00:00.000Z" }; + + // first bucket has no lower bound + List firstBucket = new ArrayList<>(3); + firstBucket.add(null); + firstBucket.add(dates[1]); + firstBucket.add(isDate ? dateTimeToLong(dates[0]) : dateNanosToLong(dates[0])); + + // last bucket has no upper bound + List lastBucket = new ArrayList<>(3); + lastBucket.add(dates[3]); + lastBucket.add(null); + lastBucket.add(isDate ? dateTimeToLong(dates[3]) : dateNanosToLong(dates[3])); + + return List.of( + firstBucket, + List.of(dates[1], dates[2], isDate ? dateTimeToLong(dates[1]) : dateNanosToLong(dates[1])), + List.of(dates[2], dates[3], isDate ? dateTimeToLong(dates[2]) : dateNanosToLong(dates[2])), + lastBucket + ); + } + + private static List> numericBuckets(List roundingPoints) { + Object p1 = roundingPoints.get(0); + Object p2 = roundingPoints.get(1); + Object p3 = roundingPoints.get(2); + Object p4 = roundingPoints.get(3); + // first bucket has no lower bound + List firstBucket = new ArrayList<>(3); + firstBucket.add(null); + firstBucket.add(p2); + firstBucket.add(p1); + // last bucket has no upper bound + List lastBucket = new ArrayList<>(3); + lastBucket.add(p4); + lastBucket.add(null); + lastBucket.add(p4); + return List.of(firstBucket, List.of(p2, p3, p2), List.of(p3, p4, p3), lastBucket); + } + + private static SearchStats searchStats() { + // create a SearchStats with min and max in milliseconds + Map minValue = Map.of("date", 1697804103360L); // 2023-10-20T12:15:03.360Z + Map maxValue = Map.of("date", 1698069301543L); // 2023-10-23T13:55:01.543Z + return new EsqlTestUtils.TestSearchStatsWithMinMax(minValue, maxValue); + } +} diff --git a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/plan/physical/EsQueryExecSerializationTests.java b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/plan/physical/EsQueryExecSerializationTests.java deleted file mode 100644 index eb53a57d3bdfb..0000000000000 --- a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/plan/physical/EsQueryExecSerializationTests.java +++ /dev/null @@ -1,91 +0,0 @@ -/* - * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one - * or more contributor license agreements. Licensed under the Elastic License - * 2.0; you may not use this file except in compliance with the Elastic License - * 2.0. - */ - -package org.elasticsearch.xpack.esql.plan.physical; - -import org.elasticsearch.index.IndexMode; -import org.elasticsearch.index.query.MatchAllQueryBuilder; -import org.elasticsearch.index.query.QueryBuilder; -import org.elasticsearch.index.query.TermQueryBuilder; -import org.elasticsearch.xpack.esql.core.expression.Attribute; -import org.elasticsearch.xpack.esql.core.expression.Expression; -import org.elasticsearch.xpack.esql.core.expression.Literal; -import org.elasticsearch.xpack.esql.core.type.DataType; -import org.elasticsearch.xpack.esql.index.EsIndexSerializationTests; - -import java.io.IOException; -import java.util.List; -import java.util.Map; - -import static org.elasticsearch.xpack.esql.index.EsIndexSerializationTests.randomIndexNameWithModes; - -public class EsQueryExecSerializationTests extends AbstractPhysicalPlanSerializationTests { - public static EsQueryExec randomEsQueryExec() { - return new EsQueryExec( - randomSource(), - randomIdentifier(), - randomFrom(IndexMode.values()), - randomIndexNameWithModes(), - randomFieldAttributes(1, 10, false), - randomQuery(), - new Literal(randomSource(), between(0, Integer.MAX_VALUE), DataType.INTEGER), - EsQueryExec.NO_SORTS, - randomEstimatedRowSize() - ); - } - - public static QueryBuilder randomQuery() { - return randomBoolean() ? new MatchAllQueryBuilder() : new TermQueryBuilder(randomAlphaOfLength(4), randomAlphaOfLength(4)); - } - - @Override - protected EsQueryExec createTestInstance() { - return randomEsQueryExec(); - } - - @Override - protected EsQueryExec mutateInstance(EsQueryExec instance) throws IOException { - String indexPattern = instance.indexPattern(); - IndexMode indexMode = instance.indexMode(); - Map indexNameWithModes = instance.indexNameWithModes(); - List attrs = instance.attrs(); - QueryBuilder query = instance.query(); - Expression limit = instance.limit(); - Integer estimatedRowSize = instance.estimatedRowSize(); - switch (between(0, 6)) { - case 0 -> indexPattern = randomValueOtherThan(indexPattern, EsIndexSerializationTests::randomIdentifier); - case 1 -> indexMode = randomValueOtherThan(indexMode, () -> randomFrom(IndexMode.values())); - case 2 -> indexNameWithModes = randomValueOtherThan(indexNameWithModes, EsIndexSerializationTests::randomIndexNameWithModes); - case 3 -> attrs = randomValueOtherThan(attrs, () -> randomFieldAttributes(1, 10, false)); - case 4 -> query = randomValueOtherThan(query, EsQueryExecSerializationTests::randomQuery); - case 5 -> limit = randomValueOtherThan( - limit, - () -> new Literal(randomSource(), between(0, Integer.MAX_VALUE), DataType.INTEGER) - ); - case 6 -> estimatedRowSize = randomValueOtherThan( - estimatedRowSize, - AbstractPhysicalPlanSerializationTests::randomEstimatedRowSize - ); - } - return new EsQueryExec( - instance.source(), - indexPattern, - indexMode, - indexNameWithModes, - attrs, - query, - limit, - EsQueryExec.NO_SORTS, - estimatedRowSize - ); - } - - @Override - protected boolean alwaysEmptySource() { - return true; - } -} diff --git a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/planner/LocalExecutionPlannerTests.java b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/planner/LocalExecutionPlannerTests.java index b56f4a3a4898b..073dd28aafbcb 100644 --- a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/planner/LocalExecutionPlannerTests.java +++ b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/planner/LocalExecutionPlannerTests.java @@ -137,8 +137,8 @@ public void testLuceneSourceOperatorHugeRowSize() throws IOException { List.of(), null, null, - null, - estimatedRowSize + estimatedRowSize, + List.of(new EsQueryExec.QueryBuilderAndTags(null, List.of())) ) ); assertThat(plan.driverFactories.size(), lessThanOrEqualTo(pragmas.taskConcurrency())); @@ -162,10 +162,10 @@ public void testLuceneTopNSourceOperator() throws IOException { IndexMode.STANDARD, index().indexNameWithModes(), List.of(), - null, limit, List.of(sort), - estimatedRowSize + estimatedRowSize, + List.of(new EsQueryExec.QueryBuilderAndTags(null, List.of())) ) ); assertThat(plan.driverFactories.size(), lessThanOrEqualTo(pragmas.taskConcurrency())); @@ -189,10 +189,10 @@ public void testLuceneTopNSourceOperatorDistanceSort() throws IOException { IndexMode.STANDARD, index().indexNameWithModes(), List.of(), - null, limit, List.of(sort), - estimatedRowSize + estimatedRowSize, + List.of(new EsQueryExec.QueryBuilderAndTags(null, List.of())) ) ); assertThat(plan.driverFactories.size(), lessThanOrEqualTo(pragmas.taskConcurrency())); @@ -215,8 +215,8 @@ public void testDriverClusterAndNodeName() throws IOException { List.of(), null, null, - null, - estimatedRowSize + estimatedRowSize, + List.of(new EsQueryExec.QueryBuilderAndTags(null, List.of())) ) ); assertThat(plan.driverFactories.size(), lessThanOrEqualTo(pragmas.taskConcurrency())); @@ -234,8 +234,8 @@ public void testParallel() throws Exception { List.of(), null, null, - null, - between(1, 1000) + between(1, 1000), + List.of(new EsQueryExec.QueryBuilderAndTags(null, List.of())) ); var limitExec = new LimitExec( Source.EMPTY, @@ -271,8 +271,8 @@ private BlockLoader constructBlockLoader() throws IOException { List.of(new FieldAttribute(Source.EMPTY, EsQueryExec.DOC_ID_FIELD.getName(), EsQueryExec.DOC_ID_FIELD)), null, null, - null, - between(1, 1000) + between(1, 1000), + List.of(new EsQueryExec.QueryBuilderAndTags(null, List.of())) ); FieldExtractExec fieldExtractExec = new FieldExtractExec( Source.EMPTY, diff --git a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/tree/EsqlNodeSubclassTests.java b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/tree/EsqlNodeSubclassTests.java index 716ae56d85d7a..7c95f0a4e2cb6 100644 --- a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/tree/EsqlNodeSubclassTests.java +++ b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/tree/EsqlNodeSubclassTests.java @@ -16,6 +16,9 @@ import org.elasticsearch.core.SuppressForbidden; import org.elasticsearch.dissect.DissectParser; import org.elasticsearch.index.IndexMode; +import org.elasticsearch.index.query.MatchAllQueryBuilder; +import org.elasticsearch.index.query.QueryBuilder; +import org.elasticsearch.index.query.TermQueryBuilder; import org.elasticsearch.test.ESTestCase; import org.elasticsearch.xpack.core.enrich.EnrichPolicy; import org.elasticsearch.xpack.esql.core.capabilities.UnresolvedException; @@ -518,6 +521,10 @@ public void accept(Page page) { ); } + if (argClass == EsQueryExec.QueryBuilderAndTags.class) { + return randomQueryBuildAndTags(); + } + try { return mock(argClass); } catch (MockitoException e) { @@ -732,6 +739,14 @@ static EsRelation randomEsRelation() { ); } + static QueryBuilder randomQuery() { + return randomBoolean() ? new MatchAllQueryBuilder() : new TermQueryBuilder(randomAlphaOfLength(4), randomAlphaOfLength(4)); + } + + static EsQueryExec.QueryBuilderAndTags randomQueryBuildAndTags() { + return new EsQueryExec.QueryBuilderAndTags(randomQuery(), List.of(randomBoolean() ? randomLong() : randomDouble())); + } + static FieldAttribute field(String name, DataType type) { return new FieldAttribute(Source.EMPTY, name, new EsField(name, type, Collections.emptyMap(), false)); }