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

feature: Added support for "search_after" in elasticsearch queries
atlanhq · Jun 13, 2024 · f9e4248 · f9e4248
2 parents 040988f + dff3191
commit f9e4248
Show file tree

Hide file tree

Showing 7 changed files with 62 additions and 7 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 sortMetadata = this.searchMetadata.getOrDefault(guid, new ElasticsearchMetadata());
+        sortMetadata.addSort(sort);
+        if (this.searchMetadata.containsKey(guid)) {
+           this.searchMetadata.replace(guid, sortMetadata);
+        } else {
+            this.searchMetadata.put(guid, sortMetadata);
+        }
+    }
+
     @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
@@ -1,15 +1,20 @@
 package org.apache.atlas.model.discovery;
 
+import com.fasterxml.jackson.annotation.JsonInclude;
 import org.apache.commons.collections.MapUtils;
 
 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;
 
+    @JsonInclude(JsonInclude.Include.NON_NULL)
+    private ArrayList<Object> sort;
+
     public Map<String, List<String>> getHighlights() {
         return highlights;
     }
@@ -23,6 +28,15 @@ public void addHighlights(Map<String, List<String>> highlights) {
         }
     }
 
+    public Object getSort() { return sort; }
+
+    public void addSort(ArrayList<Object> sort) {
+        if (sort.isEmpty()) {
+            this.sort = null;
+        } else {
+            this.sort = sort;
+        }
+    }
 
     @Override
     public String toString() {

diff --git a/intg/src/main/java/org/apache/atlas/model/discovery/SearchParams.java b/intg/src/main/java/org/apache/atlas/model/discovery/SearchParams.java
@@ -26,6 +26,8 @@ public class SearchParams {
     Async async = new Async();
     boolean showHighlights;
 
+    boolean showSearchMetadata;
+
     public String getQuery() {
         return getQuery();
     }
@@ -154,10 +156,14 @@ public String getSearchInput() {
         return this.requestMetadata.getSearchInput();
     }
 
-    public boolean isShowHighlights() {
+    public boolean getShowHighlights() {
         return showHighlights;
     }
 
+    public boolean getShowSearchMetadata() {
+        return showSearchMetadata;
+    }
+
 
     static class RequestMetadata {
         private String searchInput;

diff --git a/repository/src/main/java/org/apache/atlas/discovery/EntityDiscoveryService.java b/repository/src/main/java/org/apache/atlas/discovery/EntityDiscoveryService.java
@@ -1098,8 +1098,10 @@ private void prepareSearchResult(AtlasSearchResult ret, DirectIndexQueryResult i
                         header.setCollapse(collapse);
                     }
                 }
-
-                if (searchParams.isShowHighlights()) {
+                if (searchParams.getShowSearchMetadata()) {
+                    ret.addHighlights(header.getGuid(), result.getHighLights());
+                    ret.addSort(header.getGuid(), result.getSort());
+                } else if (searchParams.getShowHighlights()) {
                     ret.addHighlights(header.getGuid(), result.getHighLights());
                 }