Examples with RexShuttle org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rex.RexShuttle used on opensource projects

Example 11 with RexShuttle

use of org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rex.RexShuttle in project hive by apache.

the class FilterFlattenCorrelatedConditionRule method onMatch.

@Override
public void onMatch(RelOptRuleCall call) {
    Filter filter = call.rel(0);
    RelBuilder b = call.builder();
    b.push(filter.getInput());
    final int proj = b.fields().size();
    List<RexNode> projOperands = new ArrayList<>();
    RexNode newCondition = filter.getCondition().accept(new RexShuttle() {

        @Override
        public RexNode visitCall(RexCall call) {
            switch(call.getKind()) {
                case EQUALS:
                case NOT_EQUALS:
                case GREATER_THAN:
                case GREATER_THAN_OR_EQUAL:
                case LESS_THAN:
                case LESS_THAN_OR_EQUAL:
                    RexNode op0 = call.operands.get(0);
                    RexNode op1 = call.operands.get(1);
                    final int replaceIndex;
                    if (RexUtil.containsCorrelation(op1) && isUncorrelatedCall(op0)) {
                        replaceIndex = 0;
                    } else if (RexUtil.containsCorrelation(op0) && isUncorrelatedCall(op1)) {
                        replaceIndex = 1;
                    } else {
                        // Structure does not match, do not replace
                        replaceIndex = -1;
                    }
                    if (replaceIndex != -1) {
                        List<RexNode> copyOperands = new ArrayList<>(call.operands);
                        RexNode oldOp = call.operands.get(replaceIndex);
                        RexNode newOp = b.getRexBuilder().makeInputRef(oldOp.getType(), proj + projOperands.size());
                        projOperands.add(oldOp);
                        copyOperands.set(replaceIndex, newOp);
                        return call.clone(call.type, copyOperands);
                    }
                    break;
                case AND:
                case OR:
                    return super.visitCall(call);
            }
            return call;
        }
    });
    if (newCondition.equals(filter.getCondition())) {
        return;
    }
    b.projectPlus(projOperands);
    b.filter(newCondition);
    b.project(b.fields(ImmutableBitSet.range(proj).asList()));
    call.transformTo(b.build());
}

Example 12 with RexShuttle

use of org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rex.RexShuttle in project beam by apache.

the class CalcRelSplitter method createProgramForLevel.

/**
 * Creates a program containing the expressions for a given level.
 *
 * <p>The expression list of the program will consist of all entries in the expression list <code>
 * allExprs[i]</code> for which the corresponding level ordinal <code>exprLevels[i]</code> is
 * equal to <code>level</code>. Expressions are mapped according to <code>inputExprOrdinals</code>
 * .
 *
 * @param level Level ordinal
 * @param levelCount Number of levels
 * @param inputRowType Input row type
 * @param allExprs Array of all expressions
 * @param exprLevels Array of the level ordinal of each expression
 * @param inputExprOrdinals Ordinals in the expression list of input expressions. Input expression
 *     <code>i</code> will be found at position <code>inputExprOrdinals[i]</code>.
 * @param projectExprOrdinals Ordinals of the expressions to be output this level.
 * @param conditionExprOrdinal Ordinal of the expression to form the condition for this level, or
 *     -1 if there is no condition.
 * @param outputRowType Output row type
 * @return Relational expression
 */
private RexProgram createProgramForLevel(int level, int levelCount, RelDataType inputRowType, RexNode[] allExprs, int[] exprLevels, int[] inputExprOrdinals, final int[] projectExprOrdinals, int conditionExprOrdinal, @Nullable RelDataType outputRowType) {
    // Build a list of expressions to form the calc.
    List<RexNode> exprs = new ArrayList<>();
    // exprInverseOrdinals describes where an expression in allExprs comes
    // from -- from an input, from a calculated expression, or -1 if not
    // available at this level.
    int[] exprInverseOrdinals = new int[allExprs.length];
    Arrays.fill(exprInverseOrdinals, -1);
    int j = 0;
    // and are used here.
    for (int i = 0; i < inputExprOrdinals.length; i++) {
        final int inputExprOrdinal = inputExprOrdinals[i];
        exprs.add(new RexInputRef(i, allExprs[inputExprOrdinal].getType()));
        exprInverseOrdinals[inputExprOrdinal] = j;
        ++j;
    }
    // Next populate the computed expressions.
    final RexShuttle shuttle = new InputToCommonExprConverter(exprInverseOrdinals, exprLevels, level, inputExprOrdinals, allExprs);
    for (int i = 0; i < allExprs.length; i++) {
        if (exprLevels[i] == level || exprLevels[i] == -1 && level == (levelCount - 1) && allExprs[i] instanceof RexLiteral) {
            RexNode expr = allExprs[i];
            final RexNode translatedExpr = expr.accept(shuttle);
            exprs.add(translatedExpr);
            assert exprInverseOrdinals[i] == -1;
            exprInverseOrdinals[i] = j;
            ++j;
        }
    }
    // Form the projection and condition list. Project and condition
    // ordinals are offsets into allExprs, so we need to map them into
    // exprs.
    final List<RexLocalRef> projectRefs = new ArrayList<>(projectExprOrdinals.length);
    final List<String> fieldNames = new ArrayList<>(projectExprOrdinals.length);
    for (int i = 0; i < projectExprOrdinals.length; i++) {
        final int projectExprOrdinal = projectExprOrdinals[i];
        final int index = exprInverseOrdinals[projectExprOrdinal];
        assert index >= 0;
        RexNode expr = allExprs[projectExprOrdinal];
        projectRefs.add(new RexLocalRef(index, expr.getType()));
        // Inherit meaningful field name if possible.
        fieldNames.add(deriveFieldName(expr, i));
    }
    RexLocalRef conditionRef;
    if (conditionExprOrdinal >= 0) {
        final int index = exprInverseOrdinals[conditionExprOrdinal];
        conditionRef = new RexLocalRef(index, allExprs[conditionExprOrdinal].getType());
    } else {
        conditionRef = null;
    }
    if (outputRowType == null) {
        outputRowType = RexUtil.createStructType(typeFactory, projectRefs, fieldNames, null);
    }
    final RexProgram program = new RexProgram(inputRowType, exprs, projectRefs, conditionRef, outputRowType);
    // call operands), since literals should be inlined.
    return program;
}

Example 13 with RexShuttle

use of org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rex.RexShuttle in project samza by apache.

the class JoinTranslator method validateJoinQuery.

private void validateJoinQuery(LogicalJoin join, JoinInputNode.InputType inputTypeOnLeft, JoinInputNode.InputType inputTypeOnRight) {
    JoinRelType joinRelType = join.getJoinType();
    if (joinRelType.compareTo(JoinRelType.INNER) != 0 && joinRelType.compareTo(JoinRelType.LEFT) != 0 && joinRelType.compareTo(JoinRelType.RIGHT) != 0) {
        throw new SamzaException("Query with only INNER and LEFT/RIGHT OUTER join are supported.");
    }
    boolean isTablePosOnLeft = inputTypeOnLeft != JoinInputNode.InputType.STREAM;
    boolean isTablePosOnRight = inputTypeOnRight != JoinInputNode.InputType.STREAM;
    if (!isTablePosOnLeft && !isTablePosOnRight) {
        throw new SamzaException("Invalid query with both sides of join being denoted as 'stream'. " + "Stream-stream join is not yet supported. " + dumpRelPlanForNode(join));
    }
    if (isTablePosOnLeft && isTablePosOnRight) {
        throw new SamzaException("Invalid query with both sides of join being denoted as 'table'. " + dumpRelPlanForNode(join));
    }
    if (joinRelType.compareTo(JoinRelType.LEFT) == 0 && isTablePosOnLeft) {
        throw new SamzaException("Invalid query for outer left join. Left side of the join should be a 'stream' and " + "right side of join should be a 'table'. " + dumpRelPlanForNode(join));
    }
    if (joinRelType.compareTo(JoinRelType.RIGHT) == 0 && isTablePosOnRight) {
        throw new SamzaException("Invalid query for outer right join. Left side of the join should be a 'table' and " + "right side of join should be a 'stream'. " + dumpRelPlanForNode(join));
    }
    final List<RexNode> conjunctionList = new ArrayList<>();
    decomposeAndValidateConjunction(join.getCondition(), conjunctionList);
    if (conjunctionList.isEmpty()) {
        throw new SamzaException("Query results in a cross join, which is not supported. Please optimize the query." + " It is expected that the joins should include JOIN ON operator in the sql query.");
    }
    // TODO Not sure why we can not allow literal as part of the join condition will revisit this in another scope
    conjunctionList.forEach(rexNode -> rexNode.accept(new RexShuttle() {

        @Override
        public RexNode visitLiteral(RexLiteral literal) {
            throw new SamzaException("Join Condition can not allow literal " + literal.toString() + " join node" + join.getDigest());
        }
    }));
    final JoinInputNode.InputType rootTableInput = isTablePosOnRight ? inputTypeOnRight : inputTypeOnLeft;
    if (rootTableInput.compareTo(JoinInputNode.InputType.REMOTE_TABLE) != 0) {
        // it is not a remote table all is good we do not have to validate the project on key Column
        return;
    }
    /*
    For remote Table we need to validate The join Condition and The project that is above remote table scan.
     - As of today Filter need to be exactly one equi-join using the __key__ column (see SAMZA-2554)
     - The Project on the top of the remote table has to contain only simple input references to any of the column used in the join.
    */
    // First let's collect the ref of columns used by the join condition.
    List<RexInputRef> refCollector = new ArrayList<>();
    join.getCondition().accept(new RexShuttle() {

        @Override
        public RexNode visitInputRef(RexInputRef inputRef) {
            refCollector.add(inputRef);
            return inputRef;
        }
    });
    // start index of the Remote table within the Join Row
    final int tableStartIndex = isTablePosOnRight ? join.getLeft().getRowType().getFieldCount() : 0;
    // end index of the Remote table withing the Join Row
    final int tableEndIndex = isTablePosOnRight ? join.getRowType().getFieldCount() : join.getLeft().getRowType().getFieldCount();
    List<Integer> tableRefsIdx = refCollector.stream().map(x -> x.getIndex()).filter(// collect all the refs form table side
    x -> tableStartIndex <= x && x < tableEndIndex).map(// re-adjust the offset
    x -> x - tableStartIndex).sorted().collect(// we have a list with all the input from table side with 0 based index.
    Collectors.toList());
    if (conjunctionList.size() != 1 || tableRefsIdx.size() != 1) {
        // TODO We can relax this by allowing another filter to be evaluated post lookup see SAMZA-2554
        throw new SamzaException("Invalid query for join condition must contain exactly one predicate for remote table on __key__ column " + dumpRelPlanForNode(join));
    }
    // Validate the Project, follow each input and ensure that it is a simple ref with no rexCall in the way.
    if (!isValidRemoteJoinRef(tableRefsIdx.get(0), isTablePosOnRight ? join.getRight() : join.getLeft())) {
        throw new SamzaException("Invalid query for join condition can not have an expression and must be reference " + SamzaSqlRelMessage.KEY_NAME + " column " + dumpRelPlanForNode(join));
    }
}

Example 14 with RexShuttle

use of org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rex.RexShuttle in project samza by apache.

the class JoinTranslator method translate.

void translate(final LogicalJoin join, final TranslatorContext translatorContext) {
    JoinInputNode.InputType inputTypeOnLeft = JoinInputNode.getInputType(join.getLeft(), translatorContext.getExecutionContext().getSamzaSqlApplicationConfig().getInputSystemStreamConfigBySource());
    JoinInputNode.InputType inputTypeOnRight = JoinInputNode.getInputType(join.getRight(), translatorContext.getExecutionContext().getSamzaSqlApplicationConfig().getInputSystemStreamConfigBySource());
    // Do the validation of join query
    validateJoinQuery(join, inputTypeOnLeft, inputTypeOnRight);
    // At this point, one of the sides is a table. Let's figure out if it is on left or right side.
    boolean isTablePosOnRight = inputTypeOnRight != JoinInputNode.InputType.STREAM;
    // stream and table keyIds are used to extract the join condition field (key) names and values out of the stream
    // and table records.
    List<Integer> streamKeyIds = new LinkedList<>();
    List<Integer> tableKeyIds = new LinkedList<>();
    // Fetch the stream and table indices corresponding to the fields given in the join condition.
    final int leftSideSize = join.getLeft().getRowType().getFieldCount();
    final int tableStartIdx = isTablePosOnRight ? leftSideSize : 0;
    final int streamStartIdx = isTablePosOnRight ? 0 : leftSideSize;
    final int tableEndIdx = isTablePosOnRight ? join.getRowType().getFieldCount() : leftSideSize;
    join.getCondition().accept(new RexShuttle() {

        @Override
        public RexNode visitInputRef(RexInputRef inputRef) {
            // Validate the type of the input ref.
            validateJoinKeyType(inputRef);
            int index = inputRef.getIndex();
            if (index >= tableStartIdx && index < tableEndIdx) {
                tableKeyIds.add(index - tableStartIdx);
            } else {
                streamKeyIds.add(index - streamStartIdx);
            }
            return inputRef;
        }
    });
    Collections.sort(tableKeyIds);
    Collections.sort(streamKeyIds);
    // Get the two input nodes (stream and table nodes) for the join.
    JoinInputNode streamNode = new JoinInputNode(isTablePosOnRight ? join.getLeft() : join.getRight(), streamKeyIds, isTablePosOnRight ? inputTypeOnLeft : inputTypeOnRight, !isTablePosOnRight);
    JoinInputNode tableNode = new JoinInputNode(isTablePosOnRight ? join.getRight() : join.getLeft(), tableKeyIds, isTablePosOnRight ? inputTypeOnRight : inputTypeOnLeft, isTablePosOnRight);
    MessageStream<SamzaSqlRelMessage> inputStream = translatorContext.getMessageStream(streamNode.getRelNode().getId());
    Table table = getTable(tableNode, translatorContext);
    MessageStream<SamzaSqlRelMessage> outputStream = joinStreamWithTable(inputStream, table, streamNode, tableNode, join, translatorContext);
    translatorContext.registerMessageStream(join.getId(), outputStream);
    outputStream.map(outputMetricsMF);
}

Example 15 with RexShuttle

use of org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rex.RexShuttle in project calcite by apache.

the class CalcRelSplitter method createProgramForLevel.

/**
 * Creates a program containing the expressions for a given level.
 *
 * <p>The expression list of the program will consist of all entries in the
 * expression list <code>allExprs[i]</code> for which the corresponding
 * level ordinal <code>exprLevels[i]</code> is equal to <code>level</code>.
 * Expressions are mapped according to <code>inputExprOrdinals</code>.
 *
 * @param level                Level ordinal
 * @param levelCount           Number of levels
 * @param inputRowType         Input row type
 * @param allExprs             Array of all expressions
 * @param exprLevels           Array of the level ordinal of each expression
 * @param inputExprOrdinals    Ordinals in the expression list of input
 *                             expressions. Input expression <code>i</code>
 *                             will be found at position
 *                             <code>inputExprOrdinals[i]</code>.
 * @param projectExprOrdinals  Ordinals of the expressions to be output this
 *                             level.
 * @param conditionExprOrdinal Ordinal of the expression to form the
 *                             condition for this level, or -1 if there is no
 *                             condition.
 * @param outputRowType        Output row type
 * @return Relational expression
 */
private RexProgram createProgramForLevel(int level, int levelCount, RelDataType inputRowType, RexNode[] allExprs, int[] exprLevels, int[] inputExprOrdinals, final int[] projectExprOrdinals, int conditionExprOrdinal, RelDataType outputRowType) {
    // Build a list of expressions to form the calc.
    List<RexNode> exprs = new ArrayList<>();
    // exprInverseOrdinals describes where an expression in allExprs comes
    // from -- from an input, from a calculated expression, or -1 if not
    // available at this level.
    int[] exprInverseOrdinals = new int[allExprs.length];
    Arrays.fill(exprInverseOrdinals, -1);
    int j = 0;
    // and are used here.
    for (int i = 0; i < inputExprOrdinals.length; i++) {
        final int inputExprOrdinal = inputExprOrdinals[i];
        exprs.add(new RexInputRef(i, allExprs[inputExprOrdinal].getType()));
        exprInverseOrdinals[inputExprOrdinal] = j;
        ++j;
    }
    // Next populate the computed expressions.
    final RexShuttle shuttle = new InputToCommonExprConverter(exprInverseOrdinals, exprLevels, level, inputExprOrdinals, allExprs);
    for (int i = 0; i < allExprs.length; i++) {
        if (exprLevels[i] == level || exprLevels[i] == -1 && level == (levelCount - 1) && allExprs[i] instanceof RexLiteral) {
            RexNode expr = allExprs[i];
            final RexNode translatedExpr = expr.accept(shuttle);
            exprs.add(translatedExpr);
            assert exprInverseOrdinals[i] == -1;
            exprInverseOrdinals[i] = j;
            ++j;
        }
    }
    // Form the projection and condition list. Project and condition
    // ordinals are offsets into allExprs, so we need to map them into
    // exprs.
    final List<RexLocalRef> projectRefs = new ArrayList<>(projectExprOrdinals.length);
    final List<String> fieldNames = new ArrayList<>(projectExprOrdinals.length);
    for (int i = 0; i < projectExprOrdinals.length; i++) {
        final int projectExprOrdinal = projectExprOrdinals[i];
        final int index = exprInverseOrdinals[projectExprOrdinal];
        assert index >= 0;
        RexNode expr = allExprs[projectExprOrdinal];
        projectRefs.add(new RexLocalRef(index, expr.getType()));
        // Inherit meaningful field name if possible.
        fieldNames.add(deriveFieldName(expr, i));
    }
    RexLocalRef conditionRef;
    if (conditionExprOrdinal >= 0) {
        final int index = exprInverseOrdinals[conditionExprOrdinal];
        conditionRef = new RexLocalRef(index, allExprs[conditionExprOrdinal].getType());
    } else {
        conditionRef = null;
    }
    if (outputRowType == null) {
        outputRowType = RexUtil.createStructType(typeFactory, projectRefs, fieldNames, null);
    }
    final RexProgram program = new RexProgram(inputRowType, exprs, projectRefs, conditionRef, outputRowType);
    // call operands), since literals should be inlined.
    return program;
}