Merge pull request #164 from marklogic/feature/logging-tweak

Added logging of how many documents were written

Author: rjrudin

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

@@ -11,7 +11,6 @@
 import com.marklogic.client.query.QueryDefinition;
 import com.marklogic.client.query.SearchQueryDefinition;
 import com.marklogic.client.query.StructuredQueryBuilder;
-import com.marklogic.spark.Util;
 import org.apache.spark.sql.catalyst.InternalRow;
 import org.apache.spark.sql.catalyst.expressions.GenericInternalRow;
 import org.apache.spark.sql.catalyst.util.ArrayBasedMapData;
@@ -93,10 +92,9 @@ public boolean next() {
             closeCurrentDocumentPage();
             List<String> uris = getNextBatchOfUris();
             if (uris.isEmpty()) {
-                // TBD on whether this should be info/debug.
-                if (Util.MAIN_LOGGER.isInfoEnabled()) {
+                if (logger.isDebugEnabled()) {
                     long duration = System.currentTimeMillis() - startTime;
-                    Util.MAIN_LOGGER.info("Read {} documents from partition {} in {}ms", docCount, forestPartition, duration);
+                    logger.debug("Read {} documents from partition {} in {}ms", docCount, forestPartition, duration);
                 }
                 return false;
             }

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

@@ -0,0 +1,35 @@
+package com.marklogic.spark.writer;
+
+import org.apache.spark.sql.connector.write.WriterCommitMessage;
+
+public class CommitMessage implements WriterCommitMessage {
+
+    private final String action;
+    private final int docCount;
+    private final int partitionId;
+    private final long taskId;
+    private final long epochId;
+
+    public CommitMessage(String action, int docCount, int partitionId, long taskId, long epochId) {
+        this.action = action;
+        this.docCount = docCount;
+        this.partitionId = partitionId;
+        this.taskId = taskId;
+        this.epochId = epochId;
+    }
+
+    public String getAction() {
+        return action;
+    }
+
+    public int getDocCount() {
+        return docCount;
+    }
+
+    @Override
+    public String toString() {
+        return epochId != 0L ?
+            String.format("[partitionId: %d; taskId: %d; epochId: %d]; docCount: %d", partitionId, taskId, epochId, docCount) :
+            String.format("[partitionId: %d; taskId: %d]; docCount: %d", partitionId, taskId, docCount);
+    }
+}

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

@@ -53,8 +53,18 @@ public DataWriterFactory createBatchWriterFactory(PhysicalWriteInfo info) {
 
     @Override
     public void commit(WriterCommitMessage[] messages) {
-        if (messages != null && messages.length > 0 && logger.isDebugEnabled()) {
-            logger.debug("Commit messages received: {}", Arrays.asList(messages));
+        if (messages != null && messages.length > 0) {
+            if (logger.isDebugEnabled()) {
+                logger.debug("Commit messages received: {}", Arrays.asList(messages));
+            }
+            final String action = ((CommitMessage) messages[0]).getAction();
+            if (Util.MAIN_LOGGER.isInfoEnabled()) {
+                int count = 0;
+                for (WriterCommitMessage message : messages) {
+                    count += ((CommitMessage) message).getDocCount();
+                }
+                Util.MAIN_LOGGER.info("{} {} documents", action, count);
+            }
         }
     }
 

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

@@ -90,7 +90,7 @@ public void write(InternalRow row) throws IOException {
 
     @Override
     public WriterCommitMessage commit() throws IOException {
-        CommitMessage message = new CommitMessage(docCount, partitionId, taskId, epochId);
+        CommitMessage message = new CommitMessage("Wrote", docCount, partitionId, taskId, epochId);
         if (logger.isDebugEnabled()) {
             logger.debug("Committing {}", message);
         }
@@ -129,25 +129,4 @@ private void stopJobAndRelease() {
             this.databaseClient.release();
         }
     }
-
-    private static class CommitMessage implements WriterCommitMessage {
-        private int docCount;
-        private int partitionId;
-        private long taskId;
-        private long epochId;
-
-        public CommitMessage(int docCount, int partitionId, long taskId, long epochId) {
-            this.docCount = docCount;
-            this.partitionId = partitionId;
-            this.taskId = taskId;
-            this.epochId = epochId;
-        }
-
-        @Override
-        public String toString() {
-            return epochId != 0L ?
-                String.format("[partitionId: %d; taskId: %d; epochId: %d]; docCount: %d", partitionId, taskId, epochId, docCount) :
-                String.format("[partitionId: %d; taskId: %d]; docCount: %d", partitionId, taskId, docCount);
-        }
-    }
 }

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

@@ -13,6 +13,7 @@
 import com.marklogic.spark.Options;
 import com.marklogic.spark.Util;
 import com.marklogic.spark.reader.customcode.CustomCodeContext;
+import com.marklogic.spark.writer.CommitMessage;
 import org.apache.spark.sql.catalyst.InternalRow;
 import org.apache.spark.sql.catalyst.json.JacksonGenerator;
 import org.apache.spark.sql.connector.write.DataWriter;
@@ -41,6 +42,7 @@ class CustomCodeWriter implements DataWriter<InternalRow> {
     private final int partitionId;
     private final long taskId;
     private final long epochId;
+    private int itemCount;
 
     CustomCodeWriter(CustomCodeContext customCodeContext, int partitionId, long taskId, long epochId) {
         this.customCodeContext = customCodeContext;
@@ -66,6 +68,7 @@ public void write(InternalRow row) {
         String rowValue = customCodeContext.isCustomSchema() ?
             convertRowToJSONString(row) :
             row.getString(0);
+        itemCount++;
 
         this.currentBatch.add(rowValue);
 
@@ -77,7 +80,7 @@ public void write(InternalRow row) {
     @Override
     public WriterCommitMessage commit() {
         flush();
-        CommitMessage message = new CommitMessage(partitionId, taskId, epochId);
+        CommitMessage message = new CommitMessage("Processed", itemCount, partitionId, taskId, epochId);
         if (logger.isDebugEnabled()) {
             logger.debug("Committing {}", message);
         }
@@ -172,23 +175,4 @@ private void executeCall(ServerEvaluationCall call) {
             Util.MAIN_LOGGER.error(String.format("Unable to process row; cause: %s", ex.getMessage()));
         }
     }
-
-    private static class CommitMessage implements WriterCommitMessage {
-        private int partitionId;
-        private long taskId;
-        private long epochId;
-
-        public CommitMessage(int partitionId, long taskId, long epochId) {
-            this.partitionId = partitionId;
-            this.taskId = taskId;
-            this.epochId = epochId;
-        }
-
-        @Override
-        public String toString() {
-            return epochId != 0L ?
-                String.format("[partitionId: %d; taskId: %d; epochId: %d]", partitionId, taskId, epochId) :
-                String.format("[partitionId: %d; taskId: %d]", partitionId, taskId);
-        }
-    }
 }