Merge pull request #169 from marklogic/feature/2.2.0-performance

Small tweaks to avoid unnecessary object creation

Author: rjrudin

src/main/java/com/marklogic/spark/reader/customcode/CustomCodeBatch.java

@@ -9,16 +9,10 @@
 class CustomCodeBatch implements Batch {
 
     private CustomCodeContext customCodeContext;
-    private List<String> partitions;
+    private InputPartition[] inputPartitions;
 
-    public CustomCodeBatch(CustomCodeContext customCodeContext, List<String> partitions) {
+    CustomCodeBatch(CustomCodeContext customCodeContext, List<String> partitions) {
         this.customCodeContext = customCodeContext;
-        this.partitions = partitions;
-    }
-
-    @Override
-    public InputPartition[] planInputPartitions() {
-        InputPartition[] inputPartitions;
         if (partitions != null && partitions.size() > 1) {
             inputPartitions = new InputPartition[partitions.size()];
             for (int i = 0; i < partitions.size(); i++) {
@@ -27,6 +21,10 @@ public InputPartition[] planInputPartitions() {
         } else {
             inputPartitions = new InputPartition[]{new CustomCodePartition()};
         }
+    }
+
+    @Override
+    public InputPartition[] planInputPartitions() {
         return inputPartitions;
     }
 

src/main/java/com/marklogic/spark/reader/customcode/CustomCodeScan.java

@@ -15,6 +15,7 @@ class CustomCodeScan implements Scan {
 
     private CustomCodeContext customCodeContext;
     private final List<String> partitions;
+    private final CustomCodeBatch batch;
 
     public CustomCodeScan(CustomCodeContext customCodeContext) {
         this.customCodeContext = customCodeContext;
@@ -35,6 +36,8 @@ public CustomCodeScan(CustomCodeContext customCodeContext) {
                 client.release();
             }
         }
+
+        batch = new CustomCodeBatch(customCodeContext, partitions);
     }
 
     @Override
@@ -44,7 +47,7 @@ public StructType readSchema() {
 
     @Override
     public Batch toBatch() {
-        return new CustomCodeBatch(customCodeContext, partitions);
+        return batch;
     }
 
     @Override

src/main/java/com/marklogic/spark/reader/document/DocumentBatch.java

@@ -5,6 +5,7 @@
 import com.marklogic.client.io.SearchHandle;
 import com.marklogic.client.query.QueryManager;
 import com.marklogic.client.query.SearchQueryDefinition;
+import com.marklogic.spark.Util;
 import org.apache.spark.sql.connector.read.Batch;
 import org.apache.spark.sql.connector.read.InputPartition;
 import org.apache.spark.sql.connector.read.PartitionReaderFactory;
@@ -16,19 +17,15 @@ class DocumentBatch implements Batch {
     private static final Logger logger = LoggerFactory.getLogger(DocumentBatch.class);
 
     private final DocumentContext context;
-
-    DocumentBatch(DocumentContext context) {
-        this.context = context;
-    }
+    private final InputPartition[] partitions;
 
     /**
      * Reuses the DMSDK support for obtaining a list of all eligible forests. A partition reader will then be
      * created for each partition/forest.
-     *
-     * @return
      */
-    @Override
-    public InputPartition[] planInputPartitions() {
+    DocumentBatch(DocumentContext context) {
+        this.context = context;
+
         DatabaseClient client = this.context.connectToMarkLogic();
         Forest[] forests = client.newDataMovementManager().readForestConfig().listForests();
 
@@ -47,11 +44,16 @@ public InputPartition[] planInputPartitions() {
             logger.debug("Creating forest partitions; query estimate: {}; server timestamp: {}", estimate, serverTimestamp);
         }
         ForestPartitionPlanner planner = new ForestPartitionPlanner(context.getPartitionsPerForest());
-        ForestPartition[] partitions = planner.makePartitions(estimate, serverTimestamp, forests);
-        if (logger.isDebugEnabled()) {
-            logger.debug("Created {} forest partitions", partitions.length);
+        this.partitions = planner.makePartitions(estimate, serverTimestamp, forests);
+
+        if (Util.MAIN_LOGGER.isInfoEnabled()) {
+            logger.info("Created {} partitions; query estimate: {}", partitions.length, estimate);
         }
-        return partitions;
+    }
+
+    @Override
+    public InputPartition[] planInputPartitions() {
+        return this.partitions;
     }
 
     @Override

src/main/java/com/marklogic/spark/reader/document/DocumentScan.java

@@ -6,10 +6,10 @@
 
 class DocumentScan implements Scan {
 
-    private final DocumentContext context;
+    private final DocumentBatch batch;
 
     DocumentScan(DocumentContext context) {
-        this.context = context;
+        this.batch = new DocumentBatch(context);
     }
 
     @Override
@@ -19,6 +19,6 @@ public StructType readSchema() {
 
     @Override
     public Batch toBatch() {
-        return new DocumentBatch(context);
+        return this.batch;
     }
 }

src/main/java/com/marklogic/spark/reader/optic/OpticBatch.java

@@ -26,17 +26,19 @@ class OpticBatch implements Batch {
     private static final Logger logger = LoggerFactory.getLogger(OpticBatch.class);
 
     private final ReadContext readContext;
+    private final InputPartition[] partitions;
 
     OpticBatch(ReadContext readContext) {
         this.readContext = readContext;
+        PlanAnalysis planAnalysis = readContext.getPlanAnalysis();
+        partitions = planAnalysis != null ?
+            planAnalysis.getPartitionArray() :
+            new InputPartition[]{};
     }
 
     @Override
     public InputPartition[] planInputPartitions() {
-        PlanAnalysis planAnalysis = readContext.getPlanAnalysis();
-        return planAnalysis != null ?
-            planAnalysis.getPartitionArray() :
-            new InputPartition[]{};
+        return partitions;
     }
 
     @Override