[drill] 05/13: DRILL-7971: Cleanup

[vo...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

volodymyr pushed a commit to branch mongo
in repository https://gitbox.apache.org/repos/asf/drill.git

commit a4c868355a5dff829d2e4ea74bf0630a6514c742
Author: Volodymyr Vysotskyi <vv...@gmail.com>
AuthorDate: Tue Jul 13 23:33:15 2021 +0300

    DRILL-7971: Cleanup
---
 .../exec/store/mongo/MongoAggregateUtils.java      | 49 +++++++++++-----------
 1 file changed, 25 insertions(+), 24 deletions(-)

diff --git a/contrib/storage-mongo/src/main/java/org/apache/drill/exec/store/mongo/MongoAggregateUtils.java b/contrib/storage-mongo/src/main/java/org/apache/drill/exec/store/mongo/MongoAggregateUtils.java
index 505b234..f5ebb9a 100644
--- a/contrib/storage-mongo/src/main/java/org/apache/drill/exec/store/mongo/MongoAggregateUtils.java
+++ b/contrib/storage-mongo/src/main/java/org/apache/drill/exec/store/mongo/MongoAggregateUtils.java
@@ -8,7 +8,7 @@ import org.apache.calcite.rel.core.AggregateCall;
 import org.apache.calcite.rel.type.RelDataType;
 import org.apache.calcite.sql.fun.SqlStdOperatorTable;
 import org.apache.calcite.sql.validate.SqlValidatorUtil;
-import org.apache.calcite.util.Util;
+import org.apache.drill.exec.store.mongo.common.MongoCompareOp;
 import org.bson.BsonArray;
 import org.bson.BsonDocument;
 import org.bson.BsonElement;
@@ -86,13 +86,13 @@ public class MongoAggregateUtils {
     }
     List<Bson> docAggList = new ArrayList<>();
     docAggList.add(Aggregates.group(id, docList).toBsonDocument());
-    List<String> fixups;
+    List<BsonElement> fixups;
     if (aggregate.getGroupSet().cardinality() == 1) {
-      fixups = new AbstractList<String>() {
-        @Override public String get(int index) {
+      fixups = new AbstractList<BsonElement>() {
+        @Override public BsonElement get(int index) {
           String outName = outNames.get(index);
-          return maybeQuote(outName) + ": "
-              + quote("$" + (index == 0 ? "_id" : outName));
+          return new BsonElement(maybeQuote(outName),
+              new BsonString("$" + (index == 0 ? "_id" : outName)));
         }
 
         @Override public int size() {
@@ -101,22 +101,21 @@ public class MongoAggregateUtils {
       };
     } else {
       fixups = new ArrayList<>();
-      fixups.add("_id: 0");
+      fixups.add(new BsonElement("_id", new BsonInt32(0)));
       i = 0;
       for (int group : aggregate.getGroupSet()) {
         fixups.add(
-            maybeQuote(outNames.get(group))
-                + ": "
-                + quote("$_id." + outNames.get(group)));
+            new BsonElement(maybeQuote(outNames.get(group)),
+                new BsonString("$_id." + outNames.get(group))));
         ++i;
       }
       for (AggregateCall ignored : aggregate.getAggCallList()) {
         String outName = outNames.get(i++);
-        fixups.add(maybeQuote(outName) + ": " + quote("$" + outName));
+        fixups.add(new BsonElement(maybeQuote(outName), new BsonString("$" + outName)));
       }
     }
     if (!aggregate.getGroupSet().isEmpty()) {
-      docAggList.add(Aggregates.project(BsonDocument.parse(Util.toString(fixups, "{", ", ", "}") + "}")).toBsonDocument());
+      docAggList.add(Aggregates.project(new BsonDocument(fixups)).toBsonDocument());
     }
 
     return docAggList;
@@ -126,27 +125,29 @@ public class MongoAggregateUtils {
     String aggregationName = aggCall.getAggregation().getName();
     List<Integer> args = aggCall.getArgList();
     if (aggregationName.equals(SqlStdOperatorTable.COUNT.getName())) {
+      Object expr;
       if (args.size() == 0) {
-        return Accumulators.sum(maybeQuote(outName), 1);
+        // count(*) case
+        expr = 1;
       } else {
         assert args.size() == 1;
         String inName = inNames.get(args.get(0));
-        return Accumulators.sum(maybeQuote(outName),
-            new BsonDocument("$cond",
-                new BsonArray(Arrays.asList(
-                    new Document("$eq",
-                        new BsonArray(Arrays.asList(
-                            new BsonString(quote(inName)), BsonNull.VALUE))).toBsonDocument(),
-                    new BsonInt32(0),
-                    new BsonInt32(1)
-                ))
-            )
+        expr = new BsonDocument(MongoCompareOp.COND.getCompareOp(),
+            new BsonArray(Arrays.asList(
+                new Document(MongoCompareOp.EQUAL.getCompareOp(),
+                    new BsonArray(Arrays.asList(
+                        new BsonString(quote(inName)),
+                        BsonNull.VALUE))).toBsonDocument(),
+                new BsonInt32(0),
+                new BsonInt32(1)
+            ))
         );
       }
+      return Accumulators.sum(maybeQuote(outName), expr);
     } else {
       BiFunction<String, Object, BsonField> mongoAccumulator = mongoAccumulator(aggregationName);
       if (mongoAccumulator != null) {
-        return mongoAccumulator.apply(maybeQuote(outName), "$" + inNames.get(0));
+        return mongoAccumulator.apply(maybeQuote(outName), "$" + inNames.get(args.get(0)));
       }
     }
     return null;