Merge pull request #163 from marklogic/feature/12219-forest-hosts

MLE-12219 Can now create direct connections to hosts

Author: rjrudin

docs/reading-data/documents.md

@@ -210,3 +210,9 @@ with more partition readers and a higher batch size.
 You can also adjust the level of parallelism by controlling how many threads Spark uses for executing partition reads. 
 Please see your Spark distribution's documentation for further information.
 
+### Direct connections to hosts
+
+If your Spark program is able to connect to each host in your MarkLogic cluster, you can set the
+`spark.marklogic.client.connectionType` option to `direct`. Each partition reader will then connect to the
+host on which the reader's assigned forest resides. This will typically improve performance by reducing the network
+traffic, as the host that receives a request will not need to involve any other host in the processing of that request.

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

@@ -36,7 +36,19 @@ protected ContextSupport(Map<String, String> properties) {
     }
 
     public DatabaseClient connectToMarkLogic() {
+        return connectToMarkLogic(null);
+    }
+
+    /**
+     * @param host if not null, overrides the user-defined host. Used for direct connections when a load balancer is
+     *             not in front of MarkLogic.
+     * @return
+     */
+    public DatabaseClient connectToMarkLogic(String host) {
         Map<String, String> connectionProps = buildConnectionProperties();
+        if (host != null) {
+            connectionProps.put(Options.CLIENT_HOST, host);
+        }
         DatabaseClient client;
         try {
             client = DatabaseClientFactory.newClient(propertyName -> connectionProps.get("spark." + propertyName));
@@ -130,6 +142,11 @@ public final boolean isAbortOnFailure() {
         return !"false".equalsIgnoreCase(getProperties().get(Options.WRITE_ABORT_ON_FAILURE));
     }
 
+    public final boolean isDirectConnection() {
+        String value = getProperties().get(Options.CLIENT_CONNECTION_TYPE);
+        return value != null && value.equalsIgnoreCase(DatabaseClient.ConnectionType.DIRECT.name());
+    }
+
     public final boolean hasOption(String... options) {
         return Util.hasOption(this.properties, options);
     }

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

@@ -1,18 +1,21 @@
 package com.marklogic.spark.reader.document;
 
+import com.marklogic.client.datamovement.Forest;
 import org.apache.spark.sql.connector.read.InputPartition;
 
 class ForestPartition implements InputPartition {
 
     static final long serialVersionUID = 1;
 
     private final String forestName;
+    private final String host;
     private final long serverTimestamp;
     private final Long offsetStart;
     private final Long offsetEnd;
 
-    ForestPartition(String forestName, long serverTimestamp, Long offsetStart, Long offsetEnd) {
-        this.forestName = forestName;
+    ForestPartition(Forest forest, long serverTimestamp, Long offsetStart, Long offsetEnd) {
+        this.forestName = forest.getForestName();
+        this.host = forest.getHost();
         this.serverTimestamp = serverTimestamp;
         this.offsetStart = offsetStart;
         this.offsetEnd = offsetEnd;
@@ -22,6 +25,10 @@ String getForestName() {
         return forestName;
     }
 
+    String getHost() {
+        return host;
+    }
+
     long getServerTimestamp() {
         return serverTimestamp;
     }

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

@@ -28,7 +28,7 @@ ForestPartition[] makePartitions(long estimate, long serverTimestamp, Forest...
                     break;
                 }
                 Long offsetEnd = j < (partitionsPerForest - 1) ? (urisPerPartition + offset - 1) : null;
-                partitions.add(new ForestPartition(forest.getForestName(), serverTimestamp, offset, offsetEnd));
+                partitions.add(new ForestPartition(forest, serverTimestamp, offset, offsetEnd));
                 offset += urisPerPartition;
             }
         }

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

@@ -55,13 +55,17 @@ class ForestReader implements PartitionReader<InternalRow> {
     private int docCount;
 
     ForestReader(ForestPartition forestPartition, DocumentContext context) {
-        if (logger.isDebugEnabled()) {
-            logger.debug("Will read from partition: {}", forestPartition);
-        }
         this.forestPartition = forestPartition;
         this.limit = context.getLimit();
 
-        DatabaseClient client = context.connectToMarkLogic();
+        DatabaseClient client = context.isDirectConnection() ?
+            context.connectToMarkLogic(forestPartition.getHost()) :
+            context.connectToMarkLogic();
+
+        if (logger.isDebugEnabled()) {
+            logger.debug("Will read from host {} for partition: {}", client.getHost(), forestPartition);
+        }
+
         SearchQueryDefinition query = context.buildSearchQuery(client);
         this.uriBatcher = new UriBatcher(client, query, forestPartition, context.getBatchSize(), false);
 

src/test/java/com/marklogic/spark/ContextSupportTest.java

@@ -0,0 +1,45 @@
+package com.marklogic.spark;
+
+import com.marklogic.client.DatabaseClient;
+import com.marklogic.client.MarkLogicIOException;
+import org.junit.jupiter.api.Test;
+
+import java.net.UnknownHostException;
+import java.util.HashMap;
+import java.util.Map;
+
+import static org.junit.jupiter.api.Assertions.*;
+
+class ContextSupportTest extends AbstractIntegrationTest {
+
+    private Map<String, String> options = new HashMap<>();
+
+    /**
+     * For automated tests, we only have the single-node MarkLogic cluster and thus can't verify that direct
+     * connections really work. But this test at least verifies that the "host" input overrides what is in the
+     * options.
+     */
+    @Test
+    void directConnectionToHost() {
+        options.put(Options.CLIENT_URI, makeClientUri());
+
+        ContextSupport support = new ContextSupport(options);
+        DatabaseClient client = support.connectToMarkLogic(testConfig.getHost());
+        assertTrue(client.checkConnection().isConnected());
+
+        MarkLogicIOException ex = assertThrows(MarkLogicIOException.class,
+            () -> support.connectToMarkLogic("invalid-host"));
+        assertTrue(ex.getCause() instanceof UnknownHostException);
+    }
+
+    @Test
+    void isDirectConnection() {
+        assertFalse(new ContextSupport(options).isDirectConnection());
+
+        options.put(Options.CLIENT_CONNECTION_TYPE, "direct");
+        assertTrue(new ContextSupport(options).isDirectConnection());
+
+        options.put(Options.CLIENT_CONNECTION_TYPE, "gateway");
+        assertFalse(new ContextSupport(options).isDirectConnection());
+    }
+}

src/test/java/com/marklogic/spark/reader/document/ReadDocumentRowsTest.java

@@ -2,9 +2,11 @@
 
 import com.fasterxml.jackson.databind.JsonNode;
 import com.fasterxml.jackson.databind.node.ObjectNode;
+import com.marklogic.client.MarkLogicIOException;
 import com.marklogic.spark.AbstractIntegrationTest;
 import com.marklogic.spark.ConnectorException;
 import com.marklogic.spark.Options;
+import org.apache.spark.SparkException;
 import org.apache.spark.sql.DataFrameReader;
 import org.apache.spark.sql.Dataset;
 import org.apache.spark.sql.Row;
@@ -14,8 +16,7 @@
 import java.util.List;
 import java.util.stream.Collectors;
 
-import static org.junit.jupiter.api.Assertions.assertEquals;
-import static org.junit.jupiter.api.Assertions.assertTrue;
+import static org.junit.jupiter.api.Assertions.*;
 
 class ReadDocumentRowsTest extends AbstractIntegrationTest {
 
@@ -36,6 +37,19 @@ void readByCollection() {
         assertEquals("Vivianne", doc.get("ForeName").asText());
     }
 
+    @Test
+    void readViaDirectConnect() {
+        Dataset<Row> rows = startRead()
+            .option(Options.READ_DOCUMENTS_COLLECTIONS, "author")
+            .option(Options.CLIENT_CONNECTION_TYPE, "direct")
+            .load();
+
+        SparkException ex = assertThrows(SparkException.class, () -> rows.count());
+        assertTrue(ex.getCause() instanceof MarkLogicIOException, "This test is expected to fail when run against " +
+            "the cluster created by docker-compose.yaml, as the host name of the MarkLogic cluster is not expected " +
+            "to be accessible. Actual exception: " + ex.getCause());
+    }
+
     @Test
     void invalidBatchSize() {
         // Verify batch size doesn't cause an error.