Fixed bug with splitting text in text files

Author: rjrudin

marklogic-langchain4j/src/main/java/com/marklogic/langchain4j/splitter/AbstractChunkDocumentProducer.java

@@ -29,9 +29,10 @@ abstract class AbstractChunkDocumentProducer implements Iterator<DocumentWriteOp
         this.textSegments = textSegments;
         this.chunkConfig = chunkConfig;
 
-        // Chunks cannot be written to a TEXT document. So if maxChunks is zero, and we have a text document, we will
-        // instead write all the chunks to a separate document.
-        this.maxChunksPerDocument = Format.TEXT.equals(sourceDocumentFormat) && chunkConfig.getMaxChunks() == 0 ?
+        // Chunks cannot be written to the source document unless its format is JSON or XML. So if maxChunks is zero and
+        // we don't have a JSON or XML document, all chunks will be written to a separate document.
+        boolean cannotAddChunksToSourceDocument = !Format.JSON.equals(sourceDocumentFormat) && !Format.XML.equals(sourceDocumentFormat);
+        this.maxChunksPerDocument = cannotAddChunksToSourceDocument && chunkConfig.getMaxChunks() == 0 ?
             textSegments.size() :
             chunkConfig.getMaxChunks();
     }

marklogic-langchain4j/src/main/java/com/marklogic/langchain4j/splitter/DefaultChunkAssembler.java

@@ -56,13 +56,12 @@ private Format determineSourceDocumentFormat(DocumentWriteOperation sourceDocume
     }
 
     private Format determineChunkDocumentFormat(Format sourceDocumentFormat) {
-        final boolean addChunksToSourceDocument = !Format.TEXT.equals(sourceDocumentFormat) && chunkConfig.getMaxChunks() == 0;
-        if (addChunksToSourceDocument) {
+        final boolean canAddChunksToSourceDocument = Format.XML.equals(sourceDocumentFormat) || Format.JSON.equals(sourceDocumentFormat);
+        if (canAddChunksToSourceDocument && chunkConfig.getMaxChunks() == 0) {
             return sourceDocumentFormat;
         }
 
-        final String documentType = chunkConfig.getDocumentType();
-        if (documentType != null || Format.TEXT.equals(sourceDocumentFormat)) {
+        if (chunkConfig.getDocumentType() != null || !canAddChunksToSourceDocument) {
             return "xml".equalsIgnoreCase(chunkConfig.getDocumentType()) ? Format.XML : Format.JSON;
         }
 

marklogic-spark-connector/src/test/java/com/marklogic/spark/writer/splitter/SplitTextDocumentTest.java

@@ -9,6 +9,8 @@
 import com.marklogic.spark.AbstractIntegrationTest;
 import com.marklogic.spark.Options;
 import org.apache.spark.sql.DataFrameWriter;
+import org.apache.spark.sql.Dataset;
+import org.apache.spark.sql.Row;
 import org.apache.spark.sql.SaveMode;
 import org.jdom2.Namespace;
 import org.junit.jupiter.api.Test;
@@ -56,6 +58,29 @@ void xmlChunks() {
         tester.assertReadPermissionExists("spark-user-role");
     }
 
+    @Test
+    void inputDocumentHasUnknownFormat() {
+        Dataset<Row> dataset = newSparkSession().read().format(CONNECTOR_IDENTIFIER)
+            .load("src/test/resources/mixed-files/hello.txt");
+
+        assertTrue(dataset.collectAsList().get(0).isNullAt(2),
+            "The connector is not expected to determine document type when reading files.");
+
+        dataset.write().format(CONNECTOR_IDENTIFIER)
+            .option(Options.CLIENT_URI, makeClientUri())
+            .option(Options.WRITE_PERMISSIONS, DEFAULT_PERMISSIONS)
+            .option(Options.WRITE_SPLITTER_TEXT, true)
+            .option(Options.WRITE_URI_TEMPLATE, "/test/hello.txt")
+            .mode(SaveMode.Append)
+            .save();
+
+        JsonNode doc = readJsonDocument("/test/hello.txt-chunks-1.json");
+        assertEquals("hello world", doc.get("chunks").get(0).get("text").asText(),
+            "When the input document format is UNKNOWN and max chunks is zero, the connector should realize it " +
+                "cannot add chunks to a document with format=UNKNOWN and thus it should create a separate chunks " +
+                "document containing all the chunks.");
+    }
+
     @Test
     void maxChunksOfThree() {
         prepareToWriteChunkDocuments()