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

Example 51 with RexProgram

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

the class AbstractBeamCalcRel method estimateFilterSelectivity.

private static double estimateFilterSelectivity(RelNode child, RexProgram program, RelMetadataQuery mq) {
    // Similar to calcite, if the calc node is representing filter operation we estimate the filter
    // selectivity based on the number of equality conditions, number of inequality conditions, ....
    RexLocalRef programCondition = program.getCondition();
    RexNode condition;
    if (programCondition == null) {
        condition = null;
    } else {
        condition = program.expandLocalRef(programCondition);
    }
    // Currently this gets the selectivity based on Calcite's Selectivity Handler (RelMdSelectivity)
    return mq.getSelectivity(child, condition);
}

Example 52 with RexProgram

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

the class BeamIOPushDownRule method onMatch.

// ~ Methods ----------------------------------------------------------------
@Override
public void onMatch(RelOptRuleCall call) {
    final BeamIOSourceRel ioSourceRel = call.rel(1);
    final BeamSqlTable beamSqlTable = ioSourceRel.getBeamSqlTable();
    if (ioSourceRel instanceof BeamPushDownIOSourceRel) {
        return;
    }
    // Nested rows are not supported at the moment
    for (RelDataTypeField field : ioSourceRel.getRowType().getFieldList()) {
        if (field.getType() instanceof RelRecordType) {
            return;
        }
    }
    final Calc calc = call.rel(0);
    final RexProgram program = calc.getProgram();
    final Pair<ImmutableList<RexNode>, ImmutableList<RexNode>> projectFilter = program.split();
    final RelDataType calcInputRowType = program.getInputRowType();
    // When predicate push-down is not supported - all filters are unsupported.
    final BeamSqlTableFilter tableFilter = beamSqlTable.constructFilter(projectFilter.right);
    if (!beamSqlTable.supportsProjects().isSupported() && tableFilter instanceof DefaultTableFilter) {
        // Either project or filter push-down must be supported by the IO.
        return;
    }
    Set<String> usedFields = new LinkedHashSet<>();
    if (!(tableFilter instanceof DefaultTableFilter) && !beamSqlTable.supportsProjects().isSupported()) {
        // When applying standalone filter push-down all fields must be project by an IO.
        // With a single exception: Calc projects all fields (in the same order) and does nothing
        // else.
        usedFields.addAll(calcInputRowType.getFieldNames());
    } else {
        // Find all input refs used by projects
        for (RexNode project : projectFilter.left) {
            findUtilizedInputRefs(calcInputRowType, project, usedFields);
        }
        // Find all input refs used by filters
        for (RexNode filter : tableFilter.getNotSupported()) {
            findUtilizedInputRefs(calcInputRowType, filter, usedFields);
        }
    }
    if (usedFields.isEmpty()) {
        // No need to do push-down for queries like this: "select UPPER('hello')".
        return;
    }
    // IO only projects fields utilized by a calc.
    if (tableFilter.getNotSupported().containsAll(projectFilter.right) && usedFields.containsAll(ioSourceRel.getRowType().getFieldNames())) {
        return;
    }
    FieldAccessDescriptor resolved = FieldAccessDescriptor.withFieldNames(usedFields);
    resolved = resolved.resolve(beamSqlTable.getSchema());
    if (canDropCalc(program, beamSqlTable.supportsProjects(), tableFilter)) {
        call.transformTo(ioSourceRel.createPushDownRel(calc.getRowType(), resolved.getFieldsAccessed().stream().map(FieldDescriptor::getFieldName).collect(Collectors.toList()), tableFilter));
        return;
    }
    // IO only projects fields utilised by a calc.
    if (tableFilter.getNotSupported().equals(projectFilter.right) && usedFields.containsAll(ioSourceRel.getRowType().getFieldNames())) {
        return;
    }
    RelNode result = constructNodesWithPushDown(resolved, call.builder(), ioSourceRel, tableFilter, calc.getRowType(), projectFilter.left);
    if (tableFilter.getNotSupported().size() <= projectFilter.right.size() || usedFields.size() < calcInputRowType.getFieldCount()) {
        // Smaller Calc programs are indisputably better, as well as IOs with less projected fields.
        // We can consider something with the same number of filters.
        call.transformTo(result);
    }
}

Example 53 with RexProgram

use of org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rex.RexProgram 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 54 with RexProgram

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

the class BeamSqlUnparseContext method toSql.

@Override
public SqlNode toSql(RexProgram program, RexNode rex) {
    if (rex.getKind().equals(SqlKind.LITERAL)) {
        final RexLiteral literal = (RexLiteral) rex;
        SqlTypeName name = literal.getTypeName();
        SqlTypeFamily family = name.getFamily();
        if (SqlTypeName.TIMESTAMP_WITH_LOCAL_TIME_ZONE.equals(name)) {
            TimestampString timestampString = literal.getValueAs(TimestampString.class);
            return new SqlDateTimeLiteral(timestampString, POS);
        } else if (SqlTypeFamily.BINARY.equals(family)) {
            ByteString byteString = literal.getValueAs(ByteString.class);
            BitString bitString = BitString.createFromHexString(byteString.toString(16));
            return new SqlByteStringLiteral(bitString, POS);
        } else if (SqlTypeFamily.CHARACTER.equals(family)) {
            String escaped = ESCAPE_FOR_ZETA_SQL.translate(literal.getValueAs(String.class));
            return SqlLiteral.createCharString(escaped, POS);
        } else if (SqlTypeName.SYMBOL.equals(literal.getTypeName())) {
            Enum symbol = literal.getValueAs(Enum.class);
            if (TimeUnitRange.DOW.equals(symbol)) {
                return new ReplaceLiteral(literal, POS, "DAYOFWEEK");
            } else if (TimeUnitRange.DOY.equals(symbol)) {
                return new ReplaceLiteral(literal, POS, "DAYOFYEAR");
            } else if (TimeUnitRange.WEEK.equals(symbol)) {
                return new ReplaceLiteral(literal, POS, "ISOWEEK");
            }
        }
    } else if (rex.getKind().equals(SqlKind.DYNAMIC_PARAM)) {
        final RexDynamicParam param = (RexDynamicParam) rex;
        final int index = param.getIndex();
        final String name = "null_param_" + index;
        nullParams.put(name, param.getType());
        return new NamedDynamicParam(index, POS, name);
    } else if (SqlKind.SEARCH.equals(rex.getKind())) {
        // Workaround CALCITE-4716
        RexCall search = (RexCall) rex;
        RexLocalRef ref = (RexLocalRef) search.operands.get(1);
        RexLiteral literal = (RexLiteral) program.getExprList().get(ref.getIndex());
        rex = search.clone(search.getType(), ImmutableList.of(search.operands.get(0), literal));
    }
    return super.toSql(program, rex);
}