Adds support for --limit and --ion-flush-period for Ion 1.1 streams.

amazon-ion · Dec 12, 2024 · 5046e6e · 5046e6e
1 parent d9d852e
commit 5046e6e
Show file tree

Hide file tree

Showing 4 changed files with 143 additions and 22 deletions.
diff --git a/src/com/amazon/ion/benchmark/IonMeasurableWriteTask_1_1.java b/src/com/amazon/ion/benchmark/IonMeasurableWriteTask_1_1.java
@@ -1,12 +1,9 @@
 package com.amazon.ion.benchmark;
 
-import com.amazon.ion.MacroAwareIonReader;
 import com.amazon.ion.MacroAwareIonWriter;
-import com.amazon.ion.impl._Private_IonReaderBuilder;
 
 import java.io.IOException;
 import java.io.OutputStream;
-import java.nio.file.Files;
 import java.nio.file.Path;
 import java.util.function.Consumer;
 
@@ -44,16 +41,7 @@ void generateWriteInstructionsDom(Consumer<WriteInstruction<MacroAwareIonWriter>
 
     @Override
     void generateWriteInstructionsStreaming(Consumer<WriteInstruction<MacroAwareIonWriter>> instructionsSink) throws IOException {
-        if (options.limit != Integer.MAX_VALUE || options.flushPeriod != null) {
-            throw new UnsupportedOperationException("Benchmarking Ion 1.1 write using --limit or --ion-flush-period is not yet supported.");
-        }
-        // TODO support buildMacroAware from InputStream to avoid having to buffer all bytes.
-        try (
-            MacroAwareIonReader reader = ((_Private_IonReaderBuilder) IonUtilities.newReaderBuilderForInput(options))
-                .buildMacroAware(Files.readAllBytes(inputFile.toPath()))
-        ) {
-            reader.transcodeTo(new RecordingMacroAwareIonWriter(instructionsSink));
-        }
+        IonUtilities.rewriteIon11File(inputFile, options, new RecordingMacroAwareIonWriter(instructionsSink));
     }
 
     @Override

diff --git a/src/com/amazon/ion/benchmark/IonUtilities.java b/src/com/amazon/ion/benchmark/IonUtilities.java
@@ -482,18 +482,28 @@ private static void writeValuesWithOptions(
 
     /**
      * Rewrite the given Ion 1.1+ file to another Ion 1.1+ stream using the given options.
-     * @param input an ion 1.1+ file.
+     * @param inputFile an ion 1.1+ file.
      * @param options the options to use when re-writing.
      * @param writer the writer of the new stream.
      * @throws IOException if thrown when reading or writing.
      */
-    private static void rewriteIon11File(Path input, OptionsCombinationBase options, IonWriter writer) throws IOException{
-        if (options.limit != Integer.MAX_VALUE) {
-            throw new UnsupportedOperationException("Macro-aware transcoding of Ion 1.1 with the --limit option not yet supported.");
-        }
-        // TODO add a method to MacroAwareIonReader to write one value at a time so that 'limit' can be used
-        try (MacroAwareIonReader macroAwareIonReader = ((_Private_IonReaderBuilder) newReaderBuilderForInput(options)).buildMacroAware(Files.readAllBytes(input))) {
-            macroAwareIonReader.transcodeTo((MacroAwareIonWriter) writer);
+    static void rewriteIon11File(File inputFile, OptionsCombinationBase options, IonWriter writer) throws IOException {
+        try (
+            MacroAwareIonReader macroAwareIonReader = ((_Private_IonReaderBuilder) newReaderBuilderForInput(options))
+                .buildMacroAware(options.newInputStream(inputFile))
+        ) {
+            macroAwareIonReader.prepareTranscodeTo((MacroAwareIonWriter) writer);
+            int i = 0;
+            boolean isUnlimited = options.limit == Integer.MAX_VALUE;
+            while (isUnlimited || i < options.limit) {
+                if (!macroAwareIonReader.transcodeNext()) {
+                    break;
+                }
+                if (options.flushPeriod != null && i % options.flushPeriod == 0) {
+                    writer.flush();
+                }
+                i++;
+            }
         }
     }
 
@@ -517,7 +527,7 @@ static void rewriteIonFile(Format inputFormat, Path input, Path output, OptionsC
             int inputMinorVersion = getMinorVersion(inputFormat, input.toFile());
             writer = writerSupplier.get(options.newOutputStream(outputFile));
             if (inputMinorVersion > 0 && options.ionMinorVersion > 0) {
-                rewriteIon11File(input, options, writer);
+                rewriteIon11File(inputFile, options, writer);
             } else {
                 if (
                     options.flushPeriod == null &&

diff --git a/src/com/amazon/ion/benchmark/RecordingMacroAwareIonWriter.java b/src/com/amazon/ion/benchmark/RecordingMacroAwareIonWriter.java
@@ -153,6 +153,13 @@ public void writeValue(IonReader reader) {
         // Note: w -> w.writeValue(reader) is not correct because it does not capture the value at which the reader
         // is currently positioned. For now, we use IonValue to achieve this, though it would be more efficient to
         // capture a primitive and store an instruction that writes that primitive to the writer directly.
+        if (reader.isInStruct()) {
+            // The field name must be captured and written manually because the IonValue below is not considered
+            // a child of a container.
+            SymbolToken fieldName = reader.getFieldNameSymbol();
+            instructionsSink.accept(w -> w.setFieldNameSymbol(fieldName));
+        }
+        // Note: the following includes any annotations on the value.
         IonValue value = SYSTEM.newValue(reader);
         instructionsSink.accept(value::writeTo);
     }

diff --git a/tst/com/amazon/ion/benchmark/OptionsTest.java b/tst/com/amazon/ion/benchmark/OptionsTest.java
@@ -2488,6 +2488,64 @@ public void writeIon11WithMacros() throws Exception {
         assertEquals(2, countOccurrencesOfSubstring(ion11Text, "(:foo)"));
     }
 
+    @Test
+    public void writeIon11WithMacrosAndLimit() throws Exception {
+        WriteOptionsCombination optionsCombination = parseSingleOptionsCombination(
+            "write",
+            "--format",
+            "ion_text",
+            "--ion-minor-version",
+            "1",
+            "--io-type",
+            "file",
+            "--limit",
+            "2",
+            "binaryMacroInvocations.10n"
+        );
+        byte[] data = assertWriteTaskExecutesCorrectly(
+            "binaryMacroInvocations.10n",
+            optionsCombination,
+            Format.ION_TEXT,
+            IoType.FILE
+        );
+        String ion11Text = new String(data, StandardCharsets.UTF_8);
+        assertEquals(1, countOccurrencesOfSubstring(ion11Text, "add_symbols"));
+        assertEquals(1, countOccurrencesOfSubstring(ion11Text, "add_macros"));
+        assertEquals(0, countOccurrencesOfSubstring(ion11Text, "set_symbols"));
+        assertEquals(0, countOccurrencesOfSubstring(ion11Text, "set_macros"));
+        assertEquals(1, countOccurrencesOfSubstring(ion11Text, "(:Pi)"));
+        assertEquals(0, countOccurrencesOfSubstring(ion11Text, "(:foo)"));
+    }
+
+    @Test
+    public void writeIon11WithMacrosAndFlushPeriod() throws Exception {
+        WriteOptionsCombination optionsCombination = parseSingleOptionsCombination(
+            "write",
+            "--format",
+            "ion_text",
+            "--ion-minor-version",
+            "1",
+            "--io-type",
+            "file",
+            "--ion-flush-period",
+            "2",
+            "binaryMacroInvocations.10n"
+        );
+        byte[] data = assertWriteTaskExecutesCorrectly(
+            "binaryMacroInvocations.10n",
+            optionsCombination,
+            Format.ION_TEXT,
+            IoType.FILE
+        );
+        String ion11Text = new String(data, StandardCharsets.UTF_8);
+        assertEquals(1, countOccurrencesOfSubstring(ion11Text, "add_symbols"));
+        assertEquals(2, countOccurrencesOfSubstring(ion11Text, "add_macros"));
+        assertEquals(1, countOccurrencesOfSubstring(ion11Text, "set_symbols"));
+        assertEquals(1, countOccurrencesOfSubstring(ion11Text, "set_macros"));
+        assertEquals(1, countOccurrencesOfSubstring(ion11Text, "(:Pi)"));
+        assertEquals(2, countOccurrencesOfSubstring(ion11Text, "(:foo)"));
+    }
+
     @Test
     public void writeIon11WithMacrosToIon10() throws Exception {
         WriteOptionsCombination optionsCombination = parseSingleOptionsCombination(
@@ -2535,6 +2593,64 @@ public void readIon11WithMacros() throws Exception {
         assertEquals(2, countOccurrencesOfSubstring(ion11Text, "(:foo)"));
     }
 
+    @Test
+    public void readIon11WithMacrosAndLimit() throws Exception {
+        ReadOptionsCombination optionsCombination = parseSingleOptionsCombination(
+            "read",
+            "--format",
+            "ion_text",
+            "--ion-minor-version",
+            "1",
+            "--io-type",
+            "file",
+            "--limit",
+            "3",
+            "binaryMacroInvocations.10n"
+        );
+        byte[] data = assertReadTaskExecutesCorrectly(
+            "binaryMacroInvocations.10n",
+            optionsCombination,
+            Format.ION_TEXT,
+            true
+        );
+        String ion11Text = new String(data, StandardCharsets.UTF_8);
+        assertEquals(1, countOccurrencesOfSubstring(ion11Text, "add_symbols"));
+        assertEquals(2, countOccurrencesOfSubstring(ion11Text, "add_macros"));
+        assertEquals(1, countOccurrencesOfSubstring(ion11Text, "set_symbols"));
+        assertEquals(0, countOccurrencesOfSubstring(ion11Text, "set_macros"));
+        assertEquals(1, countOccurrencesOfSubstring(ion11Text, "(:Pi)"));
+        assertEquals(1, countOccurrencesOfSubstring(ion11Text, "(:foo)"));
+    }
+
+    @Test
+    public void readIon11WithMacrosAndFlushPeriod() throws Exception {
+        ReadOptionsCombination optionsCombination = parseSingleOptionsCombination(
+            "read",
+            "--format",
+            "ion_text",
+            "--ion-minor-version",
+            "1",
+            "--io-type",
+            "file",
+            "--ion-flush-period",
+            "2",
+            "binaryMacroInvocations.10n"
+        );
+        byte[] data = assertReadTaskExecutesCorrectly(
+            "binaryMacroInvocations.10n",
+            optionsCombination,
+            Format.ION_TEXT,
+            true
+        );
+        String ion11Text = new String(data, StandardCharsets.UTF_8);
+        assertEquals(1, countOccurrencesOfSubstring(ion11Text, "add_symbols"));
+        assertEquals(2, countOccurrencesOfSubstring(ion11Text, "add_macros"));
+        assertEquals(1, countOccurrencesOfSubstring(ion11Text, "set_symbols"));
+        assertEquals(1, countOccurrencesOfSubstring(ion11Text, "set_macros"));
+        assertEquals(1, countOccurrencesOfSubstring(ion11Text, "(:Pi)"));
+        assertEquals(2, countOccurrencesOfSubstring(ion11Text, "(:foo)"));
+    }
+
     @Test
     public void readIon10ConvertedFromIon11WithMacros() throws Exception {
         ReadOptionsCombination optionsCombination = parseSingleOptionsCombination(