Search in sources :

Example 1 with ArrowReader

use of org.apache.flink.table.runtime.arrow.ArrowReader in project flink by apache.

the class ArrowSourceFunction method run.

@Override
public void run(SourceContext<RowData> ctx) throws Exception {
    VectorLoader vectorLoader = new VectorLoader(root);
    while (running && !indexesToEmit.isEmpty()) {
        Tuple2<Integer, Integer> indexToEmit = indexesToEmit.peek();
        ArrowRecordBatch arrowRecordBatch = loadBatch(indexToEmit.f0);
        vectorLoader.load(arrowRecordBatch);
        arrowRecordBatch.close();
        ArrowReader arrowReader = createArrowReader(root);
        int rowCount = root.getRowCount();
        int nextRowId = indexToEmit.f1;
        while (nextRowId < rowCount) {
            RowData element = arrowReader.read(nextRowId);
            synchronized (ctx.getCheckpointLock()) {
                ctx.collect(element);
                indexToEmit.setField(++nextRowId, 1);
            }
        }
        synchronized (ctx.getCheckpointLock()) {
            indexesToEmit.pop();
        }
    }
}
Also used : VectorLoader(org.apache.arrow.vector.VectorLoader) RowData(org.apache.flink.table.data.RowData) ArrowRecordBatch(org.apache.arrow.vector.ipc.message.ArrowRecordBatch) ArrowReader(org.apache.flink.table.runtime.arrow.ArrowReader)

Aggregations

VectorLoader (org.apache.arrow.vector.VectorLoader)1 ArrowRecordBatch (org.apache.arrow.vector.ipc.message.ArrowRecordBatch)1 RowData (org.apache.flink.table.data.RowData)1 ArrowReader (org.apache.flink.table.runtime.arrow.ArrowReader)1