Support for unnesting array fields (#71)

* Add interface for From clause

* Add unnest expression

* From visitor

* Parser

* Fix integ test

* Update test

* minor change

* Add docs

* Add integration test under MongoQueryExecutorIntegrationTest

* Update docs

* address review comments

Author: Rishabh Singh

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

@@ -42,6 +42,7 @@
 import org.hypertrace.core.documentstore.expression.impl.IdentifierExpression;
 import org.hypertrace.core.documentstore.expression.impl.LogicalExpression;
 import org.hypertrace.core.documentstore.expression.impl.RelationalExpression;
+import org.hypertrace.core.documentstore.expression.impl.UnnestExpression;
 import org.hypertrace.core.documentstore.query.Filter;
 import org.hypertrace.core.documentstore.query.Pagination;
 import org.hypertrace.core.documentstore.query.Query;
@@ -477,6 +478,24 @@ public void testDistinctCount() throws IOException {
     assertDocsEqual(resultDocs, "mongo/distinct_count_response.json");
   }
 
+  @Test
+  public void testUnnestAndAggregate() throws IOException {
+    org.hypertrace.core.documentstore.query.Query query =
+        org.hypertrace.core.documentstore.query.Query.builder()
+            .addSelection(IdentifierExpression.of("sales.medium.type"))
+            .addAggregation(IdentifierExpression.of("sales.medium.type"))
+            .addSelection(
+                AggregateExpression.of(SUM, IdentifierExpression.of("sales.medium.volume")),
+                "totalSales")
+            .addFromClause(UnnestExpression.of(IdentifierExpression.of("sales")))
+            .addFromClause(UnnestExpression.of(IdentifierExpression.of("sales.medium")))
+            .addSort(IdentifierExpression.of("totalSales"), DESC)
+            .build();
+
+    Iterator<Document> iterator = collection.aggregate(query);
+    assertDocsEqual(iterator, "mongo/aggregate_on_nested_array_reponse.json");
+  }
+
   private static void assertDocsEqual(Iterator<Document> documents, String filePath)
       throws IOException {
     String fileContent = readFileFromResource(filePath).orElseThrow();

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

@@ -0,0 +1,26 @@
+[
+  {
+    "sales": {
+      "medium": {
+        "type": "online"
+      }
+    },
+    "totalSales": 9000
+  },
+  {
+    "sales": {
+      "medium": {
+        "type": "distributionChannel"
+      }
+    },
+    "totalSales": 5000
+  },
+  {
+    "sales": {
+      "medium": {
+        "type": "retail"
+      }
+    },
+    "totalSales": 1500
+  }
+]

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

@@ -15,7 +15,39 @@
           "pincode": 400004
         }
       }
-    }
+    },
+    "sales": [
+      {
+        "city": "delhi",
+        "medium": [
+          {
+            "type": "distributionChannel",
+            "volume": 1000
+          },
+          {
+            "type": "retail",
+            "volume": 500
+          },
+          {
+            "type": "online",
+            "volume": 1000
+          }
+        ]
+      },
+      {
+        "city": "pune",
+        "medium": [
+          {
+            "type": "distributionChannel",
+            "volume": 300
+          },
+          {
+            "type": "online",
+            "volume": 2000
+          }
+        ]
+      }
+    ]
   },
   {
     "_id": 2,
@@ -40,7 +72,43 @@
           "pincode": 400004
         }
       }
-    }
+    },
+    "sales": [
+      {
+        "city": "delhi",
+        "medium": [
+          {
+            "type": "distributionChannel",
+            "volume": 3000
+          },
+          {
+            "type": "retail",
+            "volume": 500
+          },
+          {
+            "type": "online",
+            "volume": 1000
+          }
+        ]
+      },
+      {
+        "city": "mumbai",
+        "medium": [
+          {
+            "type": "distributionChannel",
+            "volume": 700
+          },
+          {
+            "type": "retail",
+            "volume": 500
+          },
+          {
+            "type": "online",
+            "volume": 5000
+          }
+        ]
+      }
+    ]
   },
   {
     "_id": 4,
@@ -97,4 +165,4 @@
     "quantity": 5,
     "date": "2016-02-06T20:20:13Z"
   }
-]
+]

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

@@ -0,0 +1,31 @@
+package org.hypertrace.core.documentstore.expression.impl;
+
+import com.google.common.base.Preconditions;
+import lombok.AccessLevel;
+import lombok.AllArgsConstructor;
+import lombok.Value;
+import org.hypertrace.core.documentstore.expression.type.FromTypeExpression;
+import org.hypertrace.core.documentstore.parser.FromTypeExpressionVisitor;
+
+/**
+ * This expression allows expanding an array field to a set of rows
+ *
+ * <p><code>
+ *  UnnestExpression.of(IdentifierExpression.of("array_col")) </code>
+ */
+@Value
+@AllArgsConstructor(access = AccessLevel.PRIVATE)
+public class UnnestExpression implements FromTypeExpression {
+
+  IdentifierExpression identifierExpression;
+
+  public static UnnestExpression of(final IdentifierExpression identifierExpression) {
+    Preconditions.checkArgument(identifierExpression != null, "expression is null");
+    return new UnnestExpression(identifierExpression);
+  }
+
+  @Override
+  public <T> T accept(FromTypeExpressionVisitor visitor) {
+    return visitor.visit(this);
+  }
+}

document-store/src/main/java/org/hypertrace/core/documentstore/expression/type/FromTypeExpression.java

@@ -0,0 +1,11 @@
+package org.hypertrace.core.documentstore.expression.type;
+
+import org.hypertrace.core.documentstore.parser.FromTypeExpressionVisitor;
+
+/**
+ * Expression to retrieve rows from the referenced tables Implementations can perform table
+ * functions, join, lateral subqueries
+ */
+public interface FromTypeExpression {
+  <T> T accept(final FromTypeExpressionVisitor visitor);
+}

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

@@ -1,5 +1,6 @@
 package org.hypertrace.core.documentstore.mongo;
 
+import static java.util.Collections.singleton;
 import static java.util.function.Predicate.not;
 import static org.hypertrace.core.documentstore.mongo.MongoPaginationHelper.applyPagination;
 import static org.hypertrace.core.documentstore.mongo.MongoPaginationHelper.getLimitClause;
@@ -16,19 +17,33 @@
 import com.mongodb.client.AggregateIterable;
 import com.mongodb.client.FindIterable;
 import com.mongodb.client.MongoCursor;
+import java.util.Collection;
 import java.util.List;
+import java.util.function.Function;
 import java.util.stream.Collectors;
-import java.util.stream.Stream;
 import lombok.AllArgsConstructor;
 import lombok.extern.slf4j.Slf4j;
 import org.bson.conversions.Bson;
+import org.hypertrace.core.documentstore.mongo.parser.MongoFromTypeExpressionParser;
 import org.hypertrace.core.documentstore.mongo.query.transformer.MongoQueryTransformer;
 import org.hypertrace.core.documentstore.query.Pagination;
 import org.hypertrace.core.documentstore.query.Query;
 
 @Slf4j
 @AllArgsConstructor
 public class MongoQueryExecutor {
+  private static final List<Function<Query, Collection<BasicDBObject>>>
+      AGGREGATE_PIPELINE_FUNCTIONS =
+          List.of(
+              query -> singleton(getFilterClause(query, Query::getFilter)),
+              MongoFromTypeExpressionParser::getFromClauses,
+              query -> singleton(getGroupClause(query)),
+              query -> singleton(getProjectClause(query)),
+              query -> singleton(getFilterClause(query, Query::getAggregationFilter)),
+              query -> singleton(getSortClause(query)),
+              query -> singleton(getSkipClause(query)),
+              query -> singleton(getLimitClause(query)));
+
   final com.mongodb.client.MongoCollection<BasicDBObject> collection;
 
   public MongoCursor<BasicDBObject> find(final Query query) {
@@ -52,26 +67,9 @@ public MongoCursor<BasicDBObject> find(final Query query) {
   public MongoCursor<BasicDBObject> aggregate(final Query originalQuery) {
     Query query = transformAndLog(originalQuery);
 
-    BasicDBObject filterClause = getFilterClause(query, Query::getFilter);
-    BasicDBObject groupFilterClause = getFilterClause(query, Query::getAggregationFilter);
-
-    BasicDBObject groupClause = getGroupClause(query);
-    BasicDBObject sortClause = getSortClause(query);
-
-    BasicDBObject skipClause = getSkipClause(query);
-    BasicDBObject limitClause = getLimitClause(query);
-
-    BasicDBObject projectClause = getProjectClause(query);
-
     List<BasicDBObject> pipeline =
-        Stream.of(
-                filterClause,
-                groupClause,
-                projectClause,
-                groupFilterClause,
-                sortClause,
-                skipClause,
-                limitClause)
+        AGGREGATE_PIPELINE_FUNCTIONS.stream()
+            .flatMap(function -> function.apply(query).stream())
             .filter(not(BasicDBObject::isEmpty))
             .collect(Collectors.toList());
 

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

@@ -0,0 +1,36 @@
+package org.hypertrace.core.documentstore.mongo.parser;
+
+import com.mongodb.BasicDBObject;
+import java.util.List;
+import java.util.Map;
+import java.util.stream.Collectors;
+import org.hypertrace.core.documentstore.expression.impl.UnnestExpression;
+import org.hypertrace.core.documentstore.parser.FromTypeExpressionVisitor;
+import org.hypertrace.core.documentstore.query.Query;
+
+public class MongoFromTypeExpressionParser implements FromTypeExpressionVisitor {
+
+  private static final String PATH_KEY = "path";
+  private static final String UNWIND_OPERATOR = "$unwind";
+
+  private static final MongoIdentifierPrefixingParser mongoIdentifierPrefixingParser =
+      new MongoIdentifierPrefixingParser(new MongoIdentifierExpressionParser());
+
+  @SuppressWarnings("unchecked")
+  @Override
+  public BasicDBObject visit(UnnestExpression unnestExpression) {
+    String parsedIdentifierExpression =
+        mongoIdentifierPrefixingParser.visit(unnestExpression.getIdentifierExpression());
+    return new BasicDBObject(UNWIND_OPERATOR, Map.of(PATH_KEY, parsedIdentifierExpression));
+  }
+
+  public static List<BasicDBObject> getFromClauses(final Query query) {
+    MongoFromTypeExpressionParser mongoFromTypeExpressionParser =
+        new MongoFromTypeExpressionParser();
+    return query.getFromTypeExpressions().stream()
+        .map(
+            fromTypeExpression ->
+                (BasicDBObject) fromTypeExpression.accept(mongoFromTypeExpressionParser))
+        .collect(Collectors.toList());
+  }
+}

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

@@ -0,0 +1,7 @@
+package org.hypertrace.core.documentstore.parser;
+
+import org.hypertrace.core.documentstore.expression.impl.UnnestExpression;
+
+public interface FromTypeExpressionVisitor {
+  <T> T visit(UnnestExpression unnestExpression);
+}

document-store/src/main/java/org/hypertrace/core/documentstore/query/FromClause.java

@@ -0,0 +1,28 @@
+package org.hypertrace.core.documentstore.query;
+
+import com.google.common.base.Preconditions;
+import java.util.List;
+import java.util.Objects;
+import lombok.AccessLevel;
+import lombok.AllArgsConstructor;
+import lombok.Builder;
+import lombok.Singular;
+import lombok.Value;
+import org.hypertrace.core.documentstore.expression.type.FromTypeExpression;
+
+@Value
+@Builder(toBuilder = true)
+@AllArgsConstructor(access = AccessLevel.PRIVATE)
+public class FromClause {
+  @Singular List<FromTypeExpression> fromTypeExpressions;
+
+  public static class FromClauseBuilder {
+    public FromClause build() {
+      Preconditions.checkArgument(!fromTypeExpressions.isEmpty(), "expressions is empty");
+      Preconditions.checkArgument(
+          fromTypeExpressions.stream().noneMatch(Objects::isNull),
+          "One or more expressions is null");
+      return new FromClause(fromTypeExpressions);
+    }
+  }
+}