Merge pull request #45 from marklogic/feature/count-fix

DEVEXP-467 Fix for count(); not pushing down when groupBy exists

Author: rjrudin

src/main/java/com/marklogic/spark/reader/MarkLogicScanBuilder.java

@@ -151,9 +151,10 @@ public boolean pushAggregation(Aggregation aggregation) {
 
     @Override
     public boolean supportCompletePushDown(Aggregation aggregation) {
-        // Only a single "count()" call is supported so far. Will expand as we add support for other aggregations.
+        // Only a single "count()" call is supported so far. Will expand as we add support for other aggregations,
+        // including support for groupBy() + count().
         AggregateFunc[] expressions = aggregation.aggregateExpressions();
-        return expressions != null && expressions.length == 1 && expressions[0] instanceof CountStar;
+        return expressions.length == 1 && expressions[0] instanceof CountStar && aggregation.groupByExpressions().length == 0;
     }
 
     @Override

src/main/java/com/marklogic/spark/reader/PlanUtil.java

@@ -24,7 +24,7 @@ static ObjectNode buildGroupByCount() {
         return newOperation("group-by", args -> args
             .add(objectMapper.nullNode())
             // Using "null" is the equivalent of "count(*)" - it counts rows, not values.
-            .addObject().put("ns", "op").put("fn", "count").putArray("args").add("Count").add(objectMapper.nullNode()));
+            .addObject().put("ns", "op").put("fn", "count").putArray("args").add("count").add(objectMapper.nullNode()));
     }
 
     static ObjectNode buildLimit(int limit) {

src/main/java/com/marklogic/spark/reader/ReadContext.java

@@ -176,7 +176,7 @@ void pushDownCount() {
             // As will likely be the case for all aggregations, the schema needs to be modified. And the plan analysis is
             // rebuilt to contain a single bucket, as the assumption is that MarkLogic can efficiently determine the count
             // in a single call to /v1/rows, regardless of the number of matching rows.
-            this.schema = new StructType().add("Count", DataTypes.LongType);
+            this.schema = new StructType().add("count", DataTypes.LongType);
             this.planAnalysis = new PlanAnalysis(this.planAnalysis.boundedPlan);
         }
     }

src/test/java/com/marklogic/spark/reader/PushDownCountTest.java

@@ -1,14 +1,17 @@
 package com.marklogic.spark.reader;
 
 import com.marklogic.spark.Options;
+import org.apache.spark.sql.Row;
 import org.junit.jupiter.api.Test;
 
+import java.util.List;
+
 import static org.junit.jupiter.api.Assertions.assertEquals;
 
 public class PushDownCountTest extends AbstractPushDownTest {
 
     @Test
-    void test() {
+    void count() {
         long count = newDefaultReader()
             .option(Options.READ_NUM_PARTITIONS, 2)
             .option(Options.READ_BATCH_SIZE, 1000)
@@ -23,4 +26,24 @@ void test() {
             "that regardless of the number of matching rows, MarkLogic can efficiently determine a count in a single " +
             "request.");
     }
+
+    @Test
+    void groupByAndCount() {
+        List<Row> rows = newDefaultReader()
+            .option(Options.READ_OPTIC_DSL, QUERY_WITH_NO_QUALIFIER)
+            .load()
+            .groupBy("CitationID")
+            .count()
+            .orderBy("CitationID")
+            .collectAsList();
+
+        assertEquals(15, countOfRowsReadFromMarkLogic, "groupBy + count is not yet being pushed down to MarkLogic; " +
+            "only count() by itself is being pushed down. So expecting all rows to be read for now.");
+
+        assertEquals(4, (long) rows.get(0).getAs("count"));
+        assertEquals(4, (long) rows.get(1).getAs("count"));
+        assertEquals(4, (long) rows.get(2).getAs("count"));
+        assertEquals(1, (long) rows.get(3).getAs("count"));
+        assertEquals(2, (long) rows.get(4).getAs("count"));
+    }
 }