Repository: jtablesaw/tablesaw
Pull request #739 (2020-01-13 14:32:50)

Implement single-pass random sampling of rows while reading input
Thanks for contributing.

- [x] Tick to sign-off your agreement to the [Developer Certificate of Origin (DCO) 1.1](https://developercertificate.org)

## Description

Implemented a single-pass row random sampling while reading CSVs using the reservoir sampling algorithm. This PR was extracted from PR #738.

## Testing

Did you add a unit test?

Yes.



Patch:
diff --git a/core/src/main/java/tech/tablesaw/analytic/AnalyticQuery.java b/core/src/main/java/tech/tablesaw/analytic/AnalyticQuery.java
index 1713c4f8..6a00d013 100644
--- a/core/src/main/java/tech/tablesaw/analytic/AnalyticQuery.java
+++ b/core/src/main/java/tech/tablesaw/analytic/AnalyticQuery.java
@@ -3,8 +3,8 @@ package tech.tablesaw.analytic;
 import com.google.common.annotations.Beta;
 import com.google.common.base.Preconditions;
 import java.util.Arrays;
-import java.util.LinkedHashSet;
 import java.util.Optional;
+import java.util.Set;
 import tech.tablesaw.analytic.AnalyticQuerySteps.AddAggregateFunctions;
 import tech.tablesaw.analytic.AnalyticQuerySteps.AddAggregateFunctionsWithExecute;
 import tech.tablesaw.analytic.AnalyticQuerySteps.AddNumberingFunction;
@@ -117,7 +117,7 @@ public final class AnalyticQuery {
     return argumentList;
   }
 
-  LinkedHashSet<String> getPartitionColumns() {
+  Set<String> getPartitionColumns() {
     return windowSpecification.getPartitionColumns();
   }
 
diff --git a/core/src/main/java/tech/tablesaw/analytic/ArgumentList.java b/core/src/main/java/tech/tablesaw/analytic/ArgumentList.java
index 0720d2df..c68483c2 100644
--- a/core/src/main/java/tech/tablesaw/analytic/ArgumentList.java
+++ b/core/src/main/java/tech/tablesaw/analytic/ArgumentList.java
@@ -9,6 +9,7 @@ import java.util.LinkedHashSet;
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
+import java.util.Set;
 import java.util.stream.Stream;
 import tech.tablesaw.api.ColumnType;
 import tech.tablesaw.columns.Column;
@@ -20,12 +21,12 @@ final class ArgumentList {
   private final Map<String, FunctionCall<NumberingFunctions>> numberingFunctions;
 
   // Used to determine the order in which to add new columns.
-  private final LinkedHashSet<String> newColumnNames;
+  private final Set<String> newColumnNames;
 
-  public ArgumentList(
+  private ArgumentList(
       Map<String, FunctionCall<AggregateFunctions>> aggregateFunctions,
       Map<String, FunctionCall<NumberingFunctions>> numberingFunctions,
-      LinkedHashSet<String> newColumnNames) {
+      Set<String> newColumnNames) {
     this.aggregateFunctions = aggregateFunctions;
     this.numberingFunctions = numberingFunctions;
     this.newColumnNames = newColumnNames;
@@ -163,7 +164,7 @@ final class ArgumentList {
         new HashMap<>();
 
     // Throws if a column with the same name is registered twice.
-    private final LinkedHashSet<String> newColumnNames = new LinkedHashSet<>();
+    private final Set<String> newColumnNames = new LinkedHashSet<>();
 
     // Temporarily store analytic function data until the user calls 'as' to give the new column a
     // name
diff --git a/core/src/main/java/tech/tablesaw/analytic/WindowSpecification.java b/core/src/main/java/tech/tablesaw/analytic/WindowSpecification.java
index 9eb7d98a..188ea22e 100644
--- a/core/src/main/java/tech/tablesaw/analytic/WindowSpecification.java
+++ b/core/src/main/java/tech/tablesaw/analytic/WindowSpecification.java
@@ -7,6 +7,7 @@ import java.util.LinkedHashSet;
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
+import java.util.Set;
 import java.util.stream.Collectors;
 import tech.tablesaw.sorting.Sort;
 
@@ -14,11 +15,10 @@ import tech.tablesaw.sorting.Sort;
 final class WindowSpecification {
 
   private final String windowName;
-  private final LinkedHashSet<String> partitionColumns;
+  private final Set<String> partitionColumns;
   private final Sort sort;
 
-  private WindowSpecification(
-      String windowName, LinkedHashSet<String> partitionColumns, Sort sort) {
+  private WindowSpecification(String windowName, Set<String> partitionColumns, Sort sort) {
     this.windowName = windowName;
     this.partitionColumns = partitionColumns;
     this.sort = sort;
@@ -58,7 +58,7 @@ final class WindowSpecification {
     return windowName;
   }
 
-  public LinkedHashSet<String> getPartitionColumns() {
+  public Set<String> getPartitionColumns() {
     return partitionColumns;
   }
 
diff --git a/core/src/main/java/tech/tablesaw/api/BooleanColumn.java b/core/src/main/java/tech/tablesaw/api/BooleanColumn.java
index b131e219..5c95258f 100644
--- a/core/src/main/java/tech/tablesaw/api/BooleanColumn.java
+++ b/core/src/main/java/tech/tablesaw/api/BooleanColumn.java
@@ -447,8 +447,9 @@ public class BooleanColumn extends AbstractColumn<BooleanColumn, Boolean>
     return this;
   }
 
-  private void set(int i, byte b) {
+  private BooleanColumn set(int i, byte b) {
     data.set(i, b);
+    return this;
   }
 
   @Override
@@ -456,6 +457,11 @@ public class BooleanColumn extends AbstractColumn<BooleanColumn, Boolean>
     return set(i, val.booleanValue());
   }
 
+  @Override
+  public Column<Boolean> set(int row, String stringValue, AbstractColumnParser<?> parser) {
+    return set(row, parser.parseByte(stringValue));
+  }
+
   @Override
   public BooleanColumn lead(int n) {
     BooleanColumn column = lag(-n);
diff --git a/core/src/main/java/tech/tablesaw/api/DoubleColumn.java b/core/src/main/java/tech/tablesaw/api/DoubleColumn.java
index 4befa4fa..bb5463ba 100644
--- a/core/src/main/java/tech/tablesaw/api/DoubleColumn.java
+++ b/core/src/main/java/tech/tablesaw/api/DoubleColumn.java
@@ -317,6 +317,11 @@ public class DoubleColumn extends NumberColumn<DoubleColumn, Double>
     return this;
   }
 
+  @Override
+  public Column<Double> set(int row, String stringValue, AbstractColumnParser<?> parser) {
+    return set(row, parser.parseDouble(stringValue));
+  }
+
   @Override
   public DoubleColumn append(final Column<Double> column) {
     Preconditions.checkArgument(column.type() == this.type());
diff --git a/core/src/main/java/tech/tablesaw/api/FloatColumn.java b/core/src/main/java/tech/tablesaw/api/FloatColumn.java
index 31a0886c..c017b60d 100644
--- a/core/src/main/java/tech/tablesaw/api/FloatColumn.java
+++ b/core/src/main/java/tech/tablesaw/api/FloatColumn.java
@@ -225,6 +225,11 @@ public class FloatColumn extends NumberColumn<FloatColumn, Float> {
     return this;
   }
 
+  @Override
+  public Column<Float> set(int row, String stringValue, AbstractColumnParser<?> parser) {
+    return set(row, parser.parseFloat(stringValue));
+  }
+
   @Override
   public FloatColumn append(final Column<Float> column) {
     Preconditions.checkArgument(column.type() == this.type());
diff --git a/core/src/main/java/tech/tablesaw/api/IntColumn.java b/core/src/main/java/tech/tablesaw/api/IntColumn.java
index 01ebdf33..7045bac5 100644
--- a/core/src/main/java/tech/tablesaw/api/IntColumn.java
+++ b/core/src/main/java/tech/tablesaw/api/IntColumn.java
@@ -254,6 +254,11 @@ public class IntColumn extends NumberColumn<IntColumn, Integer>
     return set(row, ((IntColumn) column).getInt(sourceRow));
   }
 
+  @Override
+  public Column<Integer> set(int row, String stringValue, AbstractColumnParser<?> parser) {
+    return set(row, parser.parseInt(stringValue));
+  }
+
   @Override
   public IntColumn appendMissing() {
     return append(IntColumnType.missingValueIndicator());
diff --git a/core/src/main/java/tech/tablesaw/api/LongColumn.java b/core/src/main/java/tech/tablesaw/api/LongColumn.java
index bde69294..8f237416 100644
--- a/core/src/main/java/tech/tablesaw/api/LongColumn.java
+++ b/core/src/main/java/tech/tablesaw/api/LongColumn.java
@@ -260,6 +260,11 @@ public class LongColumn extends NumberColumn<LongColumn, Long> implements Catego
     return this;
   }
 
+  @Override
+  public Column<Long> set(int row, String stringValue, AbstractColumnParser<?> parser) {
+    return set(row, parser.parseLong(stringValue));
+  }
+
   @Override
   public LongColumn append(final Column<Long> column) {
     Preconditions.checkArgument(column.type() == this.type());
diff --git a/core/src/main/java/tech/tablesaw/api/ShortColumn.java b/core/src/main/java/tech/tablesaw/api/ShortColumn.java
index 698e61b8..dfe661f7 100644
--- a/core/src/main/java/tech/tablesaw/api/ShortColumn.java
+++ b/core/src/main/java/tech/tablesaw/api/ShortColumn.java
@@ -229,6 +229,11 @@ public class ShortColumn extends NumberColumn<ShortColumn, Short>
     return this;
   }
 
+  @Override
+  public Column<Short> set(int row, String stringValue, AbstractColumnParser<?> parser) {
+    return set(row, parser.parseShort(stringValue));
+  }
+
   @Override
   public ShortColumn append(final Column<Short> column) {
     Preconditions.checkArgument(column.type() == this.type());
diff --git a/core/src/main/java/tech/tablesaw/columns/Column.java b/core/src/main/java/tech/tablesaw/columns/Column.java
index e351f970..a988bdf9 100644
--- a/core/src/main/java/tech/tablesaw/columns/Column.java
+++ b/core/src/main/java/tech/tablesaw/columns/Column.java
@@ -564,6 +564,11 @@ public interface Column<T> extends Iterable<T>, Comparator<T> {
 
   Column<T> set(int row, T value);
 
+  default Column<T> set(int row, String stringValue, AbstractColumnParser<?> parser) {
+    AbstractColumnParser<T> typedParser = (AbstractColumnParser<T>) parser;
+    return set(row, typedParser.parse(stringValue));
+  }
+
   Column<T> set(int row, Column<T> sourceColumn, int sourceRow);
 
   Column<T> append(T value);
diff --git a/core/src/main/java/tech/tablesaw/io/FileReader.java b/core/src/main/java/tech/tablesaw/io/FileReader.java
index 1a9cee94..a8be58a7 100644
--- a/core/src/main/java/tech/tablesaw/io/FileReader.java
+++ b/core/src/main/java/tech/tablesaw/io/FileReader.java
@@ -12,6 +12,7 @@ import java.util.Iterator;
 import java.util.List;
 import java.util.Map;
 import java.util.NoSuchElementException;
+import java.util.Random;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 import tech.tablesaw.api.ColumnType;
@@ -22,8 +23,7 @@ import tech.tablesaw.columns.Column;
 public abstract class FileReader {
 
   private static Logger logger = LoggerFactory.getLogger(FileReader.class);
-
-  public FileReader() {}
+  private static final int UNLIMITED_SAMPLE_SIZE = -1;
 
   /**
    * Returns an array containing the inferred columnTypes for the file being read, as calculated by
@@ -97,6 +97,16 @@ public abstract class FileReader {
       Reader reader,
       ColumnType[] types,
       AbstractParser<?> parser) {
+    return parseRows(options, headerOnly, reader, types, parser, UNLIMITED_SAMPLE_SIZE);
+  }
+
+  protected Table parseRows(
+      ReadOptions options,
+      boolean headerOnly,
+      Reader reader,
+      ColumnType[] types,
+      AbstractParser<?> parser,
+      int sampleSize) {
     parser.beginParsing(reader);
     Table table = Table.create(options.tableName());
 
@@ -120,7 +130,7 @@ public abstract class FileReader {
         // get the index in the original table, which includes skipped fields
         columnIndexes[i] = headerRow.indexOf(columnNames[i]);
       }
-      addRows(options, types, parser, table, columnIndexes);
+      addRows(options, types, parser, table, columnIndexes, sampleSize);
     }
 
     return table;
@@ -131,13 +141,15 @@ public abstract class FileReader {
       ColumnType[] types,
       AbstractParser<?> reader,
       Table table,
-      int[] columnIndexes) {
-    String[] nextLine;
+      int[] columnIndexes,
+      int sampleSize) {
 
+    String[] nextLine;
     Map<String, AbstractColumnParser<?>> parserMap = getParserMap(options, table);
 
+    Random random = new Random(0);
     // Add the rows
-    for (long rowNumber = options.header() ? 1L : 0L;
+    for (int rowNumber = options.header() ? 1 : 0;
         (nextLine = reader.parseNext()) != null;
         rowNumber++) {
       // validation
@@ -168,20 +180,44 @@ public abstract class FileReader {
                 + " expected.");
       }
 
-      // append each column that we're including (not skipping)
-      int cellIndex = 0;
-      for (int columnIndex : columnIndexes) {
-        Column<?> column = table.column(cellIndex);
-        AbstractColumnParser<?> parser = parserMap.get(column.name());
-        try {
-          String value = nextLine[columnIndex];
+      int samplesCount = table.rowCount();
+      if (sampleSize < 0 || samplesCount < sampleSize) {
+        addValuesToColumns(table, columnIndexes, nextLine, parserMap, rowNumber, -1);
+      } else {
+        // find a row index to replace
+        int randomIndex = random.nextInt(samplesCount + 1);
+        // replace index if it is smaller than numSamples, otherwise ignore it.
+        if (randomIndex < sampleSize) {
+          addValuesToColumns(table, columnIndexes, nextLine, parserMap, rowNumber, randomIndex);
+        }
+      }
+    }
+  }
+
+  private void addValuesToColumns(
+      Table table,
+      int[] columnIndexes,
+      String[] nextLine,
+      Map<String, AbstractColumnParser<?>> parserMap,
+      int rowNumber,
+      int rowIndex) {
+    // append each column that we're including (not skipping)
+    int cellIndex = 0;
+    for (int columnIndex : columnIndexes) {
+      Column<?> column = table.column(cellIndex);
+      AbstractColumnParser<?> parser = parserMap.get(column.name());
+      try {
+        String value = nextLine[columnIndex];
+        if (rowIndex >= 0) {
+          column.set(rowIndex, value, parser);
+        } else {
           column.appendCell(value, parser);
-        } catch (Exception e) {
-          throw new AddCellToColumnException(
-              e, columnIndex, rowNumber, table.columnNames(), nextLine);
         }
-        cellIndex++;
+      } catch (Exception e) {
+        throw new AddCellToColumnException(
+            e, columnIndex, rowNumber, table.columnNames(), nextLine);
       }
+      cellIndex++;
     }
   }
 
diff --git a/core/src/main/java/tech/tablesaw/io/csv/CsvReadOptions.java b/core/src/main/java/tech/tablesaw/io/csv/CsvReadOptions.java
index 0821653b..df50f057 100644
--- a/core/src/main/java/tech/tablesaw/io/csv/CsvReadOptions.java
+++ b/core/src/main/java/tech/tablesaw/io/csv/CsvReadOptions.java
@@ -35,10 +35,10 @@ public class CsvReadOptions extends ReadOptions {
   private final Integer maxNumberOfColumns;
   private final Character commentPrefix;
   private final boolean lineSeparatorDetectionEnabled;
+  private final int sampleSize;
 
   private CsvReadOptions(CsvReadOptions.Builder builder) {
     super(builder);
-
     columnTypes = builder.columnTypes;
     separator = builder.separator;
     quoteChar = builder.quoteChar;
@@ -46,6 +46,7 @@ public class CsvReadOptions extends ReadOptions {
     maxNumberOfColumns = builder.maxNumberOfColumns;
     commentPrefix = builder.commentPrefix;
     lineSeparatorDetectionEnabled = builder.lineSeparatorDetectionEnabled;
+    sampleSize = builder.sampleSize;
   }
 
   public static Builder builder(Source source) {
@@ -130,6 +131,10 @@ public class CsvReadOptions extends ReadOptions {
     return maxCharsPerColumn;
   }
 
+  public int sampleSize() {
+    return sampleSize;
+  }
+
   public static class Builder extends ReadOptions.Builder {
 
     private Character separator = ',';
@@ -139,6 +144,7 @@ public class CsvReadOptions extends ReadOptions {
     private Integer maxNumberOfColumns = 10_000;
     private Character commentPrefix;
     private boolean lineSeparatorDetectionEnabled = true;
+    private int sampleSize = -1;
 
     protected Builder(Source source) {
       super(source);
@@ -196,6 +202,19 @@ public class CsvReadOptions extends ReadOptions {
       return this;
     }
 
+    /**
+     * Defines the maximum number of rows to be read from the file. Sampling is performed in a
+     * single pass using the reservoir sampling algorithm
+     * (https://en.wikipedia.org/wiki/Reservoir_sampling). Given a file with 'n' rows, if
+     * 'numSamples is smaller than 'n', than exactly 'numSamples' random samples are returned; if
+     * 'numSamples' is greater than 'n', then only 'n' samples are returned (no oversampling is
+     * performed to increase the data to match 'numSamples').
+     */
+    public Builder sampleSize(int numSamples) {
+      this.sampleSize = numSamples;
+      return this;
+    }
+
     @Override
     public CsvReadOptions build() {
       return new CsvReadOptions(this);
diff --git a/core/src/main/java/tech/tablesaw/io/csv/CsvReader.java b/core/src/main/java/tech/tablesaw/io/csv/CsvReader.java
index 29209618..cd148fab 100644
--- a/core/src/main/java/tech/tablesaw/io/csv/CsvReader.java
+++ b/core/src/main/java/tech/tablesaw/io/csv/CsvReader.java
@@ -84,7 +84,7 @@ public class CsvReader extends FileReader implements DataReader<CsvReadOptions>
     AbstractParser<?> parser = csvParser(options);
 
     try {
-      return parseRows(options, headerOnly, reader, types, parser);
+      return parseRows(options, headerOnly, reader, types, parser, options.sampleSize());
     } finally {
       if (options.source().reader() == null) {
         // if we get a reader back from options it means the client opened it, so let the client


Test patch:
diff --git a/core/src/test/java/tech/tablesaw/io/csv/CsvReaderTest.java b/core/src/test/java/tech/tablesaw/io/csv/CsvReaderTest.java
index c90aa811..8e682fb9 100644
--- a/core/src/test/java/tech/tablesaw/io/csv/CsvReaderTest.java
+++ b/core/src/test/java/tech/tablesaw/io/csv/CsvReaderTest.java
@@ -45,6 +45,8 @@ import java.time.format.DateTimeFormatterBuilder;
 import java.util.Collections;
 import java.util.List;
 import java.util.Locale;
+import java.util.Set;
+import java.util.TreeSet;
 import org.junit.jupiter.api.Test;
 import tech.tablesaw.api.ColumnType;
 import tech.tablesaw.api.DateColumn;
@@ -702,4 +704,38 @@ public class CsvReaderTest {
         Table.read().csv(CsvReadOptions.builder("../data/sacramento_real_estate_transactions.csv"));
     assertEquals(985, table.rowCount());
   }
+
+  @Test
+  public void testReadCsvWithRowSampling() throws IOException {
+    int expectedSampleSize = 10;
+    Table table =
+        Table.read().csv(CsvReadOptions.builder("../data/cake.csv").sampleSize(expectedSampleSize));
+
+    assertEquals(expectedSampleSize, table.rowCount());
+    assertEquals("[recipe, mix, temp, y]", table.columnNames().toString());
+  }
+
+  @Test
+  public void testReadCsvWithRowSamplingParsingNumbers() throws IOException {
+
+    StringBuilder csv = new StringBuilder();
+    csv.append("RandomNumbers\n");
+    Set<Integer> values =
+        new TreeSet<>(
+            asList(
+                24323, 542345, 64323, 73640, 38453, 12735, 93456, 23457, 483075, 469364, 473936));
+    values.forEach(v -> csv.append(v + "\n"));
+
+    Reader reader = new StringReader(csv.toString());
+
+    int expectedSampleSize = 5;
+    Table t = Table.read().csv(CsvReadOptions.builder(reader).sampleSize(expectedSampleSize));
+
+    assertEquals(1, t.columnCount());
+    assertEquals(expectedSampleSize, t.rowCount());
+    assertEquals(INTEGER, t.column(0).type());
+
+    List<Integer> intValues = (List<Integer>) t.column(0).asList();
+    assertEquals(true, values.containsAll(intValues));
+  }
 }