use of org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.sql.SqlKind.OR 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.org.apache.calcite.sql.SqlKind.OR in project beam by apache.
the class MongoDbTable method translateRexNodeToBson.
/**
* Recursively translates a single RexNode to MongoDB Bson filter. Supports simple comparison
* operations, negation, and nested conjunction/disjunction. Boolean fields are translated as an
* `$eq` operation with a boolean `true`.
*
* @param node {@code RexNode} to translate.
* @return {@code Bson} filter.
*/
private Bson translateRexNodeToBson(RexNode node) {
final IntFunction<String> fieldIdToName = i -> getSchema().getField(i).getName();
// Supported operations are described in MongoDbFilter#isSupported
if (node instanceof RexCall) {
RexCall compositeNode = (RexCall) node;
List<RexLiteral> literals = new ArrayList<>();
List<RexInputRef> inputRefs = new ArrayList<>();
for (RexNode operand : compositeNode.getOperands()) {
if (operand instanceof RexLiteral) {
literals.add((RexLiteral) operand);
} else if (operand instanceof RexInputRef) {
inputRefs.add((RexInputRef) operand);
}
}
// Operation is a comparison, since one of the operands in a field reference.
if (inputRefs.size() == 1) {
RexInputRef inputRef = inputRefs.get(0);
String inputFieldName = fieldIdToName.apply(inputRef.getIndex());
if (literals.size() > 0) {
// Convert literal value to the same Java type as the field we are comparing to.
Object literal = convertToExpectedType(inputRef, literals.get(0));
switch(node.getKind()) {
case IN:
return Filters.in(inputFieldName, convertToExpectedType(inputRef, literals));
case EQUALS:
return Filters.eq(inputFieldName, literal);
case NOT_EQUALS:
return Filters.not(Filters.eq(inputFieldName, literal));
case LESS_THAN:
return Filters.lt(inputFieldName, literal);
case GREATER_THAN:
return Filters.gt(inputFieldName, literal);
case GREATER_THAN_OR_EQUAL:
return Filters.gte(inputFieldName, literal);
case LESS_THAN_OR_EQUAL:
return Filters.lte(inputFieldName, literal);
default:
// Encountered an unexpected node kind, RuntimeException below.
break;
}
} else if (node.getKind().equals(SqlKind.NOT)) {
// Ex: `where not boolean_field`
return Filters.not(translateRexNodeToBson(inputRef));
} else {
throw new RuntimeException("Cannot create a filter for an unsupported node: " + node.toString());
}
} else {
// Operation is a conjunction/disjunction.
switch(node.getKind()) {
case AND:
// Recursively construct filter for each operand of conjunction.
return Filters.and(compositeNode.getOperands().stream().map(this::translateRexNodeToBson).collect(Collectors.toList()));
case OR:
// Recursively construct filter for each operand of disjunction.
return Filters.or(compositeNode.getOperands().stream().map(this::translateRexNodeToBson).collect(Collectors.toList()));
default:
// Encountered an unexpected node kind, RuntimeException below.
break;
}
}
throw new RuntimeException("Encountered an unexpected node kind: " + node.getKind().toString());
} else if (node instanceof RexInputRef && node.getType().getSqlTypeName().equals(SqlTypeName.BOOLEAN)) {
// Boolean field, must be true. Ex: `select * from table where bool_field`
return Filters.eq(fieldIdToName.apply(((RexInputRef) node).getIndex()), true);
}
throw new RuntimeException("Was expecting a RexCall or a boolean RexInputRef, but received: " + node.getClass().getSimpleName());
}
use of org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.sql.SqlKind.OR in project beam by apache.
the class BigQueryFilter method isSupported.
/**
* Check whether a {@code RexNode} is supported. As of right now BigQuery supports: 1. Complex
* predicates (both conjunction and disjunction). 2. Comparison between a column and a literal.
*
* <p>TODO: Check if comparison between two columns is supported. Also over a boolean field.
*
* @param node A node to check for predicate push-down support.
* @return A pair containing a boolean whether an expression is supported and the number of input
* references used by the expression.
*/
private Pair<Boolean, Integer> isSupported(RexNode node) {
int numberOfInputRefs = 0;
boolean isSupported = true;
if (node instanceof RexCall) {
RexCall compositeNode = (RexCall) node;
// CAST, TRIM? and REVERSE? should be supported as well.
if (!node.getKind().belongsTo(SUPPORTED_OPS)) {
isSupported = false;
} else {
for (RexNode operand : compositeNode.getOperands()) {
// All operands must be supported for a parent node to be supported.
Pair<Boolean, Integer> childSupported = isSupported(operand);
// (OR).
if (!node.getKind().belongsTo(ImmutableSet.of(AND, OR))) {
numberOfInputRefs += childSupported.getRight();
}
// Predicate functions, where more than one field is involved are unsupported.
isSupported = numberOfInputRefs < 2 && childSupported.getLeft();
}
}
} else if (node instanceof RexInputRef) {
numberOfInputRefs = 1;
} else if (node instanceof RexLiteral) {
// RexLiterals are expected, but no action is needed.
} else {
throw new UnsupportedOperationException("Encountered an unexpected node type: " + node.getClass().getSimpleName());
}
return Pair.of(isSupported, numberOfInputRefs);
}
use of org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.sql.SqlKind.OR in project beam by apache.
the class BeamAggregateProjectMergeRule method getUnderlyingIO.
/**
* Following scenarios are possible:<br>
* 1) Aggregate <- Project <- IO.<br>
* 2) Aggregate <- Project <- Chain of Project/Filter <- IO.<br>
* 3) Aggregate <- Project <- Something else.<br>
* 4) Aggregate <- Project <- Chain of Project/Filter <- Something else.
*
* @param parent project that matched this rule.
* @return {@code BeamIOSourceRel} when it is present or null when some other {@code RelNode} is
* present.
*/
private BeamIOSourceRel getUnderlyingIO(Set<RelNode> visitedNodes, SingleRel parent) {
// No need to look at the same node more than once.
if (visitedNodes.contains(parent)) {
return null;
}
visitedNodes.add(parent);
List<RelNode> nodes = ((RelSubset) parent.getInput()).getRelList();
for (RelNode node : nodes) {
if (node instanceof Filter || node instanceof Project) {
// Search node inputs for an IO.
BeamIOSourceRel child = getUnderlyingIO(visitedNodes, (SingleRel) node);
if (child != null) {
return child;
}
} else if (node instanceof BeamIOSourceRel) {
return (BeamIOSourceRel) node;
}
}
return null;
}
use of org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.sql.SqlKind.OR in project beam by apache.
the class BeamIOPushDownRule method findUtilizedInputRefs.
/**
* Given a {@code RexNode}, find all {@code RexInputRef}s a node or it's children nodes use.
*
* @param inputRowType {@code RelDataType} used for looking up names of {@code RexInputRef}.
* @param startNode A node to start at.
* @param usedFields Names of {@code RexInputRef}s are added to this list.
*/
@VisibleForTesting
void findUtilizedInputRefs(RelDataType inputRowType, RexNode startNode, Set<String> usedFields) {
Queue<RexNode> prerequisites = new ArrayDeque<>();
prerequisites.add(startNode);
// Assuming there are no cyclic nodes, traverse dependency tree until all RexInputRefs are found
while (!prerequisites.isEmpty()) {
RexNode node = prerequisites.poll();
if (node instanceof RexCall) {
// Composite expression, example: "=($t11, $t12)"
RexCall compositeNode = (RexCall) node;
// Expression from example above contains 2 operands: $t11, $t12
prerequisites.addAll(compositeNode.getOperands());
} else if (node instanceof RexInputRef) {
// Input reference
// Find a field in an inputRowType for the input reference
int inputFieldIndex = ((RexInputRef) node).getIndex();
RelDataTypeField field = inputRowType.getFieldList().get(inputFieldIndex);
// If we have not seen it before - add it to the list (hash set)
usedFields.add(field.getName());
} else if (node instanceof RexLiteral) {
// Does not contain information about columns utilized by a Calc
} else {
throw new UnsupportedOperationException("Unexpected RexNode encountered: " + node.getClass().getSimpleName());
}
}
}
Aggregations