Fix for adding embeddings to sidecar docs from text docs

Author: rjrudin

marklogic-spark-connector/src/test/java/com/marklogic/spark/writer/embedding/AddEmbeddingsFromTextTest.java

@@ -0,0 +1,76 @@
+/*
+ * Copyright © 2024 MarkLogic Corporation. All Rights Reserved.
+ */
+package com.marklogic.spark.writer.embedding;
+
+import com.fasterxml.jackson.databind.JsonNode;
+import com.fasterxml.jackson.databind.node.JsonNodeType;
+import com.marklogic.junit5.XmlNode;
+import com.marklogic.spark.AbstractIntegrationTest;
+import com.marklogic.spark.Options;
+import org.apache.spark.sql.DataFrameWriter;
+import org.apache.spark.sql.Row;
+import org.apache.spark.sql.SaveMode;
+import org.junit.jupiter.api.Test;
+
+import java.util.List;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+/**
+ * Verifies that when split text from text documents and then adding embeddings to the sidecar docs, the user doesn't
+ * need to specify the location of the chunks. The connector is expected to determine the location based on whether the
+ * sidecar docs are JSON or XML.
+ */
+class AddEmbeddingsFromTextTest extends AbstractIntegrationTest {
+
+    private static final String TEST_EMBEDDING_FUNCTION_CLASS = "com.marklogic.spark.writer.embedding.MinilmEmbeddingModelFunction";
+
+    @Test
+    void jsonSidecarDocuments() {
+        prepareToWriteChunks()
+            .mode(SaveMode.Append)
+            .save();
+
+        List<String> uris = getUrisInCollection("text-chunks", 4);
+        for (String uri : uris) {
+            assertTrue(uri.endsWith(".json"));
+            JsonNode doc = readJsonDocument(uri);
+            assertEquals(JsonNodeType.ARRAY, doc.get("chunks").getNodeType());
+        }
+    }
+
+    @Test
+    void xmlSidecarDocuments() {
+        prepareToWriteChunks()
+            .option(Options.WRITE_SPLITTER_SIDECAR_DOCUMENT_TYPE, "xml")
+            .mode(SaveMode.Append)
+            .save();
+
+        List<String> uris = getUrisInCollection("text-chunks", 4);
+        for (String uri : uris) {
+            assertTrue(uri.endsWith(".xml"));
+            XmlNode doc = readXmlDocument(uri);
+            doc.assertElementCount("/node()/chunks/chunk", 1);
+        }
+    }
+
+    private DataFrameWriter<Row> prepareToWriteChunks() {
+        return newSparkSession().read().format(CONNECTOR_IDENTIFIER)
+            .option(Options.CLIENT_URI, makeClientUri())
+            .option(Options.READ_DOCUMENTS_URIS, "/marklogic-docs/java-client-intro.txt")
+            .load()
+            .write().format(CONNECTOR_IDENTIFIER)
+            .option(Options.CLIENT_URI, makeClientUri())
+            .option(Options.WRITE_PERMISSIONS, DEFAULT_PERMISSIONS)
+            .option(Options.WRITE_URI_PREFIX, "/test")
+            .option(Options.WRITE_SPLITTER_TEXT, true)
+            .option(Options.WRITE_SPLITTER_MAX_CHUNK_SIZE, 500)
+            .option(Options.WRITE_SPLITTER_SIDECAR_MAX_CHUNKS, 1)
+            .option(Options.WRITE_SPLITTER_SIDECAR_COLLECTIONS, "text-chunks")
+            .option(Options.WRITE_EMBEDDER_MODEL_FUNCTION_CLASS_NAME, TEST_EMBEDDING_FUNCTION_CLASS);
+    }
+
+
+}

marklogic-spark-connector/src/test/java/com/marklogic/spark/writer/embedding/TestEmbeddingModel.java

@@ -32,6 +32,11 @@ public EmbeddingModel apply(Map<String, String> options) {
         return this;
     }
 
+    @Override
+    public int dimension() {
+        return 0;
+    }
+
     @Override
     public Response<List<Embedding>> embedAll(List<TextSegment> textSegments) {
         batchCounter++;

marklogic-spark-langchain4j/src/main/java/com/marklogic/spark/langchain4j/EmbeddingAdderFactory.java

@@ -8,6 +8,7 @@
 import com.marklogic.spark.ConnectorException;
 import com.marklogic.spark.Context;
 import com.marklogic.spark.Options;
+import com.marklogic.spark.Util;
 import dev.langchain4j.model.embedding.EmbeddingModel;
 
 import java.util.HashMap;
@@ -31,7 +32,11 @@ public static EmbeddingGenerator makeEmbeddingGenerator(Context context) {
         Optional<EmbeddingModel> embeddingModel = makeEmbeddingModel(context);
         if (embeddingModel.isPresent()) {
             int batchSize = context.getIntOption(Options.WRITE_EMBEDDER_BATCH_SIZE, 1, 1);
-            return new EmbeddingGenerator(embeddingModel.get(), batchSize);
+            EmbeddingModel model = embeddingModel.get();
+            if (Util.MAIN_LOGGER.isInfoEnabled()) {
+                Util.MAIN_LOGGER.info("Using embedding model with dimension: {}", model.dimension());
+            }
+            return new EmbeddingGenerator(model, batchSize);
         }
         return null;
     }
@@ -55,6 +60,10 @@ private static ChunkSelector makeChunkSelector(Context context) {
             return makeJsonChunkSelector(context);
         } else if (context.hasOption(Options.WRITE_EMBEDDER_CHUNKS_XPATH)) {
             return makeXmlChunkSelector(context);
+        } else if (context.hasOption(Options.WRITE_SPLITTER_TEXT)) {
+            return "xml".equalsIgnoreCase(context.getStringOption(Options.WRITE_SPLITTER_SIDECAR_DOCUMENT_TYPE)) ?
+                makeXmlChunkSelector(context) :
+                makeJsonChunkSelector(context);
         }
         throw new ConnectorException(String.format("To generate embeddings on documents, you must specify either " +
                 "%s or %s to define the location of chunks in documents.",