Merge pull request #74 from marklogic/feature/485-disable-aggregates

DEVEXP-485 Can now disable push down of aggregates

Author: rjrudin

docs/configuration.md

@@ -69,7 +69,7 @@ information on how data is read from MarkLogic.
 | spark.marklogic.read.opticQuery | Required; the Optic DSL query to run for retrieving rows; must use `op.fromView` as the accessor. |
 | spark.marklogic.read.numPartitions | The number of Spark partitions to create; defaults to `spark.default.parallelism` .               |
 | spark.marklogic.read.batchSize | Approximate number of rows to retrieve in each call to MarkLogic; defaults to 10000.              |
-
+| spark.marklogic.read.pushDownAggregates | Whether to push down aggregate operations to MarkLogic; defaults to `true`. Set to `false` to prevent aggregates from being pushed down to MarkLogic. |
 ## Write options
 
 These options control how the connector writes data to MarkLogic. See [the guide on writing](writing.md) for more 

docs/reading.md

@@ -160,6 +160,13 @@ The following results are returned:
 +-----+-----------+-----+
 ```
 
+### Disabling push down of aggregates
+
+If you run into any issues with aggregates being pushed down to MarkLogic, you can set the 
+`spark.marklogic.read.pushDownAggregates` option to `false`. If doing so results in what appears to be a different and 
+correct result, please [file an issue with this project](https://github.com/marklogic/marklogic-spark-connector/issues).
+
+
 ## Tuning performance
 
 The primary factor affecting how quickly the connector can retrieve rows is MarkLogic's ability to process your Optic 

src/main/java/com/marklogic/spark/Options.java

@@ -22,6 +22,7 @@ public interface Options {
     String READ_OPTIC_QUERY = "spark.marklogic.read.opticQuery";
     String READ_NUM_PARTITIONS = "spark.marklogic.read.numPartitions";
     String READ_BATCH_SIZE = "spark.marklogic.read.batchSize";
+    String READ_PUSH_DOWN_AGGREGATES = "spark.marklogic.read.pushDownAggregates";
 
     String WRITE_BATCH_SIZE = "spark.marklogic.write.batchSize";
     String WRITE_THREAD_COUNT = "spark.marklogic.write.threadCount";

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

@@ -15,9 +15,11 @@
  */
 package com.marklogic.spark.reader;
 
+import com.marklogic.spark.Options;
 import com.marklogic.spark.reader.filter.FilterFactory;
 import com.marklogic.spark.reader.filter.OpticFilter;
 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;
 import org.apache.spark.sql.connector.expressions.aggregate.Avg;
 import org.apache.spark.sql.connector.expressions.aggregate.Count;
@@ -49,10 +51,10 @@ public class MarkLogicScanBuilder implements ScanBuilder, SupportsPushDownFilter
 
     private final static Logger logger = LoggerFactory.getLogger(MarkLogicScanBuilder.class);
 
-    private ReadContext readContext;
+    private final ReadContext readContext;
     private List<Filter> pushedFilters;
 
-    private final static Set<Class> SUPPORTED_AGGREGATE_FUNCTIONS = new HashSet() {{
+    private final static Set<Class<? extends AggregateFunc>> SUPPORTED_AGGREGATE_FUNCTIONS = new HashSet() {{
         add(Avg.class);
         add(Count.class);
         add(CountStar.class);
@@ -164,7 +166,7 @@ public boolean isPartiallyPushed() {
      */
     @Override
     public boolean supportCompletePushDown(Aggregation aggregation) {
-        if (readContext.planAnalysisFoundNoRows()) {
+        if (readContext.planAnalysisFoundNoRows() || pushDownAggregatesIsDisabled()) {
             return false;
         }
 
@@ -190,6 +192,12 @@ public boolean pushAggregation(Aggregation aggregation) {
         if (readContext.planAnalysisFoundNoRows() || hasUnsupportedAggregateFunction(aggregation)) {
             return false;
         }
+
+        if (pushDownAggregatesIsDisabled()) {
+            logger.info("Push down of aggregates is disabled; Spark will handle all aggregations.");
+            return false;
+        }
+
         logger.info("Pushing down aggregation: {}", describeAggregation(aggregation));
         readContext.pushDownAggregation(aggregation);
         return true;
@@ -224,4 +232,8 @@ private String describeAggregation(Aggregation aggregation) {
             Arrays.asList(aggregation.groupByExpressions()),
             Arrays.asList(aggregation.aggregateExpressions()));
     }
+
+    private boolean pushDownAggregatesIsDisabled() {
+        return "false".equalsIgnoreCase(readContext.getProperties().get(Options.READ_PUSH_DOWN_AGGREGATES));
+    }
 }

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

@@ -0,0 +1,36 @@
+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 DisablePushDownAggregatesTest extends AbstractPushDownTest {
+
+    @Test
+    void disabled() {
+        List<Row> rows = newDefaultReader()
+            .option(Options.READ_OPTIC_QUERY, QUERY_WITH_NO_QUALIFIER)
+            .option(Options.READ_PUSH_DOWN_AGGREGATES, false)
+            .load()
+            .groupBy("CitationID")
+            .avg("LuckyNumber")
+            .orderBy("CitationID")
+            .collectAsList();
+
+        assertEquals(5, rows.size());
+        assertEquals(15, countOfRowsReadFromMarkLogic, "Because push down of aggregates is disabled, all 15 author " +
+            "rows should have been read from MarkLogic.");
+
+        // Averages should still be calculated correctly by Spark.
+        String columnName = "avg(LuckyNumber)";
+        assertEquals(2.5, (double) rows.get(0).getAs(columnName));
+        assertEquals(6.5, (double) rows.get(1).getAs(columnName));
+        assertEquals(10.5, (double) rows.get(2).getAs(columnName));
+        assertEquals(13.0, (double) rows.get(3).getAs(columnName));
+        assertEquals(14.5, (double) rows.get(4).getAs(columnName));
+    }
+}