feat: add sorting by fields to SCSI (#101)

* feat: add sorting by fields to SCSI

* deprecate finder method and use page count

* remove index column and find value type

* get value type without querying db

* create methods in RecordUtils and specify invalid criterion

* add tests and refactor to RecordUtils

* check empty sort column and return exception for invalid types

* throw unsupported exception for invalid types and support enums

* add sorting for array of records

* unsupport array and refactor construct sql query

* add default value for count

* use same finder name

dao-api/src/main/java/com/linkedin/metadata/dao/BaseLocalDAO.java

@@ -24,6 +24,7 @@
 import com.linkedin.metadata.query.IndexCriterion;
 import com.linkedin.metadata.query.IndexCriterionArray;
 import com.linkedin.metadata.query.IndexFilter;
+import com.linkedin.metadata.query.IndexSortCriterion;
 import java.time.Clock;
 import java.util.Collections;
 import java.util.ArrayList;
@@ -385,15 +386,26 @@ protected abstract <ASPECT extends RecordTemplate> void updateLocalIndex(@Nonnul
   /**
    * Returns list of urns from local secondary index that satisfy the given filter conditions.
    *
-   * <p>Results are ordered lexicographically by the string representation of the URN.
+   * <p>Results are ordered by the order criterion but defaults to sorting lexicographically by the string
+   * representation of the URN.
    *
    * @param indexFilter {@link IndexFilter} containing filter conditions to be applied
+   * @param indexSortCriterion {@link IndexSortCriterion} sorting criterion to be applied
    * @param lastUrn last urn of the previous fetched page. For the first page, this should be set as NULL
    * @param pageSize maximum number of distinct urns to return
    * @return List of urns from local secondary index that satisfy the given filter conditions
    */
   @Nonnull
-  public abstract List<URN> listUrns(@Nonnull IndexFilter indexFilter, @Nullable URN lastUrn, int pageSize);
+  public abstract List<URN> listUrns(@Nonnull IndexFilter indexFilter, @Nullable IndexSortCriterion indexSortCriterion,
+      @Nullable URN lastUrn, int pageSize);
+
+  /**
+   * Similar to {@link #listUrns(IndexFilter, IndexSortCriterion, Urn, int)} but sorts lexicographically by the URN.
+   */
+  @Nonnull
+  public List<URN> listUrns(@Nonnull IndexFilter indexFilter, @Nullable URN lastUrn, int pageSize) {
+    return listUrns(indexFilter, null, lastUrn, pageSize);
+  }
 
   /**
    * Similar to {@link #listUrns(IndexFilter, Urn, int)}. This is to get all urns with type URN.
@@ -407,21 +419,22 @@ public List<URN> listUrns(@Nonnull Class<URN> urnClazz, @Nullable URN lastUrn, i
 
   /**
    * Retrieves list of {@link UrnAspectEntry} containing latest version of aspects along with the urn for the list of urns
-   * returned from local secondary index that satisfy given filter conditions. The returned list is ordered lexicographically by the string
-   * representation of the URN.
+   * returned from local secondary index that satisfy given filter conditions. The returned list is ordered by the
+   * sort criterion but ordered lexicographically by the string representation of the URN by default.
    *
    * @param aspectClasses aspect classes whose latest versions need to be retrieved
    * @param indexFilter {@link IndexFilter} containing filter conditions to be applied
+   * @param indexSortCriterion {@link IndexSortCriterion} sort conditions to be applied
    * @param lastUrn last urn of the previous fetched page. For the first page, this should be set as NULL
    * @param pageSize maximum number of distinct urns whose aspects need to be retrieved
    * @return ordered list of latest versions of aspects along with urns returned from local secondary index
    *        satisfying given filter conditions
    */
   @Nonnull
   public List<UrnAspectEntry<URN>> getAspects(@Nonnull Set<Class<? extends RecordTemplate>> aspectClasses,
-      @Nonnull IndexFilter indexFilter, @Nullable URN lastUrn, int pageSize) {
+      @Nonnull IndexFilter indexFilter, @Nullable IndexSortCriterion indexSortCriterion, @Nullable URN lastUrn, int pageSize) {
 
-    final List<URN> urns = listUrns(indexFilter, lastUrn, pageSize);
+    final List<URN> urns = listUrns(indexFilter, indexSortCriterion, lastUrn, pageSize);
     final Map<URN, Map<Class<? extends RecordTemplate>, Optional<? extends RecordTemplate>>> urnAspectMap =
         get(aspectClasses, new HashSet<>(urns));
 
@@ -445,6 +458,16 @@ public List<UrnAspectEntry<URN>> getAspects(@Nonnull Set<Class<? extends RecordT
         .collect(Collectors.toList());
   }
 
+  /**
+   * Similar to {@link #getAspects(Set, IndexFilter, IndexSortCriterion, Urn, int)}
+   * but sorts lexicographically by the URN.
+   */
+  @Nonnull
+  public List<UrnAspectEntry<URN>> getAspects(@Nonnull Set<Class<? extends RecordTemplate>> aspectClasses,
+    @Nonnull IndexFilter indexFilter, @Nullable URN lastUrn, int pageSize) {
+    return getAspects(aspectClasses, indexFilter, null, lastUrn, pageSize);
+  }
+
   /**
    * Runs the given lambda expression in a transaction with a limited number of retries.
    *

dao-api/src/main/java/com/linkedin/metadata/dao/utils/RecordUtils.java

@@ -138,21 +138,17 @@ public static RecordTemplate toRecordTemplate(@Nonnull String className, @Nonnul
   }
 
   /**
-   * Extracts the aspect from an entity value which includes a single aspect.
+   * Gets the aspect from the aspect class.
    *
-   * @param entity the entity value.
-   * @param aspectClass the aspect class.
-   * @return the aspect which is included in the entity.
-   * */
-  @Nonnull
-  public static <ASPECT extends RecordTemplate, ENTITY extends RecordTemplate> ASPECT
-  extractAspectFromSingleAspectEntity(@Nonnull ENTITY entity, @Nonnull Class<ASPECT> aspectClass) {
-
-    // Create an empty aspect to extract it's field names
+   * @param aspectClass the aspect class
+   * @return the aspect corresponding to the aspect class
+   */
+  public static <ASPECT extends RecordTemplate> ASPECT getAspectFromClass(@Nonnull Class<ASPECT> aspectClass) {
+    // Create an empty aspect to extract its field names
     final Constructor<ASPECT> constructor;
     try {
       @SuppressWarnings("rawtypes")
-      final Class[] constructorParamArray = new Class[] {};
+      final Class[] constructorParamArray = new Class[]{};
       constructor = aspectClass.getConstructor(constructorParamArray);
     } catch (NoSuchMethodException e) {
       throw new RuntimeException("Exception occurred while trying to get the default constructor for the aspect. ", e);
@@ -164,6 +160,39 @@ public static RecordTemplate toRecordTemplate(@Nonnull String className, @Nonnul
     } catch (InstantiationException | IllegalAccessException | InvocationTargetException e) {
       throw new RuntimeException("Exception occurred while creating an instance of the aspect. ", e);
     }
+    return aspect;
+  }
+
+  /**
+   * Gets aspect from a string representing the class name.
+   *
+   * @param aspectClassString string representation of an aspect class
+   * @return the aspect corresponding to the aspect class
+   */
+  @Nonnull
+  public static <ASPECT extends RecordTemplate> ASPECT getAspectFromString(@Nonnull String aspectClassString) {
+    final Class<ASPECT> aspectClass;
+    try {
+      aspectClass = (Class<ASPECT>) Class.forName(aspectClassString);
+    } catch (ClassNotFoundException e) {
+      throw new RuntimeException("Exception occurred while trying to get the aspect class. ", e);
+    }
+
+    return getAspectFromClass(aspectClass);
+  }
+
+  /**
+   * Extracts the aspect from an entity value which includes a single aspect.
+   *
+   * @param entity the entity value.
+   * @param aspectClass the aspect class.
+   * @return the aspect which is included in the entity.
+   * */
+  @Nonnull
+  public static <ASPECT extends RecordTemplate, ENTITY extends RecordTemplate> ASPECT
+  extractAspectFromSingleAspectEntity(@Nonnull ENTITY entity, @Nonnull Class<ASPECT> aspectClass) {
+
+    final ASPECT aspect = getAspectFromClass(aspectClass);
 
     final Set<String> aspectFields = aspect.schema().getFields()
         .stream()
@@ -400,17 +429,34 @@ private static List<Object> getReferenceForAbstractArray(@Nonnull AbstractArrayT
     return Collections.emptyList();
   }
 
+  /**
+   * Returns a PathSpec as an array representation from a string representation.
+   *
+   * @param pathSpecAsString string representation of a path spec
+   * @return array representation of a path spec
+   */
+  @Nonnull
+  public static String[] getPathSpecAsArray(@Nonnull String pathSpecAsString) {
+    pathSpecAsString = LEADING_SPACESLASH_PATTERN.matcher(pathSpecAsString).replaceAll("");
+    pathSpecAsString = TRAILING_SPACESLASH_PATTERN.matcher(pathSpecAsString).replaceAll("");
+
+    if (pathSpecAsString.isEmpty()) {
+      return new String[0];
+    }
+
+    return SLASH_PATERN.split(pathSpecAsString);
+  }
+
   /**
    * Similar to {@link #getFieldValue(RecordTemplate, PathSpec)} but takes string representation of Pegasus PathSpec as
    * input.
    */
   @Nonnull
   public static Optional<Object> getFieldValue(@Nonnull RecordTemplate recordTemplate, @Nonnull String pathSpecAsString) {
-    pathSpecAsString = LEADING_SPACESLASH_PATTERN.matcher(pathSpecAsString).replaceAll("");
-    pathSpecAsString = TRAILING_SPACESLASH_PATTERN.matcher(pathSpecAsString).replaceAll("");
+    final String[] pathSpecAsArray = getPathSpecAsArray(pathSpecAsString);
 
-    if (!pathSpecAsString.isEmpty()) {
-      return getFieldValue(recordTemplate, new PathSpec(SLASH_PATERN.split(pathSpecAsString)));
+    if (pathSpecAsArray.length > 0) {
+      return getFieldValue(recordTemplate, new PathSpec(pathSpecAsArray));
     }
     return Optional.empty();
   }

dao-api/src/main/pegasus/com/linkedin/metadata/query/IndexSortCriterion.pdl

@@ -0,0 +1,22 @@
+namespace com.linkedin.metadata.query
+
+/**
+ * Sort order along with the aspect and path to sort it on, to be applied to the results.
+ */
+record IndexSortCriterion {
+
+  /**
+   * FQCN of the aspect class in the index table that this criterion refers to e.g. com.linkedin.common.Status
+   */
+  aspect: string
+
+  /**
+   * Corresponding path column of the index table that this criterion refers to e.g. /removed (corresponding to field "removed" of com.linkedin.common.Status aspect)
+   */
+  path: string
+
+  /**
+   * The order to sort the results i.e. ASCENDING or DESCENDING
+   */
+  order: SortOrder
+}

dao-api/src/test/java/com/linkedin/metadata/dao/BaseLocalDAOTest.java

@@ -7,6 +7,7 @@
 import com.linkedin.metadata.dao.retention.VersionBasedRetention;
 import com.linkedin.metadata.query.ExtraInfo;
 import com.linkedin.metadata.query.IndexFilter;
+import com.linkedin.metadata.query.IndexSortCriterion;
 import com.linkedin.testing.AspectFoo;
 import com.linkedin.testing.EntityAspectUnion;
 import com.linkedin.testing.urn.FooUrn;
@@ -100,7 +101,8 @@ public <ASPECT extends RecordTemplate> ListResult<FooUrn> listUrns(Class<ASPECT>
     }
 
     @Override
-    public List<FooUrn> listUrns(@Nonnull IndexFilter indexFilter, @Nullable FooUrn lastUrn, int pageSize) {
+    public List<FooUrn> listUrns(@Nonnull IndexFilter indexFilter, @Nullable IndexSortCriterion indexSortCriterion,
+        @Nullable FooUrn lastUrn, int pageSize) {
       return null;
     }
 

dao-api/src/test/java/com/linkedin/metadata/dao/utils/RecordUtilsTest.java

@@ -165,6 +165,47 @@ public void testGetInvalidMetadataSnapshotClassFromName() {
     ModelUtils.getMetadataSnapshotClassFromName("com.linkedin.testing.AspectInvalid");
   }
 
+  @Test
+  public void testGetPathSpecAsArray() {
+    String ps1 = "/test/path";
+    String[] psArray1 = RecordUtils.getPathSpecAsArray(ps1);
+
+    assertEquals(psArray1.length, 2);
+    assertEquals(psArray1[0], "test");
+    assertEquals(psArray1[1], "path");
+
+    String ps2 = "/";
+    String[] psArray2 = RecordUtils.getPathSpecAsArray(ps2);
+
+    assertEquals(psArray2.length, 0);
+
+    String ps3 = "/test/";
+    String[] psArray3 = RecordUtils.getPathSpecAsArray(ps3);
+
+    assertEquals(psArray3.length, 1);
+    assertEquals(psArray3[0], "test");
+
+    String ps4 = "/recordArray/*/value";
+    String[] psArray4 = RecordUtils.getPathSpecAsArray(ps4);
+
+    assertEquals(psArray4.length, 3);
+    assertEquals(psArray4[0], "recordArray");
+    assertEquals(psArray4[1], "*");
+    assertEquals(psArray4[2], "value");
+  }
+
+  @Test
+  public void testGetAspectFromString() {
+    assertEquals(RecordUtils.getAspectFromString(AspectBar.class.getCanonicalName()), new AspectBar());
+    assertThrows(RuntimeException.class,
+        () -> RecordUtils.getAspectFromString(""));
+  }
+
+  @Test
+  public void testGetAspectFromClass() {
+    assertEquals(RecordUtils.getAspectFromClass(AspectBar.class), new AspectBar());
+  }
+
   @Test
   public void testExtractAspectFromSingleAspectEntity() {
     String field1 = "foo";

dao-api/src/test/resources/baz.json

@@ -2,6 +2,7 @@
   "boolField": true,
   "stringField": "baz",
   "longField": 1234,
+  "intField": 1,
   "arrayField": [
     "1",
     "2",