Merge pull request #303 from marklogic/feature/17141-write-stream-zip-file

MLE-17141 Verifying normal zip files can be streamed

Author: rjrudin

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

@@ -43,7 +43,7 @@ class MarkLogicFileTable extends FileTable {
     @Override
     public ScanBuilder newScanBuilder(CaseInsensitiveStringMap options) {
         if ("true".equalsIgnoreCase(options.get(Options.STREAM_FILES)) && Util.MAIN_LOGGER.isInfoEnabled()) {
-            Util.MAIN_LOGGER.info("Will defer reading of file contents so they can be streamed during the writer phase.");
+            Util.MAIN_LOGGER.info("File streaming is enabled; will read files during writer phase.");
         }
         return new FileScanBuilder(options.asCaseSensitiveMap(), super.fileIndex());
     }

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

@@ -17,7 +17,7 @@ class DocumentScanBuilder implements ScanBuilder, SupportsPushDownLimit {
     DocumentScanBuilder(CaseInsensitiveStringMap options, StructType schema) {
         this.context = new DocumentContext(options, schema);
         if (this.context.isStreamingFiles() && Util.MAIN_LOGGER.isInfoEnabled()) {
-            Util.MAIN_LOGGER.info("Will defer reading documents from MarkLogic so they can be streamed to files during the writer phase.");
+            Util.MAIN_LOGGER.info("File streaming is enabled; will read documents from MarkLogic during writer phase.");
         }
     }
 

src/test/java/com/marklogic/spark/writer/file/WriteDocumentZipFilesTest.java

@@ -91,6 +91,41 @@ void opaqueURI(@TempDir @NotNull Path tempDir) throws IOException {
             "'schema-specific part', which is just example/123.xml.");
     }
 
+    /**
+     * Verifies that streaming documents to a zip file "just works" on account of supporting streaming of archive files
+     * first. The same ZipFileWriter is used for both. The only difference with archive files is that it will also
+     * check for metadata in each Spark row and include a metadata entry in the archive file.
+     *
+     * @param tempDir
+     * @throws Exception
+     */
+    @Test
+    void streamZipFile(@TempDir Path tempDir) throws Exception {
+        Dataset<Row> dataset = newSparkSession().read()
+            .format(CONNECTOR_IDENTIFIER)
+            .option(Options.READ_DOCUMENTS_PARTITIONS_PER_FOREST, 1)
+            .option(Options.CLIENT_URI, makeClientUri())
+            .option(Options.READ_DOCUMENTS_COLLECTIONS, "author")
+            .option(Options.STREAM_FILES, true)
+            .load();
+
+        assertEquals(15, dataset.count(), "Should have 1 row per author document.");
+        dataset.collectAsList().forEach(row -> {
+            assertFalse(row.isNullAt(0), "The URI column should be non-null.");
+            assertTrue(row.isNullAt(1), "The content column should be empty. The document will be read during the " +
+                "writer phase instead.");
+        });
+
+        dataset.write().format(CONNECTOR_IDENTIFIER)
+            .option(Options.STREAM_FILES, true)
+            .option(Options.CLIENT_URI, makeClientUri())
+            .option(Options.WRITE_FILES_COMPRESSION, "zip")
+            .mode(SaveMode.Append)
+            .save(tempDir.toFile().getAbsolutePath());
+
+        verifyZipFilesHaveExpectedFilenames(tempDir);
+        verifyZipFilesContainFifteenAuthors(tempDir);
+    }
 
     private Dataset<Row> readAuthorCollection() {
         return newSparkSession().read()