|
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