Search in sources :

Example 1 with PlannerContext

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

the class ExternalTableScanRuleTest method testMatchesWhenExternalScanUnsupported.

@Test
public void testMatchesWhenExternalScanUnsupported() throws ValidationException {
    final PlannerContext plannerContext = PlannerContext.create(// The actual query isn't important for this test
    "DUMMY", CalciteTests.createOperatorTable(), CalciteTests.createExprMacroTable(), CalciteTests.getJsonMapper(), new PlannerConfig(), new DruidSchemaCatalog(EasyMock.createMock(SchemaPlus.class), ImmutableMap.of("druid", new NamedDruidSchema(EasyMock.createMock(DruidSchema.class), "druid"), NamedViewSchema.NAME, new NamedViewSchema(EasyMock.createMock(ViewSchema.class)))), ImmutableMap.of());
    plannerContext.setQueryMaker(CalciteTests.createMockQueryMakerFactory(EasyMock.createMock(QuerySegmentWalker.class), EasyMock.createMock(QueryRunnerFactoryConglomerate.class)).buildForSelect(EasyMock.createMock(RelRoot.class), plannerContext));
    ExternalTableScanRule rule = new ExternalTableScanRule(plannerContext);
    rule.matches(EasyMock.createMock(RelOptRuleCall.class));
    Assert.assertEquals("SQL query requires scanning external datasources that is not suported.", plannerContext.getPlanningError());
}
Also used : QueryRunnerFactoryConglomerate(org.apache.druid.query.QueryRunnerFactoryConglomerate) PlannerContext(org.apache.druid.sql.calcite.planner.PlannerContext) NamedViewSchema(org.apache.druid.sql.calcite.schema.NamedViewSchema) PlannerConfig(org.apache.druid.sql.calcite.planner.PlannerConfig) DruidSchemaCatalog(org.apache.druid.sql.calcite.schema.DruidSchemaCatalog) RelOptRuleCall(org.apache.calcite.plan.RelOptRuleCall) NamedDruidSchema(org.apache.druid.sql.calcite.schema.NamedDruidSchema) QuerySegmentWalker(org.apache.druid.query.QuerySegmentWalker) Test(org.junit.Test)

Example 2 with PlannerContext

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

the class DruidJoinRuleTest method setup.

@Before
public void setup() {
    PlannerContext plannerContext = Mockito.mock(PlannerContext.class);
    druidJoinRule = DruidJoinRule.instance(plannerContext);
}
Also used : PlannerContext(org.apache.druid.sql.calcite.planner.PlannerContext) Before(org.junit.Before)

Example 3 with PlannerContext

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

the class ArrayConcatSqlAggregator method toDruidAggregation.

@Nullable
@Override
public Aggregation toDruidAggregation(PlannerContext plannerContext, RowSignature rowSignature, VirtualColumnRegistry virtualColumnRegistry, RexBuilder rexBuilder, String name, AggregateCall aggregateCall, Project project, List<Aggregation> existingAggregations, boolean finalizeAggregations) {
    final List<RexNode> arguments = aggregateCall.getArgList().stream().map(i -> Expressions.fromFieldAccess(rowSignature, project, i)).collect(Collectors.toList());
    Integer maxSizeBytes = null;
    if (arguments.size() > 1) {
        RexNode maxBytes = arguments.get(1);
        if (!maxBytes.isA(SqlKind.LITERAL)) {
            // maxBytes must be a literal
            return null;
        }
        maxSizeBytes = ((Number) RexLiteral.value(maxBytes)).intValue();
    }
    final DruidExpression arg = Expressions.toDruidExpression(plannerContext, rowSignature, arguments.get(0));
    final ExprMacroTable macroTable = plannerContext.getExprMacroTable();
    final String fieldName;
    final ColumnType druidType = Calcites.getValueTypeForRelDataTypeFull(aggregateCall.getType());
    if (druidType == null || !druidType.isArray()) {
        // must be an array
        return null;
    }
    final String initialvalue = ExpressionType.fromColumnTypeStrict(druidType).asTypeString() + "[]";
    if (arg.isDirectColumnAccess()) {
        fieldName = arg.getDirectColumn();
    } else {
        VirtualColumn vc = virtualColumnRegistry.getOrCreateVirtualColumnForExpression(plannerContext, arg, druidType);
        fieldName = vc.getOutputName();
    }
    if (aggregateCall.isDistinct()) {
        return Aggregation.create(new ExpressionLambdaAggregatorFactory(name, ImmutableSet.of(fieldName), null, initialvalue, null, true, false, false, StringUtils.format("array_set_add_all(\"__acc\", \"%s\")", fieldName), StringUtils.format("array_set_add_all(\"__acc\", \"%s\")", name), null, null, maxSizeBytes != null ? new HumanReadableBytes(maxSizeBytes) : null, macroTable));
    } else {
        return Aggregation.create(new ExpressionLambdaAggregatorFactory(name, ImmutableSet.of(fieldName), null, initialvalue, null, true, false, false, StringUtils.format("array_concat(\"__acc\", \"%s\")", fieldName), StringUtils.format("array_concat(\"__acc\", \"%s\")", name), null, null, maxSizeBytes != null ? new HumanReadableBytes(maxSizeBytes) : null, macroTable));
    }
}
Also used : Project(org.apache.calcite.rel.core.Project) SqlAggregator(org.apache.druid.sql.calcite.aggregation.SqlAggregator) ReturnTypes(org.apache.calcite.sql.type.ReturnTypes) DruidExpression(org.apache.druid.sql.calcite.expression.DruidExpression) HumanReadableBytes(org.apache.druid.java.util.common.HumanReadableBytes) Optionality(org.apache.calcite.util.Optionality) RexNode(org.apache.calcite.rex.RexNode) ExpressionType(org.apache.druid.math.expr.ExpressionType) VirtualColumnRegistry(org.apache.druid.sql.calcite.rel.VirtualColumnRegistry) PlannerContext(org.apache.druid.sql.calcite.planner.PlannerContext) Nullable(javax.annotation.Nullable) ImmutableSet(com.google.common.collect.ImmutableSet) SqlKind(org.apache.calcite.sql.SqlKind) ExpressionLambdaAggregatorFactory(org.apache.druid.query.aggregation.ExpressionLambdaAggregatorFactory) InferTypes(org.apache.calcite.sql.type.InferTypes) RexBuilder(org.apache.calcite.rex.RexBuilder) RexLiteral(org.apache.calcite.rex.RexLiteral) VirtualColumn(org.apache.druid.segment.VirtualColumn) SqlFunctionCategory(org.apache.calcite.sql.SqlFunctionCategory) StringUtils(org.apache.druid.java.util.common.StringUtils) Aggregation(org.apache.druid.sql.calcite.aggregation.Aggregation) Collectors(java.util.stream.Collectors) ExprMacroTable(org.apache.druid.math.expr.ExprMacroTable) List(java.util.List) RowSignature(org.apache.druid.segment.column.RowSignature) OperandTypes(org.apache.calcite.sql.type.OperandTypes) ColumnType(org.apache.druid.segment.column.ColumnType) AggregateCall(org.apache.calcite.rel.core.AggregateCall) SqlAggFunction(org.apache.calcite.sql.SqlAggFunction) Calcites(org.apache.druid.sql.calcite.planner.Calcites) Expressions(org.apache.druid.sql.calcite.expression.Expressions) ColumnType(org.apache.druid.segment.column.ColumnType) ExpressionLambdaAggregatorFactory(org.apache.druid.query.aggregation.ExpressionLambdaAggregatorFactory) DruidExpression(org.apache.druid.sql.calcite.expression.DruidExpression) VirtualColumn(org.apache.druid.segment.VirtualColumn) HumanReadableBytes(org.apache.druid.java.util.common.HumanReadableBytes) ExprMacroTable(org.apache.druid.math.expr.ExprMacroTable) RexNode(org.apache.calcite.rex.RexNode) Nullable(javax.annotation.Nullable)

Example 4 with PlannerContext

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

the class TimeArithmeticOperatorConversion method toDruidExpression.

@Override
public DruidExpression toDruidExpression(final PlannerContext plannerContext, final RowSignature rowSignature, final RexNode rexNode) {
    final RexCall call = (RexCall) rexNode;
    final List<RexNode> operands = call.getOperands();
    if (operands.size() != 2) {
        throw new IAE("Expected 2 args, got %s", operands.size());
    }
    final RexNode leftRexNode = operands.get(0);
    final RexNode rightRexNode = operands.get(1);
    final DruidExpression leftExpr = Expressions.toDruidExpression(plannerContext, rowSignature, leftRexNode);
    final DruidExpression rightExpr = Expressions.toDruidExpression(plannerContext, rowSignature, rightRexNode);
    if (leftExpr == null || rightExpr == null) {
        return null;
    }
    final ColumnType outputType = Calcites.getColumnTypeForRelDataType(rexNode.getType());
    if (rightRexNode.getType().getFamily() == SqlTypeFamily.INTERVAL_YEAR_MONTH) {
        // Period is a value in months.
        return DruidExpression.ofExpression(outputType, DruidExpression.functionCall("timestamp_shift"), ImmutableList.of(leftExpr, rightExpr.map(simpleExtraction -> null, expression -> rightRexNode.isA(SqlKind.LITERAL) ? StringUtils.format("'P%sM'", RexLiteral.value(rightRexNode)) : StringUtils.format("concat('P', %s, 'M')", expression)), DruidExpression.ofLiteral(ColumnType.LONG, DruidExpression.numberLiteral(direction > 0 ? 1 : -1)), DruidExpression.ofStringLiteral(plannerContext.getTimeZone().getID())));
    } else if (rightRexNode.getType().getFamily() == SqlTypeFamily.INTERVAL_DAY_TIME) {
        // Period is a value in milliseconds. Ignore time zone.
        return DruidExpression.ofExpression(outputType, (args) -> StringUtils.format("(%s %s %s)", args.get(0).getExpression(), direction > 0 ? "+" : "-", args.get(1).getExpression()), ImmutableList.of(leftExpr, rightExpr));
    } else if ((leftRexNode.getType().getFamily() == SqlTypeFamily.TIMESTAMP || leftRexNode.getType().getFamily() == SqlTypeFamily.DATE) && (rightRexNode.getType().getFamily() == SqlTypeFamily.TIMESTAMP || rightRexNode.getType().getFamily() == SqlTypeFamily.DATE)) {
        // Calcite represents both TIMESTAMP - INTERVAL and TIMESTAMPDIFF (TIMESTAMP - TIMESTAMP)
        // with a MINUS_DATE operator, so we must tell which case we're in by checking the type of
        // the second argument.
        Preconditions.checkState(direction < 0, "Time arithmetic require direction < 0");
        if (call.getType().getFamily() == SqlTypeFamily.INTERVAL_YEAR_MONTH) {
            return DruidExpression.ofExpression(outputType, DruidExpression.functionCall("subtract_months"), ImmutableList.of(leftExpr, rightExpr, DruidExpression.ofStringLiteral(plannerContext.getTimeZone().getID())));
        } else {
            return DruidExpression.ofExpression(outputType, (args) -> StringUtils.format("(%s %s %s)", args.get(0).getExpression(), "-", args.get(1).getExpression()), ImmutableList.of(leftExpr, rightExpr));
        }
    } else {
        // Shouldn't happen if subclasses are behaving.
        throw new ISE("Got unexpected type period type family[%s]", rightRexNode.getType().getFamily());
    }
}
Also used : RexCall(org.apache.calcite.rex.RexCall) SqlKind(org.apache.calcite.sql.SqlKind) SqlTypeFamily(org.apache.calcite.sql.type.SqlTypeFamily) SqlOperatorConversion(org.apache.druid.sql.calcite.expression.SqlOperatorConversion) RexLiteral(org.apache.calcite.rex.RexLiteral) StringUtils(org.apache.druid.java.util.common.StringUtils) ISE(org.apache.druid.java.util.common.ISE) DruidExpression(org.apache.druid.sql.calcite.expression.DruidExpression) List(java.util.List) ImmutableList(com.google.common.collect.ImmutableList) SqlStdOperatorTable(org.apache.calcite.sql.fun.SqlStdOperatorTable) RexNode(org.apache.calcite.rex.RexNode) RowSignature(org.apache.druid.segment.column.RowSignature) PlannerContext(org.apache.druid.sql.calcite.planner.PlannerContext) ColumnType(org.apache.druid.segment.column.ColumnType) Preconditions(com.google.common.base.Preconditions) SqlOperator(org.apache.calcite.sql.SqlOperator) IAE(org.apache.druid.java.util.common.IAE) Calcites(org.apache.druid.sql.calcite.planner.Calcites) RexCall(org.apache.calcite.rex.RexCall) Expressions(org.apache.druid.sql.calcite.expression.Expressions) ColumnType(org.apache.druid.segment.column.ColumnType) DruidExpression(org.apache.druid.sql.calcite.expression.DruidExpression) ISE(org.apache.druid.java.util.common.ISE) IAE(org.apache.druid.java.util.common.IAE) RexNode(org.apache.calcite.rex.RexNode)

Example 5 with PlannerContext

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

the class QuantileSqlAggregator 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 DruidExpression input = Aggregations.toDruidExpressionForNumericAggregator(plannerContext, rowSignature, Expressions.fromFieldAccess(rowSignature, project, aggregateCall.getArgList().get(0)));
    if (input == null) {
        return null;
    }
    final AggregatorFactory aggregatorFactory;
    final String histogramName = StringUtils.format("%s:agg", name);
    final RexNode probabilityArg = Expressions.fromFieldAccess(rowSignature, project, aggregateCall.getArgList().get(1));
    if (!probabilityArg.isA(SqlKind.LITERAL)) {
        // Probability must be a literal in order to plan.
        return null;
    }
    final float probability = ((Number) RexLiteral.value(probabilityArg)).floatValue();
    final int resolution;
    if (aggregateCall.getArgList().size() >= 3) {
        final RexNode resolutionArg = Expressions.fromFieldAccess(rowSignature, project, aggregateCall.getArgList().get(2));
        if (!resolutionArg.isA(SqlKind.LITERAL)) {
            // Resolution must be a literal in order to plan.
            return null;
        }
        resolution = ((Number) RexLiteral.value(resolutionArg)).intValue();
    } else {
        resolution = ApproximateHistogram.DEFAULT_HISTOGRAM_SIZE;
    }
    final int numBuckets = ApproximateHistogram.DEFAULT_BUCKET_SIZE;
    final float lowerLimit = Float.NEGATIVE_INFINITY;
    final float upperLimit = Float.POSITIVE_INFINITY;
    // Look for existing matching aggregatorFactory.
    for (final Aggregation existing : existingAggregations) {
        for (AggregatorFactory factory : existing.getAggregatorFactories()) {
            if (factory instanceof ApproximateHistogramAggregatorFactory) {
                final ApproximateHistogramAggregatorFactory theFactory = (ApproximateHistogramAggregatorFactory) factory;
                // Check input for equivalence.
                final boolean inputMatches;
                final DruidExpression virtualInput = virtualColumnRegistry.findVirtualColumnExpressions(theFactory.requiredFields()).stream().findFirst().orElse(null);
                if (virtualInput == null) {
                    inputMatches = input.isDirectColumnAccess() && input.getDirectColumn().equals(theFactory.getFieldName());
                } else {
                    inputMatches = virtualInput.equals(input);
                }
                final boolean matches = inputMatches && theFactory.getResolution() == resolution && theFactory.getNumBuckets() == numBuckets && theFactory.getLowerLimit() == lowerLimit && theFactory.getUpperLimit() == upperLimit;
                if (matches) {
                    // Found existing one. Use this.
                    return Aggregation.create(ImmutableList.of(), new QuantilePostAggregator(name, factory.getName(), probability));
                }
            }
        }
    }
    // No existing match found. Create a new one.
    if (input.isDirectColumnAccess()) {
        if (rowSignature.getColumnType(input.getDirectColumn()).map(type -> type.is(ValueType.COMPLEX)).orElse(false)) {
            aggregatorFactory = new ApproximateHistogramFoldingAggregatorFactory(histogramName, input.getDirectColumn(), resolution, numBuckets, lowerLimit, upperLimit, false);
        } else {
            aggregatorFactory = new ApproximateHistogramAggregatorFactory(histogramName, input.getDirectColumn(), resolution, numBuckets, lowerLimit, upperLimit, false);
        }
    } else {
        final String virtualColumnName = virtualColumnRegistry.getOrCreateVirtualColumnForExpression(input, ColumnType.FLOAT);
        aggregatorFactory = new ApproximateHistogramAggregatorFactory(histogramName, virtualColumnName, resolution, numBuckets, lowerLimit, upperLimit, false);
    }
    return Aggregation.create(ImmutableList.of(aggregatorFactory), new QuantilePostAggregator(name, histogramName, probability));
}
Also used : Project(org.apache.calcite.rel.core.Project) SqlAggregator(org.apache.druid.sql.calcite.aggregation.SqlAggregator) ReturnTypes(org.apache.calcite.sql.type.ReturnTypes) QuantilePostAggregator(org.apache.druid.query.aggregation.histogram.QuantilePostAggregator) DruidExpression(org.apache.druid.sql.calcite.expression.DruidExpression) ApproximateHistogram(org.apache.druid.query.aggregation.histogram.ApproximateHistogram) ImmutableList(com.google.common.collect.ImmutableList) RexNode(org.apache.calcite.rex.RexNode) VirtualColumnRegistry(org.apache.druid.sql.calcite.rel.VirtualColumnRegistry) PlannerContext(org.apache.druid.sql.calcite.planner.PlannerContext) Nullable(javax.annotation.Nullable) SqlKind(org.apache.calcite.sql.SqlKind) SqlTypeFamily(org.apache.calcite.sql.type.SqlTypeFamily) SqlTypeName(org.apache.calcite.sql.type.SqlTypeName) RexBuilder(org.apache.calcite.rex.RexBuilder) RexLiteral(org.apache.calcite.rex.RexLiteral) AggregatorFactory(org.apache.druid.query.aggregation.AggregatorFactory) SqlFunctionCategory(org.apache.calcite.sql.SqlFunctionCategory) ApproximateHistogramFoldingAggregatorFactory(org.apache.druid.query.aggregation.histogram.ApproximateHistogramFoldingAggregatorFactory) StringUtils(org.apache.druid.java.util.common.StringUtils) ValueType(org.apache.druid.segment.column.ValueType) Aggregation(org.apache.druid.sql.calcite.aggregation.Aggregation) ApproximateHistogramAggregatorFactory(org.apache.druid.query.aggregation.histogram.ApproximateHistogramAggregatorFactory) List(java.util.List) Aggregations(org.apache.druid.sql.calcite.aggregation.Aggregations) RowSignature(org.apache.druid.segment.column.RowSignature) OperandTypes(org.apache.calcite.sql.type.OperandTypes) ColumnType(org.apache.druid.segment.column.ColumnType) AggregateCall(org.apache.calcite.rel.core.AggregateCall) SqlAggFunction(org.apache.calcite.sql.SqlAggFunction) Expressions(org.apache.druid.sql.calcite.expression.Expressions) ApproximateHistogramFoldingAggregatorFactory(org.apache.druid.query.aggregation.histogram.ApproximateHistogramFoldingAggregatorFactory) QuantilePostAggregator(org.apache.druid.query.aggregation.histogram.QuantilePostAggregator) AggregatorFactory(org.apache.druid.query.aggregation.AggregatorFactory) ApproximateHistogramFoldingAggregatorFactory(org.apache.druid.query.aggregation.histogram.ApproximateHistogramFoldingAggregatorFactory) ApproximateHistogramAggregatorFactory(org.apache.druid.query.aggregation.histogram.ApproximateHistogramAggregatorFactory) Aggregation(org.apache.druid.sql.calcite.aggregation.Aggregation) DruidExpression(org.apache.druid.sql.calcite.expression.DruidExpression) ApproximateHistogramAggregatorFactory(org.apache.druid.query.aggregation.histogram.ApproximateHistogramAggregatorFactory) RexNode(org.apache.calcite.rex.RexNode) Nullable(javax.annotation.Nullable)

Aggregations

PlannerContext (org.apache.druid.sql.calcite.planner.PlannerContext)15 List (java.util.List)11 RowSignature (org.apache.druid.segment.column.RowSignature)11 Expressions (org.apache.druid.sql.calcite.expression.Expressions)10 Nullable (javax.annotation.Nullable)9 RexLiteral (org.apache.calcite.rex.RexLiteral)9 RexNode (org.apache.calcite.rex.RexNode)9 SqlKind (org.apache.calcite.sql.SqlKind)9 ColumnType (org.apache.druid.segment.column.ColumnType)9 DruidExpression (org.apache.druid.sql.calcite.expression.DruidExpression)9 Collectors (java.util.stream.Collectors)8 AggregateCall (org.apache.calcite.rel.core.AggregateCall)8 Project (org.apache.calcite.rel.core.Project)8 RexBuilder (org.apache.calcite.rex.RexBuilder)8 SqlFunctionCategory (org.apache.calcite.sql.SqlFunctionCategory)8 Aggregation (org.apache.druid.sql.calcite.aggregation.Aggregation)8 SqlAggregator (org.apache.druid.sql.calcite.aggregation.SqlAggregator)8 Calcites (org.apache.druid.sql.calcite.planner.Calcites)8 VirtualColumnRegistry (org.apache.druid.sql.calcite.rel.VirtualColumnRegistry)8 SqlAggFunction (org.apache.calcite.sql.SqlAggFunction)7