Examples with RelCollation org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rel.RelCollation used on opensource projects

Example 51 with RelCollation

use of org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rel.RelCollation in project flink by apache.

the class HiveParserCalcitePlanner method genDistSortBy.

// Generate plan for sort by, cluster by and distribute by. This is basically same as generating
// order by plan.
// Should refactor to combine them.
private Pair<RelNode, RelNode> genDistSortBy(HiveParserQB qb, RelNode srcRel, boolean outermostOB) throws SemanticException {
    RelNode res = null;
    RelNode originalInput = null;
    HiveParserQBParseInfo qbp = qb.getParseInfo();
    String destClause = qbp.getClauseNames().iterator().next();
    HiveParserASTNode sortAST = qbp.getSortByForClause(destClause);
    HiveParserASTNode distAST = qbp.getDistributeByForClause(destClause);
    HiveParserASTNode clusterAST = qbp.getClusterByForClause(destClause);
    if (sortAST != null || distAST != null || clusterAST != null) {
        List<RexNode> virtualCols = new ArrayList<>();
        List<Pair<HiveParserASTNode, TypeInfo>> vcASTAndType = new ArrayList<>();
        List<RelFieldCollation> fieldCollations = new ArrayList<>();
        List<Integer> distKeys = new ArrayList<>();
        HiveParserRowResolver inputRR = relToRowResolver.get(srcRel);
        HiveParserRexNodeConverter converter = new HiveParserRexNodeConverter(cluster, srcRel.getRowType(), relToHiveColNameCalcitePosMap.get(srcRel), 0, false, funcConverter);
        int numSrcFields = srcRel.getRowType().getFieldCount();
        // handle cluster by
        if (clusterAST != null) {
            if (sortAST != null) {
                throw new SemanticException("Cannot have both CLUSTER BY and SORT BY");
            }
            if (distAST != null) {
                throw new SemanticException("Cannot have both CLUSTER BY and DISTRIBUTE BY");
            }
            for (Node node : clusterAST.getChildren()) {
                HiveParserASTNode childAST = (HiveParserASTNode) node;
                Map<HiveParserASTNode, ExprNodeDesc> astToExprNodeDesc = semanticAnalyzer.genAllExprNodeDesc(childAST, inputRR);
                ExprNodeDesc childNodeDesc = astToExprNodeDesc.get(childAST);
                if (childNodeDesc == null) {
                    throw new SemanticException("Invalid CLUSTER BY expression: " + childAST.toString());
                }
                RexNode childRexNode = converter.convert(childNodeDesc).accept(funcConverter);
                int fieldIndex;
                if (childRexNode instanceof RexInputRef) {
                    fieldIndex = ((RexInputRef) childRexNode).getIndex();
                } else {
                    fieldIndex = numSrcFields + virtualCols.size();
                    virtualCols.add(childRexNode);
                    vcASTAndType.add(new Pair<>(childAST, childNodeDesc.getTypeInfo()));
                }
                // cluster by doesn't support specifying ASC/DESC or NULLS FIRST/LAST, so use
                // default values
                fieldCollations.add(new RelFieldCollation(fieldIndex, RelFieldCollation.Direction.ASCENDING, RelFieldCollation.NullDirection.FIRST));
                distKeys.add(fieldIndex);
            }
        } else {
            // handle sort by
            if (sortAST != null) {
                for (Node node : sortAST.getChildren()) {
                    HiveParserASTNode childAST = (HiveParserASTNode) node;
                    HiveParserASTNode nullOrderAST = (HiveParserASTNode) childAST.getChild(0);
                    HiveParserASTNode fieldAST = (HiveParserASTNode) nullOrderAST.getChild(0);
                    Map<HiveParserASTNode, ExprNodeDesc> astToExprNodeDesc = semanticAnalyzer.genAllExprNodeDesc(fieldAST, inputRR);
                    ExprNodeDesc fieldNodeDesc = astToExprNodeDesc.get(fieldAST);
                    if (fieldNodeDesc == null) {
                        throw new SemanticException("Invalid sort by expression: " + fieldAST.toString());
                    }
                    RexNode childRexNode = converter.convert(fieldNodeDesc).accept(funcConverter);
                    int fieldIndex;
                    if (childRexNode instanceof RexInputRef) {
                        fieldIndex = ((RexInputRef) childRexNode).getIndex();
                    } else {
                        fieldIndex = numSrcFields + virtualCols.size();
                        virtualCols.add(childRexNode);
                        vcASTAndType.add(new Pair<>(childAST, fieldNodeDesc.getTypeInfo()));
                    }
                    RelFieldCollation.Direction direction = RelFieldCollation.Direction.DESCENDING;
                    if (childAST.getType() == HiveASTParser.TOK_TABSORTCOLNAMEASC) {
                        direction = RelFieldCollation.Direction.ASCENDING;
                    }
                    RelFieldCollation.NullDirection nullOrder;
                    if (nullOrderAST.getType() == HiveASTParser.TOK_NULLS_FIRST) {
                        nullOrder = RelFieldCollation.NullDirection.FIRST;
                    } else if (nullOrderAST.getType() == HiveASTParser.TOK_NULLS_LAST) {
                        nullOrder = RelFieldCollation.NullDirection.LAST;
                    } else {
                        throw new SemanticException("Unexpected null ordering option: " + nullOrderAST.getType());
                    }
                    fieldCollations.add(new RelFieldCollation(fieldIndex, direction, nullOrder));
                }
            }
            // handle distribute by
            if (distAST != null) {
                for (Node node : distAST.getChildren()) {
                    HiveParserASTNode childAST = (HiveParserASTNode) node;
                    Map<HiveParserASTNode, ExprNodeDesc> astToExprNodeDesc = semanticAnalyzer.genAllExprNodeDesc(childAST, inputRR);
                    ExprNodeDesc childNodeDesc = astToExprNodeDesc.get(childAST);
                    if (childNodeDesc == null) {
                        throw new SemanticException("Invalid DISTRIBUTE BY expression: " + childAST.toString());
                    }
                    RexNode childRexNode = converter.convert(childNodeDesc).accept(funcConverter);
                    int fieldIndex;
                    if (childRexNode instanceof RexInputRef) {
                        fieldIndex = ((RexInputRef) childRexNode).getIndex();
                    } else {
                        fieldIndex = numSrcFields + virtualCols.size();
                        virtualCols.add(childRexNode);
                        vcASTAndType.add(new Pair<>(childAST, childNodeDesc.getTypeInfo()));
                    }
                    distKeys.add(fieldIndex);
                }
            }
        }
        Preconditions.checkState(!fieldCollations.isEmpty() || !distKeys.isEmpty(), "Both field collations and dist keys are empty");
        // add child SEL if needed
        RelNode realInput = srcRel;
        HiveParserRowResolver outputRR = new HiveParserRowResolver();
        if (!virtualCols.isEmpty()) {
            List<RexNode> originalInputRefs = srcRel.getRowType().getFieldList().stream().map(input -> new RexInputRef(input.getIndex(), input.getType())).collect(Collectors.toList());
            HiveParserRowResolver addedProjectRR = new HiveParserRowResolver();
            if (!HiveParserRowResolver.add(addedProjectRR, inputRR)) {
                throw new SemanticException("Duplicates detected when adding columns to RR: see previous message");
            }
            int vColPos = inputRR.getRowSchema().getSignature().size();
            for (Pair<HiveParserASTNode, TypeInfo> astTypePair : vcASTAndType) {
                addedProjectRR.putExpression(astTypePair.getKey(), new ColumnInfo(getColumnInternalName(vColPos), astTypePair.getValue(), null, false));
                vColPos++;
            }
            realInput = genSelectRelNode(CompositeList.of(originalInputRefs, virtualCols), addedProjectRR, srcRel);
            if (outermostOB) {
                if (!HiveParserRowResolver.add(outputRR, inputRR)) {
                    throw new SemanticException("Duplicates detected when adding columns to RR: see previous message");
                }
            } else {
                if (!HiveParserRowResolver.add(outputRR, addedProjectRR)) {
                    throw new SemanticException("Duplicates detected when adding columns to RR: see previous message");
                }
            }
            originalInput = srcRel;
        } else {
            if (!HiveParserRowResolver.add(outputRR, inputRR)) {
                throw new SemanticException("Duplicates detected when adding columns to RR: see previous message");
            }
        }
        // create rel node
        RelTraitSet traitSet = cluster.traitSet();
        RelCollation canonizedCollation = traitSet.canonize(RelCollationImpl.of(fieldCollations));
        res = LogicalDistribution.create(realInput, canonizedCollation, distKeys);
        Map<String, Integer> hiveColNameCalcitePosMap = buildHiveToCalciteColumnMap(outputRR);
        relToRowResolver.put(res, outputRR);
        relToHiveColNameCalcitePosMap.put(res, hiveColNameCalcitePosMap);
    }
    return (new Pair<>(res, originalInput));
}

Example 52 with RelCollation

use of org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rel.RelCollation in project calcite by apache.

the class EnumerableValues method create.

/**
 * Creates an EnumerableValues.
 */
public static EnumerableValues create(RelOptCluster cluster, final RelDataType rowType, final ImmutableList<ImmutableList<RexLiteral>> tuples) {
    final RelMetadataQuery mq = cluster.getMetadataQuery();
    final RelTraitSet traitSet = cluster.traitSetOf(EnumerableConvention.INSTANCE).replaceIfs(RelCollationTraitDef.INSTANCE, new Supplier<List<RelCollation>>() {

        public List<RelCollation> get() {
            return RelMdCollation.values(mq, rowType, tuples);
        }
    }).replaceIf(RelDistributionTraitDef.INSTANCE, new Supplier<RelDistribution>() {

        public RelDistribution get() {
            return RelMdDistribution.values(rowType, tuples);
        }
    });
    return new EnumerableValues(cluster, rowType, tuples, traitSet);
}

Example 53 with RelCollation

use of org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rel.RelCollation in project calcite by apache.

the class EnumerableFilter method create.

/**
 * Creates an EnumerableFilter.
 */
public static EnumerableFilter create(final RelNode input, RexNode condition) {
    final RelOptCluster cluster = input.getCluster();
    final RelMetadataQuery mq = cluster.getMetadataQuery();
    final RelTraitSet traitSet = cluster.traitSetOf(EnumerableConvention.INSTANCE).replaceIfs(RelCollationTraitDef.INSTANCE, new Supplier<List<RelCollation>>() {

        public List<RelCollation> get() {
            return RelMdCollation.filter(mq, input);
        }
    }).replaceIf(RelDistributionTraitDef.INSTANCE, new Supplier<RelDistribution>() {

        public RelDistribution get() {
            return RelMdDistribution.filter(mq, input);
        }
    });
    return new EnumerableFilter(cluster, traitSet, input, condition);
}

Example 54 with RelCollation

use of org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rel.RelCollation in project calcite by apache.

the class EnumerableLimit method create.

/**
 * Creates an EnumerableLimit.
 */
public static EnumerableLimit create(final RelNode input, RexNode offset, RexNode fetch) {
    final RelOptCluster cluster = input.getCluster();
    final RelMetadataQuery mq = cluster.getMetadataQuery();
    final RelTraitSet traitSet = cluster.traitSetOf(EnumerableConvention.INSTANCE).replaceIfs(RelCollationTraitDef.INSTANCE, new Supplier<List<RelCollation>>() {

        public List<RelCollation> get() {
            return RelMdCollation.limit(mq, input);
        }
    }).replaceIf(RelDistributionTraitDef.INSTANCE, new Supplier<RelDistribution>() {

        public RelDistribution get() {
            return RelMdDistribution.limit(mq, input);
        }
    });
    return new EnumerableLimit(cluster, traitSet, input, offset, fetch);
}

Example 55 with RelCollation

use of org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rel.RelCollation in project calcite by apache.

the class EnumerableMergeJoin method create.

public static EnumerableMergeJoin create(RelNode left, RelNode right, RexLiteral condition, ImmutableIntList leftKeys, ImmutableIntList rightKeys, JoinRelType joinType) throws InvalidRelException {
    final RelOptCluster cluster = right.getCluster();
    RelTraitSet traitSet = cluster.traitSet();
    if (traitSet.isEnabled(RelCollationTraitDef.INSTANCE)) {
        final RelMetadataQuery mq = cluster.getMetadataQuery();
        final List<RelCollation> collations = RelMdCollation.mergeJoin(mq, left, right, leftKeys, rightKeys);
        traitSet = traitSet.replace(collations);
    }
    return new EnumerableMergeJoin(cluster, traitSet, left, right, condition, leftKeys, rightKeys, ImmutableSet.<CorrelationId>of(), joinType);
}