Examples with PlannerContext org.apache.druid.sql.calcite.planner.PlannerContext used on opensource projects

Example 11 with PlannerContext

use of org.apache.druid.sql.calcite.planner.PlannerContext in project druid by druid-io.

the class EarliestLatestAnySqlAggregator method toDruidAggregation.

@Nullable
@Override
public Aggregation toDruidAggregation(final PlannerContext plannerContext, final RowSignature rowSignature, final VirtualColumnRegistry virtualColumnRegistry, final RexBuilder rexBuilder, final String name, final AggregateCall aggregateCall, final Project project, final List<Aggregation> existingAggregations, final boolean finalizeAggregations) {
    final List<RexNode> rexNodes = aggregateCall.getArgList().stream().map(i -> Expressions.fromFieldAccess(rowSignature, project, i)).collect(Collectors.toList());
    final List<DruidExpression> args = Expressions.toDruidExpressions(plannerContext, rowSignature, rexNodes);
    if (args == null) {
        return null;
    }
    final String aggregatorName = finalizeAggregations ? Calcites.makePrefixedName(name, "a") : name;
    final ColumnType outputType = Calcites.getColumnTypeForRelDataType(aggregateCall.getType());
    if (outputType == null) {
        throw new ISE("Cannot translate output sqlTypeName[%s] to Druid type for aggregator[%s]", aggregateCall.getType().getSqlTypeName(), aggregateCall.getName());
    }
    final String fieldName = getColumnName(plannerContext, virtualColumnRegistry, args.get(0), rexNodes.get(0));
    final AggregatorFactory theAggFactory;
    switch(args.size()) {
        case 1:
            theAggFactory = aggregatorType.createAggregatorFactory(aggregatorName, fieldName, null, outputType, -1);
            break;
        case 2:
            theAggFactory = aggregatorType.createAggregatorFactory(aggregatorName, fieldName, null, outputType, RexLiteral.intValue(rexNodes.get(1)));
            break;
        default:
            throw new IAE("aggregation[%s], Invalid number of arguments[%,d] to [%s] operator", aggregatorName, args.size(), aggregatorType.name());
    }
    return Aggregation.create(Collections.singletonList(theAggFactory), finalizeAggregations ? new FinalizingFieldAccessPostAggregator(name, aggregatorName) : null);
}

Example 12 with PlannerContext

use of org.apache.druid.sql.calcite.planner.PlannerContext in project druid by druid-io.

the class SqlLifecycleTest method testStateTransition.

@Test
public void testStateTransition() throws ValidationException, SqlParseException, RelConversionException, IOException {
    SqlLifecycle lifecycle = sqlLifecycleFactory.factorize();
    final String sql = "select 1 + ?";
    final Map<String, Object> queryContext = Collections.emptyMap();
    Assert.assertEquals(SqlLifecycle.State.NEW, lifecycle.getState());
    // test initialize
    lifecycle.initialize(sql, queryContext);
    Assert.assertEquals(SqlLifecycle.State.INITIALIZED, lifecycle.getState());
    List<TypedValue> parameters = ImmutableList.of(new SqlParameter(SqlType.BIGINT, 1L).getTypedValue());
    lifecycle.setParameters(parameters);
    // setting parameters should not change the state
    Assert.assertEquals(SqlLifecycle.State.INITIALIZED, lifecycle.getState());
    // test authorization
    DruidPlanner mockPlanner = EasyMock.createMock(DruidPlanner.class);
    PlannerContext mockPlannerContext = EasyMock.createMock(PlannerContext.class);
    ValidationResult validationResult = new ValidationResult(Collections.emptySet());
    EasyMock.expect(plannerFactory.createPlanner(EasyMock.eq(sql), EasyMock.anyObject())).andReturn(mockPlanner).once();
    EasyMock.expect(mockPlanner.getPlannerContext()).andReturn(mockPlannerContext).once();
    mockPlannerContext.setAuthenticationResult(CalciteTests.REGULAR_USER_AUTH_RESULT);
    EasyMock.expectLastCall();
    mockPlannerContext.setParameters(parameters);
    EasyMock.expectLastCall();
    EasyMock.expect(plannerFactory.getAuthorizerMapper()).andReturn(CalciteTests.TEST_AUTHORIZER_MAPPER).once();
    mockPlannerContext.setAuthorizationResult(Access.OK);
    EasyMock.expectLastCall();
    EasyMock.expect(mockPlanner.validate()).andReturn(validationResult).once();
    mockPlanner.close();
    EasyMock.expectLastCall();
    EasyMock.replay(plannerFactory, serviceEmitter, requestLogger, mockPlanner, mockPlannerContext);
    lifecycle.validateAndAuthorize(CalciteTests.REGULAR_USER_AUTH_RESULT);
    Assert.assertEquals(SqlLifecycle.State.AUTHORIZED, lifecycle.getState());
    EasyMock.verify(plannerFactory, serviceEmitter, requestLogger, mockPlanner, mockPlannerContext);
    EasyMock.reset(plannerFactory, serviceEmitter, requestLogger, mockPlanner, mockPlannerContext);
    // test prepare
    PrepareResult mockPrepareResult = EasyMock.createMock(PrepareResult.class);
    EasyMock.expect(plannerFactory.createPlannerWithContext(EasyMock.eq(mockPlannerContext))).andReturn(mockPlanner).once();
    EasyMock.expect(mockPlanner.prepare()).andReturn(mockPrepareResult).once();
    mockPlanner.close();
    EasyMock.expectLastCall();
    EasyMock.replay(plannerFactory, serviceEmitter, requestLogger, mockPlanner, mockPlannerContext, mockPrepareResult);
    lifecycle.prepare();
    // prepare doens't change lifecycle state
    Assert.assertEquals(SqlLifecycle.State.AUTHORIZED, lifecycle.getState());
    EasyMock.verify(plannerFactory, serviceEmitter, requestLogger, mockPlanner, mockPlannerContext, mockPrepareResult);
    EasyMock.reset(plannerFactory, serviceEmitter, requestLogger, mockPlanner, mockPlannerContext, mockPrepareResult);
    // test plan
    PlannerResult mockPlanResult = EasyMock.createMock(PlannerResult.class);
    EasyMock.expect(plannerFactory.createPlannerWithContext(EasyMock.eq(mockPlannerContext))).andReturn(mockPlanner).once();
    EasyMock.expect(mockPlanner.plan()).andReturn(mockPlanResult).once();
    mockPlanner.close();
    EasyMock.expectLastCall();
    EasyMock.replay(plannerFactory, serviceEmitter, requestLogger, mockPlanner, mockPlannerContext, mockPrepareResult, mockPlanResult);
    lifecycle.plan();
    Assert.assertEquals(mockPlannerContext, lifecycle.getPlannerContext());
    Assert.assertEquals(SqlLifecycle.State.PLANNED, lifecycle.getState());
    EasyMock.verify(plannerFactory, serviceEmitter, requestLogger, mockPlanner, mockPlannerContext, mockPrepareResult, mockPlanResult);
    EasyMock.reset(plannerFactory, serviceEmitter, requestLogger, mockPlanner, mockPlannerContext, mockPrepareResult, mockPlanResult);
    // test execute
    EasyMock.expect(mockPlanResult.run()).andReturn(Sequences.simple(ImmutableList.of(new Object[] { 2L }))).once();
    EasyMock.replay(plannerFactory, serviceEmitter, requestLogger, mockPlanner, mockPlannerContext, mockPrepareResult, mockPlanResult);
    lifecycle.execute();
    Assert.assertEquals(SqlLifecycle.State.EXECUTING, lifecycle.getState());
    EasyMock.verify(plannerFactory, serviceEmitter, requestLogger, mockPlanner, mockPlannerContext, mockPrepareResult, mockPlanResult);
    EasyMock.reset(plannerFactory, serviceEmitter, requestLogger, mockPlanner, mockPlannerContext, mockPrepareResult, mockPlanResult);
    // test emit
    EasyMock.expect(mockPlannerContext.getSqlQueryId()).andReturn("id").once();
    CopyOnWriteArrayList<String> nativeQueryIds = new CopyOnWriteArrayList<>(ImmutableList.of("id"));
    EasyMock.expect(mockPlannerContext.getNativeQueryIds()).andReturn(nativeQueryIds).times(2);
    EasyMock.expect(mockPlannerContext.getAuthenticationResult()).andReturn(CalciteTests.REGULAR_USER_AUTH_RESULT).once();
    serviceEmitter.emit(EasyMock.anyObject(ServiceEventBuilder.class));
    EasyMock.expectLastCall();
    serviceEmitter.emit(EasyMock.anyObject(ServiceEventBuilder.class));
    EasyMock.expectLastCall();
    requestLogger.logSqlQuery(EasyMock.anyObject());
    EasyMock.expectLastCall();
    EasyMock.replay(plannerFactory, serviceEmitter, requestLogger, mockPlanner, mockPlannerContext, mockPrepareResult, mockPlanResult);
    lifecycle.finalizeStateAndEmitLogsAndMetrics(null, null, 10);
    Assert.assertEquals(SqlLifecycle.State.DONE, lifecycle.getState());
    EasyMock.verify(plannerFactory, serviceEmitter, requestLogger, mockPlanner, mockPlannerContext, mockPrepareResult, mockPlanResult);
    EasyMock.reset(plannerFactory, serviceEmitter, requestLogger, mockPlanner, mockPlannerContext, mockPrepareResult, mockPlanResult);
}

Example 13 with PlannerContext

use of org.apache.druid.sql.calcite.planner.PlannerContext in project druid by druid-io.

the class SqlLifecycleTest method testStateTransitionHttpRequest.

@Test
public void testStateTransitionHttpRequest() throws ValidationException, SqlParseException, RelConversionException, IOException {
    // this test is a duplicate of testStateTransition except with a slight variation of how validate and authorize
    // is run
    SqlLifecycle lifecycle = sqlLifecycleFactory.factorize();
    final String sql = "select 1 + ?";
    final Map<String, Object> queryContext = Collections.emptyMap();
    Assert.assertEquals(SqlLifecycle.State.NEW, lifecycle.getState());
    // test initialize
    lifecycle.initialize(sql, queryContext);
    Assert.assertEquals(SqlLifecycle.State.INITIALIZED, lifecycle.getState());
    List<TypedValue> parameters = ImmutableList.of(new SqlParameter(SqlType.BIGINT, 1L).getTypedValue());
    lifecycle.setParameters(parameters);
    // setting parameters should not change the state
    Assert.assertEquals(SqlLifecycle.State.INITIALIZED, lifecycle.getState());
    // test authorization
    DruidPlanner mockPlanner = EasyMock.createMock(DruidPlanner.class);
    PlannerContext mockPlannerContext = EasyMock.createMock(PlannerContext.class);
    ValidationResult validationResult = new ValidationResult(Collections.emptySet());
    EasyMock.expect(plannerFactory.createPlanner(EasyMock.eq(sql), EasyMock.anyObject())).andReturn(mockPlanner).once();
    EasyMock.expect(mockPlanner.getPlannerContext()).andReturn(mockPlannerContext).once();
    mockPlannerContext.setAuthenticationResult(CalciteTests.REGULAR_USER_AUTH_RESULT);
    EasyMock.expectLastCall();
    mockPlannerContext.setParameters(parameters);
    EasyMock.expectLastCall();
    EasyMock.expect(plannerFactory.getAuthorizerMapper()).andReturn(CalciteTests.TEST_AUTHORIZER_MAPPER).once();
    mockPlannerContext.setAuthorizationResult(Access.OK);
    EasyMock.expectLastCall();
    EasyMock.expect(mockPlanner.validate()).andReturn(validationResult).once();
    mockPlanner.close();
    EasyMock.expectLastCall();
    HttpServletRequest request = EasyMock.createMock(HttpServletRequest.class);
    EasyMock.expect(request.getAttribute(AuthConfig.DRUID_AUTHENTICATION_RESULT)).andReturn(CalciteTests.REGULAR_USER_AUTH_RESULT).times(2);
    EasyMock.expect(request.getAttribute(AuthConfig.DRUID_ALLOW_UNSECURED_PATH)).andReturn(null).once();
    EasyMock.expect(request.getAttribute(AuthConfig.DRUID_AUTHORIZATION_CHECKED)).andReturn(null).once();
    request.setAttribute(AuthConfig.DRUID_AUTHORIZATION_CHECKED, true);
    EasyMock.expectLastCall();
    EasyMock.replay(plannerFactory, serviceEmitter, requestLogger, mockPlanner, mockPlannerContext, request);
    lifecycle.validateAndAuthorize(request);
    Assert.assertEquals(SqlLifecycle.State.AUTHORIZED, lifecycle.getState());
    EasyMock.verify(plannerFactory, serviceEmitter, requestLogger, mockPlanner, mockPlannerContext, request);
    EasyMock.reset(plannerFactory, serviceEmitter, requestLogger, mockPlanner, mockPlannerContext, request);
    // test prepare
    PrepareResult mockPrepareResult = EasyMock.createMock(PrepareResult.class);
    EasyMock.expect(plannerFactory.createPlannerWithContext(EasyMock.eq(mockPlannerContext))).andReturn(mockPlanner).once();
    EasyMock.expect(mockPlanner.prepare()).andReturn(mockPrepareResult).once();
    mockPlanner.close();
    EasyMock.expectLastCall();
    EasyMock.replay(plannerFactory, serviceEmitter, requestLogger, mockPlanner, mockPlannerContext, mockPrepareResult);
    lifecycle.prepare();
    // prepare doens't change lifecycle state
    Assert.assertEquals(SqlLifecycle.State.AUTHORIZED, lifecycle.getState());
    EasyMock.verify(plannerFactory, serviceEmitter, requestLogger, mockPlanner, mockPlannerContext, mockPrepareResult);
    EasyMock.reset(plannerFactory, serviceEmitter, requestLogger, mockPlanner, mockPlannerContext, mockPrepareResult);
    // test plan
    PlannerResult mockPlanResult = EasyMock.createMock(PlannerResult.class);
    EasyMock.expect(plannerFactory.createPlannerWithContext(EasyMock.eq(mockPlannerContext))).andReturn(mockPlanner).once();
    EasyMock.expect(mockPlanner.plan()).andReturn(mockPlanResult).once();
    mockPlanner.close();
    EasyMock.expectLastCall();
    EasyMock.replay(plannerFactory, serviceEmitter, requestLogger, mockPlanner, mockPlannerContext, mockPrepareResult, mockPlanResult);
    lifecycle.plan();
    Assert.assertEquals(mockPlannerContext, lifecycle.getPlannerContext());
    Assert.assertEquals(SqlLifecycle.State.PLANNED, lifecycle.getState());
    EasyMock.verify(plannerFactory, serviceEmitter, requestLogger, mockPlanner, mockPlannerContext, mockPrepareResult, mockPlanResult);
    EasyMock.reset(plannerFactory, serviceEmitter, requestLogger, mockPlanner, mockPlannerContext, mockPrepareResult, mockPlanResult);
    // test execute
    EasyMock.expect(mockPlanResult.run()).andReturn(Sequences.simple(ImmutableList.of(new Object[] { 2L }))).once();
    EasyMock.replay(plannerFactory, serviceEmitter, requestLogger, mockPlanner, mockPlannerContext, mockPrepareResult, mockPlanResult);
    lifecycle.execute();
    Assert.assertEquals(SqlLifecycle.State.EXECUTING, lifecycle.getState());
    EasyMock.verify(plannerFactory, serviceEmitter, requestLogger, mockPlanner, mockPlannerContext, mockPrepareResult, mockPlanResult);
    EasyMock.reset(plannerFactory, serviceEmitter, requestLogger, mockPlanner, mockPlannerContext, mockPrepareResult, mockPlanResult);
    // test emit
    EasyMock.expect(mockPlannerContext.getSqlQueryId()).andReturn("id").once();
    CopyOnWriteArrayList<String> nativeQueryIds = new CopyOnWriteArrayList<>(ImmutableList.of("id"));
    EasyMock.expect(mockPlannerContext.getNativeQueryIds()).andReturn(nativeQueryIds).times(2);
    EasyMock.expect(mockPlannerContext.getAuthenticationResult()).andReturn(CalciteTests.REGULAR_USER_AUTH_RESULT).once();
    serviceEmitter.emit(EasyMock.anyObject(ServiceEventBuilder.class));
    EasyMock.expectLastCall();
    serviceEmitter.emit(EasyMock.anyObject(ServiceEventBuilder.class));
    EasyMock.expectLastCall();
    requestLogger.logSqlQuery(EasyMock.anyObject());
    EasyMock.expectLastCall();
    EasyMock.replay(plannerFactory, serviceEmitter, requestLogger, mockPlanner, mockPlannerContext, mockPrepareResult, mockPlanResult);
    lifecycle.finalizeStateAndEmitLogsAndMetrics(null, null, 10);
    Assert.assertEquals(SqlLifecycle.State.DONE, lifecycle.getState());
    EasyMock.verify(plannerFactory, serviceEmitter, requestLogger, mockPlanner, mockPlannerContext, mockPrepareResult, mockPlanResult);
    EasyMock.reset(plannerFactory, serviceEmitter, requestLogger, mockPlanner, mockPlannerContext, mockPrepareResult, mockPlanResult);
}

Example 14 with PlannerContext

use of org.apache.druid.sql.calcite.planner.PlannerContext in project druid by druid-io.

the class DruidLogicalValuesRule method onMatch.

@Override
public void onMatch(RelOptRuleCall call) {
    final LogicalValues values = call.rel(0);
    final List<ImmutableList<RexLiteral>> tuples = values.getTuples();
    final List<Object[]> objectTuples = tuples.stream().map(tuple -> tuple.stream().map(v -> getValueFromLiteral(v, plannerContext)).collect(Collectors.toList()).toArray(new Object[0])).collect(Collectors.toList());
    final RowSignature rowSignature = RowSignatures.fromRelDataType(values.getRowType().getFieldNames(), values.getRowType());
    final DruidTable druidTable = new DruidTable(InlineDataSource.fromIterable(objectTuples, rowSignature), rowSignature, null, true, false);
    call.transformTo(DruidQueryRel.scanValues(values, druidTable, plannerContext));
}

Example 15 with PlannerContext

use of org.apache.druid.sql.calcite.planner.PlannerContext in project druid by druid-io.

the class GroupByRules method translateAggregateCall.

/**
 * Translate an AggregateCall to Druid equivalents.
 *
 * @return translated aggregation, or null if translation failed.
 */
public static Aggregation translateAggregateCall(final PlannerContext plannerContext, final RowSignature rowSignature, final VirtualColumnRegistry virtualColumnRegistry, final RexBuilder rexBuilder, final Project project, final List<Aggregation> existingAggregations, final String name, final AggregateCall call, final boolean finalizeAggregations) {
    final DimFilter filter;
    if (call.filterArg >= 0) {
        // AGG(xxx) FILTER(WHERE yyy)
        if (project == null) {
            // We need some kind of projection to support filtered aggregations.
            return null;
        }
        final RexNode expression = project.getChildExps().get(call.filterArg);
        final DimFilter nonOptimizedFilter = Expressions.toFilter(plannerContext, rowSignature, virtualColumnRegistry, expression);
        if (nonOptimizedFilter == null) {
            return null;
        } else {
            filter = Filtration.create(nonOptimizedFilter).optimizeFilterOnly(virtualColumnRegistry.getFullRowSignature()).getDimFilter();
        }
    } else {
        filter = null;
    }
    final SqlAggregator sqlAggregator = plannerContext.getOperatorTable().lookupAggregator(call.getAggregation());
    if (sqlAggregator == null) {
        return null;
    }
    // Compute existingAggregations for SqlAggregator impls that want it.
    final List<Aggregation> existingAggregationsWithSameFilter = new ArrayList<>();
    for (Aggregation existingAggregation : existingAggregations) {
        if (filter == null) {
            final boolean doesMatch = existingAggregation.getAggregatorFactories().stream().noneMatch(factory -> factory instanceof FilteredAggregatorFactory);
            if (doesMatch) {
                existingAggregationsWithSameFilter.add(existingAggregation);
            }
        } else {
            final boolean doesMatch = existingAggregation.getAggregatorFactories().stream().allMatch(factory -> factory instanceof FilteredAggregatorFactory && ((FilteredAggregatorFactory) factory).getFilter().equals(filter));
            if (doesMatch) {
                existingAggregationsWithSameFilter.add(Aggregation.create(existingAggregation.getAggregatorFactories().stream().map(factory -> ((FilteredAggregatorFactory) factory).getAggregator()).collect(Collectors.toList()), existingAggregation.getPostAggregator()));
            }
        }
    }
    final Aggregation retVal = sqlAggregator.toDruidAggregation(plannerContext, rowSignature, virtualColumnRegistry, rexBuilder, name, call, project, existingAggregationsWithSameFilter, finalizeAggregations);
    if (retVal == null) {
        return null;
    } else {
        // Check if this refers to the existingAggregationsWithSameFilter. If so, no need to apply the filter.
        if (isUsingExistingAggregation(retVal, existingAggregationsWithSameFilter)) {
            return retVal;
        } else {
            return retVal.filter(rowSignature, virtualColumnRegistry, filter);
        }
    }
}