Improving groupBy pushdown to work for multiple column names

Turns out this was a simple enhancement since `op.groupBy` already supports multiple column names being passed in.

Author: rjrudin

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

@@ -145,11 +145,10 @@ public boolean pushAggregation(Aggregation aggregation) {
         }
         if (supportCompletePushDown(aggregation)) {
             if (aggregation.groupByExpressions().length > 0) {
-                Expression expr = aggregation.groupByExpressions()[0];
                 if (logger.isInfoEnabled()) {
-                    logger.info("Pushing down groupBy + count on: {}", expr.describe());
+                    logger.info("Pushing down groupBy + count on: {}", Arrays.asList(aggregation.groupByExpressions()));
                 }
-                readContext.pushDownGroupByCount(expr);
+                readContext.pushDownGroupByCount(aggregation.groupByExpressions());
             } else {
                 if (logger.isInfoEnabled()) {
                     logger.info("Pushing down count()");
@@ -167,12 +166,10 @@ public boolean supportCompletePushDown(Aggregation aggregation) {
             return false;
         }
         AggregateFunc[] expressions = aggregation.aggregateExpressions();
-        if (expressions.length == 1 && expressions[0] instanceof CountStar) {
-            // If a count() is used, it's supported if there's no groupBy - i.e. just doing a count() by itself -
-            // and supported with a single groupBy - e.g. groupBy("column").count().
-            return aggregation.groupByExpressions().length < 2;
-        }
-        return false;
+        // If a count() is used, it's supported if there's no groupBy - i.e. just doing a count() by itself -
+        // and supported with 1 to many groupBy's - e.g. groupBy("column", "someOtherColumn").count().
+        // Other aggregate functions will be supported in the near future.
+        return expressions.length == 1 && expressions[0] instanceof CountStar;
     }
 
     @Override

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

@@ -48,9 +48,10 @@ static ObjectNode buildGroupByCount() {
         });
     }
 
-    static ObjectNode buildGroupByCount(String columnName) {
+    static ObjectNode buildGroupByCount(List<String> columnNames) {
         return newOperation("group-by", args -> {
-            populateSchemaCol(args.addObject(), columnName);
+            ArrayNode columns = args.addArray();
+            columnNames.forEach(columnName -> populateSchemaCol(columns.addObject(), columnName));
             addCountArg(args);
         });
     }

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

@@ -42,6 +42,8 @@
 import java.util.Iterator;
 import java.util.List;
 import java.util.Map;
+import java.util.stream.Collectors;
+import java.util.stream.Stream;
 
 /**
  * Captures state - all of which is serializable - that can be calculated at different times based on a user's inputs.
@@ -166,21 +168,30 @@ void pushDownCount() {
         modifyPlanAnalysisToUseSingleBucket();
     }
 
-    void pushDownGroupByCount(Expression groupBy) {
-        final String columnName = PlanUtil.expressionToColumnName(groupBy);
-        addOperatorToPlan(PlanUtil.buildGroupByCount(columnName));
+    void pushDownGroupByCount(Expression[] groupByExpressions) {
+        List<String> columnNames = Stream.of(groupByExpressions)
+            .map(groupBy -> PlanUtil.expressionToColumnName(groupBy))
+            .collect(Collectors.toList());
 
-        StructField columnField = null;
-        for (StructField field : this.schema.fields()) {
-            if (columnName.equals(field.name())) {
-                columnField = field;
-                break;
+        addOperatorToPlan(PlanUtil.buildGroupByCount(columnNames));
+
+        StructType newSchema = new StructType();
+
+        for (String columnName : columnNames) {
+            StructField columnField = null;
+            for (StructField field : this.schema.fields()) {
+                if (columnName.equals(field.name())) {
+                    columnField = field;
+                    break;
+                }
             }
+            if (columnField == null) {
+                throw new IllegalArgumentException("Unable to find groupBy column in schema; column name: " + columnName);
+            }
+            newSchema = newSchema.add(columnField);
         }
-        if (columnField == null) {
-            throw new IllegalArgumentException("Unable to find groupBy column in schema; groupBy expression: " + groupBy.describe());
-        }
-        this.schema = new StructType().add(columnField).add("count", DataTypes.LongType);
+
+        this.schema = newSchema.add("count", DataTypes.LongType);
         modifyPlanAnalysisToUseSingleBucket();
     }
 

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

@@ -16,6 +16,7 @@
 package com.marklogic.spark.reader;
 
 import com.marklogic.spark.Options;
+import org.apache.spark.sql.Column;
 import org.apache.spark.sql.Row;
 import org.junit.jupiter.api.Test;
 
@@ -39,6 +40,24 @@ void groupByWithNoQualifier() {
         assertEquals(1l, (long) rows.get(0).getAs("CitationID"));
     }
 
+    @Test
+    void groupByMultipleColumns() {
+        List<Row> rows = newDefaultReader()
+            .option(Options.READ_OPTIC_QUERY, QUERY_WITH_NO_QUALIFIER)
+            .load()
+            .groupBy("CitationID", "Date")
+            .count()
+            .orderBy("CitationID")
+            .collectAsList();
+
+        verifyGroupByWasPushedDown(rows);
+
+        assertEquals(1l, (long) rows.get(0).getAs("CitationID"));
+        assertEquals("2022-07-13", rows.get(0).getAs("Date").toString());
+        assertEquals(2l, (long) rows.get(1).getAs("CitationID"));
+        assertEquals("2022-05-11", rows.get(1).getAs("Date").toString());
+    }
+
     @Test
     void noRowsFound() {
         List<Row> rows = newDefaultReader()
@@ -80,6 +99,23 @@ void groupByWithSchemaAndView() {
         assertEquals(1l, (long) rows.get(0).getAs("Medical.Authors.CitationID"));
     }
 
+    @Test
+    void groupByMultipleColumnsAndSchemaAndView() {
+        List<Row> rows = newDefaultReader()
+            .load()
+            .groupBy("`Medical.Authors.CitationID`", "`Medical.Authors.Date`")
+            .count()
+            .orderBy("`Medical.Authors.CitationID`")
+            .collectAsList();
+
+        verifyGroupByWasPushedDown(rows);
+
+        verifyGroupByWasPushedDown(rows);
+        assertEquals(1l, (long) rows.get(0).getAs("Medical.Authors.CitationID"));
+        assertEquals("2022-07-13", rows.get(0).getAs("Medical.Authors.Date").toString());
+    }
+
+
     @Test
     void groupByCountLimitOrderBy() {
         List<Row> rows = newDefaultReader()