Small tweaks to avoid unnecessary object creation

Was reading through the Spark javadocs as I saw the Batch objects being created multiple times and the partitions logic being called multiple times. That's not necessary for us - the partitions will be the same for each Batch object. So having the Scan object eagerly create a Batch when possible and having that Batch eagerly calculate its partitions as well. No functional change, but does make our logging less confusing because the user won't see what seem like duplicate log entries (like for the number of partitions being used for reading documents).

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