Examples with FunctionImplementationRegistry org.apache.drill.exec.expr.fn.FunctionImplementationRegistry used on opensource projects

Example 66 with FunctionImplementationRegistry

use of org.apache.drill.exec.expr.fn.FunctionImplementationRegistry in project drill by apache.

the class TestComparisonFunctions method runTest.

public void runTest(String expression, int expectedResults) throws Throwable {
    final DrillbitContext bitContext = mockDrillbitContext();
    final UserClientConnection connection = Mockito.mock(UserClientConnection.class);
    final String planString = Resources.toString(Resources.getResource(COMPARISON_TEST_PHYSICAL_PLAN), Charsets.UTF_8).replaceAll("EXPRESSION", expression);
    if (reader == null) {
        reader = PhysicalPlanReaderTestFactory.defaultPhysicalPlanReader(c);
    }
    if (registry == null) {
        registry = new FunctionImplementationRegistry(c);
    }
    final FragmentContextImpl 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()));
    while (exec.next()) {
        assertEquals(String.format("Expression: %s;", expression), expectedResults, exec.getSelectionVector2().getCount());
    }
    exec.close();
    context.close();
    if (context.getExecutorState().getFailureCause() != null) {
        throw context.getExecutorState().getFailureCause();
    }
    assertTrue(!context.getExecutorState().isFailed());
}

Example 67 with FunctionImplementationRegistry

use of org.apache.drill.exec.expr.fn.FunctionImplementationRegistry in project drill by apache.

the class TestMathFunctions method testBasicMathFunctions.

@Test
public void testBasicMathFunctions() 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/simple_math_functions.json"), Charsets.UTF_8).read());
    final FunctionImplementationRegistry registry = new FunctionImplementationRegistry(c);
    final FragmentContextImpl context = new FragmentContextImpl(bitContext, BitControl.PlanFragment.getDefaultInstance(), connection, registry);
    final SimpleRootExec exec = new SimpleRootExec(ImplCreator.getExec(context, (FragmentRoot) plan.getSortedOperators(false).iterator().next()));
    while (exec.next()) {
        final IntVector intMulVector = exec.getValueVectorById(new SchemaPath("INTMUL", ExpressionPosition.UNKNOWN), IntVector.class);
        final Float8Vector floatMulVector = exec.getValueVectorById(new SchemaPath("FLOATMUL", ExpressionPosition.UNKNOWN), Float8Vector.class);
        final IntVector intAddVector = exec.getValueVectorById(new SchemaPath("INTADD", ExpressionPosition.UNKNOWN), IntVector.class);
        final Float8Vector floatAddVector = exec.getValueVectorById(new SchemaPath("FLOATADD", ExpressionPosition.UNKNOWN), Float8Vector.class);
        assertEquals(exec.getRecordCount(), 1);
        assertEquals(intMulVector.getAccessor().get(0), 2);
        assertEquals(floatMulVector.getAccessor().get(0), (1.1 * 2.2), 0);
        assertEquals(intAddVector.getAccessor().get(0), 3);
        assertEquals(floatAddVector.getAccessor().get(0), (1.1 + 2.2), 0);
    }
    if (context.getExecutorState().getFailureCause() != null) {
        throw context.getExecutorState().getFailureCause();
    }
    assertTrue(!context.getExecutorState().isFailed());
}

Example 68 with FunctionImplementationRegistry

use of org.apache.drill.exec.expr.fn.FunctionImplementationRegistry in project drill by apache.

the class TestTraceOutputDump method testFilter.

@Test
public void testFilter() 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("/trace/simple_trace.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()) {
    }
    exec.close();
    if (context.getExecutorState().getFailureCause() != null) {
        throw context.getExecutorState().getFailureCause();
    }
    assertTrue(!context.getExecutorState().isFailed());
    final FragmentHandle handle = context.getHandle();
    /* Form the file name to which the trace output will dump the record batches */
    final String qid = QueryIdHelper.getQueryId(handle.getQueryId());
    final int majorFragmentId = handle.getMajorFragmentId();
    final int minorFragmentId = handle.getMinorFragmentId();
    final String logLocation = c.getString(ExecConstants.TRACE_DUMP_DIRECTORY);
    final String filename = String.format("%s//%s_%d_%d_mock-scan", logLocation, qid, majorFragmentId, minorFragmentId);
    final Configuration conf = new Configuration();
    conf.set(FileSystem.FS_DEFAULT_NAME_KEY, c.getString(ExecConstants.TRACE_DUMP_FILESYSTEM));
    final FileSystem fs = FileSystem.get(conf);
    final Path path = new Path(filename);
    assertTrue("Trace file does not exist", fs.exists(path));
    final FSDataInputStream in = fs.open(path);
    final VectorAccessibleSerializable wrap = new VectorAccessibleSerializable(context.getAllocator());
    wrap.readFromStream(in);
    final VectorAccessible container = wrap.get();
    /* Assert there are no selection vectors */
    assertNull(wrap.getSv2());
    /* Assert there is only one record */
    assertEquals(1, container.getRecordCount());
    /* Read the Integer value and ASSERT its Integer.MIN_VALUE */
    final int value = (int) container.iterator().next().getValueVector().getAccessor().getObject(0);
    assertEquals(value, Integer.MIN_VALUE);
}

Example 69 with FunctionImplementationRegistry

use of org.apache.drill.exec.expr.fn.FunctionImplementationRegistry in project drill by apache.

the class TestDynamicUDFSupport method testLazyInitConcurrent.

@Test
public void testLazyInitConcurrent() throws Exception {
    FunctionImplementationRegistry functionImplementationRegistry = spyFunctionImplementationRegistry();
    copyDefaultJarsToStagingArea();
    test("create function using jar '%s'", default_binary_name);
    final CountDownLatch latch1 = new CountDownLatch(1);
    final CountDownLatch latch2 = new CountDownLatch(1);
    final String query = "select custom_lower('A') from (values(1))";
    doAnswer(new Answer<Boolean>() {

        @Override
        public Boolean answer(InvocationOnMock invocation) throws Throwable {
            latch1.await();
            boolean result = (boolean) invocation.callRealMethod();
            assertTrue("syncWithRemoteRegistry() should return true", result);
            latch2.countDown();
            return true;
        }
    }).doAnswer(new Answer() {

        @Override
        public Boolean answer(InvocationOnMock invocation) throws Throwable {
            latch1.countDown();
            latch2.await();
            boolean result = (boolean) invocation.callRealMethod();
            assertTrue("syncWithRemoteRegistry() should return true", result);
            return true;
        }
    }).when(functionImplementationRegistry).syncWithRemoteRegistry(anyLong());
    SimpleQueryRunner simpleQueryRunner = new SimpleQueryRunner(query);
    Thread thread1 = new Thread(simpleQueryRunner);
    Thread thread2 = new Thread(simpleQueryRunner);
    thread1.start();
    thread2.start();
    thread1.join();
    thread2.join();
    verify(functionImplementationRegistry, times(2)).syncWithRemoteRegistry(anyLong());
    LocalFunctionRegistry localFunctionRegistry = Deencapsulation.getField(functionImplementationRegistry, "localFunctionRegistry");
    assertEquals("Sync function registry version should match", 1L, localFunctionRegistry.getVersion());
}

Example 70 with FunctionImplementationRegistry

use of org.apache.drill.exec.expr.fn.FunctionImplementationRegistry in project drill by apache.

the class TestOptiqPlans method doLogicalTest.

private SimpleRootExec doLogicalTest(final BootStrapContext context, UserClientConnection connection, String file, ClusterCoordinator coord, DataConnectionCreator com, Controller controller, WorkEventBus workBus) throws Exception {
    new NonStrictExpectations() {

        {
            context.getMetrics();
            result = new MetricRegistry();
            context.getAllocator();
            result = RootAllocatorFactory.newRoot(config);
            context.getConfig();
            result = config;
        }
    };
    final RemoteServiceSet lss = RemoteServiceSet.getLocalServiceSet();
    final DrillbitContext bitContext = new DrillbitContext(DrillbitEndpoint.getDefaultInstance(), context, coord, controller, com, workBus, new LocalPersistentStoreProvider(config));
    final QueryContext qc = new QueryContext(UserSession.Builder.newBuilder().setSupportComplexTypes(true).build(), bitContext, QueryId.getDefaultInstance());
    final PhysicalPlanReader reader = bitContext.getPlanReader();
    final LogicalPlan plan = reader.readLogicalPlan(Files.toString(FileUtils.getResourceAsFile(file), Charsets.UTF_8));
    final PhysicalPlan pp = new BasicOptimizer(qc, connection).optimize(new BasicOptimizer.BasicOptimizationContext(qc), plan);
    final FunctionImplementationRegistry registry = new FunctionImplementationRegistry(config);
    final FragmentContext fctxt = new FragmentContext(bitContext, PlanFragment.getDefaultInstance(), connection, registry);
    final SimpleRootExec exec = new SimpleRootExec(ImplCreator.getExec(fctxt, (FragmentRoot) pp.getSortedOperators(false).iterator().next()));
    return exec;
}