Skip to content

Commit

Permalink
ARROW-474: [Java] Add initial version of streaming serialized format.
Browse files Browse the repository at this point in the history 
This patch proposes a serialized container format for streaming producer and
consumers. The goal is to allow readers and writers to produce/consume arrow
data without requiring intermediate buffering.

This is similar to the File format but reorganizes the pieces. In particular:
  - No magic header. It's likely a reader connects to a 'random' stream to read it.
  - Move footer to header. This includes similar information, including the schema.
  - ArrowRecordBatches follow one by one. Each is prefixed with an i32 length. The
    serialization is identical as the File version.
  - See Stream.fbs for more details.

This patch also implements the Java reader/writer.

Author: Nong Li <[email protected]>

Closes #288 from nongli/streaming and squashes the following commits:

554cc18 [Nong Li] Redo serialization format.
03bee58 [Nong Li] Updates from wes' comments.
7257031 [Nong Li] ARROW-474: [Java] Add initial version of streaming serialized format.
  • Loading branch information
@cerebrodata @wesm
cerebrodata authored and wesm committed Jan 19, 2017
1 parent 9b1b397 commit 6811d3f
Show file tree
Hide file tree
Showing 13 changed files with 1,100 additions and 99 deletions.
12 changes: 3 additions & 9 deletions java/vector/src/main/java/org/apache/arrow/vector/file/ArrowReader.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -31,6 +31,7 @@
import org.apache.arrow.memory.BufferAllocator;
import org.apache.arrow.vector.schema.ArrowFieldNode;
import org.apache.arrow.vector.schema.ArrowRecordBatch;
import org.apache.arrow.vector.stream.MessageSerializer;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

Expand All @@ -39,7 +40,7 @@
public class ArrowReader implements AutoCloseable {
private static final Logger LOGGER = LoggerFactory.getLogger(ArrowReader.class);

private static final byte[] MAGIC = "ARROW1".getBytes();
public static final byte[] MAGIC = "ARROW1".getBytes();

private final SeekableByteChannel in;

Expand Down Expand Up @@ -73,13 +74,6 @@ private int readFully(ByteBuffer buffer) throws IOException {
return total;
}

private static int bytesToInt(byte[] bytes) {
return ((int)(bytes[3] & 255) << 24) +
((int)(bytes[2] & 255) << 16) +
((int)(bytes[1] & 255) << 8) +
((int)(bytes[0] & 255) << 0);
}

public ArrowFooter readFooter() throws IOException {
if (footer == null) {
if (in.size() <= (MAGIC.length * 2 + 4)) {
Expand All @@ -93,7 +87,7 @@ public ArrowFooter readFooter() throws IOException {
if (!Arrays.equals(MAGIC, Arrays.copyOfRange(array, 4, array.length))) {
throw new InvalidArrowFileException("missing Magic number " + Arrays.toString(buffer.array()));
}
int footerLength = bytesToInt(array);
int footerLength = MessageSerializer.bytesToInt(array);
if (footerLength <= 0 || footerLength + MAGIC.length * 2 + 4 > in.size()) {
throw new InvalidArrowFileException("invalid footer length: " + footerLength);
}
Expand Down
102 changes: 20 additions & 82 deletions java/vector/src/main/java/org/apache/arrow/vector/file/ArrowWriter.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -18,94 +18,52 @@
package org.apache.arrow.vector.file;

import java.io.IOException;
import java.nio.ByteBuffer;
import java.nio.channels.WritableByteChannel;
import java.util.ArrayList;
import java.util.Collections;
import java.util.List;

import org.apache.arrow.vector.schema.ArrowBuffer;
import org.apache.arrow.vector.schema.ArrowRecordBatch;
import org.apache.arrow.vector.schema.FBSerializable;
import org.apache.arrow.vector.types.pojo.Schema;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import com.google.flatbuffers.FlatBufferBuilder;

import io.netty.buffer.ArrowBuf;

public class ArrowWriter implements AutoCloseable {
private static final Logger LOGGER = LoggerFactory.getLogger(ArrowWriter.class);

private static final byte[] MAGIC = "ARROW1".getBytes();

private final WritableByteChannel out;
private final WriteChannel out;

private final Schema schema;

private final List<ArrowBlock> recordBatches = new ArrayList<>();

private long currentPosition = 0;

private boolean started = false;

public ArrowWriter(WritableByteChannel out, Schema schema) {
this.out = out;
this.out = new WriteChannel(out);
this.schema = schema;
}

private void start() throws IOException {
writeMagic();
}

private long write(byte[] buffer) throws IOException {
return write(ByteBuffer.wrap(buffer));
}

private long writeZeros(int zeroCount) throws IOException {
return write(new byte[zeroCount]);
}

private long align() throws IOException {
if (currentPosition % 8 != 0) { // align on 8 byte boundaries
return writeZeros(8 - (int)(currentPosition % 8));
}
return 0;
}

private long write(ByteBuffer buffer) throws IOException {
long length = buffer.remaining();
out.write(buffer);
currentPosition += length;
return length;
}

private static byte[] intToBytes(int value) {
byte[] outBuffer = new byte[4];
outBuffer[3] = (byte)(value >>> 24);
outBuffer[2] = (byte)(value >>> 16);
outBuffer[1] = (byte)(value >>> 8);
outBuffer[0] = (byte)(value >>> 0);
return outBuffer;
}

private long writeIntLittleEndian(int v) throws IOException {
return write(intToBytes(v));
}

// TODO: write dictionaries

public void writeRecordBatch(ArrowRecordBatch recordBatch) throws IOException {
checkStarted();
align();
out.align();

// write metadata header with int32 size prefix
long offset = currentPosition;
write(recordBatch, true);
align();
long offset = out.getCurrentPosition();
out.write(recordBatch, true);
out.align();
// write body
long bodyOffset = currentPosition;
long bodyOffset = out.getCurrentPosition();
List<ArrowBuf> buffers = recordBatch.getBuffers();
List<ArrowBuffer> buffersLayout = recordBatch.getBuffersLayout();
if (buffers.size() != buffersLayout.size()) {
Expand All @@ -115,47 +73,42 @@ public void writeRecordBatch(ArrowRecordBatch recordBatch) throws IOException {
ArrowBuf buffer = buffers.get(i);
ArrowBuffer layout = buffersLayout.get(i);
long startPosition = bodyOffset + layout.getOffset();
if (startPosition != currentPosition) {
writeZeros((int)(startPosition - currentPosition));
if (startPosition != out.getCurrentPosition()) {
out.writeZeros((int)(startPosition - out.getCurrentPosition()));
}

write(buffer);
if (currentPosition != startPosition + layout.getSize()) {
throw new IllegalStateException("wrong buffer size: " + currentPosition + " != " + startPosition + layout.getSize());
out.write(buffer);
if (out.getCurrentPosition() != startPosition + layout.getSize()) {
throw new IllegalStateException("wrong buffer size: " + out.getCurrentPosition() + " != " + startPosition + layout.getSize());
}
}
int metadataLength = (int)(bodyOffset - offset);
if (metadataLength <= 0) {
throw new InvalidArrowFileException("invalid recordBatch");
}
long bodyLength = currentPosition - bodyOffset;
long bodyLength = out.getCurrentPosition() - bodyOffset;
LOGGER.debug(String.format("RecordBatch at %d, metadata: %d, body: %d", offset, metadataLength, bodyLength));
// add metadata to footer
recordBatches.add(new ArrowBlock(offset, metadataLength, bodyLength));
}

private void write(ArrowBuf buffer) throws IOException {
ByteBuffer nioBuffer = buffer.nioBuffer(buffer.readerIndex(), buffer.readableBytes());
LOGGER.debug("Writing buffer with size: " + nioBuffer.remaining());
write(nioBuffer);
}

private void checkStarted() throws IOException {
if (!started) {
started = true;
start();
}
}

@Override
public void close() throws IOException {
try {
long footerStart = currentPosition;
long footerStart = out.getCurrentPosition();
writeFooter();
int footerLength = (int)(currentPosition - footerStart);
int footerLength = (int)(out.getCurrentPosition() - footerStart);
if (footerLength <= 0 ) {
throw new InvalidArrowFileException("invalid footer");
}
writeIntLittleEndian(footerLength);
out.writeIntLittleEndian(footerLength);
LOGGER.debug(String.format("Footer starts at %d, length: %d", footerStart, footerLength));
writeMagic();
} finally {
Expand All @@ -164,27 +117,12 @@ public void close() throws IOException {
}

private void writeMagic() throws IOException {
write(MAGIC);
LOGGER.debug(String.format("magic written, now at %d", currentPosition));
out.write(ArrowReader.MAGIC);
LOGGER.debug(String.format("magic written, now at %d", out.getCurrentPosition()));
}

private void writeFooter() throws IOException {
// TODO: dictionaries
write(new ArrowFooter(schema, Collections.<ArrowBlock>emptyList(), recordBatches), false);
}

private long write(FBSerializable writer, boolean withSizePrefix) throws IOException {
FlatBufferBuilder builder = new FlatBufferBuilder();
int root = writer.writeTo(builder);
builder.finish(root);

ByteBuffer buffer = builder.dataBuffer();

if (withSizePrefix) {
writeIntLittleEndian(buffer.remaining());
}

return write(buffer);
out.write(new ArrowFooter(schema, Collections.<ArrowBlock>emptyList(), recordBatches), false);
}

}
75 changes: 75 additions & 0 deletions java/vector/src/main/java/org/apache/arrow/vector/file/ReadChannel.java
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,75 @@
/**
* Licensed to the Apache Software Foundation (ASF) under one
* or more contributor license agreements. See the NOTICE file
* distributed with this work for additional information
* regarding copyright ownership. The ASF licenses this file
* to you under the Apache License, Version 2.0 (the
* "License"); you may not use this file except in compliance
* with the License. You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.apache.arrow.vector.file;

import java.io.IOException;
import java.nio.ByteBuffer;
import java.nio.channels.ReadableByteChannel;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import io.netty.buffer.ArrowBuf;

public class ReadChannel implements AutoCloseable {

private static final Logger LOGGER = LoggerFactory.getLogger(ReadChannel.class);

private ReadableByteChannel in;
private long bytesRead = 0;

public ReadChannel(ReadableByteChannel in) {
this.in = in;
}

public long bytesRead() { return bytesRead; }

/**
* Reads bytes into buffer until it is full (buffer.remaining() == 0). Returns the
* number of bytes read which can be less than full if there are no more.
*/
public int readFully(ByteBuffer buffer) throws IOException {
LOGGER.debug("Reading buffer with size: " + buffer.remaining());
int totalRead = 0;
while (buffer.remaining() != 0) {
int read = in.read(buffer);
if (read < 0) return totalRead;
totalRead += read;
if (read == 0) break;
}
this.bytesRead += totalRead;
return totalRead;
}

/**
* Reads up to len into buffer. Returns bytes read.
*/
public int readFully(ArrowBuf buffer, int l) throws IOException {
int n = readFully(buffer.nioBuffer(buffer.writerIndex(), l));
buffer.writerIndex(n);
return n;
}

@Override
public void close() throws IOException {
if (this.in != null) {
in.close();
in = null;
}
}
}
Loading

0 comments on commit 6811d3f

Please sign in to comment.