Speed up loading test data into Elasticsearch

trinodb · Oct 4, 2019 · 36d418f · 36d418f
1 parent 956ebb7
commit 36d418f
Showing 1 changed file with 8 additions and 4 deletions.
diff --git a/presto-elasticsearch/src/test/java/io/prestosql/elasticsearch/ElasticsearchLoader.java b/presto-elasticsearch/src/test/java/io/prestosql/elasticsearch/ElasticsearchLoader.java
@@ -21,6 +21,8 @@
 import io.prestosql.spi.type.Type;
 import io.prestosql.tests.AbstractTestingPrestoClient;
 import io.prestosql.tests.ResultsSession;
+import org.elasticsearch.action.bulk.BulkRequest;
+import org.elasticsearch.action.index.IndexRequest;
 import org.elasticsearch.client.Client;
 import org.elasticsearch.common.xcontent.XContentBuilder;
 
@@ -41,7 +43,6 @@
 import static java.util.Objects.requireNonNull;
 import static org.elasticsearch.client.Requests.flushRequest;
 import static org.elasticsearch.common.xcontent.XContentFactory.jsonBuilder;
-import static org.elasticsearch.common.xcontent.XContentType.JSON;
 
 public class ElasticsearchLoader
         extends AbstractTestingPrestoClient<Void>
@@ -87,6 +88,8 @@ public void addResults(QueryStatusInfo statusInfo, QueryData data)
             }
             checkState(types.get() != null, "Type information is missing");
             List<Column> columns = statusInfo.getColumns();
+
+            BulkRequest request = new BulkRequest();
             for (List<Object> fields : data.getData()) {
                 try {
                     XContentBuilder dataBuilder = jsonBuilder().startObject();
@@ -96,14 +99,15 @@ public void addResults(QueryStatusInfo statusInfo, QueryData data)
                         dataBuilder.field(columns.get(i).getName(), value);
                     }
                     dataBuilder.endObject();
-                    client.prepareIndex(tableName, "doc")
-                            .setSource(dataBuilder.string(), JSON)
-                            .get();
+
+                    request.add(new IndexRequest(tableName, "doc").source(dataBuilder));
                 }
                 catch (IOException e) {
                     throw new UncheckedIOException("Error loading data into Elasticsearch index: " + tableName, e);
                 }
             }
+
+            client.bulk(request).actionGet();
             client.admin().indices().flush(flushRequest(tableName)).actionGet();
         }