Merge pull request #47 from marklogic/feature/groupBy

Fixes for count/orderBy on qualified column names

Author: rjrudin

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

@@ -90,8 +90,8 @@ public boolean next() {
             if (rowIterator.hasNext()) {
                 return true;
             } else {
-                if (logger.isTraceEnabled()) {
-                    logger.trace("Count of rows for partition {} and bucket {}: {}", this.partition,
+                if (logger.isDebugEnabled()) {
+                    logger.debug("Count of rows for partition {} and bucket {}: {}", this.partition,
                         this.partition.buckets.get(nextBucketIndex - 1), currentBucketRowCount);
                 }
                 currentBucketRowCount = 0;

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

@@ -17,6 +17,7 @@
 
 import com.marklogic.spark.reader.filter.FilterFactory;
 import com.marklogic.spark.reader.filter.OpticFilter;
+import org.apache.spark.sql.connector.expressions.Expression;
 import org.apache.spark.sql.connector.expressions.SortOrder;
 import org.apache.spark.sql.connector.expressions.aggregate.AggregateFunc;
 import org.apache.spark.sql.connector.expressions.aggregate.Aggregation;
@@ -140,28 +141,45 @@ public boolean pushOffset(int offset) {
     @Override
     public boolean pushAggregation(Aggregation aggregation) {
         if (supportCompletePushDown(aggregation)) {
-            if (logger.isDebugEnabled()) {
-                logger.debug("Pushing down count()");
+            if (aggregation.groupByExpressions().length > 0) {
+                Expression expr = aggregation.groupByExpressions()[0];
+                if (logger.isDebugEnabled()) {
+                    logger.debug("Pushing down by groupBy + count on: {}", expr.describe());
+                }
+                readContext.pushDownGroupByCount(expr);
+            } else {
+                if (logger.isDebugEnabled()) {
+                    logger.debug("Pushing down count()");
+                }
+                readContext.pushDownCount();
             }
-            readContext.pushDownCount();
             return true;
         }
         return false;
     }
 
     @Override
     public boolean supportCompletePushDown(Aggregation aggregation) {
-        // 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.length == 1 && expressions[0] instanceof CountStar && aggregation.groupByExpressions().length == 0;
+        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;
     }
 
     @Override
     public void pruneColumns(StructType requiredSchema) {
-        if (logger.isDebugEnabled()) {
-            logger.debug("Pushing down required schema: {}", requiredSchema.json());
+        if (requiredSchema.equals(readContext.getSchema())) {
+            if (logger.isDebugEnabled()) {
+                logger.debug("The schema to push down is equal to the existing schema, so not pushing it down.");
+            }
+        } else {
+            if (logger.isDebugEnabled()) {
+                logger.debug("Pushing down required schema: {}", requiredSchema.json());
+            }
+            readContext.pushDownRequiredSchema(requiredSchema);
         }
-        readContext.pushDownRequiredSchema(requiredSchema);
     }
 }

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

@@ -4,10 +4,14 @@
 import com.fasterxml.jackson.databind.node.ArrayNode;
 import com.fasterxml.jackson.databind.node.ObjectNode;
 import com.marklogic.spark.reader.filter.OpticFilter;
+import org.apache.spark.sql.connector.expressions.Expression;
+import org.apache.spark.sql.connector.expressions.NamedReference;
 import org.apache.spark.sql.connector.expressions.SortDirection;
 import org.apache.spark.sql.connector.expressions.SortOrder;
 import org.apache.spark.sql.types.StructField;
 import org.apache.spark.sql.types.StructType;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 
 import java.util.List;
 import java.util.function.Consumer;
@@ -18,15 +22,24 @@
  */
 abstract class PlanUtil {
 
+    private final static Logger logger = LoggerFactory.getLogger(PlanUtil.class);
+
     private final static ObjectMapper objectMapper = new ObjectMapper();
 
     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()));
     }
 
+    static ObjectNode buildGroupByCount(String columnName) {
+        return newOperation("group-by", args -> {
+            populateSchemaCol(args.addObject(), columnName);
+            // Using "null" is the equivalent of "count(*)" - it counts rows, not values.
+            args.addObject().put("ns", "op").put("fn", "count").putArray("args").add("count").add(objectMapper.nullNode());
+        });
+    }
+
     static ObjectNode buildLimit(int limit) {
         return newOperation("limit", args -> args.add(limit));
     }
@@ -37,30 +50,41 @@ static ObjectNode buildOffset(int offset) {
 
     static ObjectNode buildOrderBy(SortOrder sortOrder) {
         final String direction = SortDirection.ASCENDING.equals(sortOrder.direction()) ? "asc" : "desc";
-        final String columnName = sortOrder.expression().describe();
-        return newOperation("order-by", args -> args.addObject()
-            .put("ns", "op").put("fn", direction)
-            .putArray("args").addObject()
-            .put("ns", "op").put("fn", "col").putArray("args").add(columnName));
+        final String columnName = expressionToColumnName(sortOrder.expression());
+        return newOperation("order-by", args -> {
+            ArrayNode orderByArgs = args.addObject().put("ns", "op").put("fn", direction).putArray("args");
+            // This may be a bad hack to account for when the user does a groupBy/count/orderBy/limit, which does not
+            // seem like the correct approach - the Spark ScanBuilder javadocs indicate that it should be limit/orderBy
+            // instead. In the former scenario, we get "COUNT(*)" as the expression to order by, and we know that's not
+            // the column name.
+            if (logger.isDebugEnabled()) {
+                logger.debug("Adjusting `COUNT(*)` column to be `count`");
+            }
+            populateSchemaCol(orderByArgs.addObject(), "COUNT(*)".equals(columnName) ? "count" : columnName);
+        });
     }
 
     static ObjectNode buildSelect(StructType schema) {
         return newOperation("select", args -> {
             ArrayNode innerArgs = args.addArray();
             for (StructField field : schema.fields()) {
-                ArrayNode colArgs = innerArgs.addObject().put("ns", "op").put("fn", "schema-col").putArray("args");
-                String[] parts = field.name().split("\\.");
-                if (parts.length == 3) {
-                    colArgs.add(parts[0]).add(parts[1]).add(parts[2]);
-                } else if (parts.length == 2) {
-                    colArgs.add(objectMapper.nullNode()).add(parts[0]).add(parts[1]);
-                } else {
-                    colArgs.add(objectMapper.nullNode()).add(objectMapper.nullNode()).add(parts[0]);
-                }
+                populateSchemaCol(innerArgs.addObject(), field.name());
             }
         });
     }
 
+    private static void populateSchemaCol(ObjectNode node, String columnName) {
+        ArrayNode colArgs = node.put("ns", "op").put("fn", "schema-col").putArray("args");
+        String[] parts = columnName.split("\\.");
+        if (parts.length == 3) {
+            colArgs.add(parts[0]).add(parts[1]).add(parts[2]);
+        } else if (parts.length == 2) {
+            colArgs.add(objectMapper.nullNode()).add(parts[0]).add(parts[1]);
+        } else {
+            colArgs.add(objectMapper.nullNode()).add(objectMapper.nullNode()).add(parts[0]);
+        }
+    }
+
     static ObjectNode buildWhere(List<OpticFilter> opticFilters) {
         return newOperation("where", args -> {
             // If there's only one filter, can toss it into the "where" clause. Else, toss an "and" into the "where" and
@@ -78,4 +102,19 @@ private static ObjectNode newOperation(String name, Consumer<ArrayNode> withArgs
         withArgs.accept(operation.putArray("args"));
         return operation;
     }
+
+    static String expressionToColumnName(Expression expression) {
+        // The structure of an Expression isn't well-understood yet. But when it refers to a single column, the
+        // column name can be found in the below manner. Anything else is not supported yet.
+        NamedReference[] refs = expression.references();
+        if (refs == null || refs.length < 1) {
+            return expression.describe();
+        }
+        String[] fieldNames = refs[0].fieldNames();
+        if (fieldNames.length != 1) {
+            throw new IllegalArgumentException("Unsupported expression: " + expression + "; expecting expression " +
+                "to have exactly one field name.");
+        }
+        return fieldNames[0];
+    }
 }

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

@@ -30,8 +30,10 @@
 import com.marklogic.spark.Options;
 import com.marklogic.spark.reader.filter.OpticFilter;
 import org.apache.spark.sql.SparkSession;
+import org.apache.spark.sql.connector.expressions.Expression;
 import org.apache.spark.sql.connector.expressions.SortOrder;
 import org.apache.spark.sql.types.DataTypes;
+import org.apache.spark.sql.types.StructField;
 import org.apache.spark.sql.types.StructType;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -172,15 +174,44 @@ void pushDownTopN(SortOrder[] orders, int limit) {
     void pushDownCount() {
         if (planAnalysisFoundAtLeastOneRow()) {
             addOperatorToPlan(PlanUtil.buildGroupByCount());
-
-            // 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.
+            // As will likely be the case for all aggregations, the schema needs to be modified.
             this.schema = new StructType().add("count", DataTypes.LongType);
-            this.planAnalysis = new PlanAnalysis(this.planAnalysis.boundedPlan);
+            modifyPlanAnalysisToUseSingleBucket();
+        }
+    }
+
+    void pushDownGroupByCount(Expression groupBy) {
+        if (planAnalysisFoundAtLeastOneRow()) {
+            final String columnName = PlanUtil.expressionToColumnName(groupBy);
+            addOperatorToPlan(PlanUtil.buildGroupByCount(columnName));
+
+            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; groupBy expression: " + groupBy.describe());
+            }
+            this.schema = new StructType().add(columnField).add("count", DataTypes.LongType);
+            modifyPlanAnalysisToUseSingleBucket();
         }
     }
 
+    /**
+     * Used when the assumption is that MarkLogic can efficiently execute a plan in a single call to /v1/rows. This is
+     * typically done for "count()" operations. In such a scenario, returning 2 or more rows may produce an incorrect
+     * result as well - for example, for a "count()" call, only the first row will be reported as the count.
+     */
+    private void modifyPlanAnalysisToUseSingleBucket() {
+        if (logger.isDebugEnabled()) {
+            logger.debug("Modifying plan analysis to use a single bucket");
+        }
+        this.planAnalysis = new PlanAnalysis(this.planAnalysis.boundedPlan);
+    }
+
     void pushDownRequiredSchema(StructType requiredSchema) {
         if (planAnalysisFoundAtLeastOneRow()) {
             this.schema = requiredSchema;
@@ -202,6 +233,9 @@ private boolean planAnalysisFoundAtLeastOneRow() {
      * @param operator
      */
     private void addOperatorToPlan(ObjectNode operator) {
+        if (logger.isDebugEnabled()) {
+            logger.debug("Adding operator to plan: {}", operator);
+        }
         ArrayNode operators = (ArrayNode) planAnalysis.boundedPlan.get("$optic").get("args");
         operators.insert(operators.size() - 1, operator);
     }

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

@@ -26,24 +26,4 @@ void count() {
             "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"));
-    }
 }

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

@@ -0,0 +1,103 @@
+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 PushDownGroupByCountTest extends AbstractPushDownTest {
+
+    @Test
+    void groupByWithNoQualifier() {
+        List<Row> rows = newDefaultReader()
+            .option(Options.READ_OPTIC_DSL, QUERY_WITH_NO_QUALIFIER)
+            .load()
+            .groupBy("CitationID")
+            .count()
+            .orderBy("CitationID")
+            .collectAsList();
+
+        verifyGroupByWasPushedDown(rows);
+        assertEquals(1l, (long) rows.get(0).getAs("CitationID"));
+    }
+
+    @Test
+    void groupByWithView() {
+        List<Row> rows = newDefaultReader()
+            .option(Options.READ_OPTIC_DSL, "op.fromView('Medical', 'Authors', 'example')")
+            .load()
+            .groupBy("`example.CitationID`")
+            .count()
+            .orderBy("`example.CitationID`")
+            .collectAsList();
+
+        verifyGroupByWasPushedDown(rows);
+        assertEquals(1l, (long) rows.get(0).getAs("example.CitationID"));
+    }
+
+    @Test
+    void groupByWithSchemaAndView() {
+        List<Row> rows = newDefaultReader()
+            .load()
+            .groupBy("`Medical.Authors.CitationID`")
+            .count()
+            .orderBy("`Medical.Authors.CitationID`")
+            .collectAsList();
+
+        verifyGroupByWasPushedDown(rows);
+        assertEquals(1l, (long) rows.get(0).getAs("Medical.Authors.CitationID"));
+    }
+
+    @Test
+    void groupByCountLimitOrderBy() {
+        List<Row> rows = newDefaultReader()
+            .option(Options.READ_OPTIC_DSL, QUERY_WITH_NO_QUALIFIER)
+            .load()
+            .groupBy("CitationID")
+            .count()
+            .limit(4)
+            // When the user puts the orderBy after limit, Spark doesn't push the orderBy down. Spark will instead
+            // apply the orderBy itself.
+            .orderBy("count")
+            .collectAsList();
+
+        assertEquals(4, rows.size());
+        assertEquals(4, countOfRowsReadFromMarkLogic);
+        assertEquals(4l, (long) rows.get(0).getAs("CitationID"));
+        assertEquals(1l, (long) rows.get(0).getAs("count"));
+    }
+
+    @Test
+    void groupByCountOrderByLimit() {
+        List<Row> rows = newDefaultReader()
+            .option(Options.READ_OPTIC_DSL, QUERY_WITH_NO_QUALIFIER)
+            .load()
+            .groupBy("CitationID")
+            .count()
+            // If the user puts orderBy before limit, Spark will send "COUNT(*)" as the column name for the orderBy.
+            // The connector is expected to translate that into "count"; not sure how it should work otherwise. Spark
+            // is expected to push down the limit as well.
+            .orderBy("count")
+            .limit(4)
+            .collectAsList();
+
+        assertEquals(4, rows.size());
+        assertEquals(4, countOfRowsReadFromMarkLogic);
+        assertEquals(4l, (long) rows.get(0).getAs("CitationID"));
+        assertEquals(1l, (long) rows.get(0).getAs("count"));
+    }
+
+    private void verifyGroupByWasPushedDown(List<Row> rows) {
+        assertEquals(5, countOfRowsReadFromMarkLogic, "groupBy should be pushed down to MarkLogic when used with " +
+            "count, and since there are 5 CitationID values, 5 rows should be returned.");
+
+        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"));
+    }
+}