CSVIngester extends DataIngester { @Override public MBStream<Datum> getStream() throws Exception { if(!loaded) { long st = System.currentTimeMillis(); filename = conf.getString(MacroBaseConf.CSV_INPUT_FILE); Compression compression = conf.getCsvCompression(); if (compression == Compression.GZIP) { InputStream fileStream = new FileInputStream(filename); InputStream gzipStream = new GZIPInputStream(fileStream); Reader decoder = new InputStreamReader(gzipStream); csvParser = new CSVParser(decoder, CSVFormat.DEFAULT.withHeader()); } else { File csvFile = new File(conf.getString(MacroBaseConf.CSV_INPUT_FILE)); csvParser = CSVParser.parse(csvFile, Charset.defaultCharset(), CSVFormat.DEFAULT.withHeader()); } schema = csvParser.getHeaderMap(); for (Map.Entry<String, Integer> se : schema.entrySet()) { conf.getEncoder().recordAttributeName(se.getValue() + 1, se.getKey()); } Iterator<CSVRecord> rawIterator = csvParser.iterator(); int numRows = 0; while (rawIterator.hasNext()) { try { CSVRecord record = rawIterator.next(); Datum curRow = parseRecord(record); dataStream.add(curRow); numRows++; } catch (NumberFormatException e) { badRows++; } } log.info("{}/{} rows successfully parsed ({} malformed rows)", numRows, numRows + badRows, badRows); } return dataStream; } CSVIngester(MacroBaseConf conf); @Override MBStream<Datum> getStream(); @Override RowSet getRows(String baseQuery,
                          Map<String, String> preds,
                          int limit,
                          int offset);  }
@Test public void testSimple() throws Exception { MacroBaseConf conf = new MacroBaseConf(); conf.set(MacroBaseConf.CSV_INPUT_FILE, "src/test/resources/data/simple.csv"); conf.set(MacroBaseConf.ATTRIBUTES, Lists.newArrayList("A2","A5")); conf.set(MacroBaseConf.METRICS, Lists.newArrayList("A1","A3","A4")); CSVIngester ingester = new CSVIngester(conf); List<Datum> data = ingester.getStream().drain(); Datum datum = data.get(0); assertEquals(43, datum.metrics().getEntry(0), 0.0); assertEquals(45, datum.metrics().getEntry(1), 0.0); }
@Test public void testMissingData() throws Exception { MacroBaseConf conf = new MacroBaseConf(); conf.set(MacroBaseConf.CSV_INPUT_FILE, "src/test/resources/data/missingdata.csv"); conf.set(MacroBaseConf.ATTRIBUTES, Lists.newArrayList("a1")); conf.set(MacroBaseConf.METRICS, Lists.newArrayList("m1","m2")); CSVIngester ingester = new CSVIngester(conf); List<Datum> data = ingester.getStream().drain(); int count = 0; for(Datum d : data) { if (count == 3) { assertEquals(5, d.metrics().getEntry(0), 0.0); } count++; } assertEquals(4, count); }