Iterate over Elasticsearch documents in index order

This can improve performance of scroll queries. Cherry-pick of trinodb/trino@7f4e39d Co-authored-by: Martin Traverso mtraverso@gmail.com
prestodb · Jun 10, 2020 · 4d436fd · 4d436fd
1 parent 4c7acd9
commit 4d436fd
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 2 deletions.
diff --git a/...to-elasticsearch/src/main/java/com/facebook/presto/elasticsearch/ElasticsearchClient.java b/...to-elasticsearch/src/main/java/com/facebook/presto/elasticsearch/ElasticsearchClient.java
@@ -483,12 +483,14 @@ else if (value.has("properties")) {
         return new IndexMetadata.ObjectType(result.build());
     }
 
-    public SearchResponse beginSearch(String index, int shard, QueryBuilder query, Optional<List<String>> fields, List<String> documentFields)
+    public SearchResponse beginSearch(String index, int shard, QueryBuilder query, Optional<List<String>> fields, List<String> documentFields, Optional<String> sort)
     {
         SearchSourceBuilder sourceBuilder = SearchSourceBuilder.searchSource()
                 .query(query)
                 .size(scrollSize);
 
+        sort.ifPresent(sourceBuilder::sort);
+
         fields.ifPresent(values -> {
             if (values.isEmpty()) {
                 sourceBuilder.fetchSource(false);

diff --git a/...lasticsearch/src/main/java/com/facebook/presto/elasticsearch/ElasticsearchPageSource.java b/...lasticsearch/src/main/java/com/facebook/presto/elasticsearch/ElasticsearchPageSource.java
@@ -122,13 +122,23 @@ public ElasticsearchPageSource(
                 .filter(name -> !BuiltinColumns.NAMES.contains(name))
                 .collect(toList());
 
+        // sorting by _doc (index order) get special treatment in Elasticsearch and is more efficient
+        Optional<String> sort = Optional.of("_doc");
+
+        if (table.getQuery().isPresent()) {
+            // However, if we're using a custom Elasticsearch query, use default sorting.
+            // Documents will be scored and returned based on relevance
+            sort = Optional.empty();
+        }
+
         long start = System.nanoTime();
         SearchResponse searchResponse = client.beginSearch(
                 split.getIndex(),
                 split.getShard(),
                 buildSearchQuery(session, split.getTupleDomain().transform(ElasticsearchColumnHandle.class::cast), table.getQuery()),
                 needAllFields ? Optional.empty() : Optional.of(requiredFields),
-                documentFields);
+                documentFields,
+                sort);
         readTimeNanos += System.nanoTime() - start;
         this.iterator = new SearchHitIterator(client, () -> searchResponse);
     }