Search in sources :

Example 6 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 BeamZetaSqlCatalog method addTableToLeafCatalog.

/**
 * Assume last element in tablePath is a table name, and everything before is catalogs. So the
 * logic is to create nested catalogs until the last level, then add a table at the last level.
 *
 * <p>Table schema is extracted from Calcite schema based on the table name resolution strategy,
 * e.g. either by drilling down the schema.getSubschema() path or joining the table name with dots
 * to construct a single compound identifier (e.g. Data Catalog use case).
 */
private void addTableToLeafCatalog(List<String> tablePath, QueryTrait queryTrait) {
    SimpleCatalog leafCatalog = createNestedCatalogs(zetaSqlCatalog, tablePath);
    org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.schema.Table calciteTable = TableResolution.resolveCalciteTable(calciteSchema, tablePath);
    if (calciteTable == null) {
        throw new ZetaSqlException("Wasn't able to resolve the path " + tablePath + " in schema: " + calciteSchema.getName());
    }
    RelDataType rowType = calciteTable.getRowType(typeFactory);
    TableResolution.SimpleTableWithPath tableWithPath = TableResolution.SimpleTableWithPath.of(tablePath);
    queryTrait.addResolvedTable(tableWithPath);
    addFieldsToTable(tableWithPath, rowType);
    leafCatalog.addSimpleTable(tableWithPath.getTable());
}
Also used : SimpleCatalog(com.google.zetasql.SimpleCatalog) RelDataType(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rel.type.RelDataType)

Example 7 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 canDropCalc.

/**
 * Perform a series of checks to determine whether a Calc can be dropped. Following conditions
 * need to be met in order for that to happen (logical AND):<br>
 * 1. Program should do simple projects, project each field once, and project fields in the same
 * order when field reordering is not supported.<br>
 * 2. Predicate can be completely pushed-down.<br>
 * 3. Project push-down is supported by the IO or all fields are projected by a Calc.
 *
 * @param program A {@code RexProgram} of a {@code Calc}.
 * @param projectSupport An enum containing information about IO project push-down capabilities.
 * @param tableFilter A class containing information about IO predicate push-down capabilities.
 * @return True when Calc can be dropped, false otherwise.
 */
private boolean canDropCalc(RexProgram program, ProjectSupport projectSupport, BeamSqlTableFilter tableFilter) {
    RelDataType calcInputRowType = program.getInputRowType();
    // Program should do simple projects, project each field once, and project fields in the same
    // order when field reordering is not supported.
    boolean fieldReorderingSupported = projectSupport.withFieldReordering();
    if (!isProjectRenameOnlyProgram(program, fieldReorderingSupported)) {
        return false;
    }
    // Predicate can be completely pushed-down
    if (!tableFilter.getNotSupported().isEmpty()) {
        return false;
    }
    // Project push-down is supported by the IO or all fields are projected by a Calc.
    boolean isProjectSupported = projectSupport.isSupported();
    boolean allFieldsProjected = program.getProjectList().stream().map(ref -> program.getInputRowType().getFieldList().get(ref.getIndex()).getName()).collect(Collectors.toList()).equals(calcInputRowType.getFieldNames());
    return isProjectSupported || allFieldsProjected;
}
Also used : RelDataType(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rel.type.RelDataType)

Example 8 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 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 9 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 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 10 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 chooseLevels.

/**
 * Figures out which expressions to calculate at which level.
 *
 * @param exprs Array of expressions
 * @param conditionOrdinal Ordinal of the condition expression, or -1 if no condition
 * @param exprLevels Level ordinal for each expression (output)
 * @param levelTypeOrdinals The type of each level (output)
 * @return Number of levels required
 */
private int chooseLevels(final RexNode[] exprs, int conditionOrdinal, int[] exprLevels, int[] levelTypeOrdinals) {
    final int inputFieldCount = program.getInputRowType().getFieldCount();
    int levelCount = 0;
    final MaxInputFinder maxInputFinder = new MaxInputFinder(exprLevels);
    boolean[] relTypesPossibleForTopLevel = new boolean[relTypes.length];
    Arrays.fill(relTypesPossibleForTopLevel, true);
    // Compute the order in which to visit expressions.
    final List<Set<Integer>> cohorts = getCohorts();
    final List<Integer> permutation = computeTopologicalOrdering(exprs, cohorts);
    for (int i : permutation) {
        RexNode expr = exprs[i];
        final boolean condition = i == conditionOrdinal;
        if (i < inputFieldCount) {
            assert expr instanceof RexInputRef;
            exprLevels[i] = -1;
            continue;
        }
        // Deduce the minimum level of the expression. An expression must
        // be at a level greater than or equal to all of its inputs.
        int level = maxInputFinder.maxInputFor(expr);
        // If the expression is in a cohort, it can occur no lower than the
        // levels of other expressions in the same cohort.
        Set<Integer> cohort = findCohort(cohorts, i);
        if (cohort != null) {
            for (Integer exprOrdinal : cohort) {
                if (exprOrdinal == i) {
                    // of effort to repeat.
                    continue;
                }
                final RexNode cohortExpr = exprs[exprOrdinal];
                int cohortLevel = maxInputFinder.maxInputFor(cohortExpr);
                if (cohortLevel > level) {
                    level = cohortLevel;
                }
            }
        }
        // If that is not possible, try to implement it at higher levels.
        levelLoop: for (; ; ++level) {
            if (level >= levelCount) {
                // This is a new level. We can use any type we like.
                for (int relTypeOrdinal = 0; relTypeOrdinal < relTypes.length; relTypeOrdinal++) {
                    if (!relTypesPossibleForTopLevel[relTypeOrdinal]) {
                        continue;
                    }
                    if (relTypes[relTypeOrdinal].canImplement(expr, condition)) {
                        // Success. We have found a type where we can
                        // implement this expression.
                        exprLevels[i] = level;
                        levelTypeOrdinals[level] = relTypeOrdinal;
                        assert (level == 0) || (levelTypeOrdinals[level - 1] != levelTypeOrdinals[level]) : "successive levels of same type";
                        // Previous reltypes are not possible.
                        for (int j = 0; j < relTypeOrdinal; ++j) {
                            relTypesPossibleForTopLevel[j] = false;
                        }
                        // Successive reltypes may be possible.
                        for (int j = relTypeOrdinal + 1; j < relTypes.length; ++j) {
                            if (relTypesPossibleForTopLevel[j]) {
                                relTypesPossibleForTopLevel[j] = relTypes[j].canImplement(expr, condition);
                            }
                        }
                        // Move to next level.
                        levelTypeOrdinals[levelCount] = firstSet(relTypesPossibleForTopLevel);
                        ++levelCount;
                        Arrays.fill(relTypesPossibleForTopLevel, true);
                        break levelLoop;
                    }
                }
                // level, with all options open?
                if (count(relTypesPossibleForTopLevel) >= relTypes.length) {
                    // Cannot implement for any type.
                    throw new AssertionError("cannot implement " + expr);
                }
                levelTypeOrdinals[levelCount] = firstSet(relTypesPossibleForTopLevel);
                ++levelCount;
                Arrays.fill(relTypesPossibleForTopLevel, true);
            } else {
                final int levelTypeOrdinal = levelTypeOrdinals[level];
                if (!relTypes[levelTypeOrdinal].canImplement(expr, condition)) {
                    // continue to next level.
                    continue;
                }
                exprLevels[i] = level;
                break;
            }
        }
    }
    if (levelCount == 0) {
        // At least one level is always required.
        levelCount = 1;
    }
    return levelCount;
}
Also used : Set(java.util.Set) RelTraitSet(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.plan.RelTraitSet) 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)

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