|
17 | 17 | import org.apache.http.client.CredentialsProvider; |
18 | 18 | import org.apache.http.impl.client.BasicCredentialsProvider; |
19 | 19 | import org.apache.logging.log4j.core.config.plugins.util.PluginManager; |
20 | | -import org.apache.lucene.util.IOConsumer; |
21 | 20 | import org.elasticsearch.client.Request; |
22 | 21 | import org.elasticsearch.client.Response; |
23 | 22 | import org.elasticsearch.client.ResponseException; |
|
31 | 30 | import org.elasticsearch.inference.TaskType; |
32 | 31 | import org.elasticsearch.logging.LogManager; |
33 | 32 | import org.elasticsearch.logging.Logger; |
34 | | -import org.elasticsearch.test.ESTestCase; |
35 | 33 | import org.elasticsearch.test.rest.ESRestTestCase; |
36 | 34 | import org.elasticsearch.xcontent.XContentType; |
37 | 35 |
|
|
45 | 43 | import java.util.List; |
46 | 44 | import java.util.Map; |
47 | 45 | import java.util.Set; |
48 | | -import java.util.concurrent.Semaphore; |
49 | 46 | import java.util.regex.Matcher; |
50 | 47 | import java.util.regex.Pattern; |
51 | 48 | import java.util.stream.Collectors; |
|
58 | 55 | import static org.elasticsearch.xpack.esql.EsqlTestUtils.reader; |
59 | 56 |
|
60 | 57 | public class CsvTestsDataLoader { |
61 | | - private static final int PARALLEL_THREADS = 10; |
62 | 58 | private static final int BULK_DATA_SIZE = 100_000; |
63 | 59 | private static final TestDataset EMPLOYEES = new TestDataset("employees", "mapping-default.json", "employees.csv").noSubfields(); |
64 | 60 | private static final TestDataset EMPLOYEES_INCOMPATIBLE = new TestDataset( |
@@ -404,35 +400,16 @@ private static void loadDataSetIntoEs( |
404 | 400 | IndexCreator indexCreator |
405 | 401 | ) throws IOException { |
406 | 402 | Logger logger = LogManager.getLogger(CsvTestsDataLoader.class); |
407 | | - List<TestDataset> datasets = availableDatasetsForEs(supportsIndexModeLookup, supportsSourceFieldMapping, inferenceEnabled).stream() |
408 | | - .toList(); |
409 | 403 |
|
410 | | - executeInParallel(datasets, dataset -> createIndex(client, dataset, indexCreator), "Failed to create indices in parallel"); |
411 | | - |
412 | | - executeInParallel(datasets, dataset -> loadData(client, dataset, logger), "Failed to load data in parallel"); |
413 | | - |
414 | | - forceMerge(client, datasets.stream().map(d -> d.indexName).collect(Collectors.toSet()), logger); |
415 | | - |
416 | | - executeInParallel( |
417 | | - ENRICH_POLICIES, |
418 | | - policy -> loadEnrichPolicy(client, policy.policyName, policy.policyFileName, logger), |
419 | | - "Failed to load enrich policies in parallel" |
420 | | - ); |
421 | | - |
422 | | - } |
423 | | - |
424 | | - private static <T> void executeInParallel(List<T> items, IOConsumer<T> consumer, String errorMessage) { |
425 | | - Semaphore semaphore = new Semaphore(PARALLEL_THREADS); |
426 | | - ESTestCase.runInParallel(items.size(), i -> { |
427 | | - try { |
428 | | - semaphore.acquire(); |
429 | | - consumer.accept(items.get(i)); |
430 | | - } catch (IOException | InterruptedException e) { |
431 | | - throw new RuntimeException(errorMessage, e); |
432 | | - } finally { |
433 | | - semaphore.release(); |
434 | | - } |
435 | | - }); |
| 404 | + Set<String> loadedDatasets = new HashSet<>(); |
| 405 | + for (var dataset : availableDatasetsForEs(supportsIndexModeLookup, supportsSourceFieldMapping, inferenceEnabled)) { |
| 406 | + load(client, dataset, logger, indexCreator); |
| 407 | + loadedDatasets.add(dataset.indexName); |
| 408 | + } |
| 409 | + forceMerge(client, loadedDatasets, logger); |
| 410 | + for (var policy : ENRICH_POLICIES) { |
| 411 | + loadEnrichPolicy(client, policy.policyName, policy.policyFileName, logger); |
| 412 | + } |
436 | 413 | } |
437 | 414 |
|
438 | 415 | public static void createInferenceEndpoints(RestClient client) throws IOException { |
@@ -589,13 +566,11 @@ private static URL getResource(String name) { |
589 | 566 | return result; |
590 | 567 | } |
591 | 568 |
|
592 | | - private static void createIndex(RestClient client, TestDataset dataset, IndexCreator indexCreator) throws IOException { |
| 569 | + private static void load(RestClient client, TestDataset dataset, Logger logger, IndexCreator indexCreator) throws IOException { |
593 | 570 | URL mapping = getResource("/" + dataset.mappingFileName); |
594 | 571 | Settings indexSettings = dataset.readSettingsFile(); |
595 | 572 | indexCreator.createIndex(client, dataset.indexName, readMappingFile(mapping, dataset.typeMapping), indexSettings); |
596 | | - } |
597 | 573 |
|
598 | | - private static void loadData(RestClient client, TestDataset dataset, Logger logger) throws IOException { |
599 | 574 | // Some examples only test that the query and mappings are valid, and don't need example data. Use .noData() for those |
600 | 575 | if (dataset.dataFileName != null) { |
601 | 576 | URL data = getResource("/data/" + dataset.dataFileName); |
|
0 commit comments