From 1c2c81dac727615ffe91bd1aa7b745d5fda28eb0 Mon Sep 17 00:00:00 2001 From: Hitesh Khandelwal Date: Mon, 10 Jun 2024 19:27:13 +0530 Subject: [PATCH 1/5] feature: Added support for "search_after" in elasticsearch queries --- .../repository/graphdb/AtlasIndexQuery.java | 4 ++-- .../graphdb/janus/AtlasElasticsearchQuery.java | 14 ++++++++++++++ .../graphdb/janus/AtlasJanusIndexQuery.java | 5 +++++ .../model/discovery/AtlasSearchResult.java | 18 ++++++++++++++++-- .../model/discovery/ElasticsearchMetadata.java | 13 +++++++++++++ .../discovery/EntityDiscoveryService.java | 1 + 6 files changed, 51 insertions(+), 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 index 09fa0d6d06..d31f721b47 100644 --- 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 { DirectIndexQueryResult getCollapseVertices(String key); Map> getHighLights(); - + ArrayList getSort(); } } diff --git a/graphdb/janus/src/main/java/org/apache/atlas/repository/graphdb/janus/AtlasElasticsearchQuery.java b/graphdb/janus/src/main/java/org/apache/atlas/repository/graphdb/janus/AtlasElasticsearchQuery.java index 6d3ee807f7..9aa7cfe8ba 100644 --- a/graphdb/janus/src/main/java/org/apache/atlas/repository/graphdb/janus/AtlasElasticsearchQuery.java +++ b/graphdb/janus/src/main/java/org/apache/atlas/repository/graphdb/janus/AtlasElasticsearchQuery.java @@ -544,6 +544,11 @@ public DirectIndexQueryResult getCollapseVerti public Map> getHighLights() { return new HashMap<>(); } + + @Override + public ArrayList getSort() { + return new ArrayList<>(); + } } @@ -609,6 +614,15 @@ public Map> getHighLights() { } return new HashMap<>(); } + + @Override + public ArrayList getSort() { + Object sort = this.hit.get("sort"); + if (Objects.nonNull(sort) && sort instanceof List) { + return (ArrayList) sort; + } + return new ArrayList<>(); + } } public class AsyncQueryResult { diff --git a/graphdb/janus/src/main/java/org/apache/atlas/repository/graphdb/janus/AtlasJanusIndexQuery.java b/graphdb/janus/src/main/java/org/apache/atlas/repository/graphdb/janus/AtlasJanusIndexQuery.java index d4417e8f6e..f0959e48c4 100644 --- a/graphdb/janus/src/main/java/org/apache/atlas/repository/graphdb/janus/AtlasJanusIndexQuery.java +++ b/graphdb/janus/src/main/java/org/apache/atlas/repository/graphdb/janus/AtlasJanusIndexQuery.java @@ -153,5 +153,10 @@ public DirectIndexQueryResult getCollapseVerti public Map> getHighLights() { return new HashMap<>(); } + + @Override + public ArrayList 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 index 30d31b28a9..015e32a0b9 100644 --- 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 aggregations; private Map searchScore; - private Map searchMetadata; + private LinkedHashMap searchMetadata; @@ -162,13 +163,26 @@ public Map getSearchMetadata() { public void addHighlights(String guid, Map> 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 index 270ea5e8d8..270b9c27b4 100644 --- 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> highlights; + private ArrayList sort; public Map> getHighlights() { return highlights; @@ -23,6 +25,17 @@ public void addHighlights(Map> highlights) { } } + public Object getSort() { return sort; } + + public void addSort(ArrayList 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 index 66e8d58909..f7a530cc41 100644 --- a/repository/src/main/java/org/apache/atlas/discovery/EntityDiscoveryService.java +++ b/repository/src/main/java/org/apache/atlas/discovery/EntityDiscoveryService.java @@ -1101,6 +1101,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); From 8e9b3eb1392e375e23117c2064f24739eda904cd Mon Sep 17 00:00:00 2001 From: Hitesh Khandelwal Date: Wed, 12 Jun 2024 12:18:53 +0530 Subject: [PATCH 2/5] feature: Added support for "search_after" in elasticsearch queries --- .../atlas/model/discovery/ElasticsearchMetadata.java | 3 +++ .../org/apache/atlas/model/discovery/SearchParams.java | 8 +++++++- .../apache/atlas/discovery/EntityDiscoveryService.java | 5 +++-- 3 files changed, 13 insertions(+), 3 deletions(-) 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 index 270b9c27b4..66bd40f038 100644 --- 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,5 +1,6 @@ package org.apache.atlas.model.discovery; +import com.fasterxml.jackson.annotation.JsonInclude; import org.apache.commons.collections.MapUtils; import java.util.HashMap; @@ -10,6 +11,8 @@ public class ElasticsearchMetadata { private Map> highlights; + + @JsonInclude(JsonInclude.Include.NON_NULL) private ArrayList sort; public Map> getHighlights() { 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 index 93e6ea6b88..6179f291a5 100644 --- 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 index f7a530cc41..2a3390cfc9 100644 --- a/repository/src/main/java/org/apache/atlas/discovery/EntityDiscoveryService.java +++ b/repository/src/main/java/org/apache/atlas/discovery/EntityDiscoveryService.java @@ -1098,10 +1098,11 @@ 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()); } ret.addEntity(header); From 8724ab29509d625a9934c0bdb39f96239e422beb Mon Sep 17 00:00:00 2001 From: Hitesh Khandelwal Date: Wed, 12 Jun 2024 12:26:27 +0530 Subject: [PATCH 3/5] fix names --- .../apache/atlas/model/discovery/AtlasSearchResult.java | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) 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 index 015e32a0b9..96bc0dc587 100644 --- a/intg/src/main/java/org/apache/atlas/model/discovery/AtlasSearchResult.java +++ b/intg/src/main/java/org/apache/atlas/model/discovery/AtlasSearchResult.java @@ -174,12 +174,12 @@ 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); + ElasticsearchMetadata sortMetadata = this.searchMetadata.getOrDefault(guid, new ElasticsearchMetadata()); + sortMetadata.addSort(sort); if (this.searchMetadata.containsKey(guid)) { - this.searchMetadata.replace(guid, v); + this.searchMetadata.replace(guid, sortMetadata); } else { - this.searchMetadata.put(guid, v); + this.searchMetadata.put(guid, sortMetadata); } } From 3b2a6f14fd1d7fea64c29ee562a559f76cc9f4da Mon Sep 17 00:00:00 2001 From: Hitesh Khandelwal Date: Wed, 12 Jun 2024 16:03:19 +0530 Subject: [PATCH 4/5] remove unnecessary if condition --- .../atlas/model/discovery/ElasticsearchMetadata.java | 8 +------- 1 file changed, 1 insertion(+), 7 deletions(-) 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 index 66bd40f038..8f6c9e9b3b 100644 --- a/intg/src/main/java/org/apache/atlas/model/discovery/ElasticsearchMetadata.java +++ b/intg/src/main/java/org/apache/atlas/model/discovery/ElasticsearchMetadata.java @@ -31,13 +31,7 @@ public void addHighlights(Map> highlights) { public Object getSort() { return sort; } public void addSort(ArrayList sort) { - - if(!sort.isEmpty()) { - if (MapUtils.isEmpty(this.highlights)) { - this.sort = new ArrayList<>(); - } - this.sort = sort; - } + this.sort = sort; } @Override From dff31911bb8c86d2caaf160b45920fe4196572c9 Mon Sep 17 00:00:00 2001 From: Hitesh Khandelwal Date: Wed, 12 Jun 2024 18:00:38 +0530 Subject: [PATCH 5/5] when sort is empty return null --- .../apache/atlas/model/discovery/ElasticsearchMetadata.java | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) 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 index 8f6c9e9b3b..096ff82b83 100644 --- a/intg/src/main/java/org/apache/atlas/model/discovery/ElasticsearchMetadata.java +++ b/intg/src/main/java/org/apache/atlas/model/discovery/ElasticsearchMetadata.java @@ -31,7 +31,11 @@ public void addHighlights(Map> highlights) { public Object getSort() { return sort; } public void addSort(ArrayList sort) { - this.sort = sort; + if (sort.isEmpty()) { + this.sort = null; + } else { + this.sort = sort; + } } @Override