Search in sources :

Example 16 with AND

use of org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.sql.SqlKind.AND in project beam by apache.

the class BeamIOPushDownRule method constructNodesWithPushDown.

/**
 * Construct a new {@link BeamIOSourceRel} with predicate and/or project pushed-down and a new
 * {@code Calc} to do field reordering/field duplication/complex projects.
 *
 * @param resolved A descriptor of fields used by a {@code Calc}.
 * @param relBuilder A {@code RelBuilder} for constructing {@code Project} and {@code Filter} Rel
 *     nodes with operations unsupported by the IO.
 * @param ioSourceRel Original {@code BeamIOSourceRel} we are attempting to perform push-down for.
 * @param tableFilter A class containing information about IO predicate push-down capabilities.
 * @param calcDataType A Calcite output schema of an original {@code Calc}.
 * @param calcProjects A list of projected {@code RexNode}s by a {@code Calc}.
 * @return An alternative {@code RelNode} with supported filters/projects pushed-down to IO Rel.
 */
private RelNode constructNodesWithPushDown(FieldAccessDescriptor resolved, RelBuilder relBuilder, BeamIOSourceRel ioSourceRel, BeamSqlTableFilter tableFilter, RelDataType calcDataType, List<RexNode> calcProjects) {
    Schema newSchema = SelectHelpers.getOutputSchema(ioSourceRel.getBeamSqlTable().getSchema(), resolved);
    RelDataType calcInputType = CalciteUtils.toCalciteRowType(newSchema, ioSourceRel.getCluster().getTypeFactory());
    BeamIOSourceRel newIoSourceRel = ioSourceRel.createPushDownRel(calcInputType, newSchema.getFieldNames(), tableFilter);
    relBuilder.push(newIoSourceRel);
    List<RexNode> newProjects = new ArrayList<>();
    List<RexNode> newFilter = new ArrayList<>();
    // Ex: let's say the original fields are (number before each element is the index):
    // {0:unused1, 1:id, 2:name, 3:unused2},
    // where only 'id' and 'name' are being used. Then the new calcInputType should be as follows:
    // {0:id, 1:name}.
    // A mapping list will contain 2 entries: {0:1, 1:2},
    // showing how used field names map to the original fields.
    List<Integer> mapping = resolved.getFieldsAccessed().stream().map(FieldDescriptor::getFieldId).collect(Collectors.toList());
    // Map filters to new RexInputRef.
    for (RexNode filter : tableFilter.getNotSupported()) {
        newFilter.add(reMapRexNodeToNewInputs(filter, mapping));
    }
    // Map projects to new RexInputRef.
    for (RexNode project : calcProjects) {
        newProjects.add(reMapRexNodeToNewInputs(project, mapping));
    }
    if (RexUtil.isIdentity(newProjects, newIoSourceRel.getRowType())) {
        // Force a rename prior to filter for identity function.
        relBuilder.project(newProjects, calcDataType.getFieldNames(), true);
    }
    relBuilder.filter(newFilter);
    relBuilder.project(newProjects, calcDataType.getFieldNames());
    return relBuilder.build();
}
Also used : Schema(org.apache.beam.sdk.schemas.Schema) BeamIOSourceRel(org.apache.beam.sdk.extensions.sql.impl.rel.BeamIOSourceRel) ArrayList(java.util.ArrayList) RelDataType(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rel.type.RelDataType) RexNode(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rex.RexNode)

Example 17 with AND

use of org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.sql.SqlKind.AND in project beam by apache.

the class BeamIOPushDownRule method isProjectRenameOnlyProgram.

/**
 * Determine whether a program only performs renames and/or projects. RexProgram#isTrivial is not
 * sufficient in this case, because number of projects does not need to be the same as inputs.
 * Calc should NOT be dropped in the following cases:<br>
 * 1. Projected fields are manipulated (ex: 'select field1+10').<br>
 * 2. When the same field projected more than once.<br>
 * 3. When an IO does not supports field reordering and projects fields in a different (from
 * schema) order.
 *
 * @param program A program to check.
 * @param projectReorderingSupported Whether project push-down supports field reordering.
 * @return True when program performs only projects (w/o any modifications), false otherwise.
 */
@VisibleForTesting
boolean isProjectRenameOnlyProgram(RexProgram program, boolean projectReorderingSupported) {
    int fieldCount = program.getInputRowType().getFieldCount();
    Set<Integer> projectIndex = new HashSet<>();
    int previousIndex = -1;
    for (RexLocalRef ref : program.getProjectList()) {
        int index = ref.getIndex();
        if (// Projected values are InputRefs.
        index >= fieldCount || // Each field projected once.
        !projectIndex.add(ref.getIndex()) || (!projectReorderingSupported && index <= previousIndex)) {
            // In the same order.
            return false;
        }
        previousIndex = index;
    }
    return true;
}
Also used : RexLocalRef(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rex.RexLocalRef) HashSet(java.util.HashSet) LinkedHashSet(java.util.LinkedHashSet) VisibleForTesting(org.apache.beam.vendor.guava.v26_0_jre.com.google.common.annotations.VisibleForTesting)

Example 18 with AND

use of org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.sql.SqlKind.AND 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 19 with AND

use of org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.sql.SqlKind.AND 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 20 with AND

use of org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.sql.SqlKind.AND in project beam by apache.

the class BeamBigQuerySqlDialect method unparseCall.

@Override
public void unparseCall(final SqlWriter writer, final SqlCall call, final int leftPrec, final int rightPrec) {
    switch(call.getKind()) {
        case ROW:
            final SqlWriter.Frame structFrame = writer.startFunCall("STRUCT");
            for (SqlNode operand : call.getOperandList()) {
                writer.sep(",");
                operand.unparse(writer, leftPrec, rightPrec);
            }
            writer.endFunCall(structFrame);
            break;
        case OTHER_FUNCTION:
            String funName = call.getOperator().getName();
            if (DOUBLE_LITERAL_WRAPPERS.containsKey(funName)) {
                // self-designed function dealing with the unparsing of ZetaSQL DOUBLE positive
                // infinity, negative infinity and NaN
                unparseDoubleLiteralWrapperFunction(writer, funName);
                break;
            } else if (NUMERIC_LITERAL_WRAPPER.equals(funName)) {
                // self-designed function dealing with the unparsing of ZetaSQL NUMERIC literal
                unparseNumericLiteralWrapperFunction(writer, call, leftPrec, rightPrec);
                break;
            } else if (FUNCTIONS_USING_INTERVAL.contains(funName)) {
                unparseFunctionsUsingInterval(writer, call, leftPrec, rightPrec);
                break;
            } else if (EXTRACT_FUNCTIONS.containsKey(funName)) {
                unparseExtractFunctions(writer, call, leftPrec, rightPrec);
                break;
            } else if (IN_ARRAY_OPERATOR.equals(funName)) {
                unparseInArrayOperator(writer, call, leftPrec, rightPrec);
                break;
            }
        // fall through
        default:
            super.unparseCall(writer, call, leftPrec, rightPrec);
    }
}
Also used : SqlWriter(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.sql.SqlWriter) SqlNode(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.sql.SqlNode)

Aggregations

RexNode (org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rex.RexNode)16 ArrayList (java.util.ArrayList)11 Test (org.junit.Test)10 RelNode (org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rel.RelNode)9 List (java.util.List)8 Collectors (java.util.stream.Collectors)8 Serializable (java.io.Serializable)7 Arrays (java.util.Arrays)7 RelDataType (org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rel.type.RelDataType)7 IOException (java.io.IOException)6 Statement (java.sql.Statement)6 HashMap (java.util.HashMap)6 Map (java.util.Map)6 Schema (org.apache.beam.sdk.schemas.Schema)6 PCollection (org.apache.beam.sdk.values.PCollection)6 JsonProcessingException (com.fasterxml.jackson.core.JsonProcessingException)5 ObjectMapper (com.fasterxml.jackson.databind.ObjectMapper)5 ByteArrayOutputStream (java.io.ByteArrayOutputStream)5 StandardCharsets (java.nio.charset.StandardCharsets)5 UTF_8 (java.nio.charset.StandardCharsets.UTF_8)5