Merge pull request #3229 from atlanhq/plt-1651-search-after

feature: Added support for "search_after" in elasticsearch queries
atlanhq · Jun 10, 2024 · b9173f8 · b9173f8
2 parents e0f6bcf + 1c2c81d
commit b9173f8
Show file tree

Hide file tree

Showing 6 changed files with 51 additions and 4 deletions.
diff --git a/graphdb/api/src/main/java/org/apache/atlas/repository/graphdb/AtlasIndexQuery.java b/graphdb/api/src/main/java/org/apache/atlas/repository/graphdb/AtlasIndexQuery.java
@@ -26,7 +26,7 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
-
+import java.util.ArrayList;
 /**
  * A graph query that runs directly against a particular index.
  *
@@ -103,7 +103,7 @@ interface Result<V, E> {
         DirectIndexQueryResult<V, E> getCollapseVertices(String key);
 
         Map<String, List<String>> getHighLights();
-
+        ArrayList<Object> getSort();
     }
 
 }
diff --git a/...anus/src/main/java/org/apache/atlas/repository/graphdb/janus/AtlasElasticsearchQuery.java b/...anus/src/main/java/org/apache/atlas/repository/graphdb/janus/AtlasElasticsearchQuery.java
@@ -544,6 +544,11 @@ public DirectIndexQueryResult<AtlasJanusVertex, AtlasJanusEdge> getCollapseVerti
         public Map<String, List<String>> getHighLights() {
             return new HashMap<>();
         }
+
+        @Override
+        public ArrayList<Object> getSort() {
+            return new ArrayList<>();
+        }
     }
 
 
@@ -609,6 +614,15 @@ public Map<String, List<String>> getHighLights() {
             }
             return new HashMap<>();
         }
+
+        @Override
+        public ArrayList<Object> getSort() {
+            Object sort = this.hit.get("sort");
+            if (Objects.nonNull(sort) && sort instanceof List) {
+                return (ArrayList<Object>) sort;
+            }
+            return new ArrayList<>();
+        }
     }
 
     public class AsyncQueryResult {

diff --git a/...b/janus/src/main/java/org/apache/atlas/repository/graphdb/janus/AtlasJanusIndexQuery.java b/...b/janus/src/main/java/org/apache/atlas/repository/graphdb/janus/AtlasJanusIndexQuery.java
@@ -153,5 +153,10 @@ public DirectIndexQueryResult<AtlasJanusVertex, AtlasJanusEdge> getCollapseVerti
         public Map<String, List<String>> getHighLights() {
             return new HashMap<>();
         }
+
+        @Override
+        public ArrayList<Object> getSort() {
+            return new ArrayList<>();
+        }
     }
 }
diff --git a/intg/src/main/java/org/apache/atlas/model/discovery/AtlasSearchResult.java b/intg/src/main/java/org/apache/atlas/model/discovery/AtlasSearchResult.java
@@ -37,6 +37,7 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
+import java.util.LinkedHashMap;
 
 import static com.fasterxml.jackson.annotation.JsonAutoDetect.Visibility.NONE;
 import static com.fasterxml.jackson.annotation.JsonAutoDetect.Visibility.PUBLIC_ONLY;
@@ -59,7 +60,7 @@ public class AtlasSearchResult implements Serializable {
     private Map<String, Object>            aggregations;
     private Map<String,Double>             searchScore;
 
-    private Map<String, ElasticsearchMetadata>   searchMetadata;
+    private LinkedHashMap<String, ElasticsearchMetadata> searchMetadata;
 
 
 
@@ -162,13 +163,26 @@ public Map<String, ElasticsearchMetadata> getSearchMetadata() {
 
     public void addHighlights(String guid, Map<String, List<String>> highlights) {
         if(MapUtils.isEmpty(this.searchMetadata)) {
-            this.searchMetadata = new HashMap<>();
+            this.searchMetadata = new LinkedHashMap<>();
         }
         ElasticsearchMetadata v = this.searchMetadata.getOrDefault(guid, new ElasticsearchMetadata());
         v.addHighlights(highlights);
         this.searchMetadata.put(guid, v);
     }
 
+    public void addSort(String guid, ArrayList sort) {
+        if(MapUtils.isEmpty(this.searchMetadata)) {
+            this.searchMetadata = new LinkedHashMap<>();
+        }
+        ElasticsearchMetadata v = this.searchMetadata.getOrDefault(guid, new ElasticsearchMetadata());
+        v.addSort(sort);
+        if (this.searchMetadata.containsKey(guid)) {
+           this.searchMetadata.replace(guid, v);
+        } else {
+            this.searchMetadata.put(guid, v);
+        }
+    }
+
     @Override
     public int hashCode() { return Objects.hash(queryType, searchParameters, queryText, type, classification, entities, attributes, fullTextResult, referredEntities, nextMarker); }
 

diff --git a/intg/src/main/java/org/apache/atlas/model/discovery/ElasticsearchMetadata.java b/intg/src/main/java/org/apache/atlas/model/discovery/ElasticsearchMetadata.java
@@ -5,10 +5,12 @@
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.ArrayList;
 
 public class ElasticsearchMetadata {
 
     private Map<String, List<String>> highlights;
+    private ArrayList<Object> sort;
 
     public Map<String, List<String>> getHighlights() {
         return highlights;
@@ -23,6 +25,17 @@ public void addHighlights(Map<String, List<String>> highlights) {
         }
     }
 
+    public Object getSort() { return sort; }
+
+    public void addSort(ArrayList<Object> sort) {
+
+        if(!sort.isEmpty()) {
+            if (MapUtils.isEmpty(this.highlights)) {
+                this.sort = new ArrayList<>();
+            }
+            this.sort = sort;
+        }
+    }
 
     @Override
     public String toString() {

diff --git a/repository/src/main/java/org/apache/atlas/discovery/EntityDiscoveryService.java b/repository/src/main/java/org/apache/atlas/discovery/EntityDiscoveryService.java
@@ -1113,6 +1113,7 @@ private void prepareSearchResult(AtlasSearchResult ret, DirectIndexQueryResult i
 
                 if (searchParams.isShowHighlights()) {
                     ret.addHighlights(header.getGuid(), result.getHighLights());
+                    ret.addSort(header.getGuid(), result.getSort());
                 }
 
                 ret.addEntity(header);