Merge pull request #2761 from atlanhq/feat/highlights-DG-721

Add highlights to search query
atlanhq · Jan 18, 2024 · bfcad2a · bfcad2a
2 parents 7194492 + 294ecda
commit bfcad2a
Show file tree

Hide file tree

Showing 8 changed files with 86 additions and 3 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
@@ -23,6 +23,7 @@
 import org.apache.tinkerpop.gremlin.process.traversal.Order;
 
 import java.util.Iterator;
+import java.util.List;
 import java.util.Map;
 import java.util.Set;
 
@@ -101,6 +102,8 @@ interface Result<V, E> {
 
         DirectIndexQueryResult<V, E> getCollapseVertices(String key);
 
+        Map<String, List<String>> getHighLights();
+
     }
 
 }
diff --git a/graphdb/api/src/main/java/org/apache/atlas/repository/graphdb/DirectIndexQueryResult.java b/graphdb/api/src/main/java/org/apache/atlas/repository/graphdb/DirectIndexQueryResult.java
@@ -33,4 +33,5 @@ public Integer getApproximateCount() {
     public void setApproximateCount(Integer approximateCount) {
         this.approximateCount = approximateCount;
     }
+
 }
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
@@ -280,6 +280,11 @@ public Set<String> getCollapseKeys() {
         public DirectIndexQueryResult<AtlasJanusVertex, AtlasJanusEdge> getCollapseVertices(String key) {
             return null;
         }
+
+        @Override
+        public Map<String, List<String>> getHighLights() {
+            return new HashMap<>();
+        }
     }
 
 
@@ -336,5 +341,14 @@ public double getScore() {
             }
             return Double.parseDouble(String.valueOf(hit.get("_score")));
         }
+
+        @Override
+        public Map<String, List<String>> getHighLights() {
+            Object highlight = this.hit.get("highlight");
+            if(Objects.nonNull(highlight)) {
+                return (Map<String, List<String>>) highlight;
+            }
+            return new HashMap<>();
+        }
     }
 }
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
@@ -17,9 +17,7 @@
  */
 package org.apache.atlas.repository.graphdb.janus;
 
-import java.util.Iterator;
-import java.util.Map;
-import java.util.Set;
+import java.util.*;
 
 import com.google.common.base.Preconditions;
 import org.apache.atlas.exception.AtlasBaseException;
@@ -150,5 +148,10 @@ public Set<String> getCollapseKeys() {
         public DirectIndexQueryResult<AtlasJanusVertex, AtlasJanusEdge> getCollapseVertices(String key) {
             return null;
         }
+
+        @Override
+        public Map<String, List<String>> getHighLights() {
+            return new HashMap<>();
+        }
     }
 }
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
@@ -24,6 +24,7 @@
 
 import org.apache.atlas.model.instance.AtlasEntityHeader;
 import org.apache.commons.collections.CollectionUtils;
+import org.apache.commons.collections.MapUtils;
 import org.apache.commons.lang.StringUtils;
 
 import javax.xml.bind.annotation.XmlAccessType;
@@ -58,6 +59,10 @@ public class AtlasSearchResult implements Serializable {
     private Map<String, Object>            aggregations;
     private Map<String,Double>             searchScore;
 
+    private Map<String, ElasticsearchMetadata>   searchMetadata;
+
+
+
     public AtlasSearchResult() {}
 
     public AtlasSearchResult(AtlasQueryType queryType) {
@@ -151,6 +156,19 @@ public void setSearchScore(HashMap<String, Double> searchScore) {
 
     public void setNextMarker(String nextMarker) { this.nextMarker = nextMarker; }
 
+    public Map<String, ElasticsearchMetadata> getSearchMetadata() {
+        return searchMetadata;
+    }
+
+    public void addHighlights(String guid, Map<String, List<String>> highlights) {
+        if(MapUtils.isEmpty(this.searchMetadata)) {
+            this.searchMetadata = new HashMap<>();
+        }
+        ElasticsearchMetadata v = this.searchMetadata.getOrDefault(guid, new ElasticsearchMetadata());
+        v.addHighlights(highlights);
+        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
@@ -0,0 +1,33 @@
+package org.apache.atlas.model.discovery;
+
+import org.apache.commons.collections.MapUtils;
+
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+
+public class ElasticsearchMetadata {
+
+    private Map<String, List<String>> highlights;
+
+    public Map<String, List<String>> getHighlights() {
+        return highlights;
+    }
+
+    public void addHighlights(Map<String, List<String>> highlights) {
+        if(MapUtils.isNotEmpty(highlights)) {
+            if (MapUtils.isEmpty(this.highlights)) {
+                this.highlights = new HashMap<>();
+            }
+            this.highlights.putAll(highlights);
+        }
+    }
+
+
+    @Override
+    public String toString() {
+        return "SearchMetadata{" +
+                "highlights=" + highlights +
+                '}';
+    }
+}
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
@@ -20,6 +20,7 @@ public class SearchParams {
     boolean excludeClassifications;
 
     RequestMetadata requestMetadata = new RequestMetadata();
+    boolean showHighlights;
 
     public String getQuery() {
         return getQuery();
@@ -117,6 +118,10 @@ public String getSearchInput() {
         return this.requestMetadata.getSearchInput();
     }
 
+    public boolean isShowHighlights() {
+        return showHighlights;
+    }
+
     static class RequestMetadata {
         private String searchInput;
         private Set<String> utmTags;
@@ -146,4 +151,6 @@ public void setSaveSearchLog(boolean saveSearchLog) {
             this.saveSearchLog = saveSearchLog;
         }
     }
+
+
 }
diff --git a/repository/src/main/java/org/apache/atlas/discovery/EntityDiscoveryService.java b/repository/src/main/java/org/apache/atlas/discovery/EntityDiscoveryService.java
@@ -1093,6 +1093,10 @@ private void prepareSearchResult(AtlasSearchResult ret, DirectIndexQueryResult i
                     }
                 }
 
+                if (searchParams.isShowHighlights()) {
+                    ret.addHighlights(header.getGuid(), result.getHighLights());
+                }
+
                 ret.addEntity(header);
             }
         } catch (Exception e) {
-Original file line number
+Diff line change
@@ Expand Up / @@ -33,4 +33,5 @@ public Integer getApproximateCount() { @@
         public void setApproximateCount(Integer approximateCount) {
             this.approximateCount = approximateCount;
         }
     }