Examples with TFRecordCodec org.apache.beam.sdk.io.TFRecordIO.TFRecordCodec used on opensource projects

Example 1 with TFRecordCodec

use of org.apache.beam.sdk.io.TFRecordIO.TFRecordCodec in project beam by apache.

the class TFRecordIOTest method testTFRecordCodec.

@Test
public void testTFRecordCodec() throws IOException {
    Decoder b64 = Base64.getDecoder();
    TFRecordCodec codec = new TFRecordCodec();
    ByteArrayOutputStream baos = new ByteArrayOutputStream();
    PickyWriteChannel outChan = new PickyWriteChannel(baos);
    codec.write(outChan, "foo".getBytes(StandardCharsets.UTF_8));
    assertArrayEquals(b64.decode(FOO_RECORD_BASE64), baos.toByteArray());
    codec.write(outChan, "bar".getBytes(StandardCharsets.UTF_8));
    assertArrayEquals(b64.decode(FOO_BAR_RECORD_BASE64), baos.toByteArray());
    PickyReadChannel inChan = new PickyReadChannel(new ByteArrayInputStream(baos.toByteArray()));
    byte[] foo = codec.read(inChan);
    byte[] bar = codec.read(inChan);
    assertNull(codec.read(inChan));
    assertEquals("foo", new String(foo, StandardCharsets.UTF_8));
    assertEquals("bar", new String(bar, StandardCharsets.UTF_8));
}