diff --git a/x-pack/plugin/esql/qa/server/src/main/java/org/elasticsearch/xpack/esql/qa/rest/generative/EsqlQueryGenerator.java b/x-pack/plugin/esql/qa/server/src/main/java/org/elasticsearch/xpack/esql/qa/rest/generative/EsqlQueryGenerator.java index ef02d4a1f8c98..a8d4b9ca677fc 100644 --- a/x-pack/plugin/esql/qa/server/src/main/java/org/elasticsearch/xpack/esql/qa/rest/generative/EsqlQueryGenerator.java +++ b/x-pack/plugin/esql/qa/server/src/main/java/org/elasticsearch/xpack/esql/qa/rest/generative/EsqlQueryGenerator.java @@ -38,7 +38,11 @@ public class EsqlQueryGenerator { - public record Column(String name, String type) {} + public static final String COLUMN_NAME = "name"; + public static final String COLUMN_TYPE = "type"; + public static final String COLUMN_ORIGINAL_TYPES = "original_types"; + + public record Column(String name, String type, List originalTypes) {} public record QueryExecuted(String query, int depth, List outputSchema, List> result, Exception exception) {} @@ -290,7 +294,8 @@ public static boolean fieldCanBeUsed(Column field) { // this is a known pathological case, no need to test it for now || field.name().equals("") // no dense vectors for now, they are not supported in most commands - || field.type().contains("vector")) == false; + || field.type().contains("vector") + || field.originalTypes.stream().anyMatch(x -> x.contains("vector"))) == false; } public static String unquote(String colName) { 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 bf4f1e76f1860..f9d0b20a1e742 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 @@ -29,6 +29,9 @@ import static org.elasticsearch.xpack.esql.CsvTestsDataLoader.ENRICH_POLICIES; import static org.elasticsearch.xpack.esql.CsvTestsDataLoader.availableDatasetsForEs; import static org.elasticsearch.xpack.esql.CsvTestsDataLoader.loadDataSetIntoEs; +import static org.elasticsearch.xpack.esql.qa.rest.generative.EsqlQueryGenerator.COLUMN_NAME; +import static org.elasticsearch.xpack.esql.qa.rest.generative.EsqlQueryGenerator.COLUMN_ORIGINAL_TYPES; +import static org.elasticsearch.xpack.esql.qa.rest.generative.EsqlQueryGenerator.COLUMN_TYPE; public abstract class GenerativeRestTest extends ESRestTestCase { @@ -194,12 +197,23 @@ private EsqlQueryGenerator.QueryExecuted execute(String command, int depth) { } @SuppressWarnings("unchecked") - private List outputSchema(Map a) { - List> cols = (List>) a.get("columns"); + private static List outputSchema(Map a) { + List> cols = (List>) a.get("columns"); if (cols == null) { return null; } - return cols.stream().map(x -> new EsqlQueryGenerator.Column(x.get("name"), x.get("type"))).collect(Collectors.toList()); + return cols.stream() + .map(x -> new EsqlQueryGenerator.Column((String) x.get(COLUMN_NAME), (String) x.get(COLUMN_TYPE), originalTypes(x))) + .collect(Collectors.toList()); + } + + @SuppressWarnings("unchecked") + private static List originalTypes(Map x) { + List originalTypes = (List) x.get(COLUMN_ORIGINAL_TYPES); + if (originalTypes == null) { + return List.of(); + } + return originalTypes; } private List availableIndices() throws IOException {