Search in sources :

Example 1 with Preconditions.checkArgument

use of org.apache.beam.vendor.calcite.v1_28_0.com.google.common.base.Preconditions.checkArgument in project beam by apache.

the class SqlInOperatorRewriter method apply.

@Override
public RexNode apply(RexBuilder rexBuilder, List<RexNode> operands) {
    Preconditions.checkArgument(operands.size() >= 2, "IN should have at least two arguments in function call.");
    final RexNode arg = operands.get(0);
    final List<RexNode> ranges = ImmutableList.copyOf(operands.subList(1, operands.size()));
    // ZetaSQL has weird behavior for NULL...
    for (RexNode node : ranges) {
        if (node instanceof RexLiteral && ((RexLiteral) node).isNull()) {
            throw new UnsupportedOperationException("IN NULL unsupported");
        }
    }
    return rexBuilder.makeIn(arg, ranges);
}
Also used : RexLiteral(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rex.RexLiteral) RexNode(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rex.RexNode)

Example 2 with Preconditions.checkArgument

use of org.apache.beam.vendor.calcite.v1_28_0.com.google.common.base.Preconditions.checkArgument in project beam by apache.

the class SqlCaseWithValueOperatorRewriter method apply.

@Override
public RexNode apply(RexBuilder rexBuilder, List<RexNode> operands) {
    Preconditions.checkArgument(operands.size() % 2 == 0 && !operands.isEmpty(), "$case_with_value should have an even number of arguments greater than 0 in function call" + " (The value operand, the else operand, and paired when/then operands).");
    SqlOperator op = SqlStdOperatorTable.CASE;
    List<RexNode> newOperands = new ArrayList<>();
    RexNode value = operands.get(0);
    for (int i = 1; i < operands.size() - 2; i += 2) {
        RexNode when = operands.get(i);
        RexNode then = operands.get(i + 1);
        newOperands.add(rexBuilder.makeCall(SqlStdOperatorTable.EQUALS, ImmutableList.of(value, when)));
        newOperands.add(then);
    }
    RexNode elseOperand = Iterables.getLast(operands);
    newOperands.add(elseOperand);
    return rexBuilder.makeCall(op, newOperands);
}
Also used : SqlOperator(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.sql.SqlOperator) ArrayList(java.util.ArrayList) RexNode(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rex.RexNode)

Example 3 with Preconditions.checkArgument

use of org.apache.beam.vendor.calcite.v1_28_0.com.google.common.base.Preconditions.checkArgument in project beam by apache.

the class CalciteQueryPlanner method convertToBeamRel.

/**
 * It parses and validate the input query, then convert into a {@link BeamRelNode} tree. Note that
 * query parameters are not yet supported.
 */
@Override
public BeamRelNode convertToBeamRel(String sqlStatement, QueryParameters queryParameters) throws ParseException, SqlConversionException {
    Preconditions.checkArgument(queryParameters.getKind() == Kind.NONE, "Beam SQL Calcite dialect does not yet support query parameters.");
    BeamRelNode beamRelNode;
    try {
        SqlNode parsed = planner.parse(sqlStatement);
        TableResolutionUtils.setupCustomTableResolution(connection, parsed);
        SqlNode validated = planner.validate(parsed);
        LOG.info("SQL:\n" + validated);
        // root of original logical plan
        RelRoot root = planner.rel(validated);
        LOG.info("SQLPlan>\n" + RelOptUtil.toString(root.rel));
        RelTraitSet desiredTraits = root.rel.getTraitSet().replace(BeamLogicalConvention.INSTANCE).replace(root.collation).simplify();
        // beam physical plan
        root.rel.getCluster().setMetadataProvider(ChainedRelMetadataProvider.of(ImmutableList.of(NonCumulativeCostImpl.SOURCE, RelMdNodeStats.SOURCE, root.rel.getCluster().getMetadataProvider())));
        root.rel.getCluster().setMetadataQuerySupplier(BeamRelMetadataQuery::instance);
        RelMetadataQuery.THREAD_PROVIDERS.set(JaninoRelMetadataProvider.of(root.rel.getCluster().getMetadataProvider()));
        root.rel.getCluster().invalidateMetadataQuery();
        beamRelNode = (BeamRelNode) planner.transform(0, desiredTraits, root.rel);
        LOG.info("BEAMPlan>\n" + RelOptUtil.toString(beamRelNode));
    } catch (RelConversionException | CannotPlanException e) {
        throw new SqlConversionException(String.format("Unable to convert query %s", sqlStatement), e);
    } catch (SqlParseException | ValidationException e) {
        throw new ParseException(String.format("Unable to parse query %s", sqlStatement), e);
    } finally {
        planner.close();
    }
    return beamRelNode;
}
Also used : BeamRelNode(org.apache.beam.sdk.extensions.sql.impl.rel.BeamRelNode) BeamRelMetadataQuery(org.apache.beam.sdk.extensions.sql.impl.planner.BeamRelMetadataQuery) ValidationException(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.tools.ValidationException) CannotPlanException(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.plan.RelOptPlanner.CannotPlanException) SqlParseException(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.sql.parser.SqlParseException) RelRoot(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rel.RelRoot) RelTraitSet(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.plan.RelTraitSet) RelConversionException(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.tools.RelConversionException) SqlParseException(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.sql.parser.SqlParseException) SqlNode(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.sql.SqlNode)

Example 4 with Preconditions.checkArgument

use of org.apache.beam.vendor.calcite.v1_28_0.com.google.common.base.Preconditions.checkArgument in project beam by apache.

the class CalcRelSplitter method execute.

// ~ Methods ----------------------------------------------------------------
public RelNode execute() {
    // expressions to the left.
    assert program.isValid(Litmus.THROW, null);
    final List<RexNode> exprList = program.getExprList();
    final RexNode[] exprs = exprList.toArray(new RexNode[0]);
    assert !RexUtil.containComplexExprs(exprList);
    // Figure out what level each expression belongs to.
    int[] exprLevels = new int[exprs.length];
    // The type of a level is given by
    // relTypes[levelTypeOrdinals[level]].
    int[] levelTypeOrdinals = new int[exprs.length];
    int levelCount = chooseLevels(exprs, -1, exprLevels, levelTypeOrdinals);
    // For each expression, figure out which is the highest level where it
    // is used.
    int[] exprMaxUsingLevelOrdinals = new HighestUsageFinder(exprs, exprLevels).getMaxUsingLevelOrdinals();
    // If expressions are used as outputs, mark them as higher than that.
    final List<RexLocalRef> projectRefList = program.getProjectList();
    final RexLocalRef conditionRef = program.getCondition();
    for (RexLocalRef projectRef : projectRefList) {
        exprMaxUsingLevelOrdinals[projectRef.getIndex()] = levelCount;
    }
    if (conditionRef != null) {
        exprMaxUsingLevelOrdinals[conditionRef.getIndex()] = levelCount;
    }
    // Print out what we've got.
    if (RULE_LOGGER.isTraceEnabled()) {
        traceLevelExpressions(exprs, exprLevels, levelTypeOrdinals, levelCount);
    }
    // Now build the calcs.
    RelNode rel = child;
    final int inputFieldCount = program.getInputRowType().getFieldCount();
    int[] inputExprOrdinals = identityArray(inputFieldCount);
    boolean doneCondition = false;
    for (int level = 0; level < levelCount; level++) {
        final int[] projectExprOrdinals;
        final RelDataType outputRowType;
        if (level == (levelCount - 1)) {
            outputRowType = program.getOutputRowType();
            projectExprOrdinals = new int[projectRefList.size()];
            for (int i = 0; i < projectExprOrdinals.length; i++) {
                projectExprOrdinals[i] = projectRefList.get(i).getIndex();
            }
        } else {
            outputRowType = null;
            // Project the expressions which are computed at this level or
            // before, and will be used at later levels.
            List<Integer> projectExprOrdinalList = new ArrayList<>();
            for (int i = 0; i < exprs.length; i++) {
                RexNode expr = exprs[i];
                if (expr instanceof RexLiteral) {
                    // Don't project literals. They are always created in
                    // the level where they are used.
                    exprLevels[i] = -1;
                    continue;
                }
                if ((exprLevels[i] <= level) && (exprMaxUsingLevelOrdinals[i] > level)) {
                    projectExprOrdinalList.add(i);
                }
            }
            projectExprOrdinals = Ints.toArray(projectExprOrdinalList);
        }
        final RelType relType = relTypes[levelTypeOrdinals[level]];
        // Can we do the condition this level?
        int conditionExprOrdinal = -1;
        if ((conditionRef != null) && !doneCondition) {
            conditionExprOrdinal = conditionRef.getIndex();
            if ((exprLevels[conditionExprOrdinal] > level) || !relType.supportsCondition()) {
                // stand down -- we're not ready to do the condition yet
                conditionExprOrdinal = -1;
            } else {
                doneCondition = true;
            }
        }
        RexProgram program1 = createProgramForLevel(level, levelCount, rel.getRowType(), exprs, exprLevels, inputExprOrdinals, projectExprOrdinals, conditionExprOrdinal, outputRowType);
        rel = relType.makeRel(cluster, traits, relBuilder, rel, program1);
        rel = handle(rel);
        // The outputs of this level will be the inputs to the next level.
        inputExprOrdinals = projectExprOrdinals;
    }
    Preconditions.checkArgument(doneCondition || (conditionRef == null), "unhandled condition");
    return rel;
}
Also used : RexLiteral(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rex.RexLiteral) RexProgram(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rex.RexProgram) ArrayList(java.util.ArrayList) RelDataType(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rel.type.RelDataType) RelNode(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rel.RelNode) RexLocalRef(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rex.RexLocalRef) RexNode(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rex.RexNode)

Example 5 with Preconditions.checkArgument

use of org.apache.beam.vendor.calcite.v1_28_0.com.google.common.base.Preconditions.checkArgument in project beam by apache.

the class SqlCoalesceOperatorRewriter method apply.

@Override
public RexNode apply(RexBuilder rexBuilder, List<RexNode> operands) {
    Preconditions.checkArgument(operands.size() >= 1, "COALESCE should have at least one argument in function call.");
    // No need for a case operator if there's only one operand
    if (operands.size() == 1) {
        return operands.get(0);
    }
    SqlOperator op = SqlStdOperatorTable.CASE;
    List<RexNode> newOperands = new ArrayList<>();
    for (RexNode operand : Util.skipLast(operands)) {
        newOperands.add(rexBuilder.makeCall(SqlStdOperatorTable.IS_NOT_NULL, ImmutableList.of(operand)));
        newOperands.add(operand);
    }
    newOperands.add(Util.last(operands));
    return rexBuilder.makeCall(op, newOperands);
}
Also used : SqlOperator(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.sql.SqlOperator) ArrayList(java.util.ArrayList) RexNode(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rex.RexNode)

Aggregations

RexNode (org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rex.RexNode)6 SqlOperator (org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.sql.SqlOperator)4 ArrayList (java.util.ArrayList)3 RexLiteral (org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rex.RexLiteral)2 BeamRelMetadataQuery (org.apache.beam.sdk.extensions.sql.impl.planner.BeamRelMetadataQuery)1 BeamRelNode (org.apache.beam.sdk.extensions.sql.impl.rel.BeamRelNode)1 FieldAccessDescriptor (org.apache.beam.sdk.schemas.FieldAccessDescriptor)1 Schema (org.apache.beam.sdk.schemas.Schema)1 Row (org.apache.beam.sdk.values.Row)1 CannotPlanException (org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.plan.RelOptPlanner.CannotPlanException)1 RelTraitSet (org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.plan.RelTraitSet)1 RelNode (org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rel.RelNode)1 RelRoot (org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rel.RelRoot)1 JoinRelType (org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rel.core.JoinRelType)1 RelDataType (org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rel.type.RelDataType)1 RexLocalRef (org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rex.RexLocalRef)1 RexProgram (org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rex.RexProgram)1 SqlNode (org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.sql.SqlNode)1 SqlParseException (org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.sql.parser.SqlParseException)1 RelConversionException (org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.tools.RelConversionException)1