Merge pull request #272 from marklogic/feature/reprocess-logging

Added progress logging for reading items

Author: rjrudin

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

@@ -101,8 +101,7 @@ public Table getTable(StructType schema, Transform[] partitioning, Map<String, S
             ReadProgressLogger.initialize(readProgressInterval, "Rows read: {}");
             return new MarkLogicTable(schema, properties);
         } else if (Util.isReadWithCustomCodeOperation(properties)) {
-            // Not yet logging progress for reading with custom code, as it's assumed the user will then write with
-            // custom code.
+            ReadProgressLogger.initialize(readProgressInterval, "Items read: {}");
             return new MarkLogicTable(schema, properties);
         }
 

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

@@ -4,6 +4,7 @@
 import com.marklogic.client.eval.EvalResultIterator;
 import com.marklogic.client.eval.ServerEvaluationCall;
 import com.marklogic.spark.Options;
+import com.marklogic.spark.ReadProgressLogger;
 import com.marklogic.spark.reader.JsonRowDeserializer;
 import org.apache.spark.sql.catalyst.InternalRow;
 import org.apache.spark.sql.catalyst.expressions.GenericInternalRow;
@@ -19,6 +20,10 @@ class CustomCodePartitionReader implements PartitionReader<InternalRow> {
     private final JsonRowDeserializer jsonRowDeserializer;
     private final DatabaseClient databaseClient;
 
+    // Only needed for logging progress.
+    private final long batchSize;
+    private long progressCounter;
+
     public CustomCodePartitionReader(CustomCodeContext customCodeContext, String partition) {
         this.databaseClient = customCodeContext.connectToMarkLogic();
         this.serverEvaluationCall = customCodeContext.buildCall(
@@ -31,6 +36,8 @@ public CustomCodePartitionReader(CustomCodeContext customCodeContext, String par
             this.serverEvaluationCall.addVariable("PARTITION", partition);
         }
 
+        this.batchSize = customCodeContext.getNumericOption(Options.READ_BATCH_SIZE, 1, 1);
+
         this.isCustomSchema = customCodeContext.isCustomSchema();
         this.jsonRowDeserializer = new JsonRowDeserializer(customCodeContext.getSchema());
     }
@@ -49,6 +56,11 @@ public InternalRow get() {
         if (this.isCustomSchema) {
             return this.jsonRowDeserializer.deserializeJson(val);
         }
+        progressCounter++;
+        if (progressCounter >= batchSize) {
+            ReadProgressLogger.logProgressIfNecessary(progressCounter);
+            progressCounter = 0;
+        }
         return new GenericInternalRow(new Object[]{UTF8String.fromString(val)});
     }
 

src/test/java/com/marklogic/spark/reader/customcode/ReadWithCustomCodeTest.java

@@ -197,6 +197,8 @@ void verifyTimeoutWorks() {
     private List<Row> readRows(String option, String value) {
         return startRead()
             .option(option, value)
+            // Adding these only for manual inspection of logging and to ensure they don't cause errors.
+            .option(Options.READ_LOG_PROGRESS, "1")
             .load()
             .collectAsList();
     }