| 41 | import static org.junit.jupiter.api.Assertions.assertTrue; |
| 42 | |
| 43 | public class TestRowCount implements TestConf { |
| 44 | private Path workDir = new Path( |
| 45 | Paths.get(System.getProperty("test.tmp.dir"), "orc-test-count").toString()); |
| 46 | private FileSystem fs; |
| 47 | |
| 48 | @BeforeEach |
| 49 | public void openFileSystem() throws Exception { |
| 50 | fs = FileSystem.getLocal(conf); |
| 51 | fs.mkdirs(workDir); |
| 52 | fs.deleteOnExit(workDir); |
| 53 | } |
| 54 | |
| 55 | @Test |
| 56 | public void testCount() throws Exception { |
| 57 | TypeDescription schema = TypeDescription.fromString("struct<x:int>"); |
| 58 | Map<String, Integer> fileToRowCountMap = new LinkedHashMap<>(); |
| 59 | fileToRowCountMap.put(workDir + File.separator + "test-count-1.orc", 10000); |
| 60 | fileToRowCountMap.put(workDir + File.separator + "test-count-2.orc", 20000); |
| 61 | for (Map.Entry<String, Integer> fileToRowCount : fileToRowCountMap.entrySet()) { |
| 62 | Writer writer = OrcFile.createWriter(new Path(fileToRowCount.getKey()), |
| 63 | OrcFile.writerOptions(conf) |
| 64 | .setSchema(schema)); |
| 65 | VectorizedRowBatch batch = schema.createRowBatch(); |
| 66 | LongColumnVector x = (LongColumnVector) batch.cols[0]; |
| 67 | for (int r = 0; r < fileToRowCount.getValue(); ++r) { |
| 68 | int row = batch.size++; |
| 69 | x.vector[row] = r; |
| 70 | if (batch.size == batch.getMaxSize()) { |
| 71 | writer.addRowBatch(batch); |
| 72 | batch.reset(); |
| 73 | } |
| 74 | } |
| 75 | if (batch.size != 0) { |
| 76 | writer.addRowBatch(batch); |
| 77 | } |
| 78 | writer.close(); |
| 79 | } |
| 80 | |
| 81 | PrintStream origOut = System.out; |
| 82 | ByteArrayOutputStream myOut = new ByteArrayOutputStream(); |
| 83 | // replace stdout and run command |
| 84 | System.setOut(new PrintStream(myOut, false, StandardCharsets.UTF_8)); |
| 85 | RowCount.main(conf, new String[]{workDir.toString()}); |
| 86 | System.out.flush(); |
| 87 | System.setOut(origOut); |
| 88 | String output = myOut.toString(StandardCharsets.UTF_8); |
| 89 | assertTrue(output.contains(" 10000")); |
| 90 | assertTrue(output.contains(" 20000")); |
| 91 | } |
| 92 | } |