ATLAS-4681: Relationship Search : SortBy is not working

Signed-off-by: Pinal Shah <[email protected]> (cherry picked from commit 66ed196)
apache · Oct 17, 2022 · 36fd6c5 · 36fd6c5
1 parent f58fa0e
commit 36fd6c5
Show file tree

Hide file tree

Showing 3 changed files with 37 additions and 3 deletions.
diff --git a/repository/src/main/java/org/apache/atlas/discovery/RelationshipSearchProcessor.java b/repository/src/main/java/org/apache/atlas/discovery/RelationshipSearchProcessor.java
@@ -17,6 +17,7 @@
  */
 package org.apache.atlas.discovery;
 
+import org.apache.atlas.SortOrder;
 import org.apache.atlas.model.discovery.SearchParameters;
 import org.apache.atlas.repository.Constants;
 import org.apache.atlas.repository.graphdb.AtlasEdge;
@@ -30,6 +31,7 @@
 import org.apache.commons.collections.CollectionUtils;
 import org.apache.commons.collections.Predicate;
 import org.apache.commons.collections.PredicateUtils;
+import org.apache.commons.lang.StringUtils;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -43,6 +45,8 @@
 import java.util.stream.Collectors;
 
 import static org.apache.atlas.repository.Constants.RELATIONSHIP_TYPE_PROPERTY_KEY;
+import static org.apache.atlas.repository.graphdb.AtlasGraphQuery.SortOrder.ASC;
+import static org.apache.atlas.repository.graphdb.AtlasGraphQuery.SortOrder.DESC;
 
 public class RelationshipSearchProcessor extends SearchProcessor {
     private static final Logger LOG      = LoggerFactory.getLogger(RelationshipSearchProcessor.class);
@@ -63,6 +67,8 @@ public class RelationshipSearchProcessor extends SearchProcessor {
         final Set<String> typeNames                          = CollectionUtils.isNotEmpty(types) ? types.stream().map(AtlasRelationshipType::getTypeName).collect(Collectors.toSet()) : null;
         final String      typeAndSubTypesQryStr              = AtlasStructType.AtlasAttribute.escapeIndexQueryValue(typeNames, true);
         final Predicate   typeNamePredicate                  = SearchPredicateUtil.generateIsRelationshipEdgePredicate(context.getTypeRegistry());
+        final String      sortBy                             = context.getSearchParameters().getSortBy();
+        final SortOrder   sortOrder                          = context.getSearchParameters().getSortOrder();
         inMemoryPredicate = typeNamePredicate;
 
         processSearchAttributes(types, filterCriteria, indexAttributes, graphAttributes, allAttributes);
@@ -112,6 +118,17 @@ public class RelationshipSearchProcessor extends SearchProcessor {
             if (attributePredicate != null) {
                 inMemoryPredicate = PredicateUtils.andPredicate(inMemoryPredicate, attributePredicate);
             }
+
+            if (StringUtils.isNotEmpty(sortBy)) {
+                final AtlasRelationshipType relationshipType   = types.iterator().next();
+                AtlasStructType.AtlasAttribute sortByAttribute = relationshipType.getAttribute(sortBy);
+
+                if (sortByAttribute != null && StringUtils.isNotEmpty(sortByAttribute.getVertexPropertyName())) {
+                    AtlasGraphQuery.SortOrder qrySortOrder = sortOrder == SortOrder.ASCENDING ? ASC : DESC;
+
+                    graphQuery.orderBy(sortByAttribute.getVertexPropertyName(), qrySortOrder);
+                }
+            }
         } else {
             graphQuery = null;
         }

diff --git a/repository/src/main/java/org/apache/atlas/discovery/SearchProcessor.java b/repository/src/main/java/org/apache/atlas/discovery/SearchProcessor.java
@@ -1366,10 +1366,12 @@ protected static Iterator<AtlasIndexQuery.Result> executeIndexQuery(SearchContex
     }
 
     protected static Iterator<AtlasIndexQuery.Result> executeIndexQueryForEdge(SearchContext context, AtlasIndexQuery indexQuery, int qryOffset, int limit) {
-        String sortBy = getSortByAttribute(context);
-        if (sortBy != null && !sortBy.isEmpty()) {
+        final AtlasRelationshipType relationshipType   = context.getRelationshipTypes().iterator().next();
+        AtlasStructType.AtlasAttribute sortByAttribute = relationshipType.getAttribute(context.getSearchParameters().getSortBy());
+
+        if (sortByAttribute != null && StringUtils.isNotEmpty(sortByAttribute.getVertexPropertyName())) {
             Order sortOrder = getSortOrderAttribute(context);
-            return indexQuery.edges(qryOffset, limit, sortBy, sortOrder);
+            return indexQuery.edges(qryOffset, limit, sortByAttribute.getVertexPropertyName(), sortOrder);
         }
         return indexQuery.edges(qryOffset, limit);
     }

diff --git a/repository/src/test/java/org/apache/atlas/discovery/RelationshipSearchProcessorTest.java b/repository/src/test/java/org/apache/atlas/discovery/RelationshipSearchProcessorTest.java
@@ -131,6 +131,21 @@ public void sortByPostName() throws AtlasBaseException {
         assertEquals("ganeshchaturthi-post@Divya",edges.get(0).getProperty("user_post.post_name", String.class));
     }
 
+    @Test
+    public void sortByReaction() throws AtlasBaseException {
+        SearchParameters params = new SearchParameters();
+        params.setRelationshipName("user_post");
+        params.setRelationshipFilters(getSingleFilterCondition("user_name", SearchParameters.Operator.CONTAINS, "Ajay"));
+        params.setLimit(20);
+        params.setSortBy("reaction");
+        params.setSortOrder(SortOrder.DESCENDING);
+
+        List<AtlasEdge> edges = executeAndAssert(params, 7);
+
+        assertEquals("wow", edges.get(0).getProperty("user_post.reaction", String.class));
+        assertEquals("create", edges.get(6).getProperty("user_post.reaction", String.class));
+
+    }
     @Test
     public void searchBymultipleTypes() throws AtlasBaseException {
         SearchParameters params = new SearchParameters();