MLE-26953 Refactoring: DocumentInputs now has List<ChunkInputs>

This combines the 3 separate lists of text, classifications, and embeddings. Will make it much easier to add a model name. No change in functionality, just moving things around in the implementation.

Author: rjrudin

marklogic-spark-connector/src/main/java/com/marklogic/spark/core/ChunkInputs.java

@@ -0,0 +1,39 @@
+/*
+ * Copyright (c) 2023-2026 Progress Software Corporation and/or its subsidiaries or affiliates. All Rights Reserved.
+ */
+package com.marklogic.spark.core;
+
+/**
+ * Encapsulates the data associated with a chunk of text, including its embedding and classification. Note there's
+ * some naming issues to work out with this class and the Chunk interface.
+ */
+public class ChunkInputs {
+
+    private final String text;
+    private float[] embedding;
+    private byte[] classification;
+
+    public ChunkInputs(String text) {
+        this.text = text;
+    }
+
+    public String getText() {
+        return text;
+    }
+
+    public float[] getEmbedding() {
+        return embedding;
+    }
+
+    public void setEmbedding(float[] embedding) {
+        this.embedding = embedding;
+    }
+
+    public byte[] getClassification() {
+        return classification;
+    }
+
+    public void setClassification(byte[] classification) {
+        this.classification = classification;
+    }
+}

marklogic-spark-connector/src/main/java/com/marklogic/spark/core/DocumentInputs.java

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2023-2025 Progress Software Corporation and/or its subsidiaries or affiliates. All Rights Reserved.
+ * Copyright (c) 2023-2026 Progress Software Corporation and/or its subsidiaries or affiliates. All Rights Reserved.
  */
 package com.marklogic.spark.core;
 
@@ -30,11 +30,7 @@ public class DocumentInputs {
     private Map<String, String> extractedMetadata;
 
     private byte[] documentClassification;
-    private List<byte[]> chunkClassifications;
-    private List<float[]> embeddings;
-
-    // These will be created via a splitter.
-    private List<String> chunks;
+    private List<ChunkInputs> chunkInputsList;
 
     public DocumentInputs(String initialUri, AbstractWriteHandle content, JsonNode columnValuesForUriTemplate,
                           DocumentMetadataHandle initialMetadata) {
@@ -78,17 +74,31 @@ public AbstractWriteHandle getContent() {
     }
 
     public void addChunkClassification(byte[] classification) {
-        if (chunkClassifications == null) {
-            chunkClassifications = new ArrayList<>();
+        if (chunkInputsList == null || chunkInputsList.isEmpty()) {
+            throw new IllegalStateException("Cannot add classification: no chunks exist");
+        }
+        // Find the next chunk without a classification
+        for (ChunkInputs chunk : chunkInputsList) {
+            if (chunk.getClassification() == null) {
+                chunk.setClassification(classification);
+                return;
+            }
         }
-        chunkClassifications.add(classification);
+        throw new IllegalStateException("Cannot add classification: all chunks already have classifications");
     }
 
     public void addEmbedding(float[] embedding) {
-        if (embeddings == null) {
-            embeddings = new ArrayList<>();
+        if (chunkInputsList == null || chunkInputsList.isEmpty()) {
+            throw new IllegalStateException("Cannot add embedding: no chunks exist");
         }
-        embeddings.add(embedding);
+        // Find the next chunk without an embedding
+        for (ChunkInputs chunk : chunkInputsList) {
+            if (chunk.getEmbedding() == null) {
+                chunk.setEmbedding(embedding);
+                return;
+            }
+        }
+        throw new IllegalStateException("Cannot add embedding: all chunks already have embeddings");
     }
 
     public String getInitialUri() {
@@ -124,15 +134,25 @@ public void setExtractedMetadata(Map<String, String> extractedMetadata) {
     }
 
     public List<String> getChunks() {
-        return chunks;
+        if (chunkInputsList == null) {
+            return null;
+        }
+        List<String> texts = new ArrayList<>(chunkInputsList.size());
+        for (ChunkInputs chunk : chunkInputsList) {
+            texts.add(chunk.getText());
+        }
+        return texts;
     }
 
     public void setChunks(List<String> chunks) {
-        this.chunks = chunks;
-    }
-
-    public List<byte[]> getClassifications() {
-        return chunkClassifications;
+        if (chunks == null) {
+            this.chunkInputsList = null;
+        } else {
+            this.chunkInputsList = new ArrayList<>(chunks.size());
+            for (String text : chunks) {
+                this.chunkInputsList.add(new ChunkInputs(text));
+            }
+        }
     }
 
     public byte[] getDocumentClassification() {
@@ -143,7 +163,7 @@ public void setDocumentClassification(byte[] documentClassification) {
         this.documentClassification = documentClassification;
     }
 
-    public List<float[]> getEmbeddings() {
-        return embeddings;
+    public List<ChunkInputs> getChunkInputsList() {
+        return chunkInputsList;
     }
 }

marklogic-spark-connector/src/main/java/com/marklogic/spark/core/splitter/AbstractChunkDocumentProducer.java

@@ -1,10 +1,11 @@
 /*
- * Copyright (c) 2023-2025 Progress Software Corporation and/or its subsidiaries or affiliates. All Rights Reserved.
+ * Copyright (c) 2023-2026 Progress Software Corporation and/or its subsidiaries or affiliates. All Rights Reserved.
  */
 package com.marklogic.spark.core.splitter;
 
 import com.marklogic.client.document.DocumentWriteOperation;
 import com.marklogic.client.io.Format;
+import com.marklogic.spark.core.ChunkInputs;
 
 import java.util.Iterator;
 import java.util.List;
@@ -16,27 +17,23 @@
 abstract class AbstractChunkDocumentProducer implements Iterator<DocumentWriteOperation> {
 
     protected final DocumentWriteOperation sourceDocument;
-    protected final List<String> textSegments;
+    protected final List<ChunkInputs> chunkInputsList;
     protected final ChunkConfig chunkConfig;
-    protected final List<byte[]> classifications;
-    protected final List<float[]> embeddings;
     protected final int maxChunksPerDocument;
 
     protected int listIndex = -1;
     private int chunkDocumentCounter = 1;
 
-    AbstractChunkDocumentProducer(DocumentWriteOperation sourceDocument, Format sourceDocumentFormat, List<String> textSegments, ChunkConfig chunkConfig, List<byte[]> classifications, List<float[]> embeddings) {
+    AbstractChunkDocumentProducer(DocumentWriteOperation sourceDocument, Format sourceDocumentFormat, List<ChunkInputs> chunkInputsList, ChunkConfig chunkConfig) {
         this.sourceDocument = sourceDocument;
-        this.textSegments = textSegments;
+        this.chunkInputsList = chunkInputsList;
         this.chunkConfig = chunkConfig;
-        this.classifications = classifications;
-        this.embeddings = embeddings;
 
         // 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() :
+            chunkInputsList.size() :
             chunkConfig.getMaxChunks();
     }
 
@@ -47,7 +44,7 @@ abstract class AbstractChunkDocumentProducer implements Iterator<DocumentWriteOp
 
     @Override
     public final boolean hasNext() {
-        return listIndex < textSegments.size();
+        return listIndex < chunkInputsList.size();
     }
 
     // Sonar complains that a NoSuchElementException should be thrown here, but that would only occur if the
@@ -58,7 +55,7 @@ public DocumentWriteOperation next() {
         if (listIndex == -1) {
             listIndex++;
             if (this.maxChunksPerDocument == 0) {
-                listIndex = textSegments.size();
+                listIndex = chunkInputsList.size();
                 return addChunksToSourceDocument();
             }
             return sourceDocument;
@@ -83,15 +80,4 @@ protected final String makeChunkDocumentUri(DocumentWriteOperation sourceDocumen
         }
         return uri;
     }
-
-    /**
-     * Return the embedding at position n if it exists.
-     * @param embeddings the embeddings list
-     * @param n the position for the embedding requests
-     * @return the embedding float array
-     */
-    protected float[] getEmbeddingIfExists(List<float[]> embeddings, int n) {
-        return (embeddings != null && n < embeddings.size() ? embeddings.get(n) : null);
-    }
-
 }

marklogic-spark-connector/src/main/java/com/marklogic/spark/core/splitter/ChunkAssembler.java

@@ -1,9 +1,10 @@
 /*
- * Copyright (c) 2023-2025 Progress Software Corporation and/or its subsidiaries or affiliates. All Rights Reserved.
+ * Copyright (c) 2023-2026 Progress Software Corporation and/or its subsidiaries or affiliates. All Rights Reserved.
  */
 package com.marklogic.spark.core.splitter;
 
 import com.marklogic.client.document.DocumentWriteOperation;
+import com.marklogic.spark.core.ChunkInputs;
 
 import java.util.Iterator;
 import java.util.List;
@@ -15,11 +16,8 @@ public interface ChunkAssembler {
 
     /**
      * @param sourceDocument
-     * @param chunks
-     * @param classifications
-     * @param embeddings
+     * @param chunkInputsList
      * @return an iterator, which allows for an implementation to lazily construct documents if necessary.
      */
-    Iterator<DocumentWriteOperation> assembleChunks(DocumentWriteOperation sourceDocument, List<String> chunks,
-                                                    List<byte[]> classifications, List<float[]> embeddings);
+    Iterator<DocumentWriteOperation> assembleChunks(DocumentWriteOperation sourceDocument, List<ChunkInputs> chunkInputsList);
 }

marklogic-spark-connector/src/main/java/com/marklogic/spark/core/splitter/DefaultChunkAssembler.java

@@ -1,11 +1,12 @@
 /*
- * Copyright (c) 2023-2025 Progress Software Corporation and/or its subsidiaries or affiliates. All Rights Reserved.
+ * Copyright (c) 2023-2026 Progress Software Corporation and/or its subsidiaries or affiliates. All Rights Reserved.
  */
 package com.marklogic.spark.core.splitter;
 
 import com.marklogic.client.document.DocumentWriteOperation;
 import com.marklogic.client.io.Format;
 import com.marklogic.spark.Util;
+import com.marklogic.spark.core.ChunkInputs;
 
 import java.util.Iterator;
 import java.util.List;
@@ -20,7 +21,7 @@ public DefaultChunkAssembler(ChunkConfig chunkConfig) {
     }
 
     @Override
-    public Iterator<DocumentWriteOperation> assembleChunks(DocumentWriteOperation sourceDocument, List<String> textSegments, List<byte[]> classifications, List<float[]> embeddings) {
+    public Iterator<DocumentWriteOperation> assembleChunks(DocumentWriteOperation sourceDocument, List<ChunkInputs> chunkInputsList) {
         final Format sourceDocumentFormat = Util.determineSourceDocumentFormat(sourceDocument.getContent(), sourceDocument.getUri());
         if (sourceDocumentFormat == null) {
             Util.MAIN_LOGGER.warn("Cannot split document with URI {}; cannot determine the document format.", sourceDocument.getUri());
@@ -30,8 +31,8 @@ public Iterator<DocumentWriteOperation> assembleChunks(DocumentWriteOperation so
         final Format chunkDocumentFormat = determineChunkDocumentFormat(sourceDocumentFormat);
 
         return Format.XML.equals(chunkDocumentFormat) ?
-            new XmlChunkDocumentProducer(sourceDocument, sourceDocumentFormat, textSegments, chunkConfig, classifications, embeddings) :
-            new JsonChunkDocumentProducer(sourceDocument, sourceDocumentFormat, textSegments, chunkConfig, classifications, embeddings);
+            new XmlChunkDocumentProducer(sourceDocument, sourceDocumentFormat, chunkInputsList, chunkConfig) :
+            new JsonChunkDocumentProducer(sourceDocument, sourceDocumentFormat, chunkInputsList, chunkConfig);
     }
 
     private Format determineChunkDocumentFormat(Format sourceDocumentFormat) {

marklogic-spark-connector/src/main/java/com/marklogic/spark/core/splitter/JsonChunkDocumentProducer.java

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2023-2025 Progress Software Corporation and/or its subsidiaries or affiliates. All Rights Reserved.
+ * Copyright (c) 2023-2026 Progress Software Corporation and/or its subsidiaries or affiliates. All Rights Reserved.
  */
 package com.marklogic.spark.core.splitter;
 
@@ -14,6 +14,7 @@
 import com.marklogic.client.io.JacksonHandle;
 import com.marklogic.client.io.marker.AbstractWriteHandle;
 import com.marklogic.spark.ConnectorException;
+import com.marklogic.spark.core.ChunkInputs;
 import com.marklogic.spark.core.embedding.Chunk;
 import com.marklogic.spark.core.embedding.DocumentAndChunks;
 import com.marklogic.spark.core.embedding.JsonChunk;
@@ -30,8 +31,8 @@ class JsonChunkDocumentProducer extends AbstractChunkDocumentProducer {
     private final XmlMapper xmlMapper;
 
     JsonChunkDocumentProducer(DocumentWriteOperation sourceDocument, Format sourceDocumentFormat,
-                              List<String> textSegments, ChunkConfig chunkConfig, List<byte[]> classifications, List<float[]> embeddings) {
-        super(sourceDocument, sourceDocumentFormat, textSegments, chunkConfig, classifications, embeddings);
+                              List<ChunkInputs> chunkInputsList, ChunkConfig chunkConfig) {
+        super(sourceDocument, sourceDocumentFormat, chunkInputsList, chunkConfig);
         xmlMapper = new XmlMapper();
     }
 
@@ -42,24 +43,21 @@ protected DocumentWriteOperation addChunksToSourceDocument() {
 
         ArrayNode chunksArray = doc.putArray(determineChunksArrayName(doc));
         List<Chunk> chunks = new ArrayList<>();
-        int chunksCounter = 0;
-        for (String text : textSegments) {
+        for (ChunkInputs chunkInputs : chunkInputsList) {
             ObjectNode chunk = chunksArray.addObject();
-            chunk.put("text", text);
-            if (classifications != null && classifications.size() > chunksCounter) {
+            chunk.put("text", chunkInputs.getText());
+            if (chunkInputs.getClassification() != null) {
                 try {
-                    JsonNode classification = xmlMapper.readTree(classifications.get(chunksCounter));
+                    JsonNode classification = xmlMapper.readTree(chunkInputs.getClassification());
                     chunk.set("classification", classification);
                 } catch (IOException e) {
                     throw new ConnectorException(String.format("Unable to classify data from document with URI: %s; cause: %s", sourceDocument.getUri(), e.getMessage()), e);
                 }
             }
-            float[] embedding = getEmbeddingIfExists(embeddings, chunksCounter);
             var jsonChunk = new JsonChunk(chunk, null, chunkConfig.getEmbeddingName(), chunkConfig.isBase64EncodeVectors());
-            if (embedding != null) {
-                jsonChunk.addEmbedding(embedding);
+            if (chunkInputs.getEmbedding() != null) {
+                jsonChunk.addEmbedding(chunkInputs.getEmbedding());
             }
-            chunksCounter++;
             chunks.add(jsonChunk);
         }
 
@@ -81,26 +79,23 @@ protected DocumentWriteOperation makeChunkDocument() {
 
         ArrayNode chunksArray = rootField.putArray(DEFAULT_CHUNKS_ARRAY_NAME);
         List<Chunk> chunks = new ArrayList<>();
-        int chunksCounter = 0;
         for (int i = 0; i < this.maxChunksPerDocument && hasNext(); i++) {
-            String text = textSegments.get(listIndex);
+            ChunkInputs chunkInputs = chunkInputsList.get(listIndex);
             ObjectNode chunk = chunksArray.addObject();
-            chunk.put("text", text);
-            if (classifications != null && classifications.size() > chunksCounter) {
+            chunk.put("text", chunkInputs.getText());
+            if (chunkInputs.getClassification() != null) {
                 try {
-                    JsonNode classification = xmlMapper.readTree(classifications.get(chunksCounter));
+                    JsonNode classification = xmlMapper.readTree(chunkInputs.getClassification());
                     chunk.set("classification", classification);
                 } catch (IOException e) {
                     throw new ConnectorException(String.format("Unable to classify data from document with URI: %s; cause: %s", uri, e.getMessage()), e);
                 }
             }
-            float[] embedding = getEmbeddingIfExists(embeddings, listIndex);
             var jsonChunk = new JsonChunk(chunk, null, chunkConfig.getEmbeddingName(), chunkConfig.isBase64EncodeVectors());
-            if (embedding != null) {
-                jsonChunk.addEmbedding(embedding);
+            if (chunkInputs.getEmbedding() != null) {
+                jsonChunk.addEmbedding(chunkInputs.getEmbedding());
             }
             chunks.add(jsonChunk);
-            chunksCounter++;
             listIndex++;
         }