Merge pull request #297 from marklogic/feature/17095-add-iterator

rjrudin · web-flow · commit 5aaf88326e1c · 2024-09-26T12:48:38.000-04:00
MLE-17095 Refactoring; changed row converter to return iterator
diff --git a/src/main/java/com/marklogic/spark/reader/file/ArchiveFileReader.java b/src/main/java/com/marklogic/spark/reader/file/ArchiveFileReader.java
@@ -56,27 +56,8 @@ public boolean next() {
             if (isLegacyFormat == null) {
                 isLegacyFormat = !nextZipEntry.getName().endsWith(".metadata");
             }
-            if (!isLegacyFormat) {
-                return readMetadataFollowedByContentEntry();
-            }
-
-            byte[] content = fileContext.readBytes(currentZipInputStream);
-            if (content == null || content.length == 0) {
-                return openNextFileAndReadNextEntry();
-            }
-            final String zipEntryName = nextZipEntry.getName();
-
-            byte[] metadataBytes = readMetadataEntry(zipEntryName);
-            if (metadataBytes == null || metadataBytes.length == 0) {
-                return openNextFileAndReadNextEntry();
-            }
 
-            DocumentMetadataHandle metadata = new DocumentMetadataHandle();
-            metadata.fromBuffer(metadataBytes);
-            this.nextRowToReturn = new DocumentRowBuilder(this.metadataCategories)
-                .withUri(zipEntryName).withContent(content).withMetadata(metadata)
-                .buildRow();
-            return true;
+            return isLegacyFormat ? readContentFollowedByMetadata(nextZipEntry) : readMetadataFollowedByContent();
         } catch (IOException e) {
             String message = String.format("Unable to read archive file at %s; cause: %s", this.currentFilePath, e.getMessage());
             if (fileContext.isReadAbortOnFailure()) {
@@ -97,7 +78,35 @@ public void close() {
         IOUtils.closeQuietly(this.currentZipInputStream);
     }
 
-    private boolean readMetadataFollowedByContentEntry() throws IOException {
+    /**
+     * This is the Flux 1.0 "legacy" approach, where content was written first, followed by metadata. This does not
+     * support streaming.
+     */
+    private boolean readContentFollowedByMetadata(ZipEntry contentZipEntry) throws IOException {
+        byte[] content = fileContext.readBytes(currentZipInputStream);
+        if (content == null || content.length == 0) {
+            return openNextFileAndReadNextEntry();
+        }
+        final String zipEntryName = contentZipEntry.getName();
+
+        byte[] metadataBytes = readMetadataEntry(zipEntryName);
+        if (metadataBytes == null || metadataBytes.length == 0) {
+            return openNextFileAndReadNextEntry();
+        }
+
+        DocumentMetadataHandle metadata = new DocumentMetadataHandle();
+        metadata.fromBuffer(metadataBytes);
+
+        this.nextRowToReturn = new DocumentRowBuilder(this.metadataCategories)
+            .withUri(zipEntryName).withContent(content).withMetadata(metadata)
+            .buildRow();
+        return true;
+    }
+
+    /**
+     * This is the Flux 1.1+ approach, where the metadata entry is written first. This supports streaming.
+     */
+    private boolean readMetadataFollowedByContent() throws IOException {
         byte[] metadataBytes = fileContext.readBytes(currentZipInputStream);
         if (metadataBytes == null || metadataBytes.length == 0) {
             return openNextFileAndReadNextEntry();
diff --git a/src/main/java/com/marklogic/spark/writer/ArbitraryRowConverter.java b/src/main/java/com/marklogic/spark/writer/ArbitraryRowConverter.java
@@ -20,9 +20,10 @@
 import org.apache.spark.sql.types.StructType;
 
 import java.util.ArrayList;
+import java.util.Iterator;
 import java.util.List;
-import java.util.Optional;
 import java.util.UUID;
+import java.util.stream.Stream;
 
 /**
  * Handles building a document from an "arbitrary" row - i.e. one with an unknown schema, where the row will be
@@ -53,7 +54,7 @@ class ArbitraryRowConverter implements RowConverter {
     }
 
     @Override
-    public Optional<DocBuilder.DocumentInputs> convertRow(InternalRow row) {
+    public Iterator<DocBuilder.DocumentInputs> convertRow(InternalRow row) {
         String initialUri = null;
         if (this.filePathIndex > -1) {
             initialUri = row.getString(this.filePathIndex) + "/" + UUID.randomUUID();
@@ -103,7 +104,7 @@ else if (deserializedJson != null) {
             }
         }
 
-        return Optional.of(new DocBuilder.DocumentInputs(initialUri, contentHandle, uriTemplateValues, null));
+        return Stream.of(new DocBuilder.DocumentInputs(initialUri, contentHandle, uriTemplateValues, null)).iterator();
     }
 
     @Override
diff --git a/src/main/java/com/marklogic/spark/writer/DocumentRowConverter.java b/src/main/java/com/marklogic/spark/writer/DocumentRowConverter.java
@@ -10,7 +10,6 @@
 import com.marklogic.client.io.DocumentMetadataHandle;
 import com.marklogic.client.io.Format;
 import com.marklogic.client.io.InputStreamHandle;
-import com.marklogic.client.io.marker.AbstractWriteHandle;
 import com.marklogic.spark.ConnectorException;
 import com.marklogic.spark.Options;
 import com.marklogic.spark.reader.document.DocumentRowSchema;
@@ -21,8 +20,9 @@
 import java.io.IOException;
 import java.io.ObjectInputStream;
 import java.util.ArrayList;
+import java.util.Iterator;
 import java.util.List;
-import java.util.Optional;
+import java.util.stream.Stream;
 
 /**
  * Knows how to build a document from a row corresponding to our {@code DocumentRowSchema}.
@@ -42,31 +42,24 @@ class DocumentRowConverter implements RowConverter {
     }
 
     @Override
-    public Optional<DocBuilder.DocumentInputs> convertRow(InternalRow row) {
+    public Iterator<DocBuilder.DocumentInputs> convertRow(InternalRow row) {
         final String uri = row.getString(0);
 
         final boolean isNakedProperties = row.isNullAt(1);
         if (isNakedProperties) {
             DocumentMetadataHandle metadata = DocumentRowSchema.makeDocumentMetadata(row);
-            return Optional.of(new DocBuilder.DocumentInputs(uri, null, null, metadata));
+            return Stream.of(new DocBuilder.DocumentInputs(uri, null, null, metadata)).iterator();
         }
 
-        Content content = this.isStreamingFromFiles ?
-            readContentFromFile(uri, row) :
-            readContentFromRow(uri, row);
-
-        DocumentMetadataHandle metadata = DocumentRowSchema.makeDocumentMetadata(row);
-        return Optional.of(new DocBuilder.DocumentInputs(
-            uri, content.contentHandle, content.uriTemplateValues, metadata)
-        );
+        return this.isStreamingFromFiles ? readContentFromFile(uri, row) : readContentFromRow(uri, row);
     }
 
     @Override
     public List<DocBuilder.DocumentInputs> getRemainingDocumentInputs() {
         return new ArrayList<>();
     }
 
-    private Content readContentFromRow(String uri, InternalRow row) {
+    private Iterator<DocBuilder.DocumentInputs> readContentFromRow(String uri, InternalRow row) {
         BytesHandle bytesHandle = new BytesHandle(row.getBinary(1));
         if (this.documentFormat != null) {
             bytesHandle.withFormat(this.documentFormat);
@@ -76,7 +69,8 @@ private Content readContentFromRow(String uri, InternalRow row) {
             String format = row.isNullAt(2) ? null : row.getString(2);
             uriTemplateValues = deserializeContentToJson(uri, bytesHandle, format);
         }
-        return new Content(bytesHandle, uriTemplateValues);
+        DocumentMetadataHandle metadata = DocumentRowSchema.makeDocumentMetadata(row);
+        return Stream.of(new DocBuilder.DocumentInputs(uri, bytesHandle, uriTemplateValues, metadata)).iterator();
     }
 
     private JsonNode deserializeContentToJson(String initialUri, BytesHandle contentHandle, String format) {
@@ -97,7 +91,7 @@ private JsonNode deserializeContentToJson(String initialUri, BytesHandle content
      * In a scenario where the user wants to stream a file into MarkLogic, the content column will contain a serialized
      * instance of {@code FileContext}, which is used to stream the file into a {@code InputStreamHandle}.
      */
-    private Content readContentFromFile(String filePath, InternalRow row) {
+    private Iterator<DocBuilder.DocumentInputs> readContentFromFile(String filePath, InternalRow row) {
         byte[] bytes = row.getBinary(1);
         String filePathInErrorMessage = filePath;
         try {
@@ -109,27 +103,10 @@ private Content readContentFromFile(String filePath, InternalRow row) {
             if (this.documentFormat != null) {
                 streamHandle.withFormat(this.documentFormat);
             }
-            return new Content(streamHandle, null);
+            DocumentMetadataHandle metadata = DocumentRowSchema.makeDocumentMetadata(row);
+            return Stream.of(new DocBuilder.DocumentInputs(filePath, streamHandle, null, metadata)).iterator();
         } catch (Exception e) {
             throw new ConnectorException(String.format("Unable to read from file %s; cause: %s", filePathInErrorMessage, e.getMessage()));
         }
     }
-
-    private static class Content {
-        private final AbstractWriteHandle contentHandle;
-        private final JsonNode uriTemplateValues;
-
-        public Content(AbstractWriteHandle contentHandle, JsonNode uriTemplateValues) {
-            this.contentHandle = contentHandle;
-            this.uriTemplateValues = uriTemplateValues;
-        }
-
-        AbstractWriteHandle getContentHandle() {
-            return contentHandle;
-        }
-
-        JsonNode getUriTemplateValues() {
-            return uriTemplateValues;
-        }
-    }
 }
diff --git a/src/main/java/com/marklogic/spark/writer/FileRowConverter.java b/src/main/java/com/marklogic/spark/writer/FileRowConverter.java
@@ -14,8 +14,10 @@
 
 import java.io.IOException;
 import java.util.ArrayList;
+import java.util.Iterator;
 import java.util.List;
 import java.util.Optional;
+import java.util.stream.Stream;
 
 /**
  * Knows how to build a document from a row corresponding to our {@code FileRowSchema}.
@@ -33,12 +35,12 @@ class FileRowConverter implements RowConverter {
     }
 
     @Override
-    public Optional<DocBuilder.DocumentInputs> convertRow(InternalRow row) {
+    public Iterator<DocBuilder.DocumentInputs> convertRow(InternalRow row) {
         final String path = row.getString(writeContext.getFileSchemaPathPosition());
         BytesHandle contentHandle = new BytesHandle(row.getBinary(writeContext.getFileSchemaContentPosition()));
         forceFormatIfNecessary(contentHandle);
         Optional<JsonNode> uriTemplateValues = deserializeContentToJson(path, contentHandle, row);
-        return Optional.of(new DocBuilder.DocumentInputs(path, contentHandle, uriTemplateValues.orElse(null), null));
+        return Stream.of(new DocBuilder.DocumentInputs(path, contentHandle, uriTemplateValues.orElse(null), null)).iterator();
     }
 
     @Override
diff --git a/src/main/java/com/marklogic/spark/writer/RowConverter.java b/src/main/java/com/marklogic/spark/writer/RowConverter.java
@@ -5,22 +5,20 @@
 
 import org.apache.spark.sql.catalyst.InternalRow;
 
+import java.util.Iterator;
 import java.util.List;
-import java.util.Optional;
 
 /**
  * Strategy interface for how a Spark row is converted into a set of inputs for writing a document to MarkLogic.
  */
 public interface RowConverter {
 
     /**
-     * An implementation can return an empty Optional, which will happen when the row will be used with other rows to
-     * form a document.
-     *
      * @param row
-     * @return
+     * @return an iterator of inputs for creating documents to write to MarkLogic. An iterator is used to allow the
+     * implementor to return multiple documents if necessary.
      */
-    Optional<DocBuilder.DocumentInputs> convertRow(InternalRow row);
+    Iterator<DocBuilder.DocumentInputs> convertRow(InternalRow row);
 
     /**
      * Called when {@code WriteBatcherDataWriter} has no more rows to send, but the implementation may have one or
diff --git a/src/main/java/com/marklogic/spark/writer/WriteBatcherDataWriter.java b/src/main/java/com/marklogic/spark/writer/WriteBatcherDataWriter.java
@@ -30,8 +30,8 @@
 import org.slf4j.LoggerFactory;
 
 import java.util.ArrayList;
+import java.util.Iterator;
 import java.util.List;
-import java.util.Optional;
 import java.util.Set;
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.concurrent.atomic.AtomicReference;
@@ -92,9 +92,10 @@ class WriteBatcherDataWriter implements DataWriter<InternalRow> {
     @Override
     public void write(InternalRow row) {
         throwWriteFailureIfExists();
-        Optional<DocBuilder.DocumentInputs> document = rowConverter.convertRow(row);
-        if (document.isPresent()) {
-            DocumentWriteOperation writeOp = this.docBuilder.build(document.get());
+
+        Iterator<DocBuilder.DocumentInputs> iterator = rowConverter.convertRow(row);
+        while (iterator.hasNext()) {
+            DocumentWriteOperation writeOp = this.docBuilder.build(iterator.next());
             if (this.isStreamingFiles) {
                 writeDocumentViaPutOperation(writeOp);
             } else {
diff --git a/src/main/java/com/marklogic/spark/writer/rdf/RdfRowConverter.java b/src/main/java/com/marklogic/spark/writer/rdf/RdfRowConverter.java
@@ -14,6 +14,7 @@
 
 import java.util.*;
 import java.util.stream.Collectors;
+import java.util.stream.Stream;
 
 /**
  * Converts each row into a sem:triple element, which is then added to a sem:triples XML document associated with a
@@ -60,7 +61,7 @@ public RdfRowConverter(WriteContext writeContext) {
     }
 
     @Override
-    public Optional<DocBuilder.DocumentInputs> convertRow(InternalRow row) {
+    public Iterator<DocBuilder.DocumentInputs> convertRow(InternalRow row) {
         final String graph = determineGraph(row);
         graphs.add(graph);
 
@@ -75,9 +76,9 @@ public Optional<DocBuilder.DocumentInputs> convertRow(InternalRow row) {
         triplesDocument.addTriple(row);
         if (triplesDocument.hasMaxTriples()) {
             triplesDocuments.remove(graph);
-            return Optional.of(triplesDocument.buildDocument());
+            return Stream.of(triplesDocument.buildDocument()).iterator();
         }
-        return Optional.empty();
+        return Stream.<DocBuilder.DocumentInputs>empty().iterator();
     }
 
     /**

Original file line number	Diff line number	Diff line change
`@@ -20,9 +20,10 @@`
`20`	`20`	`import org.apache.spark.sql.types.StructType;`
`21`	`21`
`22`	`22`	`import java.util.ArrayList;`
	`23`	`+import java.util.Iterator;`
`23`	`24`	`import java.util.List;`
`24`		`-import java.util.Optional;`
`25`	`25`	`import java.util.UUID;`
	`26`	`+import java.util.stream.Stream;`
`26`	`27`
`27`	`28`	`/**`
`28`	`29`	`* Handles building a document from an "arbitrary" row - i.e. one with an unknown schema, where the row will be`
`@@ -53,7 +54,7 @@ class ArbitraryRowConverter implements RowConverter {`
`53`	`54`	`}`
`54`	`55`
`55`	`56`	`@Override`
`56`		`- public Optional<DocBuilder.DocumentInputs> convertRow(InternalRow row) {`
	`57`	`+ public Iterator<DocBuilder.DocumentInputs> convertRow(InternalRow row) {`
`57`	`58`	`String initialUri = null;`
`58`	`59`	`if (this.filePathIndex > -1) {`
`59`	`60`	`initialUri = row.getString(this.filePathIndex) + "/" + UUID.randomUUID();`
`@@ -103,7 +104,7 @@ else if (deserializedJson != null) {`
`103`	`104`	`}`
`104`	`105`	`}`
`105`	`106`
`106`		`- return Optional.of(new DocBuilder.DocumentInputs(initialUri, contentHandle, uriTemplateValues, null));`
	`107`	`+ return Stream.of(new DocBuilder.DocumentInputs(initialUri, contentHandle, uriTemplateValues, null)).iterator();`
`107`	`108`	`}`
`108`	`109`
`109`	`110`	`@Override`