private void loadRepeatedOffsetAddress() {
DrillBuf buf = vector.getOffsetVector().getBuffer();
checkBuf(buf);
this.repeatedOffset = buf.memoryAddress() + 4;
this.repeatedOffsetOriginal = buf.memoryAddress() + 4;
this.repeatedOffsetMax = buf.memoryAddress() + buf.capacity();
}
private void loadVarCharDataAddress() {
DrillBuf buf = vector.getDataVector().getBuffer();
checkBuf(buf);
this.characterData = buf.memoryAddress();
this.characterDataOriginal = buf.memoryAddress();
this.characterDataMax = buf.memoryAddress() + buf.capacity();
}
/**
* Reads a zero-compressed encoded long from input stream and returns it.
*
* @param buffer DrillBuf to read from
* @return deserialized long from stream.
*/
public static long readVLong(DrillBuf buffer, int start, int end) {
buffer.readerIndex(start);
byte firstByte = buffer.readByte();
int len = decodeVIntSize(firstByte);
int availableBytes = (end - start);
if (len == 1) {
return firstByte;
} else if (availableBytes < len) {
throw new NumberFormatException(
"Expected "
+ len
+ " bytes but the buffer '"
+ DrillStringUtils.toBinaryString(buffer, start, end)
+ "' has "
+ availableBytes
+ " bytes.");
}
long longValue = 0;
for (int idx = 0; idx < len - 1; idx++) {
byte byteValue = buffer.readByte();
longValue = longValue << 8;
longValue = longValue | (byteValue & 0xFF);
}
return (isNegativeVInt(firstByte) ? (longValue ^ -1L) : longValue);
}
public DrillBuf replace(DrillBuf old, int newSize) {
if (managedBuffers.remove(old.memoryAddress()) == null) {
throw new IllegalStateException("Tried to remove unmanaged buffer.");
}
old.release(1);
return getManagedBuffer(newSize);
}
private DrillBuf(
BufferAllocator allocator,
Accountor a,
ByteBuf replacement,
DrillBuf buffer,
int index,
int length,
boolean root) {
super(length);
if (index < 0 || index > buffer.capacity() - length) {
throw new IndexOutOfBoundsException(
buffer.toString() + ".slice(" + index + ", " + length + ')');
}
this.length = length;
writerIndex(length);
this.b = replacement;
this.addr = buffer.memoryAddress() + index;
this.offset = index;
this.acct = a;
this.length = length;
this.rootBuffer = root;
this.allocator = allocator;
}
@Override
public void addBinary(Binary value) {
holder.buffer = buf = buf.reallocIfNeeded(value.length());
buf.setBytes(0, value.toByteBuffer());
holder.start = 0;
holder.end = value.length();
writer.write(holder);
}
@Test
public void testNullableVarCharVectorLoad() {
final MaterializedField field =
MaterializedField.create(EMPTY_SCHEMA_PATH, NullableVarCharHolder.TYPE);
// Create a new value vector for 1024 nullable variable length strings.
final NullableVarCharVector vector1 = new NullableVarCharVector(field, allocator);
final NullableVarCharVector.Mutator mutator = vector1.getMutator();
vector1.allocateNew(1024 * 10, 1024);
// Populate the vector.
final StringBuilder stringBuilder = new StringBuilder();
final int valueCount = 10;
for (int i = 0; i < valueCount; ++i) {
stringBuilder.append('x');
mutator.set(i, stringBuilder.toString().getBytes(utf8Charset));
}
// Check the contents.
final NullableVarCharVector.Accessor accessor1 = vector1.getAccessor();
stringBuilder.setLength(0);
for (int i = 0; i < valueCount; ++i) {
stringBuilder.append('x');
final Object object = accessor1.getObject(i);
assertEquals(stringBuilder.toString(), object.toString());
}
mutator.setValueCount(valueCount);
assertEquals(valueCount, vector1.getAccessor().getValueCount());
// Still ok after setting value count?
stringBuilder.setLength(0);
for (int i = 0; i < valueCount; ++i) {
stringBuilder.append('x');
final Object object = accessor1.getObject(i);
assertEquals(stringBuilder.toString(), object.toString());
}
// Combine into a single buffer so we can load it into a new vector.
final DrillBuf[] buffers1 = vector1.getBuffers(false);
final DrillBuf buffer1 = combineBuffers(allocator, buffers1);
final NullableVarCharVector vector2 = new NullableVarCharVector(field, allocator);
vector2.load(vector1.getMetadata(), buffer1);
// Check the vector's contents.
final NullableVarCharVector.Accessor accessor2 = vector2.getAccessor();
stringBuilder.setLength(0);
for (int i = 0; i < valueCount; ++i) {
stringBuilder.append('x');
final Object object = accessor2.getObject(i);
assertEquals(stringBuilder.toString(), object.toString());
}
vector1.close();
vector2.close();
buffer1.release();
}
private void loadVarCharOffsetAddress() {
DrillBuf buf = vector.getDataVector().getOffsetVector().getBuffer();
checkBuf(buf);
this.charLengthOffset = buf.memoryAddress() + 4;
this.charLengthOffsetOriginal =
buf.memoryAddress()
+ 4; // add four as offsets conceptually start at 1. (first item is 0..1)
this.charLengthOffsetMax = buf.memoryAddress() + buf.capacity();
}
public static boolean checkBufRefs(final ValueVector vv) {
for (final DrillBuf buffer : vv.getBuffers(false)) {
if (buffer.refCnt() <= 0) {
throw new IllegalStateException("zero refcount");
}
}
return true;
}
public static long hash64(long start, long end, DrillBuf buffer, long seed) {
if (BoundsChecking.BOUNDS_CHECKING_ENABLED) {
buffer.checkBytes((int) start, (int) end);
}
long s = buffer.memoryAddress() + start;
long e = buffer.memoryAddress() + end;
return hash64bytes(s, e, seed);
}
public DrillBinaryToDecimal38Converter(
Decimal38SparseWriter writer, int precision, int scale, DrillBuf buf) {
this.writer = writer;
this.buf = buf.reallocIfNeeded(38);
holder.precision = precision;
holder.scale = scale;
}
public static VarCharHolder getVarCharHolder(DrillBuf buf, String s) {
VarCharHolder vch = new VarCharHolder();
byte[] b = s.getBytes(Charsets.UTF_8);
vch.start = 0;
vch.end = b.length;
vch.buffer = buf.reallocIfNeeded(b.length);
vch.buffer.setBytes(0, b);
return vch;
}
/**
* Serializes a long to a binary stream with zero-compressed encoding. For -112 <= i <= 127,
* only one byte is used with the actual value. For other values of i, the first byte value
* indicates whether the long is positive or negative, and the number of bytes that follow. If
* the first byte value v is between -113 and -120, the following long is positive, with number
* of bytes that follow are -(v+112). If the first byte value v is between -121 and -128, the
* following long is negative, with number of bytes that follow are -(v+120). Bytes are stored
* in the high-non-zero-byte-first order.
*
* @param buffer DrillBuf to write to
* @param i Long to be serialized
*/
public static void writeVLong(DrillBuf buffer, int start, int end, long i) {
int availableBytes = (end - start);
if (availableBytes < getVIntSize(i)) {
throw new NumberFormatException(
"Expected "
+ getVIntSize(i)
+ " bytes but the buffer '"
+ DrillStringUtils.toBinaryString(buffer, start, end)
+ "' has only "
+ availableBytes
+ " bytes.");
}
buffer.writerIndex(start);
if (i >= -112 && i <= 127) {
buffer.writeByte((byte) i);
return;
}
int len = -112;
if (i < 0) {
i ^= -1L; // take one's complement'
len = -120;
}
long tmp = i;
while (tmp != 0) {
tmp = tmp >> 8;
len--;
}
buffer.writeByte((byte) len);
len = (len < -120) ? -(len + 120) : -(len + 112);
for (int idx = len; idx != 0; idx--) {
int shiftbits = (idx - 1) * 8;
long mask = 0xFFL << shiftbits;
buffer.writeByte((byte) ((i & mask) >> shiftbits));
}
}
public RawFragmentBatch(
RemoteConnection connection,
FragmentRecordBatch header,
DrillBuf body,
ResponseSender sender) {
super();
this.header = header;
this.body = body;
this.connection = connection;
this.sender = sender;
if (body != null) {
body.retain();
}
}
private static DrillBuf combineBuffers(
final BufferAllocator allocator, final DrillBuf[] buffers) {
// find the total size we'll need
int size = 0;
for (final DrillBuf buffer : buffers) {
size += buffer.readableBytes();
}
// create the new buffer
final DrillBuf newBuf = allocator.buffer(size);
final DrillBuf writeBuf = newBuf;
for (final DrillBuf buffer : buffers) {
final DrillBuf readBuf = (DrillBuf) buffer.slice();
final int nBytes = readBuf.readableBytes();
final byte[] bytes = new byte[nBytes];
readBuf.readBytes(bytes);
writeBuf.writeBytes(bytes);
}
return newBuf;
}
public static Decimal38SparseHolder getDecimal38Holder(DrillBuf buf, String decimal) {
Decimal38SparseHolder dch = new Decimal38SparseHolder();
BigDecimal bigDecimal = new BigDecimal(decimal);
dch.scale = bigDecimal.scale();
dch.precision = bigDecimal.precision();
Decimal38SparseHolder.setSign(bigDecimal.signum() == -1, dch.start, dch.buffer);
dch.start = 0;
dch.buffer = buf.reallocIfNeeded(dch.maxPrecision * DecimalUtility.INTEGER_SIZE);
DecimalUtility.getSparseFromBigDecimal(
bigDecimal, dch.buffer, dch.start, dch.scale, dch.precision, dch.nDecimalDigits);
return dch;
}
private void writeBinary(
final MapOrListWriterImpl writer,
String fieldName,
boolean isList,
final VarBinaryHolder vb,
final byte[] bytes) {
ensure(bytes.length);
workBuf.setBytes(0, bytes);
vb.buffer = workBuf;
vb.start = 0;
vb.end = bytes.length;
if (isList == false) {
writer.binary(fieldName).write(vb);
} else {
writer.list.varBinary().write(vb);
}
}
private void writeString(
String readString, final MapOrListWriterImpl writer, String fieldName, boolean isList) {
int length;
byte[] strBytes;
try {
strBytes = readString.getBytes("UTF-8");
} catch (UnsupportedEncodingException e) {
throw new DrillRuntimeException("Unable to read string value for field: " + fieldName, e);
}
length = strBytes.length;
ensure(length);
workBuf.setBytes(0, strBytes);
final VarCharHolder vh = new VarCharHolder();
vh.buffer = workBuf;
vh.start = 0;
vh.end = length;
if (isList == false) {
writer.varChar(fieldName).write(vh);
} else {
writer.list.varChar().write(vh);
}
}
public long getByteCount() {
return body == null ? 0 : body.readableBytes();
}
public void release() {
if (body != null) {
body.release();
}
}
@Override
public DrillBuf slice(int index, int length) {
DrillBuf buf = new DrillBuf(this, index, length);
buf.writerIndex = length;
return buf;
}
/**
* Helper method to check if the buffer we are accessing has a minimum reference count and has not
* been deallocated
*
* @param b working drill buffer
*/
private void checkBuf(DrillBuf b) {
if (b.refCnt() < 1) {
throw new IllegalStateException("Cannot access a dereferenced buffer.");
}
}
@Test
public void testAllocators() throws Exception {
// Setup a drillbit (initializes a root allocator)
final DrillConfig config = DrillConfig.create(TEST_CONFIGURATIONS);
final RemoteServiceSet serviceSet = RemoteServiceSet.getLocalServiceSet();
final Drillbit bit = new Drillbit(config, serviceSet);
bit.run();
final DrillbitContext bitContext = bit.getContext();
FunctionImplementationRegistry functionRegistry =
bitContext.getFunctionImplementationRegistry();
StoragePluginRegistry storageRegistry = new StoragePluginRegistry(bitContext);
// Create a few Fragment Contexts
BitControl.PlanFragment.Builder pfBuilder1 = BitControl.PlanFragment.newBuilder();
pfBuilder1.setMemInitial(1500000);
BitControl.PlanFragment pf1 = pfBuilder1.build();
BitControl.PlanFragment.Builder pfBuilder2 = BitControl.PlanFragment.newBuilder();
pfBuilder2.setMemInitial(500000);
BitControl.PlanFragment pf2 = pfBuilder1.build();
FragmentContext fragmentContext1 = new FragmentContext(bitContext, pf1, null, functionRegistry);
FragmentContext fragmentContext2 = new FragmentContext(bitContext, pf2, null, functionRegistry);
// Get a few physical operators. Easiest way is to read a physical plan.
PhysicalPlanReader planReader =
new PhysicalPlanReader(
config,
config.getMapper(),
CoordinationProtos.DrillbitEndpoint.getDefaultInstance(),
storageRegistry);
PhysicalPlan plan =
planReader.readPhysicalPlan(
Files.toString(FileUtils.getResourceAsFile(planFile), Charsets.UTF_8));
List<PhysicalOperator> physicalOperators = plan.getSortedOperators();
Iterator<PhysicalOperator> physicalOperatorIterator = physicalOperators.iterator();
PhysicalOperator physicalOperator1 = physicalOperatorIterator.next();
PhysicalOperator physicalOperator2 = physicalOperatorIterator.next();
PhysicalOperator physicalOperator3 = physicalOperatorIterator.next();
PhysicalOperator physicalOperator4 = physicalOperatorIterator.next();
PhysicalOperator physicalOperator5 = physicalOperatorIterator.next();
PhysicalOperator physicalOperator6 = physicalOperatorIterator.next();
// Create some bogus Operator profile defs and stats to create operator contexts
OpProfileDef def;
OperatorStats stats;
// Use some bogus operator type to create a new operator context.
def =
new OpProfileDef(
physicalOperator1.getOperatorId(),
UserBitShared.CoreOperatorType.MOCK_SUB_SCAN_VALUE,
OperatorContext.getChildCount(physicalOperator1));
stats = fragmentContext1.getStats().getOperatorStats(def, fragmentContext1.getAllocator());
// Add a couple of Operator Contexts
// Initial allocation = 1000000 bytes for all operators
OperatorContext oContext11 = fragmentContext1.newOperatorContext(physicalOperator1, true);
DrillBuf b11 = oContext11.getAllocator().buffer(1000000);
OperatorContext oContext12 =
fragmentContext1.newOperatorContext(physicalOperator2, stats, true);
DrillBuf b12 = oContext12.getAllocator().buffer(500000);
OperatorContext oContext21 = fragmentContext1.newOperatorContext(physicalOperator3, true);
def =
new OpProfileDef(
physicalOperator4.getOperatorId(),
UserBitShared.CoreOperatorType.TEXT_WRITER_VALUE,
OperatorContext.getChildCount(physicalOperator4));
stats = fragmentContext2.getStats().getOperatorStats(def, fragmentContext2.getAllocator());
OperatorContext oContext22 =
fragmentContext2.newOperatorContext(physicalOperator4, stats, true);
DrillBuf b22 = oContext22.getAllocator().buffer(2000000);
// New Fragment begins
BitControl.PlanFragment.Builder pfBuilder3 = BitControl.PlanFragment.newBuilder();
pfBuilder3.setMemInitial(1000000);
BitControl.PlanFragment pf3 = pfBuilder3.build();
FragmentContext fragmentContext3 = new FragmentContext(bitContext, pf3, null, functionRegistry);
// New fragment starts an operator that allocates an amount within the limit
def =
new OpProfileDef(
physicalOperator5.getOperatorId(),
UserBitShared.CoreOperatorType.UNION_VALUE,
OperatorContext.getChildCount(physicalOperator5));
stats = fragmentContext3.getStats().getOperatorStats(def, fragmentContext3.getAllocator());
OperatorContext oContext31 =
fragmentContext3.newOperatorContext(physicalOperator5, stats, true);
DrillBuf b31a = oContext31.getAllocator().buffer(200000);
// Previously running operator completes
b22.release();
((AutoCloseable) oContext22).close();
// Fragment 3 asks for more and fails
boolean outOfMem = false;
try {
DrillBuf b31b = oContext31.getAllocator().buffer(4400000);
if (b31b != null) {
b31b.release();
} else {
outOfMem = true;
}
} catch (Exception e) {
outOfMem = true;
}
assertEquals(true, (boolean) outOfMem);
// Operator is Exempt from Fragment limits. Fragment 3 asks for more and succeeds
outOfMem = false;
OperatorContext oContext32 = fragmentContext3.newOperatorContext(physicalOperator6, false);
DrillBuf b32 = null;
try {
b32 = oContext32.getAllocator().buffer(4400000);
} catch (Exception e) {
outOfMem = true;
} finally {
if (b32 != null) {
b32.release();
} else {
outOfMem = true;
}
closeOp(oContext32);
}
assertEquals(false, (boolean) outOfMem);
b11.release();
closeOp(oContext11);
b12.release();
closeOp(oContext12);
closeOp(oContext21);
b31a.release();
closeOp(oContext31);
fragmentContext1.close();
fragmentContext2.close();
fragmentContext3.close();
bit.close();
serviceSet.close();
}
/**
* Special constructor used for RPC ownership transfer. Takes a snapshot slice of the current buf
* but points directly to the underlying UnsafeLittleEndian buffer. Does this by calling unwrap()
* twice on the provided DrillBuf and expecting an UnsafeDirectLittleEndian buffer. This operation
* includes taking a new reference count on the underlying buffer and maintaining returning with a
* current reference count for itself (masking the underlying reference count).
*
* @param allocator
* @param a Allocator used when users try to receive allocator from buffer.
* @param b Accountor used for accounting purposes.
*/
public DrillBuf(BufferAllocator allocator, Accountor a, DrillBuf b) {
this(allocator, a, getUnderlying(b), b, 0, b.length, true);
assert b.unwrap().unwrap() instanceof UnsafeDirectLittleEndian;
b.unwrap().unwrap().retain();
}
public DrillFixedBinaryToVarbinaryConverter(VarBinaryWriter writer, int length, DrillBuf buf) {
this.writer = writer;
holder.buffer = buf = buf.reallocIfNeeded(length);
holder.start = 0;
holder.end = length;
}
private void ensure(final int length) {
workBuf = workBuf.reallocIfNeeded(length);
}
public DrillBuf getManagedBuffer(int size) {
DrillBuf newBuf = allocator.buffer(size, this);
managedBuffers.put(newBuf.memoryAddress(), newBuf);
return newBuf;
}
private static ByteBuf getUnderlying(DrillBuf b) {
ByteBuf underlying = b.unwrap().unwrap();
return underlying.slice((int) (b.memoryAddress() - underlying.memoryAddress()), b.length);
}
