Merge pull request #266 from marklogic/feature/read-docs-progress

Fixed progress logging bug

Author: rjrudin

CONTRIBUTING.md

@@ -156,6 +156,20 @@ df2.head()
 json.loads(df2.head()['content'])
 ```
 
+For a quick test of writing documents, use the following:
+
+```
+
+spark.read.option("header", True).csv("src/test/resources/data.csv")\
+    .repartition(2)\
+    .write.format("marklogic")\
+    .option("spark.marklogic.client.uri", "spark-test-user:spark@localhost:8000")\
+    .option("spark.marklogic.write.permissions", "spark-user-role,read,spark-user-role,update")\
+    .option("spark.marklogic.write.logProgress", 50)\
+    .option("spark.marklogic.write.batchSize", 10)\
+    .mode("append")\
+    .save()
+```
 
 # Testing against a local Spark cluster
 

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

@@ -93,6 +93,7 @@ public Table getTable(StructType schema, Transform[] partitioning, Map<String, S
             return new MarkLogicTable(schema, properties);
         }
 
+        WriteProgressLogger.progressCounter.set(0);
         return new MarkLogicTable(new WriteContext(schema, properties));
     }
 

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

@@ -90,7 +90,9 @@ public abstract class Options {
     public static final String WRITE_THREAD_COUNT_PER_PARTITION = "spark.marklogic.write.threadCountPerPartition";
     public static final String WRITE_ABORT_ON_FAILURE = "spark.marklogic.write.abortOnFailure";
 
-    // For logging progress when writing documents or processing with custom code.
+    // For logging progress when writing documents or processing with custom code. Defines the interval at which
+    // progress should be logged - e.g. a value of 10,000 will result in a message being logged on every 10,000 items
+    // being written/processed.
     public static final String WRITE_LOG_PROGRESS = "spark.marklogic.write.logProgress";
 
     // For writing via custom code.

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

@@ -0,0 +1,36 @@
+/*
+ * Copyright © 2024 Progress Software Corporation and/or its subsidiaries or affiliates. All Rights Reserved.
+ */
+package com.marklogic.spark;
+
+import java.io.Serializable;
+
+public abstract class ProgressLogger implements Serializable {
+
+    static final long serialVersionUID = 1;
+
+    private final long progressInterval;
+    private final int batchSize;
+    private final String message;
+
+    protected ProgressLogger(long progressInterval, int batchSize, String message) {
+        this.progressInterval = progressInterval;
+        this.batchSize = batchSize;
+        this.message = message;
+    }
+
+    protected abstract long getNewSum(long itemCount);
+
+    public void logProgressIfNecessary(long itemCount) {
+        if (this.progressInterval > 0) {
+            long sum = getNewSum(itemCount);
+            if (sum >= progressInterval) {
+                long lowerBound = sum / (this.progressInterval);
+                long upperBound = (lowerBound * this.progressInterval) + this.batchSize;
+                if (sum >= lowerBound && sum < upperBound) {
+                    Util.MAIN_LOGGER.info(message, sum);
+                }
+            }
+        }
+    }
+}

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

@@ -0,0 +1,21 @@
+/*
+ * Copyright © 2024 Progress Software Corporation and/or its subsidiaries or affiliates. All Rights Reserved.
+ */
+package com.marklogic.spark;
+
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.atomic.AtomicLong;
+
+public class WriteProgressLogger extends ProgressLogger {
+
+    public static final AtomicLong progressCounter = new AtomicLong(0);
+
+    public WriteProgressLogger(long progressInterval, int batchSize, String message) {
+        super(progressInterval, batchSize, message);
+    }
+
+    @Override
+    protected long getNewSum(long itemCount) {
+        return progressCounter.addAndGet(itemCount);
+    }
+}

src/main/java/com/marklogic/spark/writer/WriteContext.java

@@ -24,10 +24,7 @@
 import com.marklogic.client.document.ServerTransform;
 import com.marklogic.client.impl.GenericDocumentImpl;
 import com.marklogic.client.io.Format;
-import com.marklogic.spark.ConnectorException;
-import com.marklogic.spark.ContextSupport;
-import com.marklogic.spark.Options;
-import com.marklogic.spark.Util;
+import com.marklogic.spark.*;
 import com.marklogic.spark.reader.document.DocumentRowSchema;
 import com.marklogic.spark.reader.file.TripleRowSchema;
 import org.apache.spark.sql.types.StructType;
@@ -36,18 +33,16 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
-import java.util.concurrent.atomic.AtomicInteger;
 import java.util.stream.Stream;
 
 public class WriteContext extends ContextSupport {
 
     static final long serialVersionUID = 1;
-    private static final AtomicInteger progressTracker = new AtomicInteger();
 
     private final StructType schema;
     private final boolean usingFileSchema;
     private final int batchSize;
-    private final int logProgress;
+    private final ProgressLogger progressLogger;
 
     private int fileSchemaContentPosition;
     private int fileSchemaPathPosition;
@@ -59,7 +54,8 @@ public WriteContext(StructType schema, Map<String, String> properties) {
         super(properties);
         this.schema = schema;
         this.batchSize = (int) getNumericOption(Options.WRITE_BATCH_SIZE, 100, 1);
-        this.logProgress = (int) getNumericOption(Options.WRITE_LOG_PROGRESS, 10000, 0);
+        this.progressLogger = new WriteProgressLogger(getNumericOption(Options.WRITE_LOG_PROGRESS, 0, 0),
+            batchSize, "Documents written: {}");
 
         // We support the Spark binaryFile schema - https://spark.apache.org/docs/latest/sql-data-sources-binaryFile.html -
         // so that reader can be reused for loading files as-is.
@@ -274,22 +270,9 @@ private void logBatchOnSuccess(WriteBatch batch) {
                 docCount--;
             }
         }
-        if (this.logProgress > 0) {
-            logProgressIfNecessary(docCount);
-        }
-        if (logger.isDebugEnabled()) {
-            logger.debug("Wrote batch; length: {}; job batch number: {}", docCount, batch.getJobBatchNumber());
-        }
-    }
-
-    private void logProgressIfNecessary(int docCount) {
-        int sum = progressTracker.addAndGet(docCount);
-        if (sum >= logProgress) {
-            int lowerBound = sum / (this.logProgress);
-            int upperBound = (lowerBound * this.logProgress) + this.batchSize;
-            if (sum >= lowerBound && sum < upperBound) {
-                Util.MAIN_LOGGER.info("Documents written: {}", sum);
-            }
+        progressLogger.logProgressIfNecessary(docCount);
+        if (logger.isTraceEnabled()) {
+            logger.trace("Wrote batch; length: {}; job batch number: {}", docCount, batch.getJobBatchNumber());
         }
     }
 

src/main/java/com/marklogic/spark/writer/customcode/CustomCodeWriter.java

@@ -9,10 +9,7 @@
 import com.marklogic.client.io.JacksonHandle;
 import com.marklogic.client.io.StringHandle;
 import com.marklogic.client.io.marker.AbstractWriteHandle;
-import com.marklogic.spark.ConnectorException;
-import com.marklogic.spark.JsonRowSerializer;
-import com.marklogic.spark.Options;
-import com.marklogic.spark.Util;
+import com.marklogic.spark.*;
 import com.marklogic.spark.reader.customcode.CustomCodeContext;
 import com.marklogic.spark.writer.CommitMessage;
 import org.apache.spark.sql.catalyst.InternalRow;
@@ -23,19 +20,18 @@
 
 import java.util.ArrayList;
 import java.util.List;
-import java.util.concurrent.atomic.AtomicInteger;
 import java.util.stream.Collectors;
 
 class CustomCodeWriter implements DataWriter<InternalRow> {
 
     private static final Logger logger = LoggerFactory.getLogger(CustomCodeWriter.class);
-    private static final AtomicInteger progressTracker = new AtomicInteger();
 
     private final DatabaseClient databaseClient;
     private final CustomCodeContext customCodeContext;
     private final JsonRowSerializer jsonRowSerializer;
     private final int batchSize;
-    private final int logProgress;
+    private final ProgressLogger progressLogger;
+
     private final List<String> currentBatch = new ArrayList<>();
     private final String externalVariableDelimiter;
     private ObjectMapper objectMapper;
@@ -50,7 +46,8 @@ class CustomCodeWriter implements DataWriter<InternalRow> {
         this.jsonRowSerializer = new JsonRowSerializer(customCodeContext.getSchema(), customCodeContext.getProperties());
 
         this.batchSize = (int) customCodeContext.getNumericOption(Options.WRITE_BATCH_SIZE, 1, 1);
-        this.logProgress = (int) customCodeContext.getNumericOption(Options.WRITE_LOG_PROGRESS, 10000, 0);
+        this.progressLogger = new WriteProgressLogger(customCodeContext.getNumericOption(Options.WRITE_LOG_PROGRESS, 0, 0),
+            this.batchSize, "Items processed: {}");
 
         this.externalVariableDelimiter = customCodeContext.optionExists(Options.WRITE_EXTERNAL_VARIABLE_DELIMITER) ?
             customCodeContext.getProperties().get(Options.WRITE_EXTERNAL_VARIABLE_DELIMITER) : ",";
@@ -154,9 +151,7 @@ private void executeCall(ServerEvaluationCall call, int itemCount) {
         try {
             call.evalAs(String.class);
             this.successItemCount += itemCount;
-            if (this.logProgress > 0) {
-                logProgressIfNecessary(itemCount);
-            }
+            this.progressLogger.logProgressIfNecessary(itemCount);
         } catch (RuntimeException ex) {
             if (customCodeContext.isAbortOnFailure()) {
                 throw ex;
@@ -165,13 +160,4 @@ private void executeCall(ServerEvaluationCall call, int itemCount) {
             Util.MAIN_LOGGER.error(String.format("Unable to process row; cause: %s", ex.getMessage()));
         }
     }
-
-    private void logProgressIfNecessary(int itemCount) {
-        int sum = progressTracker.addAndGet(itemCount);
-        int lowerBound = sum / (this.logProgress);
-        int upperBound = (lowerBound * this.logProgress) + this.batchSize;
-        if (sum >= lowerBound && sum < upperBound) {
-            Util.MAIN_LOGGER.info("Items processed: {}", sum);
-        }
-    }
 }

src/test/java/com/marklogic/spark/writer/WriteRowsTest.java

@@ -48,6 +48,12 @@ void logProgressTest() {
             .save();
 
         verifyTwoHundredDocsWereWritten();
+
+        // For manual inspection, run it again to ensure that the progress counter was reset.
+        newWriter(2)
+            .option(Options.WRITE_BATCH_SIZE, 10)
+            .option(Options.WRITE_LOG_PROGRESS, 40)
+            .save();
     }
 
     @Test