Search in sources :

Example 11 with OR

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;
}
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 12 with OR

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());
}
Also used : Document(org.bson.Document) PBegin(org.apache.beam.sdk.values.PBegin) LoggerFactory(org.slf4j.LoggerFactory) MongoDbIO(org.apache.beam.sdk.io.mongodb.MongoDbIO) IsBounded(org.apache.beam.sdk.values.PCollection.IsBounded) JsonToRow(org.apache.beam.sdk.transforms.JsonToRow) SimpleFunction(org.apache.beam.sdk.transforms.SimpleFunction) ProjectSupport(org.apache.beam.sdk.extensions.sql.meta.ProjectSupport) Table(org.apache.beam.sdk.extensions.sql.meta.Table) Matcher(java.util.regex.Matcher) AND(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.sql.SqlKind.AND) RexLiteral(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rex.RexLiteral) MapElements(org.apache.beam.sdk.transforms.MapElements) SchemaBaseBeamTable(org.apache.beam.sdk.extensions.sql.meta.SchemaBaseBeamTable) FieldType(org.apache.beam.sdk.schemas.Schema.FieldType) Collectors(java.util.stream.Collectors) FindQuery(org.apache.beam.sdk.io.mongodb.FindQuery) Serializable(java.io.Serializable) InvalidTableException(org.apache.beam.sdk.extensions.sql.meta.provider.InvalidTableException) VisibleForTesting(org.apache.beam.vendor.calcite.v1_28_0.com.google.common.annotations.VisibleForTesting) COMPARISON(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.sql.SqlKind.COMPARISON) POutput(org.apache.beam.sdk.values.POutput) List(java.util.List) ParDo(org.apache.beam.sdk.transforms.ParDo) BeamTableStatistics(org.apache.beam.sdk.extensions.sql.impl.BeamTableStatistics) ImmutableList(org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList) Pattern(java.util.regex.Pattern) Experimental(org.apache.beam.sdk.annotations.Experimental) RexNode(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rex.RexNode) JsonWriterSettings(org.bson.json.JsonWriterSettings) SqlKind(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.sql.SqlKind) SelectHelpers(org.apache.beam.sdk.schemas.utils.SelectHelpers) ArrayList(java.util.ArrayList) PTransform(org.apache.beam.sdk.transforms.PTransform) Filters(com.mongodb.client.model.Filters) DefaultTableFilter(org.apache.beam.sdk.extensions.sql.meta.DefaultTableFilter) Bson(org.bson.conversions.Bson) ToJson(org.apache.beam.sdk.transforms.ToJson) JsonMode(org.bson.json.JsonMode) FieldAccessDescriptor(org.apache.beam.sdk.schemas.FieldAccessDescriptor) Row(org.apache.beam.sdk.values.Row) PipelineOptions(org.apache.beam.sdk.options.PipelineOptions) IntFunction(java.util.function.IntFunction) FieldTypeDescriptors(org.apache.beam.sdk.schemas.FieldTypeDescriptors) DoFn(org.apache.beam.sdk.transforms.DoFn) Logger(org.slf4j.Logger) BeamSqlTableFilter(org.apache.beam.sdk.extensions.sql.meta.BeamSqlTableFilter) PCollection(org.apache.beam.sdk.values.PCollection) Schema(org.apache.beam.sdk.schemas.Schema) RexCall(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rex.RexCall) SqlTypeName(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.sql.type.SqlTypeName) RexInputRef(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rex.RexInputRef) OR(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.sql.SqlKind.OR) RexCall(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rex.RexCall) RexLiteral(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rex.RexLiteral) ArrayList(java.util.ArrayList) RexInputRef(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rex.RexInputRef) RexNode(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rex.RexNode)

Example 13 with OR

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);
}
Also used : RexCall(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rex.RexCall) RexLiteral(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rex.RexLiteral) RexInputRef(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rex.RexInputRef) RexNode(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rex.RexNode)

Example 14 with OR

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;
}
Also used : Project(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rel.core.Project) RelNode(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rel.RelNode) Filter(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rel.core.Filter) BeamIOSourceRel(org.apache.beam.sdk.extensions.sql.impl.rel.BeamIOSourceRel) RelSubset(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.plan.volcano.RelSubset)

Example 15 with OR

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());
        }
    }
}
Also used : RexCall(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rex.RexCall) RexLiteral(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rex.RexLiteral) RelDataTypeField(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rel.type.RelDataTypeField) RexInputRef(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rex.RexInputRef) ArrayDeque(java.util.ArrayDeque) RexNode(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rex.RexNode) VisibleForTesting(org.apache.beam.vendor.guava.v26_0_jre.com.google.common.annotations.VisibleForTesting)

Aggregations

RexNode (org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rex.RexNode)13 ArrayList (java.util.ArrayList)8 RelDataType (org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rel.type.RelDataType)7 RexLiteral (org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rex.RexLiteral)6 RelNode (org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rel.RelNode)5 RexInputRef (org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rex.RexInputRef)5 RexCall (org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rex.RexCall)4 ResolvedExpr (com.google.zetasql.resolvedast.ResolvedNodes.ResolvedExpr)3 RelDataTypeField (org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rel.type.RelDataTypeField)3 RexLocalRef (org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rex.RexLocalRef)3 ImmutableList (org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList)3 ResolvedAggregateFunctionCall (com.google.zetasql.resolvedast.ResolvedNodes.ResolvedAggregateFunctionCall)2 ResolvedComputedColumn (com.google.zetasql.resolvedast.ResolvedNodes.ResolvedComputedColumn)2 HashSet (java.util.HashSet)2 LinkedHashSet (java.util.LinkedHashSet)2 List (java.util.List)2 Collectors (java.util.stream.Collectors)2 BeamIOSourceRel (org.apache.beam.sdk.extensions.sql.impl.rel.BeamIOSourceRel)2 Schema (org.apache.beam.sdk.schemas.Schema)2 CallImplementor (org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.adapter.enumerable.CallImplementor)2