Examples with ArrowStreamWriter org.apache.arrow.vector.ipc.ArrowStreamWriter used on opensource projects

Example 6 with ArrowStreamWriter

use of org.apache.arrow.vector.ipc.ArrowStreamWriter in project flink by apache.

the class ArrowUtilsTest method testReadArrowBatches.

@Test
public void testReadArrowBatches() throws IOException {
    VectorSchemaRoot root = VectorSchemaRoot.create(ArrowUtils.toArrowSchema(rowType), allocator);
    ArrowWriter<RowData> arrowWriter = ArrowUtils.createRowDataArrowWriter(root, rowType);
    ByteArrayOutputStream baos = new ByteArrayOutputStream();
    ArrowStreamWriter arrowStreamWriter = new ArrowStreamWriter(root, null, baos);
    arrowStreamWriter.start();
    List<RowData> testData = Arrays.asList(new GenericRowData(rowType.getFieldCount()), new GenericRowData(rowType.getFieldCount()), new GenericRowData(rowType.getFieldCount()), new GenericRowData(rowType.getFieldCount()), new GenericRowData(rowType.getFieldCount()));
    int batches = 3;
    List<List<RowData>> subLists = Lists.partition(testData, testData.size() / batches + 1);
    for (List<RowData> subList : subLists) {
        for (RowData value : subList) {
            arrowWriter.write(value);
        }
        arrowWriter.finish();
        arrowStreamWriter.writeBatch();
        arrowWriter.reset();
    }
    assertEquals(batches, ArrowUtils.readArrowBatches(Channels.newChannel(new ByteArrayInputStream(baos.toByteArray()))).length);
}

Example 7 with ArrowStreamWriter

use of org.apache.arrow.vector.ipc.ArrowStreamWriter in project flink by apache.

the class ArrowSourceFunctionTestBase method createTestArrowSourceFunction.

/**
 * Create continuous monitoring function with 1 reader-parallelism and interval.
 */
private ArrowSourceFunction createTestArrowSourceFunction(List<RowData> testData, int batches) throws IOException {
    ArrowWriter<RowData> arrowWriter = createArrowWriter();
    ByteArrayOutputStream baos = new ByteArrayOutputStream();
    ArrowStreamWriter arrowStreamWriter = new ArrowStreamWriter(root, null, baos);
    arrowStreamWriter.start();
    List<List<RowData>> subLists = Lists.partition(testData, testData.size() / batches + 1);
    for (List<RowData> subList : subLists) {
        for (RowData value : subList) {
            arrowWriter.write(value);
        }
        arrowWriter.finish();
        arrowStreamWriter.writeBatch();
        arrowWriter.reset();
    }
    ArrowSourceFunction arrowSourceFunction = createArrowSourceFunction(ArrowUtils.readArrowBatches(Channels.newChannel(new ByteArrayInputStream(baos.toByteArray()))));
    arrowSourceFunction.setRuntimeContext(Mockito.mock(RuntimeContext.class));
    return arrowSourceFunction;
}

Example 8 with ArrowStreamWriter

use of org.apache.arrow.vector.ipc.ArrowStreamWriter in project flink by apache.

the class ArrowReaderWriterTestBase method testBasicFunctionality.

@Test
public void testBasicFunctionality() {
    try {
        ByteArrayOutputStream baos = new ByteArrayOutputStream();
        Tuple2<ArrowWriter<T>, ArrowStreamWriter> tuple2 = createArrowWriter(baos);
        ArrowWriter<T> arrowWriter = tuple2.f0;
        ArrowStreamWriter arrowStreamWriter = tuple2.f1;
        T[] testData = getTestData();
        for (T value : testData) {
            arrowWriter.write(value);
        }
        arrowWriter.finish();
        arrowStreamWriter.writeBatch();
        ArrowReader arrowReader = createArrowReader(new ByteArrayInputStream(baos.toByteArray()));
        for (int i = 0; i < testData.length; i++) {
            RowData deserialized = arrowReader.read(i);
            assertThat("Deserialized value is wrong.", deserialized, CustomEqualityMatcher.deeplyEquals(testData[i]).withChecker(checker));
        }
    } catch (Exception e) {
        System.err.println(e.getMessage());
        e.printStackTrace();
        fail("Exception in test: " + e.getMessage());
    }
}