|
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