diff --git a/x-pack/plugin/esql/qa/server/src/main/java/org/elasticsearch/xpack/esql/qa/rest/generative/GenerativeRestTest.java b/x-pack/plugin/esql/qa/server/src/main/java/org/elasticsearch/xpack/esql/qa/rest/generative/GenerativeRestTest.java index a5cb60dbf2681..22f7d02408f9b 100644 --- a/x-pack/plugin/esql/qa/server/src/main/java/org/elasticsearch/xpack/esql/qa/rest/generative/GenerativeRestTest.java +++ b/x-pack/plugin/esql/qa/server/src/main/java/org/elasticsearch/xpack/esql/qa/rest/generative/GenerativeRestTest.java @@ -20,9 +20,17 @@ import org.elasticsearch.xpack.esql.generator.QueryExecuted; import org.elasticsearch.xpack.esql.generator.QueryExecutor; import org.elasticsearch.xpack.esql.generator.command.CommandGenerator; +import org.elasticsearch.xpack.esql.generator.command.pipe.DissectGenerator; import org.elasticsearch.xpack.esql.generator.command.pipe.EnrichGenerator; import org.elasticsearch.xpack.esql.generator.command.pipe.EvalGenerator; +import org.elasticsearch.xpack.esql.generator.command.pipe.GrokGenerator; +import org.elasticsearch.xpack.esql.generator.command.pipe.InlineStatsGenerator; import org.elasticsearch.xpack.esql.generator.command.pipe.LookupJoinGenerator; +import org.elasticsearch.xpack.esql.generator.command.pipe.MvExpandGenerator; +import org.elasticsearch.xpack.esql.generator.command.pipe.RegisteredDomainGenerator; +import org.elasticsearch.xpack.esql.generator.command.pipe.RenameGenerator; +import org.elasticsearch.xpack.esql.generator.command.pipe.StatsGenerator; +import org.elasticsearch.xpack.esql.generator.command.pipe.UriPartsGenerator; import org.elasticsearch.xpack.esql.generator.command.source.FromGenerator; import org.elasticsearch.xpack.esql.qa.rest.ProfileLogger; import org.elasticsearch.xpack.esql.qa.rest.RestEsqlTestCase; @@ -588,19 +596,19 @@ static List updateIndexMapped( Set createdColumns = new HashSet<>(); switch (commandName) { - case "eval" -> { + case EvalGenerator.EVAL -> { Object newCols = command.context().get(EvalGenerator.NEW_COLUMNS); if (newCols instanceof List list) { list.forEach(name -> createdColumns.add((String) name)); } } - case "grok" -> { + case GrokGenerator.GROK -> { Matcher gm = GROK_GENERATED_FIELD_PATTERN.matcher(command.commandString()); while (gm.find()) { createdColumns.add(unquote(gm.group(1))); } } - case "dissect" -> { + case DissectGenerator.DISSECT -> { Matcher dm = DISSECT_GENERATED_FIELD_PATTERN.matcher(command.commandString()); while (dm.find()) { String generated = dm.group(1); @@ -609,19 +617,19 @@ static List updateIndexMapped( } } } - case "mv_expand" -> { + case MvExpandGenerator.MV_EXPAND -> { String expanded = command.commandString().replaceFirst("(?i)^\\s*\\|\\s*mv_expand\\s+", "").trim(); // Not truly a newly created column, but we need to override the indexMapped flag so that full-text functions don't use it. // https://github.com/elastic/elasticsearch/issues/142713 createdColumns.add(unquote(expanded)); } - case "stats", "inline stats" -> { + case StatsGenerator.STATS, InlineStatsGenerator.INLINE_STATS -> { return newSchema.stream().map(col -> new Column(col.name(), col.type(), col.originalTypes(), false)).toList(); } - case "rename" -> { + case RenameGenerator.RENAME -> { return handleRenameIndexMapped(newSchema, prevMapped, command.commandString()); } - case "registered_domain" -> { + case RegisteredDomainGenerator.REGISTERED_DOMAIN -> { String prefix = (String) command.context().get("prefix"); if (prefix != null) { for (String subField : List.of("domain", "registered_domain", "top_level_domain", "subdomain")) { @@ -629,7 +637,7 @@ static List updateIndexMapped( } } } - case "uri_parts" -> { + case UriPartsGenerator.URI_PARTS -> { String prefix = (String) command.context().get("prefix"); if (prefix != null) { for (Column col : newSchema) { @@ -639,7 +647,7 @@ static List updateIndexMapped( } } } - case "enrich" -> { + case EnrichGenerator.ENRICH -> { // Enrich fields can shadow existing index columns, so we use the policy's declared enrich_fields // from the context to ensure they are marked as non-index-mapped even when names collide. Object enrichFieldsObj = command.context().get(EnrichGenerator.ENRICH_FIELDS); @@ -647,7 +655,7 @@ static List updateIndexMapped( enrichFieldsList.forEach(name -> createdColumns.add((String) name)); } } - case "lookup join" -> { + case LookupJoinGenerator.LOOKUP_JOIN -> { // LookupJoinGenerator embeds RENAME commands before the actual LOOKUP JOIN to align // left-side key columns with lookup index key names. Process these renames so that // fields renamed from non-index-mapped sources correctly inherit indexMapped=false diff --git a/x-pack/plugin/esql/qa/testFixtures/src/main/java/org/elasticsearch/xpack/esql/generator/EsqlQueryGenerator.java b/x-pack/plugin/esql/qa/testFixtures/src/main/java/org/elasticsearch/xpack/esql/generator/EsqlQueryGenerator.java index 80be48cdbe7e7..e6cce8eaa7aa4 100644 --- a/x-pack/plugin/esql/qa/testFixtures/src/main/java/org/elasticsearch/xpack/esql/generator/EsqlQueryGenerator.java +++ b/x-pack/plugin/esql/qa/testFixtures/src/main/java/org/elasticsearch/xpack/esql/generator/EsqlQueryGenerator.java @@ -19,6 +19,7 @@ import org.elasticsearch.xpack.esql.generator.command.pipe.GrokGenerator; import org.elasticsearch.xpack.esql.generator.command.pipe.InlineStatsGenerator; import org.elasticsearch.xpack.esql.generator.command.pipe.KeepGenerator; +import org.elasticsearch.xpack.esql.generator.command.pipe.LimitByGenerator; import org.elasticsearch.xpack.esql.generator.command.pipe.LimitGenerator; import org.elasticsearch.xpack.esql.generator.command.pipe.LookupJoinGenerator; import org.elasticsearch.xpack.esql.generator.command.pipe.MvExpandGenerator; @@ -114,6 +115,7 @@ public class EsqlQueryGenerator { GrokGenerator.INSTANCE, KeepGenerator.INSTANCE, InlineStatsGenerator.INSTANCE, + LimitByGenerator.INSTANCE, LimitGenerator.INSTANCE, LookupJoinGenerator.INSTANCE, MvExpandGenerator.INSTANCE, diff --git a/x-pack/plugin/esql/qa/testFixtures/src/main/java/org/elasticsearch/xpack/esql/generator/command/pipe/InlineStatsGenerator.java b/x-pack/plugin/esql/qa/testFixtures/src/main/java/org/elasticsearch/xpack/esql/generator/command/pipe/InlineStatsGenerator.java index 973f459777ecf..28153e37c104b 100644 --- a/x-pack/plugin/esql/qa/testFixtures/src/main/java/org/elasticsearch/xpack/esql/generator/command/pipe/InlineStatsGenerator.java +++ b/x-pack/plugin/esql/qa/testFixtures/src/main/java/org/elasticsearch/xpack/esql/generator/command/pipe/InlineStatsGenerator.java @@ -13,7 +13,7 @@ import java.util.List; public class InlineStatsGenerator extends StatsGenerator { - public static final String INLINE_STATS = "inline stats"; + public static final String INLINE_STATS = "inline_stats"; public static final CommandGenerator INSTANCE = new InlineStatsGenerator(); @Override diff --git a/x-pack/plugin/esql/qa/testFixtures/src/main/java/org/elasticsearch/xpack/esql/generator/command/pipe/LimitByGenerator.java b/x-pack/plugin/esql/qa/testFixtures/src/main/java/org/elasticsearch/xpack/esql/generator/command/pipe/LimitByGenerator.java new file mode 100644 index 0000000000000..a13cc7954f7dd --- /dev/null +++ b/x-pack/plugin/esql/qa/testFixtures/src/main/java/org/elasticsearch/xpack/esql/generator/command/pipe/LimitByGenerator.java @@ -0,0 +1,148 @@ +/* + * 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.generator.command.pipe; + +import org.elasticsearch.xpack.esql.generator.Column; +import org.elasticsearch.xpack.esql.generator.EsqlQueryGenerator; +import org.elasticsearch.xpack.esql.generator.QueryExecutor; +import org.elasticsearch.xpack.esql.generator.command.CommandGenerator; + +import java.util.ArrayList; +import java.util.HashMap; +import java.util.LinkedHashSet; +import java.util.List; +import java.util.Map; +import java.util.Set; + +import static org.elasticsearch.test.ESTestCase.randomIntBetween; + +public class LimitByGenerator implements CommandGenerator { + public static final CommandGenerator INSTANCE = new LimitByGenerator(); + public static final String LIMIT_BY = "limit_by"; + + private static final String LIMIT_CONTEXT = "limit"; + private static final String GROUPINGS_CONTEXT = "groupings"; + + @Override + public CommandDescription generate( + List previousCommands, + List previousOutput, + QuerySchema schema, + QueryExecutor executor + ) { + List groupable = previousOutput.stream() + .filter(EsqlQueryGenerator::groupable) + .filter(EsqlQueryGenerator::fieldCanBeUsed) + .toList(); + if (groupable.isEmpty()) { + return EMPTY_DESCRIPTION; + } + + int limit = randomIntBetween(0, 100); + int groupingCount = randomIntBetween(1, Math.min(3, groupable.size())); + Set groupings = new LinkedHashSet<>(); + for (int i = 0; i < groupingCount; i++) { + String col = EsqlQueryGenerator.randomGroupableName(groupable); + if (col != null) { + groupings.add(col); + } + } + if (groupings.isEmpty()) { + return EMPTY_DESCRIPTION; + } + + String cmd = " | LIMIT " + limit + " BY " + String.join(", ", groupings); + return new CommandDescription(LIMIT_BY, this, cmd, Map.of(LIMIT_CONTEXT, limit, GROUPINGS_CONTEXT, List.copyOf(groupings))); + } + + @Override + public ValidationResult validateOutput( + List previousCommands, + CommandDescription commandDescription, + List previousColumns, + List> previousOutput, + List columns, + List> output + ) { + int limit = (int) commandDescription.context().get(LIMIT_CONTEXT); + + if (limit == 0 && output.isEmpty() == false) { + return new ValidationResult(false, "LIMIT 0 BY should return no rows, got [" + output.size() + "]"); + } + + ValidationResult columnsResult = CommandGenerator.expectSameColumns(previousCommands, previousColumns, columns); + if (columnsResult.success() == false) { + return columnsResult; + } + + return validatePerGroupRowCounts(commandDescription, columns, output, limit); + } + + @SuppressWarnings("unchecked") + private static ValidationResult validatePerGroupRowCounts( + CommandDescription commandDescription, + List columns, + List> output, + int limit + ) { + List groupings = (List) commandDescription.context().get(GROUPINGS_CONTEXT); + + List groupingIndices = new ArrayList<>(groupings.size()); + for (String grouping : groupings) { + String rawName = EsqlQueryGenerator.unquote(grouping); + int idx = -1; + for (int i = 0; i < columns.size(); i++) { + if (columns.get(i).name().equals(rawName)) { + idx = i; + break; + } + } + if (idx == -1) { + return new ValidationResult( + false, + "LIMIT " + + limit + + " BY: grouping column [" + + rawName + + "] was not found in the output schema. Available columns: " + + columns.stream().map(Column::name).toList() + ); + } + groupingIndices.add(idx); + } + + Map, Integer> groupCounts = new HashMap<>(); + for (List row : output) { + List key = new ArrayList<>(groupingIndices.size()); + for (int idx : groupingIndices) { + Object value = row.get(idx); + key.add(value); + } + groupCounts.merge(key, 1, Integer::sum); + } + + for (var entry : groupCounts.entrySet()) { + if (entry.getValue() > limit) { + return new ValidationResult( + false, + "LIMIT " + + limit + + " BY: group " + + entry.getKey() + + " has [" + + entry.getValue() + + "] rows, expected at most [" + + limit + + "]" + ); + } + } + + return VALIDATION_OK; + } +} diff --git a/x-pack/plugin/esql/qa/testFixtures/src/main/java/org/elasticsearch/xpack/esql/generator/command/pipe/StatsGenerator.java b/x-pack/plugin/esql/qa/testFixtures/src/main/java/org/elasticsearch/xpack/esql/generator/command/pipe/StatsGenerator.java index b7cc4bef22f54..25896cc9740fb 100644 --- a/x-pack/plugin/esql/qa/testFixtures/src/main/java/org/elasticsearch/xpack/esql/generator/command/pipe/StatsGenerator.java +++ b/x-pack/plugin/esql/qa/testFixtures/src/main/java/org/elasticsearch/xpack/esql/generator/command/pipe/StatsGenerator.java @@ -43,7 +43,7 @@ public CommandDescription generate( return EMPTY_DESCRIPTION; } StringBuilder cmd = new StringBuilder(" | "); - cmd.append(commandName()); + cmd.append(commandName().replace("_", " ")); cmd.append(" "); int nStats = randomIntBetween(1, 5); for (int i = 0; i < nStats; i++) { diff --git a/x-pack/plugin/esql/qa/testFixtures/src/main/java/org/elasticsearch/xpack/esql/generator/function/FullTextFunctionGenerator.java b/x-pack/plugin/esql/qa/testFixtures/src/main/java/org/elasticsearch/xpack/esql/generator/function/FullTextFunctionGenerator.java index 68e146911423a..b2f88c8ffed08 100644 --- a/x-pack/plugin/esql/qa/testFixtures/src/main/java/org/elasticsearch/xpack/esql/generator/function/FullTextFunctionGenerator.java +++ b/x-pack/plugin/esql/qa/testFixtures/src/main/java/org/elasticsearch/xpack/esql/generator/function/FullTextFunctionGenerator.java @@ -9,6 +9,12 @@ import org.elasticsearch.xpack.esql.generator.Column; import org.elasticsearch.xpack.esql.generator.command.CommandGenerator; +import org.elasticsearch.xpack.esql.generator.command.pipe.ChangePointGenerator; +import org.elasticsearch.xpack.esql.generator.command.pipe.InlineStatsGenerator; +import org.elasticsearch.xpack.esql.generator.command.pipe.LimitByGenerator; +import org.elasticsearch.xpack.esql.generator.command.pipe.LimitGenerator; +import org.elasticsearch.xpack.esql.generator.command.pipe.MvExpandGenerator; +import org.elasticsearch.xpack.esql.generator.command.pipe.StatsGenerator; import java.util.HashSet; import java.util.List; @@ -29,6 +35,18 @@ private FullTextFunctionGenerator() {} private static final Set QSTR_KQL_SAFE_COMMANDS = Set.of("from", "where", "sort"); + /** + * Commands after which full-text expressions (match, qstr, kql, etc.) are not allowed. + */ + private static final Set FULL_TEXT_FORBIDDEN_AFTER_COMMANDS = Set.of( + LimitGenerator.LIMIT, + LimitByGenerator.LIMIT_BY, + StatsGenerator.STATS, + InlineStatsGenerator.INLINE_STATS, + ChangePointGenerator.CHANGE_POINT, + MvExpandGenerator.MV_EXPAND + ); + private static boolean isFullTextAllowed(List previousCommands) { if (previousCommands == null || previousCommands.isEmpty()) { return false; @@ -37,11 +55,7 @@ private static boolean isFullTextAllowed(List