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);
}
@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();
}
public void release() {
if (body != null) {
body.release();
}
}
@Override
public void close() {
release();
}
@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();
}