Examples with RecordAndPosition org.apache.flink.connector.file.src.util.RecordAndPosition used on opensource projects

Example 1 with RecordAndPosition

use of org.apache.flink.connector.file.src.util.RecordAndPosition in project flink by apache.

the class OrcColumnarRowInputFormatTest method innerTestRestore.

private void innerTestRestore(OrcColumnarRowInputFormat<?, FileSourceSplit> format, FileSourceSplit split, int breakCnt, int expectedCnt, long expectedTotalF0) throws IOException {
    AtomicInteger cnt = new AtomicInteger(0);
    AtomicLong totalF0 = new AtomicLong(0);
    Consumer<RowData> consumer = row -> {
        Assert.assertFalse(row.isNullAt(0));
        Assert.assertFalse(row.isNullAt(1));
        totalF0.addAndGet(row.getInt(0));
        assertNotNull(row.getString(1).toString());
        cnt.incrementAndGet();
    };
    // ---------- restore reading ---------------
    long offset = -1;
    long recordSkipCount = -1;
    try (BulkFormat.Reader<RowData> reader = createReader(format, split)) {
        while (cnt.get() < breakCnt) {
            BulkFormat.RecordIterator<RowData> batch = reader.readBatch();
            Assert.assertNotNull(batch);
            RecordAndPosition<RowData> record;
            while ((record = batch.next()) != null && cnt.get() < breakCnt) {
                consumer.accept(record.getRecord());
                offset = record.getOffset();
                recordSkipCount = record.getRecordSkipCount();
            }
            batch.releaseBatch();
        }
    }
    Utils.forEachRemaining(restoreReader(format, split, offset, recordSkipCount), consumer);
    // ---------- end restore reading ---------------
    // the results should be the same as:
    // forEach(format, split, consumer);
    // check that all rows have been read
    assertEquals(expectedCnt, cnt.get());
    assertEquals(expectedTotalF0, totalF0.get());
}