|
|
@@ -29,6 +29,9 @@ import static org.elasticsearch.xpack.esql.CsvTestsDataLoader.CSV_DATASET_MAP;
|
|
|
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 @@ public abstract class GenerativeRestTest extends ESRestTestCase {
|
|
|
}
|
|
|
|
|
|
@SuppressWarnings("unchecked")
|
|
|
- private List<EsqlQueryGenerator.Column> outputSchema(Map<String, Object> a) {
|
|
|
- List<Map<String, String>> cols = (List<Map<String, String>>) a.get("columns");
|
|
|
+ private static List<EsqlQueryGenerator.Column> outputSchema(Map<String, Object> a) {
|
|
|
+ List<Map<String, ?>> cols = (List<Map<String, ?>>) 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<String> originalTypes(Map<String, ?> x) {
|
|
|
+ List<String> originalTypes = (List<String>) x.get(COLUMN_ORIGINAL_TYPES);
|
|
|
+ if (originalTypes == null) {
|
|
|
+ return List.of();
|
|
|
+ }
|
|
|
+ return originalTypes;
|
|
|
}
|
|
|
|
|
|
private List<String> availableIndices() throws IOException {
|