|
32 | 32 | import static org.elasticsearch.xpack.esql.CsvTestsDataLoader.ENRICH_POLICIES; |
33 | 33 | import static org.elasticsearch.xpack.esql.CsvTestsDataLoader.availableDatasetsForEs; |
34 | 34 | import static org.elasticsearch.xpack.esql.CsvTestsDataLoader.loadDataSetIntoEs; |
| 35 | +import static org.elasticsearch.xpack.esql.qa.rest.generative.EsqlQueryGenerator.COLUMN_NAME; |
| 36 | +import static org.elasticsearch.xpack.esql.qa.rest.generative.EsqlQueryGenerator.COLUMN_ORIGINAL_TYPES; |
| 37 | +import static org.elasticsearch.xpack.esql.qa.rest.generative.EsqlQueryGenerator.COLUMN_TYPE; |
35 | 38 |
|
36 | 39 | public abstract class GenerativeRestTest extends ESRestTestCase { |
37 | 40 |
|
@@ -212,11 +215,22 @@ public static EsqlQueryGenerator.QueryExecuted execute(String command, int depth |
212 | 215 |
|
213 | 216 | @SuppressWarnings("unchecked") |
214 | 217 | private static List<EsqlQueryGenerator.Column> outputSchema(Map<String, Object> a) { |
215 | | - List<Map<String, String>> cols = (List<Map<String, String>>) a.get("columns"); |
| 218 | + List<Map<String, ?>> cols = (List<Map<String, ?>>) a.get("columns"); |
216 | 219 | if (cols == null) { |
217 | 220 | return null; |
218 | 221 | } |
219 | | - return cols.stream().map(x -> new EsqlQueryGenerator.Column(x.get("name"), x.get("type"))).collect(Collectors.toList()); |
| 222 | + return cols.stream() |
| 223 | + .map(x -> new EsqlQueryGenerator.Column((String) x.get(COLUMN_NAME), (String) x.get(COLUMN_TYPE), originalTypes(x))) |
| 224 | + .collect(Collectors.toList()); |
| 225 | + } |
| 226 | + |
| 227 | + @SuppressWarnings("unchecked") |
| 228 | + private static List<String> originalTypes(Map<String, ?> x) { |
| 229 | + List<String> originalTypes = (List<String>) x.get(COLUMN_ORIGINAL_TYPES); |
| 230 | + if (originalTypes == null) { |
| 231 | + return List.of(); |
| 232 | + } |
| 233 | + return originalTypes; |
220 | 234 | } |
221 | 235 |
|
222 | 236 | private List<String> availableIndices() throws IOException { |
|
0 commit comments