Search in sources :

Example 56 with Project

use of org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rel.core.Project in project druid by druid-io.

the class DruidQuery method computeGrouping.

@Nonnull
private static Grouping computeGrouping(final PartialDruidQuery partialQuery, final PlannerContext plannerContext, final RowSignature rowSignature, final VirtualColumnRegistry virtualColumnRegistry, final RexBuilder rexBuilder, final boolean finalizeAggregations) {
    final Aggregate aggregate = Preconditions.checkNotNull(partialQuery.getAggregate(), "aggregate");
    final Project aggregateProject = partialQuery.getAggregateProject();
    final List<DimensionExpression> dimensions = computeDimensions(partialQuery, plannerContext, rowSignature, virtualColumnRegistry);
    final Subtotals subtotals = computeSubtotals(partialQuery, rowSignature);
    final List<Aggregation> aggregations = computeAggregations(partialQuery, plannerContext, rowSignature, virtualColumnRegistry, rexBuilder, finalizeAggregations);
    final RowSignature aggregateRowSignature = RowSignatures.fromRelDataType(ImmutableList.copyOf(Iterators.concat(dimensions.stream().map(DimensionExpression::getOutputName).iterator(), aggregations.stream().map(Aggregation::getOutputName).iterator())), aggregate.getRowType());
    final DimFilter havingFilter = computeHavingFilter(partialQuery, plannerContext, aggregateRowSignature);
    final Grouping grouping = Grouping.create(dimensions, subtotals, aggregations, havingFilter, aggregateRowSignature);
    if (aggregateProject == null) {
        return grouping;
    } else {
        return grouping.applyProject(plannerContext, aggregateProject);
    }
}
Also used : Aggregation(org.apache.druid.sql.calcite.aggregation.Aggregation) Project(org.apache.calcite.rel.core.Project) DimensionExpression(org.apache.druid.sql.calcite.aggregation.DimensionExpression) Aggregate(org.apache.calcite.rel.core.Aggregate) RowSignature(org.apache.druid.segment.column.RowSignature) DimFilter(org.apache.druid.query.filter.DimFilter) Nonnull(javax.annotation.Nonnull)

Example 57 with Project

use of org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rel.core.Project in project druid by druid-io.

the class DruidQuery method computeSorting.

@Nonnull
private static Sorting computeSorting(final PartialDruidQuery partialQuery, final PlannerContext plannerContext, final RowSignature rowSignature, @Nullable final VirtualColumnRegistry virtualColumnRegistry) {
    final Sort sort = Preconditions.checkNotNull(partialQuery.getSort(), "sort");
    final Project sortProject = partialQuery.getSortProject();
    // Extract limit and offset.
    final OffsetLimit offsetLimit = OffsetLimit.fromSort(sort);
    // Extract orderBy column specs.
    final List<OrderByColumnSpec> orderBys = new ArrayList<>(sort.getChildExps().size());
    for (int sortKey = 0; sortKey < sort.getChildExps().size(); sortKey++) {
        final RexNode sortExpression = sort.getChildExps().get(sortKey);
        final RelFieldCollation collation = sort.getCollation().getFieldCollations().get(sortKey);
        final OrderByColumnSpec.Direction direction;
        final StringComparator comparator;
        if (collation.getDirection() == RelFieldCollation.Direction.ASCENDING) {
            direction = OrderByColumnSpec.Direction.ASCENDING;
        } else if (collation.getDirection() == RelFieldCollation.Direction.DESCENDING) {
            direction = OrderByColumnSpec.Direction.DESCENDING;
        } else {
            throw new ISE("Don't know what to do with direction[%s]", collation.getDirection());
        }
        final SqlTypeName sortExpressionType = sortExpression.getType().getSqlTypeName();
        if (SqlTypeName.NUMERIC_TYPES.contains(sortExpressionType) || SqlTypeName.TIMESTAMP == sortExpressionType || SqlTypeName.DATE == sortExpressionType) {
            comparator = StringComparators.NUMERIC;
        } else {
            comparator = StringComparators.LEXICOGRAPHIC;
        }
        if (sortExpression.isA(SqlKind.INPUT_REF)) {
            final RexInputRef ref = (RexInputRef) sortExpression;
            final String fieldName = rowSignature.getColumnName(ref.getIndex());
            orderBys.add(new OrderByColumnSpec(fieldName, direction, comparator));
        } else {
            // We don't support sorting by anything other than refs which actually appear in the query result.
            throw new CannotBuildQueryException(sort, sortExpression);
        }
    }
    // Extract any post-sort Projection.
    final Projection projection;
    if (sortProject == null) {
        projection = null;
    } else if (partialQuery.getAggregate() == null) {
        if (virtualColumnRegistry == null) {
            throw new ISE("Must provide 'virtualColumnRegistry' for pre-aggregation Projection!");
        }
        projection = Projection.preAggregation(sortProject, plannerContext, rowSignature, virtualColumnRegistry);
    } else {
        projection = Projection.postAggregation(sortProject, plannerContext, rowSignature, "s");
    }
    return Sorting.create(orderBys, offsetLimit, projection);
}
Also used : OffsetLimit(org.apache.druid.sql.calcite.planner.OffsetLimit) SqlTypeName(org.apache.calcite.sql.type.SqlTypeName) IntArrayList(it.unimi.dsi.fastutil.ints.IntArrayList) ArrayList(java.util.ArrayList) StringComparator(org.apache.druid.query.ordering.StringComparator) OrderByColumnSpec(org.apache.druid.query.groupby.orderby.OrderByColumnSpec) Project(org.apache.calcite.rel.core.Project) RelFieldCollation(org.apache.calcite.rel.RelFieldCollation) Sort(org.apache.calcite.rel.core.Sort) RexInputRef(org.apache.calcite.rex.RexInputRef) ISE(org.apache.druid.java.util.common.ISE) RexNode(org.apache.calcite.rex.RexNode) Nonnull(javax.annotation.Nonnull)

Example 58 with Project

use of org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rel.core.Project in project druid by druid-io.

the class PartialDruidQuery method withSelectProject.

public PartialDruidQuery withSelectProject(final Project newSelectProject) {
    validateStage(Stage.SELECT_PROJECT);
    // Possibly merge together two projections.
    final Project theProject;
    if (selectProject == null) {
        theProject = newSelectProject;
    } else {
        final List<RexNode> newProjectRexNodes = RelOptUtil.pushPastProject(newSelectProject.getProjects(), selectProject);
        if (RexUtil.isIdentity(newProjectRexNodes, selectProject.getInput().getRowType())) {
            // The projection is gone.
            theProject = null;
        } else {
            final RelBuilder relBuilder = builderSupplier.get();
            relBuilder.push(selectProject.getInput());
            relBuilder.project(newProjectRexNodes, newSelectProject.getRowType().getFieldNames(), true);
            theProject = (Project) relBuilder.build();
        }
    }
    return new PartialDruidQuery(builderSupplier, scan, whereFilter, theProject, aggregate, aggregateProject, havingFilter, sort, sortProject);
}
Also used : Project(org.apache.calcite.rel.core.Project) RelBuilder(org.apache.calcite.tools.RelBuilder) RexNode(org.apache.calcite.rex.RexNode)

Example 59 with Project

use of org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rel.core.Project in project druid by druid-io.

the class DruidJoinRule method onMatch.

@Override
public void onMatch(RelOptRuleCall call) {
    final Join join = call.rel(0);
    final DruidRel<?> left = call.rel(1);
    final DruidRel<?> right = call.rel(2);
    final RexBuilder rexBuilder = join.getCluster().getRexBuilder();
    final DruidRel<?> newLeft;
    final DruidRel<?> newRight;
    final Filter leftFilter;
    final List<RexNode> newProjectExprs = new ArrayList<>();
    // Already verified to be present in "matches", so just call "get".
    // Can't be final, because we're going to reassign it up to a couple of times.
    ConditionAnalysis conditionAnalysis = analyzeCondition(join.getCondition(), join.getLeft().getRowType(), right).get();
    final boolean isLeftDirectAccessPossible = enableLeftScanDirect && (left instanceof DruidQueryRel);
    if (left.getPartialDruidQuery().stage() == PartialDruidQuery.Stage.SELECT_PROJECT && (isLeftDirectAccessPossible || left.getPartialDruidQuery().getWhereFilter() == null)) {
        // Swap the left-side projection above the join, so the left side is a simple scan or mapping. This helps us
        // avoid subqueries.
        final RelNode leftScan = left.getPartialDruidQuery().getScan();
        final Project leftProject = left.getPartialDruidQuery().getSelectProject();
        leftFilter = left.getPartialDruidQuery().getWhereFilter();
        // Left-side projection expressions rewritten to be on top of the join.
        newProjectExprs.addAll(leftProject.getProjects());
        newLeft = left.withPartialQuery(PartialDruidQuery.create(leftScan));
        conditionAnalysis = conditionAnalysis.pushThroughLeftProject(leftProject);
    } else {
        // Leave left as-is. Write input refs that do nothing.
        for (int i = 0; i < left.getRowType().getFieldCount(); i++) {
            newProjectExprs.add(rexBuilder.makeInputRef(join.getRowType().getFieldList().get(i).getType(), i));
        }
        newLeft = left;
        leftFilter = null;
    }
    if (right.getPartialDruidQuery().stage() == PartialDruidQuery.Stage.SELECT_PROJECT && right.getPartialDruidQuery().getWhereFilter() == null && !right.getPartialDruidQuery().getSelectProject().isMapping() && conditionAnalysis.onlyUsesMappingsFromRightProject(right.getPartialDruidQuery().getSelectProject())) {
        // Swap the right-side projection above the join, so the right side is a simple scan or mapping. This helps us
        // avoid subqueries.
        final RelNode rightScan = right.getPartialDruidQuery().getScan();
        final Project rightProject = right.getPartialDruidQuery().getSelectProject();
        // Right-side projection expressions rewritten to be on top of the join.
        for (final RexNode rexNode : RexUtil.shift(rightProject.getProjects(), newLeft.getRowType().getFieldCount())) {
            if (join.getJoinType().generatesNullsOnRight()) {
                newProjectExprs.add(makeNullableIfLiteral(rexNode, rexBuilder));
            } else {
                newProjectExprs.add(rexNode);
            }
        }
        newRight = right.withPartialQuery(PartialDruidQuery.create(rightScan));
        conditionAnalysis = conditionAnalysis.pushThroughRightProject(rightProject);
    } else {
        // Leave right as-is. Write input refs that do nothing.
        for (int i = 0; i < right.getRowType().getFieldCount(); i++) {
            newProjectExprs.add(rexBuilder.makeInputRef(join.getRowType().getFieldList().get(left.getRowType().getFieldCount() + i).getType(), newLeft.getRowType().getFieldCount() + i));
        }
        newRight = right;
    }
    // Druid join written on top of the new left and right sides.
    final DruidJoinQueryRel druidJoin = DruidJoinQueryRel.create(join.copy(join.getTraitSet(), conditionAnalysis.getCondition(rexBuilder), newLeft, newRight, join.getJoinType(), join.isSemiJoinDone()), leftFilter, left.getPlannerContext());
    final RelBuilder relBuilder = call.builder().push(druidJoin).project(RexUtil.fixUp(rexBuilder, newProjectExprs, RelOptUtil.getFieldTypeList(druidJoin.getRowType())));
    call.transformTo(relBuilder.build());
}
Also used : DruidJoinQueryRel(org.apache.druid.sql.calcite.rel.DruidJoinQueryRel) RelBuilder(org.apache.calcite.tools.RelBuilder) ArrayList(java.util.ArrayList) Join(org.apache.calcite.rel.core.Join) DruidQueryRel(org.apache.druid.sql.calcite.rel.DruidQueryRel) Project(org.apache.calcite.rel.core.Project) RelNode(org.apache.calcite.rel.RelNode) Filter(org.apache.calcite.rel.core.Filter) RexBuilder(org.apache.calcite.rex.RexBuilder) RexNode(org.apache.calcite.rex.RexNode)

Example 60 with Project

use of org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rel.core.Project 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)

Aggregations

Project (org.apache.calcite.rel.core.Project)143 RexNode (org.apache.calcite.rex.RexNode)77 RelNode (org.apache.calcite.rel.RelNode)71 ArrayList (java.util.ArrayList)51 LogicalProject (org.apache.calcite.rel.logical.LogicalProject)35 RexBuilder (org.apache.calcite.rex.RexBuilder)28 RelDataType (org.apache.calcite.rel.type.RelDataType)26 Test (org.junit.Test)23 Aggregate (org.apache.calcite.rel.core.Aggregate)22 Filter (org.apache.calcite.rel.core.Filter)22 Join (org.apache.calcite.rel.core.Join)22 List (java.util.List)20 RexLiteral (org.apache.calcite.rex.RexLiteral)19 AggregateCall (org.apache.calcite.rel.core.AggregateCall)18 Sort (org.apache.calcite.rel.core.Sort)18 RelBuilder (org.apache.calcite.tools.RelBuilder)17 ImmutableBitSet (org.apache.calcite.util.ImmutableBitSet)16 HiveProject (org.apache.hadoop.hive.ql.optimizer.calcite.reloperators.HiveProject)16 Collectors (java.util.stream.Collectors)15 RelDataTypeField (org.apache.calcite.rel.type.RelDataTypeField)15