public void writeInt(BaseWriter.ComplexWriter writer, RepeatedIntHolder... in) {
int n = -1;
for (RepeatedIntHolder input : in) {
if (n != -1 && n != input.end - input.start) {
throw new IllegalArgumentException("Arguments to zip must all be same length");
}
n = input.end - input.start;
}
if (n == -1) {
throw new IllegalArgumentException("Must have at least one argument to zip");
}
IntVector.Accessor[] v = new IntVector.Accessor[in.length];
for (int i = 0; i < in.length; i++) {
v[i] = in[i].vector.getAccessor();
}
writer.setValueCount(n);
BaseWriter.ListWriter outer = writer.rootAsList();
outer.start(); // [ outer list
for (int i = 0; i < n; i++) {
BaseWriter.ListWriter inner = outer.list();
inner.start();
for (IntVector.Accessor accessor : v) {
inner.integer().writeInt(accessor.get(i));
}
inner.end(); // ] inner list
}
outer.end(); // ] outer list
}
@Test
public void sortOneKeyAscending(
@Injectable final DrillbitContext bitContext, @Injectable UserClientConnection connection)
throws Throwable {
new NonStrictExpectations() {
{
bitContext.getMetrics();
result = new MetricRegistry();
bitContext.getAllocator();
result = new TopLevelAllocator();
bitContext.getOperatorCreatorRegistry();
result = new OperatorCreatorRegistry(c);
bitContext.getConfig();
result = c;
bitContext.getCompiler();
result = CodeCompiler.getTestCompiler(c);
}
};
final PhysicalPlanReader reader =
new PhysicalPlanReader(
c, c.getMapper(), CoordinationProtos.DrillbitEndpoint.getDefaultInstance());
final PhysicalPlan plan =
reader.readPhysicalPlan(
Files.toString(FileUtils.getResourceAsFile("/sort/one_key_sort.json"), Charsets.UTF_8));
final FunctionImplementationRegistry registry = new FunctionImplementationRegistry(c);
final FragmentContext context =
new FragmentContext(bitContext, PlanFragment.getDefaultInstance(), connection, registry);
final SimpleRootExec exec =
new SimpleRootExec(
ImplCreator.getExec(
context, (FragmentRoot) plan.getSortedOperators(false).iterator().next()));
int previousInt = Integer.MIN_VALUE;
int recordCount = 0;
int batchCount = 0;
while (exec.next()) {
batchCount++;
final IntVector c1 =
exec.getValueVectorById(
new SchemaPath("blue", ExpressionPosition.UNKNOWN), IntVector.class);
final IntVector c2 =
exec.getValueVectorById(
new SchemaPath("green", ExpressionPosition.UNKNOWN), IntVector.class);
final IntVector.Accessor a1 = c1.getAccessor();
final IntVector.Accessor a2 = c2.getAccessor();
for (int i = 0; i < c1.getAccessor().getValueCount(); i++) {
recordCount++;
assertTrue(previousInt <= a1.get(i));
previousInt = a1.get(i);
assertEquals(previousInt, a2.get(i));
}
}
System.out.println(String.format("Sorted %,d records in %d batches.", recordCount, batchCount));
if (context.getFailureCause() != null) {
throw context.getFailureCause();
}
assertTrue(!context.isFailed());
}