Fixes for count/orderBy on qualified column names

And... I implemented groupBy + count while I was in here fixing things as it was trivial to do. Also tweaked some logging based on testing. But the main thing is the addition of tests to verify that column names work, regardless of whether they have no qualifier, only a view, or a schema + view.

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"));
+    }
+}