CSVParser.getRecords() knows how to use CSVFormat's maxRows

CSVParser.stream() knows how to use CSVFormat's maxRows

Author: garydgregory

src/changes/changes.xml

@@ -51,6 +51,8 @@
       <action type="add" dev="ggregory" due-to="Gary Gregory">CSVPrinter.printRecords(ResultSet) knows how to use CSVFormat's maxRows.</action>
       <action type="add" dev="ggregory" due-to="Gary Gregory">CSVPrinter.printRecords(Iterable) knows how to use CSVFormat's maxRows.</action>
       <action type="add" dev="ggregory" due-to="Gary Gregory">CSVPrinter.printRecords(Stream) knows how to use CSVFormat's maxRows.</action>
+      <action type="add" dev="ggregory" due-to="Gary Gregory">CSVParser.stream() knows how to use CSVFormat's maxRows.</action>
+      <action type="add" dev="ggregory" due-to="Gary Gregory">CSVParser.getRecords() knows how to use CSVFormat's maxRows.</action>
       <!-- UPDATE -->
       <action type="update" dev="ggregory" due-to="Gary Gregory">Bump com.opencsv:opencsv from 5.9 to 5.10.</action>
       <action type="update" dev="ggregory" due-to="Gary Gregory">Bump commons-codec:commons-codec from 1.17.2 to 1.18.0 #522.</action>

src/main/java/org/apache/commons/csv/CSVFormat.java

@@ -40,9 +40,11 @@
 import java.util.Objects;
 import java.util.Set;
 import java.util.function.Supplier;
+import java.util.stream.Stream;
 
 import org.apache.commons.codec.binary.Base64OutputStream;
 import org.apache.commons.io.IOUtils;
+import org.apache.commons.io.function.IOStream;
 import org.apache.commons.io.function.Uncheck;
 import org.apache.commons.io.output.AppendableOutputStream;
 
@@ -2088,6 +2090,14 @@ public boolean isQuoteCharacterSet() {
         return quoteCharacter != null;
     }
 
+    <T> IOStream<T> limit(final IOStream<T> stream) {
+        return getMaxRows() > 0 ? stream.limit(getMaxRows()) : stream;
+    }
+
+    <T> Stream<T> limit(final Stream<T> stream) {
+        return getMaxRows() > 0 ? stream.limit(getMaxRows()) : stream;
+    }
+
     /**
      * Parses the specified content.
      *
@@ -3178,4 +3188,5 @@ public CSVFormat withTrim() {
     public CSVFormat withTrim(final boolean trim) {
         return builder().setTrim(trim).get();
     }
+
 }

src/main/java/org/apache/commons/csv/CSVParser.java

@@ -938,7 +938,7 @@ CSVRecord nextRecord() throws IOException {
      * @since 1.9.0
      */
     public Stream<CSVRecord> stream() {
-        return StreamSupport.stream(Spliterators.spliteratorUnknownSize(iterator(), Spliterator.ORDERED), false);
+        return format.limit(StreamSupport.stream(Spliterators.spliteratorUnknownSize(iterator(), Spliterator.ORDERED), false));
     }
 
 }

src/main/java/org/apache/commons/csv/CSVPrinter.java

@@ -345,7 +345,7 @@ private void printRecordObject(final Object value) throws IOException {
 
     @SuppressWarnings("resource")
     private void printRecords(final IOStream<?> stream) throws IOException {
-        (format.getMaxRows() > 0 ? stream.limit(format.getMaxRows()) : stream).forEachOrdered(this::printRecordObject);
+        format.limit(stream).forEachOrdered(this::printRecordObject);
     }
 
     /**

src/test/java/org/apache/commons/csv/CSVParserTest.java

@@ -62,6 +62,7 @@
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.params.ParameterizedTest;
 import org.junit.jupiter.params.provider.EnumSource;
+import org.junit.jupiter.params.provider.ValueSource;
 
 /**
  * CSVParserTest
@@ -887,6 +888,20 @@ public void testGetRecordsFromBrokenInputStream() throws IOException {
 
     }
 
+    @ParameterizedTest
+    @ValueSource(longs = { -1, 0, 1, 2, 3, 4, Long.MAX_VALUE })
+    public void testGetRecordsMaxRows(final long maxRows) throws IOException {
+        try (CSVParser parser = CSVParser.parse(CSV_INPUT, CSVFormat.DEFAULT.builder().setIgnoreSurroundingSpaces(true).setMaxRows(maxRows).get())) {
+            final List<CSVRecord> records = parser.getRecords();
+            final long expectedLength = maxRows <= 0 || maxRows > RESULT.length ? RESULT.length : maxRows;
+            assertEquals(expectedLength, records.size());
+            assertFalse(records.isEmpty());
+            for (int i = 0; i < expectedLength; i++) {
+                assertArrayEquals(RESULT[i], records.get(i).values());
+            }
+        }
+    }
+
     @Test
     public void testGetRecordThreeBytesRead() throws Exception {
         final String code = "id,date,val5,val4\n" +
@@ -1656,6 +1671,23 @@ public void testStream() throws Exception {
         }
     }
 
+    @ParameterizedTest
+    @ValueSource(longs = { -1, 0, 1, 2, 3, 4, Long.MAX_VALUE })
+    public void testStreamMaxRows(final long maxRows) throws Exception {
+        final Reader in = new StringReader("a,b,c\n1,2,3\nx,y,z");
+        try (CSVParser parser = CSVFormat.DEFAULT.builder().setMaxRows(maxRows).get().parse(in)) {
+            final List<CSVRecord> list = parser.stream().collect(Collectors.toList());
+            assertFalse(list.isEmpty());
+            assertArrayEquals(new String[] { "a", "b", "c" }, list.get(0).values());
+            if (maxRows <= 0 || maxRows > 1) {
+                assertArrayEquals(new String[] { "1", "2", "3" }, list.get(1).values());
+            }
+            if (maxRows <= 0 || maxRows > 2) {
+                assertArrayEquals(new String[] { "x", "y", "z" }, list.get(2).values());
+            }
+        }
+    }
+
     @Test
     public void testThrowExceptionWithLineAndPosition() throws IOException {
         final String csvContent = "col1,col2,col3,col4,col5,col6,col7,col8,col9,col10\nrec1,rec2,rec3,rec4,rec5,rec6,rec7,rec8,\"\"rec9\"\",rec10";

src/test/java/org/apache/commons/csv/CSVPrinterTest.java

@@ -710,7 +710,7 @@ public void testExcelPrintAllIterableOfLists() throws IOException {
     }
 
     @ParameterizedTest
-    @ValueSource(longs = { -1, 0, 1, 2, Integer.MAX_VALUE })
+    @ValueSource(longs = { -1, 0, 1, 2, Long.MAX_VALUE })
     public void testExcelPrintAllStreamOfArrays(final long maxRows) throws IOException {
         final StringWriter sw = new StringWriter();
         final CSVFormat format = CSVFormat.EXCEL.builder().setMaxRows(maxRows).get();
@@ -840,7 +840,7 @@ public void testJdbcPrinterWithFirstEmptyValue2() throws IOException, ClassNotFo
     }
 
     @ParameterizedTest
-    @ValueSource(longs = { -1, 0, 1, 2, 3, 4, Integer.MAX_VALUE })
+    @ValueSource(longs = { -1, 0, 1, 2, 3, 4, Long.MAX_VALUE })
     public void testJdbcPrinterWithResultSet(final long maxRows) throws IOException, ClassNotFoundException, SQLException {
         final StringWriter sw = new StringWriter();
         final CSVFormat format = CSVFormat.DEFAULT.builder().setMaxRows(maxRows).get();
@@ -864,11 +864,11 @@ public void testJdbcPrinterWithResultSet(final long maxRows) throws IOException,
             assertEquals(allRows, resultString);
             expectedRowsWithHeader = TABLE_AND_HEADER_RECORD_COUNT;
         }
-        assertRowCount(format, resultString, expectedRowsWithHeader);
+        assertRowCount(CSVFormat.DEFAULT, resultString, expectedRowsWithHeader);
     }
 
     @ParameterizedTest
-    @ValueSource(longs = { -1, 0, 3, 4, Integer.MAX_VALUE })
+    @ValueSource(longs = { -1, 0, 3, 4, Long.MAX_VALUE })
     public void testJdbcPrinterWithResultSetHeader(final long maxRows) throws IOException, ClassNotFoundException, SQLException {
         final StringWriter sw = new StringWriter();
         try (Connection connection = getH2Connection()) {
@@ -887,13 +887,13 @@ public void testJdbcPrinterWithResultSetHeader(final long maxRows) throws IOExce
                     assertEquals(TABLE_RECORD_COUNT * 2, printer.getRecordCount());
                     assertNotEquals("ID,NAME" + RECORD_SEPARATOR + "1,r1" + RECORD_SEPARATOR + "2,r2" + RECORD_SEPARATOR, sw.toString());
                 }
-                assertRowCount(format, sw.toString(), TABLE_AND_HEADER_RECORD_COUNT + TABLE_RECORD_COUNT);
+                assertRowCount(CSVFormat.DEFAULT, sw.toString(), TABLE_AND_HEADER_RECORD_COUNT + TABLE_RECORD_COUNT);
             }
         }
     }
 
     @ParameterizedTest
-    @ValueSource(longs = { -1, 0, 3, 4, Integer.MAX_VALUE })
+    @ValueSource(longs = { -1, 0, 3, 4, Long.MAX_VALUE })
     public void testJdbcPrinterWithResultSetMetaData(final long maxRows) throws IOException, ClassNotFoundException, SQLException {
         final StringWriter sw = new StringWriter();
         try (Connection connection = getH2Connection()) {
@@ -1509,7 +1509,7 @@ public void testPrintCSVRecord() throws IOException {
     }
 
     @ParameterizedTest
-    @ValueSource(longs = { -1, 0, 3, 4, Integer.MAX_VALUE })
+    @ValueSource(longs = { -1, 0, 3, 4, Long.MAX_VALUE })
     public void testPrintCSVRecords(final long maxRows) throws IOException {
         // @formatter:off
         final String code = "a1,b1\n" + // 1)