Add config preserving empty array while unwinding (#75)

Author: Rishabh Singh

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

@@ -498,15 +498,44 @@ public void testUnnestAndAggregate() throws IOException {
             .addSelection(
                 AggregateExpression.of(SUM, IdentifierExpression.of("sales.medium.volume")),
                 "totalSales")
-            .addFromClause(UnnestExpression.of(IdentifierExpression.of("sales")))
-            .addFromClause(UnnestExpression.of(IdentifierExpression.of("sales.medium")))
+            // we don't want to consider entries where sales data is missing
+            .addFromClause(UnnestExpression.of(IdentifierExpression.of("sales"), false))
+            .addFromClause(UnnestExpression.of(IdentifierExpression.of("sales.medium"), false))
             .addSort(IdentifierExpression.of("totalSales"), DESC)
             .build();
 
     Iterator<Document> iterator = collection.aggregate(query);
     assertDocsEqual(iterator, "mongo/aggregate_on_nested_array_reponse.json");
   }
 
+  @Test
+  public void testUnnestAndAggregate_preserveEmptyTrue() throws IOException {
+    // include all documents in the result irrespective of `sales` field
+    org.hypertrace.core.documentstore.query.Query query =
+        org.hypertrace.core.documentstore.query.Query.builder()
+            .addSelection(AggregateExpression.of(COUNT, IdentifierExpression.of("item")), "count")
+            .addFromClause(UnnestExpression.of(IdentifierExpression.of("sales"), true))
+            .addFromClause(UnnestExpression.of(IdentifierExpression.of("sales.medium"), true))
+            .build();
+
+    Iterator<Document> iterator = collection.aggregate(query);
+    assertDocsEqual(iterator, "mongo/unwind_preserving_empty_array_response.json");
+  }
+
+  @Test
+  public void testUnnestAndAggregate_preserveEmptyFalse() throws IOException {
+    // consider only those documents where sales field is missing
+    org.hypertrace.core.documentstore.query.Query query =
+        org.hypertrace.core.documentstore.query.Query.builder()
+            .addSelection(AggregateExpression.of(COUNT, IdentifierExpression.of("item")), "count")
+            .addFromClause(UnnestExpression.of(IdentifierExpression.of("sales"), false))
+            .addFromClause(UnnestExpression.of(IdentifierExpression.of("sales.medium"), true))
+            .build();
+
+    Iterator<Document> iterator = collection.aggregate(query);
+    assertDocsEqual(iterator, "mongo/unwind_not_preserving_empty_array_response.json");
+  }
+
   private static void assertDocsEqual(Iterator<Document> documents, String filePath)
       throws IOException {
     String fileContent = readFileFromResource(filePath).orElseThrow();

document-store/src/integrationTest/resources/mongo/collection_data.json

@@ -54,7 +54,13 @@
     "item": "Mirror",
     "price": 20,
     "quantity": 1,
-    "date": "2014-03-01T09:00:00Z"
+    "date": "2014-03-01T09:00:00Z",
+    "sales": [
+      {
+        "city": "delhi",
+        "medium": []
+      }
+    ]
   },
   {
     "_id": 3,
@@ -115,7 +121,8 @@
     "item": "Shampoo",
     "price": 5,
     "quantity": 20,
-    "date": "2014-04-04T11:21:39.736Z"
+    "date": "2014-04-04T11:21:39.736Z",
+    "sales": []
   },
   {
     "_id": 5,

document-store/src/integrationTest/resources/mongo/unwind_not_preserving_empty_array_response.json

@@ -0,0 +1,5 @@
+[
+  {
+  "count": 12
+  }
+]

document-store/src/integrationTest/resources/mongo/unwind_preserving_empty_array_response.json

@@ -0,0 +1,5 @@
+[
+  {
+    "count": 17
+  }
+]

document-store/src/main/java/org/hypertrace/core/documentstore/expression/impl/UnnestExpression.java

@@ -19,9 +19,12 @@ public class UnnestExpression implements FromTypeExpression {
 
   IdentifierExpression identifierExpression;
 
-  public static UnnestExpression of(final IdentifierExpression identifierExpression) {
+  boolean preserveNullAndEmptyArrays;
+
+  public static UnnestExpression of(
+      final IdentifierExpression identifierExpression, boolean preserveNullAndEmptyArrays) {
     Preconditions.checkArgument(identifierExpression != null, "expression is null");
-    return new UnnestExpression(identifierExpression);
+    return new UnnestExpression(identifierExpression, preserveNullAndEmptyArrays);
   }
 
   @Override

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

@@ -11,6 +11,7 @@
 public class MongoFromTypeExpressionParser implements FromTypeExpressionVisitor {
 
   private static final String PATH_KEY = "path";
+  public static final String PRESERVE_NULL_AND_EMPTY_ARRAYS = "preserveNullAndEmptyArrays";
   private static final String UNWIND_OPERATOR = "$unwind";
 
   private static final MongoIdentifierPrefixingParser mongoIdentifierPrefixingParser =
@@ -21,7 +22,13 @@ public class MongoFromTypeExpressionParser implements FromTypeExpressionVisitor
   public BasicDBObject visit(UnnestExpression unnestExpression) {
     String parsedIdentifierExpression =
         mongoIdentifierPrefixingParser.visit(unnestExpression.getIdentifierExpression());
-    return new BasicDBObject(UNWIND_OPERATOR, Map.of(PATH_KEY, parsedIdentifierExpression));
+    return new BasicDBObject(
+        UNWIND_OPERATOR,
+        Map.of(
+            PATH_KEY,
+            parsedIdentifierExpression,
+            PRESERVE_NULL_AND_EMPTY_ARRAYS,
+            unnestExpression.isPreserveNullAndEmptyArrays()));
   }
 
   public static List<BasicDBObject> getFromClauses(final Query query) {

document-store/src/test/java/org/hypertrace/core/documentstore/mongo/MongoQueryExecutorTest.java

@@ -603,15 +603,18 @@ public void testUnwindAndGroup() {
                 RelationalExpression.of(
                     IdentifierExpression.of("class"), LTE, ConstantExpression.of(10)))
             .addAggregation(IdentifierExpression.of("class.students.courses"))
-            .addFromClause(UnnestExpression.of(IdentifierExpression.of("class.students")))
-            .addFromClause(UnnestExpression.of(IdentifierExpression.of("class.students.courses")))
+            .addFromClause(UnnestExpression.of(IdentifierExpression.of("class.students"), true))
+            .addFromClause(
+                UnnestExpression.of(IdentifierExpression.of("class.students.courses"), true))
             .build();
 
     List<BasicDBObject> pipeline =
         List.of(
             BasicDBObject.parse("{\"$match\": {\"class\": {\"$lte\": 10}}}"),
-            BasicDBObject.parse("{\"$unwind\": {\"path\": \"$class.students\"}}"),
-            BasicDBObject.parse("{\"$unwind\": {\"path\": \"$class.students.courses\"}}"),
+            BasicDBObject.parse(
+                "{\"$unwind\": {\"path\": \"$class.students\", \"preserveNullAndEmptyArrays\": true}}"),
+            BasicDBObject.parse(
+                "{\"$unwind\": {\"path\": \"$class.students.courses\", \"preserveNullAndEmptyArrays\": true}}"),
             BasicDBObject.parse(
                 "{\"$group\": {\"_id\": {\"class\\\\u002estudents\\\\u002ecourses\": \"$class.students.courses\"}}}"));