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);
}
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);
}
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;
}
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;
}
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);
}
Aggregations