Examples with FragmentContextImpl org.apache.drill.exec.ops.FragmentContextImpl used on opensource projects

Example 86 with FragmentContextImpl

use of org.apache.drill.exec.ops.FragmentContextImpl in project drill by apache.

the class TopNBatchTest method priorityQueueOrderingTest.

/**
 * Priority queue unit test.
 * @throws Exception
 */
@Test
public void priorityQueueOrderingTest() throws Exception {
    Properties properties = new Properties();
    DrillConfig drillConfig = DrillConfig.create(properties);
    DrillbitContext drillbitContext = mockDrillbitContext();
    when(drillbitContext.getFunctionImplementationRegistry()).thenReturn(new FunctionImplementationRegistry(drillConfig));
    FieldReference expr = FieldReference.getWithQuotedRef("colA");
    Order.Ordering ordering = new Order.Ordering(Order.Ordering.ORDER_DESC, expr, Order.Ordering.NULLS_FIRST);
    List<Order.Ordering> orderings = Lists.newArrayList(ordering);
    MaterializedField colA = MaterializedField.create("colA", Types.required(TypeProtos.MinorType.INT));
    MaterializedField colB = MaterializedField.create("colB", Types.required(TypeProtos.MinorType.INT));
    List<MaterializedField> cols = Lists.newArrayList(colA, colB);
    BatchSchema batchSchema = new BatchSchema(BatchSchema.SelectionVectorMode.NONE, cols);
    FragmentContextImpl context = new FragmentContextImpl(drillbitContext, BitControl.PlanFragment.getDefaultInstance(), null, drillbitContext.getFunctionImplementationRegistry());
    RowSet expectedRowSet;
    try (RootAllocator allocator = new RootAllocator(100_000_000)) {
        expectedRowSet = new RowSetBuilder(allocator, batchSchema).addRow(110, 10).addRow(109, 9).addRow(108, 8).addRow(107, 7).addRow(106, 6).addRow(105, 5).addRow(104, 4).addRow(103, 3).addRow(102, 2).addRow(101, 1).build();
        PriorityQueue queue;
        ExpandableHyperContainer hyperContainer;
        {
            VectorContainer container = new RowSetBuilder(allocator, batchSchema).build().container();
            hyperContainer = new ExpandableHyperContainer(container);
            queue = TopNBatch.createNewPriorityQueue(TopNBatch.createMainMappingSet(), TopNBatch.createLeftMappingSet(), TopNBatch.createRightMappingSet(), orderings, hyperContainer, false, true, 10, allocator, batchSchema.getSelectionVectorMode(), context);
        }
        List<RecordBatchData> testBatches = Lists.newArrayList();
        try {
            final Random random = new Random();
            final int bound = 100;
            final int numBatches = 11;
            final int numRecordsPerBatch = 100;
            for (int batchCounter = 0; batchCounter < numBatches; batchCounter++) {
                RowSetBuilder rowSetBuilder = new RowSetBuilder(allocator, batchSchema);
                rowSetBuilder.addRow((batchCounter + bound), batchCounter);
                for (int recordCounter = 0; recordCounter < numRecordsPerBatch; recordCounter++) {
                    rowSetBuilder.addRow(random.nextInt(bound), random.nextInt(bound));
                }
                VectorContainer vectorContainer = rowSetBuilder.build().container();
                queue.add(new RecordBatchData(vectorContainer, allocator));
            }
            queue.generate();
            VectorContainer resultContainer = queue.getHyperBatch();
            resultContainer.buildSchema(BatchSchema.SelectionVectorMode.NONE);
            RowSet.HyperRowSet actualHyperSet = HyperRowSetImpl.fromContainer(resultContainer, queue.getFinalSv4());
            new RowSetComparison(expectedRowSet).verify(actualHyperSet);
        } finally {
            if (expectedRowSet != null) {
                expectedRowSet.clear();
            }
            queue.cleanup();
            hyperContainer.clear();
            for (RecordBatchData testBatch : testBatches) {
                testBatch.clear();
            }
        }
    }
}

Example 87 with FragmentContextImpl

use of org.apache.drill.exec.ops.FragmentContextImpl in project drill by apache.

the class TestAgg method doTest.

private SimpleRootExec doTest(String file) throws Exception {
    final DrillbitContext bitContext = mockDrillbitContext();
    final UserClientConnection connection = Mockito.mock(UserClientConnection.class);
    final PhysicalPlanReader reader = PhysicalPlanReaderTestFactory.defaultPhysicalPlanReader(c);
    final PhysicalPlan plan = reader.readPhysicalPlan(Files.asCharSource(DrillFileUtils.getResourceAsFile(file), Charsets.UTF_8).read());
    final FunctionImplementationRegistry registry = new FunctionImplementationRegistry(c);
    final FragmentContextImpl context = new FragmentContextImpl(bitContext, PlanFragment.getDefaultInstance(), connection, registry);
    final SimpleRootExec exec = new SimpleRootExec(ImplCreator.getExec(context, (FragmentRoot) plan.getSortedOperators(false).iterator().next()));
    return exec;
}

Example 88 with FragmentContextImpl

use of org.apache.drill.exec.ops.FragmentContextImpl in project drill by apache.

the class TestSimpleFunctions method testSubstringNegative.

@Test
public void testSubstringNegative() throws Throwable {
    final DrillbitContext bitContext = mockDrillbitContext();
    final UserClientConnection connection = Mockito.mock(UserClientConnection.class);
    final PhysicalPlanReader reader = PhysicalPlanReaderTestFactory.defaultPhysicalPlanReader(c);
    final PhysicalPlan plan = reader.readPhysicalPlan(Files.asCharSource(DrillFileUtils.getResourceAsFile("/functions/testSubstringNegative.json"), Charsets.UTF_8).read());
    final FunctionImplementationRegistry registry = new FunctionImplementationRegistry(c);
    final FragmentContextImpl context = new FragmentContextImpl(bitContext, PlanFragment.getDefaultInstance(), connection, registry);
    final SimpleRootExec exec = new SimpleRootExec(ImplCreator.getExec(context, (FragmentRoot) plan.getSortedOperators(false).iterator().next()));
    while (exec.next()) {
        final NullableVarCharVector c1 = exec.getValueVectorById(new SchemaPath("col3", ExpressionPosition.UNKNOWN), NullableVarCharVector.class);
        final NullableVarCharVector.Accessor a1 = c1.getAccessor();
        int count = 0;
        for (int i = 0; i < c1.getAccessor().getValueCount(); i++) {
            if (!a1.isNull(i)) {
                final NullableVarCharHolder holder = new NullableVarCharHolder();
                a1.get(i, holder);
                // when offset is negative, substring return empty string.
                assertEquals("", StringFunctionHelpers.toStringFromUTF8(holder.start, holder.end, holder.buffer));
                ++count;
            }
        }
        assertEquals(50, count);
    }
    if (context.getExecutorState().getFailureCause() != null) {
        throw context.getExecutorState().getFailureCause();
    }
    assertTrue(!context.getExecutorState().isFailed());
}

Example 89 with FragmentContextImpl

use of org.apache.drill.exec.ops.FragmentContextImpl in project drill by apache.

the class TestSimpleFunctions method testSubstring.

@Test
public void testSubstring() throws Throwable {
    final DrillbitContext bitContext = mockDrillbitContext();
    final UserClientConnection connection = Mockito.mock(UserClientConnection.class);
    final PhysicalPlanReader reader = PhysicalPlanReaderTestFactory.defaultPhysicalPlanReader(c);
    final PhysicalPlan plan = reader.readPhysicalPlan(Files.asCharSource(DrillFileUtils.getResourceAsFile("/functions/testSubstring.json"), Charsets.UTF_8).read());
    final FunctionImplementationRegistry registry = new FunctionImplementationRegistry(c);
    final FragmentContextImpl context = new FragmentContextImpl(bitContext, PlanFragment.getDefaultInstance(), connection, registry);
    final SimpleRootExec exec = new SimpleRootExec(ImplCreator.getExec(context, (FragmentRoot) plan.getSortedOperators(false).iterator().next()));
    while (exec.next()) {
        final NullableVarCharVector c1 = exec.getValueVectorById(new SchemaPath("col3", ExpressionPosition.UNKNOWN), NullableVarCharVector.class);
        final NullableVarCharVector.Accessor a1 = c1.getAccessor();
        int count = 0;
        for (int i = 0; i < c1.getAccessor().getValueCount(); i++) {
            if (!a1.isNull(i)) {
                final NullableVarCharHolder holder = new NullableVarCharHolder();
                a1.get(i, holder);
                assertEquals("aaaa", StringFunctionHelpers.toStringFromUTF8(holder.start, holder.end, holder.buffer));
                ++count;
            }
        }
        assertEquals(50, count);
    }
    if (context.getExecutorState().getFailureCause() != null) {
        throw context.getExecutorState().getFailureCause();
    }
    assertTrue(!context.getExecutorState().isFailed());
}

Example 90 with FragmentContextImpl

use of org.apache.drill.exec.ops.FragmentContextImpl in project drill by apache.

the class TestStringFunctions method runTest.

public void runTest(Object[] expectedResults, String planPath) throws Throwable {
    final DrillbitContext bitContext = mockDrillbitContext();
    final UserClientConnection connection = Mockito.mock(UserClientConnection.class);
    final String planString = Resources.toString(Resources.getResource(planPath), Charsets.UTF_8);
    if (reader == null) {
        reader = PhysicalPlanReaderTestFactory.defaultPhysicalPlanReader(c);
    }
    if (registry == null) {
        registry = new FunctionImplementationRegistry(c);
    }
    if (context == null) {
        context = new FragmentContextImpl(bitContext, PlanFragment.getDefaultInstance(), connection, registry);
    }
    final PhysicalPlan plan = reader.readPhysicalPlan(planString);
    final SimpleRootExec exec = new SimpleRootExec(ImplCreator.getExec(context, (FragmentRoot) plan.getSortedOperators(false).iterator().next()));
    // skip schema batch
    exec.next();
    while (exec.next()) {
        final Object[] res = getRunResult(exec);
        assertEquals("return count does not match", expectedResults.length, res.length);
        for (int i = 0; i < res.length; i++) {
            assertEquals(String.format("column %s does not match", i), expectedResults[i], res[i]);
        }
    }
    if (context.getExecutorState().getFailureCause() != null) {
        throw context.getExecutorState().getFailureCause();
    }
    assertTrue(!context.getExecutorState().isFailed());
}