/**
* 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.avro.io;
import java.io.IOException;
import java.io.OutputStream;
import java.util.Arrays;
import org.apache.avro.AvroTypeException;
import org.apache.avro.Schema;
/** A {@link BinaryEncoder} implementation that writes large arrays and maps as a
* sequence of blocks. So long as individual primitive values fit in memory,
* arbitrarily long arrays and maps may be written and subsequently read without
* exhausting memory. Values are buffered until the specified block size would
* be exceeded, minimizing block overhead.
* <p/>
* Use {@link EncoderFactory#blockingBinaryEncoder(OutputStream, BinaryEncoder)}
* to construct and configure.
* <p/>
* BlockingBinaryEncoder buffers writes, data may not appear on the output until
* {@link #flush()} is called.
* <p/>
* BlockingBinaryEncoder is not thread-safe
*
* @see BinaryEncoder
* @see EncoderFactory
* @see Encoder
*/
public class BlockingBinaryEncoder extends BufferedBinaryEncoder {
/* Implementation note:
*
* Blocking is complicated because of nesting. If a large, nested
* value overflows your buffer, you've got to do a lot of dancing
* around to output the blocks correctly.
*
* To handle this complexity, this class keeps a stack of blocked
* values: each time a new block is started (e.g., by a call to
* {@link #writeArrayStart}), an entry is pushed onto this stack.
*
* In this stack, we keep track of the state of a block. Blocks can
* be in two states. "Regular" blocks have a non-zero byte count.
* "Overflow" blocks help us deal with the case where a block
* contains a value that's too big to buffer. In this case, the
* block contains only one item, and we give it an unknown
* byte-count. Because these values (1,unknown) are fixed, we're
* able to write the header for these overflow blocks to the
* underlying stream without seeing the entire block. After writing
* this header, we've freed our buffer space to be fully devoted to
* blocking the large, inner value.
*/
private static class BlockedValue {
public enum State {
/**
* The bottom element of our stack represents being _outside_
* of a blocked value.
*/
ROOT,
/**
* Represents the "regular" case, i.e., a blocked-value whose
* current block is fully contained in the buffer. In this
* case, {@link BlockedValue#start} points to the start of the
* blocks _data_ -- but no room has been left for a header!
* When this block is terminated, it's data will have to be
* moved over a bit to make room for the header. */
REGULAR,
/**
* Represents a blocked-value whose current block is in the
* overflow state. In this case, {@link BlockedValue#start} is zero. The
* header for such a block has _already been written_ (we've
* written out a header indicating that the block has a single
* item, and we put a "zero" down for the byte-count to indicate
* that we don't know the physical length of the buffer. Any blocks
* _containing_ this block must be in the {@link #OVERFLOW}
* state. */
OVERFLOW
}
/** The type of this blocked value (ARRAY or MAP). */
public Schema.Type type;
/** The state of this BlockedValue */
public State state;
/** The location in the buffer where this blocked value starts */
public int start;
/**
* The index one past the last byte for the previous item. If this
* is the first item, this is same as {@link #start}.
*/
public int lastFullItem;
/**
* Number of items in this blocked value that are stored
* in the buffer.
*/
public int items;
/** Number of items left to write*/
public long itemsLeftToWrite;
/** Create a ROOT instance. */
public BlockedValue() {
this.type = null;
this.state = BlockedValue.State.ROOT;
this.start = this.lastFullItem = 0;
this.items = 1; // Makes various assertions work out
}
/**
* Check invariants of <code>this</code> and also the
* <code>BlockedValue</code> containing <code>this</code>.
*/
public boolean check(BlockedValue prev, int pos) {
assert state != State.ROOT || type == null;
assert (state == State.ROOT ||
type == Schema.Type.ARRAY || type == Schema.Type.MAP);
assert 0 <= items;
assert 0 != items || start == pos; // 0==itms ==> start==pos
assert 1 < items || start == lastFullItem; // 1<=itms ==> start==lFI
assert items <= 1 || start <= lastFullItem; // 1<itms ==> start<=lFI
assert lastFullItem <= pos;
switch (state) {
case ROOT:
assert start == 0;
assert prev == null;
break;
case REGULAR:
assert start >= 0;
assert prev.lastFullItem <= start;
assert 1 <= prev.items;
break;
case OVERFLOW:
assert start == 0;
assert items == 1;
assert prev.state == State.ROOT || prev.state == State.OVERFLOW;
break;
}
return false;
}
}
/**
* The buffer to hold the bytes before being written into the underlying
* stream.
*/
private byte[] buf;
/**
* Index into the location in {@link #buf}, where next byte can be written.
*/
private int pos;
/**
* The state stack.
*/
private BlockedValue[] blockStack;
private int stackTop = -1;
private static final int STACK_STEP = 10;
//buffer large enough for up to two ints for a block header
//rounded up to a multiple of 4 bytes.
private byte[] headerBuffer = new byte[12];
private boolean check() {
assert buf != null;
assert 0 <= pos;
assert pos <= buf.length : pos + " " + buf.length;
assert blockStack != null;
BlockedValue prev = null;
for (int i = 0; i <= stackTop; i++) {
BlockedValue v = blockStack[i];
v.check(prev, pos);
prev = v;
}
return true;
}
BlockingBinaryEncoder(OutputStream out,
int blockBufferSize, int binaryEncoderBufferSize) {
super(out, binaryEncoderBufferSize);
this.buf = new byte[blockBufferSize];
this.pos = 0;
blockStack = new BlockedValue[0];
expandStack();
BlockedValue bv = blockStack[++stackTop];
bv.type = null;
bv.state = BlockedValue.State.ROOT;
bv.start = bv.lastFullItem = 0;
bv.items = 1;
assert check();
}
private void expandStack() {
int oldLength = blockStack.length;
blockStack = Arrays.copyOf(blockStack,
blockStack.length + STACK_STEP);
for (int i = oldLength; i < blockStack.length; i++) {
blockStack[i] = new BlockedValue();
}
}
BlockingBinaryEncoder configure(OutputStream out, int blockBufferSize,
int binaryEncoderBufferSize) {
super.configure(out, binaryEncoderBufferSize);
pos = 0;
stackTop = 0;
if (null == buf || buf.length != blockBufferSize) {
buf = new byte[blockBufferSize];
}
assert check();
return this;
}
@Override
public void flush() throws IOException {
BlockedValue bv = blockStack[stackTop];
if (bv.state == BlockedValue.State.ROOT) {
super.writeFixed(buf, 0, pos);
pos = 0;
} else {
while (bv.state != BlockedValue.State.OVERFLOW) {
compact();
}
}
super.flush();
assert check();
}
@Override
public void writeBoolean(boolean b) throws IOException {
ensureBounds(1);
pos += BinaryData.encodeBoolean(b, buf, pos);
}
@Override
public void writeInt(int n) throws IOException {
ensureBounds(5);
pos += BinaryData.encodeInt(n, buf, pos);
}
@Override
public void writeLong(long n) throws IOException {
ensureBounds(10);
pos += BinaryData.encodeLong(n, buf, pos);
}
@Override
public void writeFloat(float f) throws IOException {
ensureBounds(4);
pos += BinaryData.encodeFloat(f, buf, pos);
}
@Override
public void writeDouble(double d) throws IOException {
ensureBounds(8);
pos += BinaryData.encodeDouble(d, buf, pos);
}
@Override
public void writeFixed(byte[] bytes, int start, int len) throws IOException {
doWriteBytes(bytes, start, len);
}
@Override
protected void writeZero() throws IOException {
ensureBounds(1);
buf[pos++] = (byte) 0;
}
@Override
public void writeArrayStart() throws IOException {
if (stackTop + 1 == blockStack.length) {
expandStack();
}
BlockedValue bv = blockStack[++stackTop];
bv.type = Schema.Type.ARRAY;
bv.state = BlockedValue.State.REGULAR;
bv.start = bv.lastFullItem = pos;
bv.items = 0;
assert check();
}
@Override
public void setItemCount(long itemCount) throws IOException {
BlockedValue v = blockStack[stackTop];
assert v.type == Schema.Type.ARRAY || v.type == Schema.Type.MAP;
assert v.itemsLeftToWrite == 0;
v.itemsLeftToWrite = itemCount;
assert check();
}
@Override
public void startItem() throws IOException {
if (blockStack[stackTop].state == BlockedValue.State.OVERFLOW) {
finishOverflow();
}
BlockedValue t = blockStack[stackTop];
t.items++;
t.lastFullItem = pos;
t.itemsLeftToWrite--;
assert check();
}
@Override
public void writeArrayEnd() throws IOException {
BlockedValue top = blockStack[stackTop];
if (top.type != Schema.Type.ARRAY) {
throw new AvroTypeException("Called writeArrayEnd outside of an array.");
}
if (top.itemsLeftToWrite != 0) {
throw new AvroTypeException("Failed to write expected number of array elements.");
}
endBlockedValue();
assert check();
}
@Override
public void writeMapStart() throws IOException {
if (stackTop + 1 == blockStack.length) {
expandStack();
}
BlockedValue bv = blockStack[++stackTop];
bv.type = Schema.Type.MAP;
bv.state = BlockedValue.State.REGULAR;
bv.start = bv.lastFullItem = pos;
bv.items = 0;
assert check();
}
@Override
public void writeMapEnd() throws IOException {
BlockedValue top = blockStack[stackTop];
if (top.type != Schema.Type.MAP) {
throw new AvroTypeException("Called writeMapEnd outside of a map.");
}
if (top.itemsLeftToWrite != 0) {
throw new AvroTypeException("Failed to read write expected number of array elements.");
}
endBlockedValue();
assert check();
}
@Override
public void writeIndex(int unionIndex) throws IOException {
ensureBounds(5);
pos += BinaryData.encodeInt(unionIndex, buf, pos);
}
@Override
public int bytesBuffered() {
return pos + super.bytesBuffered();
}
private void endBlockedValue() throws IOException {
for (; ;) {
assert check();
BlockedValue t = blockStack[stackTop];
assert t.state != BlockedValue.State.ROOT;
if (t.state == BlockedValue.State.OVERFLOW) {
finishOverflow();
}
assert t.state == BlockedValue.State.REGULAR;
if (0 < t.items) {
int byteCount = pos - t.start;
if (t.start == 0 &&
blockStack[stackTop - 1].state
!= BlockedValue.State.REGULAR) { // Lucky us -- don't have to move
super.writeInt(-t.items);
super.writeInt(byteCount);
} else {
int headerSize = 0;
headerSize += BinaryData.encodeInt(-t.items, headerBuffer, headerSize);
headerSize += BinaryData.encodeInt(byteCount, headerBuffer, headerSize);
if (buf.length >= pos + headerSize) {
pos += headerSize;
final int m = t.start;
System.arraycopy(buf, m, buf, m + headerSize, byteCount);
System.arraycopy(headerBuffer, 0, buf, m, headerSize);
} else {
compact();
continue;
}
}
}
stackTop--;
ensureBounds(1);
buf[pos++] = 0; // Sentinel for last block in a blocked value
assert check();
if (blockStack[stackTop].state == BlockedValue.State.ROOT) {
flush();
}
return;
}
}
/**
* Called when we've finished writing the last item in an overflow
* buffer. When this is finished, the top of the stack will be
* an empty block in the "regular" state.
* @throws IOException
*/
private void finishOverflow() throws IOException {
BlockedValue s = blockStack[stackTop];
if (s.state != BlockedValue.State.OVERFLOW) {
throw new IllegalStateException("Not an overflow block");
}
assert check();
// Flush any remaining data for this block
super.writeFixed(buf, 0, pos);
pos = 0;
// Reset top of stack to be in REGULAR mode
s.state = BlockedValue.State.REGULAR;
s.start = s.lastFullItem = 0;
s.items = 0;
assert check();
}
private void ensureBounds(int l) throws IOException {
while (buf.length < (pos + l)) {
if (blockStack[stackTop].state == BlockedValue.State.REGULAR) {
compact();
} else {
super.writeFixed(buf, 0, pos);
pos = 0;
}
}
}
private void doWriteBytes(byte[] bytes, int start, int len)
throws IOException {
if (len < buf.length) {
ensureBounds(len);
System.arraycopy(bytes, start, buf, pos, len);
pos += len;
} else {
ensureBounds(buf.length);
assert blockStack[stackTop].state == BlockedValue.State.ROOT ||
blockStack[stackTop].state == BlockedValue.State.OVERFLOW;
write(bytes, start, len);
}
}
private void write(byte[] b, int off, int len) throws IOException {
if (blockStack[stackTop].state == BlockedValue.State.ROOT) {
super.writeFixed(b, off, len);
} else {
assert check();
while (buf.length < (pos + len)) {
if (blockStack[stackTop].state == BlockedValue.State.REGULAR) {
compact();
} else {
super.writeFixed(buf, 0, pos);
pos = 0;
if (buf.length <= len) {
super.writeFixed(b, off, len);
len = 0;
}
}
}
System.arraycopy(b, off, buf, pos, len);
pos += len;
}
assert check();
}
/** Only call if you're there are REGULAR-state values on the stack. */
private void compact() throws IOException {
assert check();
// Find first REGULAR-state value
BlockedValue s = null;
int i;
for (i = 1; i <= stackTop; i++) {
s = blockStack[i];
if (s.state == BlockedValue.State.REGULAR) break;
}
assert s != null;
// We're going to transition "s" into the overflow state. To do
// this, We're going to flush any bytes prior to "s", then write
// any full items of "s" into a block, start an overflow
// block, write any remaining bytes of "s" up to the start of the
// next more deeply-nested blocked-value, and finally move over
// any remaining bytes (which will be from more deeply-nested
// blocked values).
// Flush any bytes prios to "s"
super.writeFixed(buf, 0, s.start);
// Write any full items of "s"
if (1 < s.items) {
super.writeInt(-(s.items - 1));
super.writeInt(s.lastFullItem - s.start);
super.writeFixed(buf, s.start, s.lastFullItem - s.start);
s.start = s.lastFullItem;
s.items = 1;
}
// Start an overflow block for s
super.writeInt(1);
// Write any remaining bytes for "s", up to the next-most
// deeply-nested value
BlockedValue n = ((i + 1) <= stackTop ?
blockStack[i + 1] : null);
int end = (n == null ? pos : n.start);
super.writeFixed(buf, s.lastFullItem, end - s.lastFullItem);
// Move over any bytes that remain (and adjust indices)
System.arraycopy(buf, end, buf, 0, pos - end);
for (int j = i + 1; j <= stackTop; j++) {
n = blockStack[j];
n.start -= end;
n.lastFullItem -= end;
}
pos -= end;
assert s.items == 1;
s.start = s.lastFullItem = 0;
s.state = BlockedValue.State.OVERFLOW;
assert check();
}
}