APIST-1653 : Changes to enhance existing length op

document-store/src/integrationTest/java/org/hypertrace/core/documentstore/DocStoreQueryV1Test.java

@@ -919,6 +919,58 @@ public void testAggregateWithMultipleGroupingLevels(String dataStoreName) throws
     testCountApi(dataStoreName, query, "query/multi_level_grouping_response.json");
   }
 
+  @ParameterizedTest
+  @ArgumentsSource(AllProvider.class)
+  public void testSortByListSizeWithMissingField(String dataStoreName) throws IOException {
+    Datastore datastore = datastoreMap.get(dataStoreName);
+    String collectionName = "list_size_sort_collection";
+    datastore.deleteCollection(collectionName);
+    datastore.createCollection(collectionName, null);
+    Collection collection = datastore.getCollection(collectionName);
+
+    try {
+      collection.upsert(
+          new SingleValueKey(TENANT_ID, "three"),
+          new JSONDocument("{\"item\":\"three\",\"tags\":[\"a\",\"b\",\"c\"]}"));
+      collection.upsert(
+          new SingleValueKey(TENANT_ID, "one"),
+          new JSONDocument("{\"item\":\"one\",\"tags\":[\"x\"]}"));
+      // Document intentionally missing the "tags" field; LENGTH must resolve to 0 instead of
+      // failing
+      collection.upsert(
+          new SingleValueKey(TENANT_ID, "none"), new JSONDocument("{\"item\":\"none\"}"));
+
+      Query query =
+          Query.builder()
+              .addSelection(IdentifierExpression.of("item"))
+              .addSelection(
+                  FunctionExpression.builder()
+                      .operator(LENGTH)
+                      .operand(IdentifierExpression.of("tags"))
+                      .build(),
+                  "tag_count")
+              .addSort(IdentifierExpression.of("tag_count"), ASC)
+              .build();
+
+      Iterator<Document> resultDocs = collection.aggregate(query);
+      List<Map<String, Object>> results = new ArrayList<>();
+      while (resultDocs.hasNext()) {
+        results.add(Utils.convertDocumentToMap(resultDocs.next()));
+      }
+
+      assertEquals(3, results.size());
+      // Document without the "tags" field counts as 0 and sorts first in ascending order
+      assertEquals("none", results.get(0).get("item"));
+      assertEquals(0, ((Number) results.get(0).get("tag_count")).intValue());
+      assertEquals("one", results.get(1).get("item"));
+      assertEquals(1, ((Number) results.get(1).get("tag_count")).intValue());
+      assertEquals("three", results.get(2).get("item"));
+      assertEquals(3, ((Number) results.get(2).get("tag_count")).intValue());
+    } finally {
+      datastore.deleteCollection(collectionName);
+    }
+  }
+
   @ParameterizedTest
   @ArgumentsSource(AllProvider.class)
   public void testAggregateWithFunctionalLeftHandSideFilter(final String dataStoreName)

document-store/src/main/java/org/hypertrace/core/documentstore/mongo/query/parser/MongoFunctionExpressionParser.java

@@ -69,6 +69,12 @@ Map<String, Object> parse(final FunctionExpression expression) {
 
     if (numArgs == 1) {
       Object value = expression.getOperands().get(0).accept(parser);
+      // $size fails when the operand resolves to a missing/absent (or null) field. Default such a
+      // value to an empty array so that LENGTH of an absent field is 0 instead of throwing an
+      // error.
+      if (operator == LENGTH) {
+        value = Map.of("$ifNull", List.of(value, List.of()));
+      }
       return Map.of(key, value);
     }
 

document-store/src/main/java/org/hypertrace/core/documentstore/postgres/query/v1/vistors/PostgresFunctionExpressionVisitor.java

@@ -6,6 +6,7 @@
 import java.util.stream.Collectors;
 import lombok.NoArgsConstructor;
 import org.apache.commons.lang3.StringUtils;
+import org.hypertrace.core.documentstore.DocumentType;
 import org.hypertrace.core.documentstore.expression.impl.FunctionExpression;
 import org.hypertrace.core.documentstore.expression.operators.FunctionOperator;
 import org.hypertrace.core.documentstore.expression.type.SelectTypeExpression;
@@ -49,10 +50,11 @@ public String visit(final FunctionExpression expression) {
     }
 
     if (numArgs == 1) {
+      if (expression.getOperator().equals(FunctionOperator.LENGTH)) {
+        return buildLengthExpression(expression.getOperands().get(0));
+      }
       String parsedExpression = getParsedExpression(expression.getOperands().get(0));
-      return expression.getOperator().equals(FunctionOperator.LENGTH)
-          ? String.format("ARRAY_LENGTH( %s, %s )", parsedExpression, ARRAY_DIMENSION)
-          : String.format("%s( %s )", expression.getOperator(), parsedExpression);
+      return String.format("%s( %s )", expression.getOperator(), parsedExpression);
     }
 
     Collector<String, ?, String> collector =
@@ -83,6 +85,27 @@ private Collector getCollectorForFunctionOperator(FunctionOperator operator) {
         String.format("Query operation:%s not supported", operator));
   }
 
+  private String buildLengthExpression(final SelectTypeExpression operand) {
+    Optional<String> identifier = Optional.ofNullable(operand.accept(identifierExpressionVisitor));
+    Optional<String> resolvedSelection =
+        identifier.map(v -> getPostgresQueryParser().getPgSelections().get(v));
+    if (resolvedSelection.isPresent()) {
+      return String.format(
+          "COALESCE( ARRAY_LENGTH( %s, %s ), 0 )", resolvedSelection.get(), ARRAY_DIMENSION);
+    }
+    PostgresFieldIdentifierExpressionVisitor fieldVisitor =
+        new PostgresFieldIdentifierExpressionVisitor(getPostgresQueryParser());
+    String parsedExpression = operand.accept(fieldVisitor);
+    if (getPostgresQueryParser().getPgColTransformer().getDocumentType() == DocumentType.FLAT) {
+      return String.format(
+          "COALESCE( ARRAY_LENGTH( %s, %s ), 0 )", parsedExpression, ARRAY_DIMENSION);
+    }
+    return String.format(
+        "jsonb_array_length( CASE WHEN jsonb_typeof( %s ) = 'array' THEN %s"
+            + " ELSE '[]'::jsonb END )",
+        parsedExpression, parsedExpression);
+  }
+
   private String getParsedExpression(final SelectTypeExpression expression) {
     Optional<String> identifier =
         Optional.ofNullable(expression.accept(identifierExpressionVisitor));

document-store/src/test/java/org/hypertrace/core/documentstore/postgres/query/v1/PostgresQueryParserTest.java

@@ -395,7 +395,7 @@ void testAggregationExpressionDistinctCount() {
 
     assertEquals(
         "SELECT ARRAY_AGG(DISTINCT CAST (document->>'quantity' AS NUMERIC)) AS \"qty_distinct\", "
-            + "ARRAY_LENGTH( ARRAY_AGG(DISTINCT CAST (document->>'quantity' AS NUMERIC)), 1 ) AS \"qty_distinct_length\" "
+            + "COALESCE( ARRAY_LENGTH( ARRAY_AGG(DISTINCT CAST (document->>'quantity' AS NUMERIC)), 1 ), 0 ) AS \"qty_distinct_length\" "
             + "FROM \"testCollection\" "
             + "WHERE CAST (document->>'price' AS NUMERIC) = ? "
             + "GROUP BY document->'item'",
@@ -435,10 +435,10 @@ void testAggregateWithMultipleGroupingLevels() {
     assertEquals(
         "SELECT document->'item' AS \"item\", document->'price' AS \"price\", "
             + "ARRAY_AGG(DISTINCT CAST (document->>'quantity' AS NUMERIC)) AS \"quantities\", "
-            + "ARRAY_LENGTH( ARRAY_AGG(DISTINCT CAST (document->>'quantity' AS NUMERIC)), 1 ) AS \"num_quantities\" "
+            + "COALESCE( ARRAY_LENGTH( ARRAY_AGG(DISTINCT CAST (document->>'quantity' AS NUMERIC)), 1 ), 0 ) AS \"num_quantities\" "
             + "FROM \"testCollection\" "
             + "GROUP BY document->'item',document->'price' "
-            + "HAVING ARRAY_LENGTH( ARRAY_AGG(DISTINCT CAST (document->>'quantity' AS NUMERIC)), 1 ) = ? "
+            + "HAVING COALESCE( ARRAY_LENGTH( ARRAY_AGG(DISTINCT CAST (document->>'quantity' AS NUMERIC)), 1 ), 0 ) = ? "
             + "ORDER BY document->'item' DESC NULLS LAST",
         sql);
 

document-store/src/test/resources/mongo/pipeline/distinct_count.json

@@ -19,7 +19,12 @@
   {
     "$project": {
       "section_count": {
-        "$size": "$section_count"
+        "$size": {
+          "$ifNull": [
+            "$section_count",
+            []
+          ]
+        }
       }
     }
   }

document-store/src/test/resources/mongo/pipeline/field_count.json

@@ -10,7 +10,12 @@
   {
     "$project": {
       "total": {
-        "$size": "$total"
+        "$size": {
+          "$ifNull": [
+            "$total",
+            []
+          ]
+        }
       }
     }
   }

document-store/src/test/resources/mongo/pipeline/optimize_sorts_simple_sort_with_aggregation_selection.json

@@ -10,7 +10,12 @@
   {
     "$project": {
       "total": {
-        "$size": "$total"
+        "$size": {
+          "$ifNull": [
+            "$total",
+            []
+          ]
+        }
       }
     }
   },

document-store/src/test/resources/mongo/pipeline/simple.json

@@ -10,7 +10,12 @@
   {
     "$project": {
       "total": {
-        "$size": "$total"
+        "$size": {
+          "$ifNull": [
+            "$total",
+            []
+          ]
+        }
       }
     }
   }

document-store/src/test/resources/mongo/pipeline/with_projections.json

@@ -11,7 +11,12 @@
     "$project": {
       "name": 1,
       "total": {
-        "$size": "$total"
+        "$size": {
+          "$ifNull": [
+            "$total",
+            []
+          ]
+        }
       }
     }
   }