Merge pull request #483 from marklogic/feature/more-polaris

Fixing a few high Polaris issues

Author: rjrudin

docker-compose.yaml

@@ -1,4 +1,4 @@
-name: marklogic_spark
+name: docker-tests-marklogic-spark
 
 services:
 

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

@@ -0,0 +1,36 @@
+/*
+ * Copyright © 2025 MarkLogic Corporation. All Rights Reserved.
+ */
+package com.marklogic.spark;
+
+import java.io.Serializable;
+
+/**
+ * Stateful class that is intended to be used in a singleton manner with synchronized access to its one method.
+ */
+class ProgressLogger implements Serializable {
+
+    static final long serialVersionUID = 1L;
+
+    private final long progressInterval;
+    private final String message;
+
+    private long progressCounter;
+    private long nextProgressInterval;
+
+    ProgressLogger(long progressInterval, String message) {
+        this.progressInterval = progressInterval;
+        this.message = message;
+        this.nextProgressInterval = progressInterval;
+    }
+
+    void logProgressIfNecessary(long itemCount) {
+        if (Util.MAIN_LOGGER.isInfoEnabled() && progressInterval > 0) {
+            this.progressCounter += itemCount;
+            if (progressCounter >= nextProgressInterval) {
+                Util.MAIN_LOGGER.info(message, nextProgressInterval);
+                nextProgressInterval += progressInterval;
+            }
+        }
+    }
+}

marklogic-spark-connector/src/main/java/com/marklogic/spark/ReadProgressLogger.java

@@ -4,7 +4,6 @@
 package com.marklogic.spark;
 
 import java.io.Serializable;
-import java.util.concurrent.atomic.AtomicLong;
 
 /**
  * Handles the progress counter for any operation involving reading from MarkLogic. A Spark job/application can only have
@@ -15,25 +14,17 @@ public class ReadProgressLogger implements Serializable {
 
     static final long serialVersionUID = 1L;
 
-    private static final AtomicLong progressCounter = new AtomicLong(0);
     private static final Object lock = new Object();
-    private static long progressInterval;
-    private static long nextProgressInterval;
-    private static String message;
+    private static ProgressLogger progressLogger;
 
     public static void initialize(long progressInterval, String message) {
-        progressCounter.set(0);
-        ReadProgressLogger.progressInterval = progressInterval;
-        nextProgressInterval = progressInterval;
-        ReadProgressLogger.message = message;
+        progressLogger = new ProgressLogger(progressInterval, message);
     }
 
     public static void logProgressIfNecessary(long itemCount) {
-        if (Util.MAIN_LOGGER.isInfoEnabled() && progressInterval > 0
-            && progressCounter.addAndGet(itemCount) >= nextProgressInterval) {
+        if (Util.MAIN_LOGGER.isInfoEnabled() && progressLogger != null) {
             synchronized (lock) {
-                Util.MAIN_LOGGER.info(message, nextProgressInterval);
-                nextProgressInterval += progressInterval;
+                progressLogger.logProgressIfNecessary(itemCount);
             }
         }
     }

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

@@ -4,7 +4,6 @@
 package com.marklogic.spark;
 
 import java.io.Serializable;
-import java.util.concurrent.atomic.AtomicLong;
 
 /**
  * Handles the progress counter for any operation involving writing to MarkLogic. A Spark job/application can only have
@@ -15,25 +14,17 @@ public class WriteProgressLogger implements Serializable {
 
     static final long serialVersionUID = 1L;
 
-    private static final AtomicLong progressCounter = new AtomicLong(0);
     private static final Object lock = new Object();
-    private static long progressInterval;
-    private static long nextProgressInterval;
-    private static String message;
+    private static ProgressLogger progressLogger;
 
     public static void initialize(long progressInterval, String message) {
-        progressCounter.set(0);
-        WriteProgressLogger.progressInterval = progressInterval;
-        nextProgressInterval = progressInterval;
-        WriteProgressLogger.message = message;
+        progressLogger = new ProgressLogger(progressInterval, message);
     }
 
     public static void logProgressIfNecessary(long itemCount) {
-        if (Util.MAIN_LOGGER.isInfoEnabled() && progressInterval > 0
-            && progressCounter.addAndGet(itemCount) >= nextProgressInterval) {
+        if (Util.MAIN_LOGGER.isInfoEnabled() && progressLogger != null) {
             synchronized (lock) {
-                Util.MAIN_LOGGER.info(message, nextProgressInterval);
-                nextProgressInterval += progressInterval;
+                progressLogger.logProgressIfNecessary(itemCount);
             }
         }
     }

marklogic-spark-connector/src/main/java/com/marklogic/spark/reader/document/OpticTriplesReader.java

@@ -8,9 +8,11 @@
 import com.marklogic.client.query.SearchQueryDefinition;
 import com.marklogic.client.row.RowManager;
 import com.marklogic.client.row.RowRecord;
+import com.marklogic.client.row.RowSet;
 import com.marklogic.client.type.PlanColumn;
 import com.marklogic.spark.Options;
 import com.marklogic.spark.ReadProgressLogger;
+import org.apache.commons.io.IOUtils;
 import org.apache.spark.sql.catalyst.InternalRow;
 import org.apache.spark.sql.catalyst.expressions.GenericInternalRow;
 import org.apache.spark.sql.connector.read.PartitionReader;
@@ -46,6 +48,7 @@ class OpticTriplesReader implements PartitionReader<InternalRow> {
     private final long batchSize;
     private long progressCounter;
 
+    private RowSet<RowRecord> currentRowSet;
     private Iterator<RowRecord> currentRowIterator;
 
     public OpticTriplesReader(ForestPartition forestPartition, DocumentContext context) {
@@ -98,7 +101,7 @@ public InternalRow get() {
 
     @Override
     public void close() {
-        // Nothing to close.
+        IOUtils.closeQuietly(this.currentRowSet);
     }
 
     private void readNextBatchOfTriples(List<String> uris) {
@@ -112,8 +115,8 @@ private void readNextBatchOfTriples(List<String> uris) {
         }
 
         plan = bindDatatypeAndLang(plan);
-
-        currentRowIterator = rowManager.resultRows(plan).iterator();
+        this.currentRowSet = rowManager.resultRows(plan);
+        this.currentRowIterator = this.currentRowSet.iterator();
     }
 
     /**