Search in sources :

Example 16 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 onMatch.

// ~ Methods ----------------------------------------------------------------
@Override
public void onMatch(RelOptRuleCall call) {
    final BeamIOSourceRel ioSourceRel = call.rel(1);
    final BeamSqlTable beamSqlTable = ioSourceRel.getBeamSqlTable();
    if (ioSourceRel instanceof BeamPushDownIOSourceRel) {
        return;
    }
    // Nested rows are not supported at the moment
    for (RelDataTypeField field : ioSourceRel.getRowType().getFieldList()) {
        if (field.getType() instanceof RelRecordType) {
            return;
        }
    }
    final Calc calc = call.rel(0);
    final RexProgram program = calc.getProgram();
    final Pair<ImmutableList<RexNode>, ImmutableList<RexNode>> projectFilter = program.split();
    final RelDataType calcInputRowType = program.getInputRowType();
    // When predicate push-down is not supported - all filters are unsupported.
    final BeamSqlTableFilter tableFilter = beamSqlTable.constructFilter(projectFilter.right);
    if (!beamSqlTable.supportsProjects().isSupported() && tableFilter instanceof DefaultTableFilter) {
        // Either project or filter push-down must be supported by the IO.
        return;
    }
    Set<String> usedFields = new LinkedHashSet<>();
    if (!(tableFilter instanceof DefaultTableFilter) && !beamSqlTable.supportsProjects().isSupported()) {
        // When applying standalone filter push-down all fields must be project by an IO.
        // With a single exception: Calc projects all fields (in the same order) and does nothing
        // else.
        usedFields.addAll(calcInputRowType.getFieldNames());
    } else {
        // Find all input refs used by projects
        for (RexNode project : projectFilter.left) {
            findUtilizedInputRefs(calcInputRowType, project, usedFields);
        }
        // Find all input refs used by filters
        for (RexNode filter : tableFilter.getNotSupported()) {
            findUtilizedInputRefs(calcInputRowType, filter, usedFields);
        }
    }
    if (usedFields.isEmpty()) {
        // No need to do push-down for queries like this: "select UPPER('hello')".
        return;
    }
    // IO only projects fields utilized by a calc.
    if (tableFilter.getNotSupported().containsAll(projectFilter.right) && usedFields.containsAll(ioSourceRel.getRowType().getFieldNames())) {
        return;
    }
    FieldAccessDescriptor resolved = FieldAccessDescriptor.withFieldNames(usedFields);
    resolved = resolved.resolve(beamSqlTable.getSchema());
    if (canDropCalc(program, beamSqlTable.supportsProjects(), tableFilter)) {
        call.transformTo(ioSourceRel.createPushDownRel(calc.getRowType(), resolved.getFieldsAccessed().stream().map(FieldDescriptor::getFieldName).collect(Collectors.toList()), tableFilter));
        return;
    }
    // IO only projects fields utilised by a calc.
    if (tableFilter.getNotSupported().equals(projectFilter.right) && usedFields.containsAll(ioSourceRel.getRowType().getFieldNames())) {
        return;
    }
    RelNode result = constructNodesWithPushDown(resolved, call.builder(), ioSourceRel, tableFilter, calc.getRowType(), projectFilter.left);
    if (tableFilter.getNotSupported().size() <= projectFilter.right.size() || usedFields.size() < calcInputRowType.getFieldCount()) {
        // Smaller Calc programs are indisputably better, as well as IOs with less projected fields.
        // We can consider something with the same number of filters.
        call.transformTo(result);
    }
}
Also used : LinkedHashSet(java.util.LinkedHashSet) BeamPushDownIOSourceRel(org.apache.beam.sdk.extensions.sql.impl.rel.BeamPushDownIOSourceRel) FieldAccessDescriptor(org.apache.beam.sdk.schemas.FieldAccessDescriptor) RexProgram(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rex.RexProgram) ImmutableList(org.apache.beam.vendor.calcite.v1_28_0.com.google.common.collect.ImmutableList) Calc(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rel.core.Calc) RelDataType(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rel.type.RelDataType) RelRecordType(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rel.type.RelRecordType) RelDataTypeField(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rel.type.RelDataTypeField) BeamSqlTableFilter(org.apache.beam.sdk.extensions.sql.meta.BeamSqlTableFilter) RelNode(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rel.RelNode) BeamSqlTable(org.apache.beam.sdk.extensions.sql.meta.BeamSqlTable) BeamIOSourceRel(org.apache.beam.sdk.extensions.sql.impl.rel.BeamIOSourceRel) DefaultTableFilter(org.apache.beam.sdk.extensions.sql.meta.DefaultTableFilter) RexNode(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rex.RexNode)

Example 17 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 BeamJoinRel method extractJoinRexNodes.

static List<Pair<RexNode, RexNode>> extractJoinRexNodes(RexNode condition) {
    // or it's a JOIN ON false because: condition == false
    if (condition instanceof RexLiteral) {
        throw new UnsupportedOperationException("CROSS JOIN, JOIN ON FALSE is not supported!");
    }
    RexCall call = (RexCall) condition;
    List<Pair<RexNode, RexNode>> pairs = new ArrayList<>();
    if ("AND".equals(call.getOperator().getName())) {
        List<RexNode> operands = call.getOperands();
        for (RexNode rexNode : operands) {
            Pair<RexNode, RexNode> pair = extractJoinPairOfRexNodes((RexCall) rexNode);
            pairs.add(pair);
        }
    } else if ("=".equals(call.getOperator().getName())) {
        pairs.add(extractJoinPairOfRexNodes(call));
    } else {
        throw new UnsupportedOperationException("Operator " + call.getOperator().getName() + " is not supported in join condition");
    }
    return pairs;
}
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) ArrayList(java.util.ArrayList) Pair(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.util.Pair) RexNode(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rex.RexNode)

Example 18 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 AggregateScanConverter method convertAggregateScanInputScanToLogicalProject.

private LogicalProject convertAggregateScanInputScanToLogicalProject(ResolvedAggregateScan node, RelNode input) {
    // AggregateScan's input is the source of data (e.g. TableScan), which is different from the
    // design of CalciteSQL, in which the LogicalAggregate's input is a LogicalProject, whose input
    // is a LogicalTableScan. When AggregateScan's input is WithRefScan, the WithRefScan is
    // ebullient to a LogicalTableScan. So it's still required to build another LogicalProject as
    // the input of LogicalAggregate.
    List<RexNode> projects = new ArrayList<>();
    List<String> fieldNames = new ArrayList<>();
    // LogicalAggregate.
    for (ResolvedComputedColumn computedColumn : node.getGroupByList()) {
        projects.add(getExpressionConverter().convertRexNodeFromResolvedExpr(computedColumn.getExpr(), node.getInputScan().getColumnList(), input.getRowType().getFieldList(), ImmutableMap.of()));
        fieldNames.add(getTrait().resolveAlias(computedColumn.getColumn()));
    }
    // TODO: remove duplicate columns in projects.
    for (ResolvedComputedColumn resolvedComputedColumn : node.getAggregateList()) {
        // Should create Calcite's RexInputRef from ResolvedColumn from ResolvedComputedColumn.
        // TODO: handle aggregate function with more than one argument and handle OVER
        // TODO: is there is general way for column reference tracking and deduplication for
        // aggregation?
        ResolvedAggregateFunctionCall aggregateFunctionCall = ((ResolvedAggregateFunctionCall) resolvedComputedColumn.getExpr());
        if (aggregateFunctionCall.getArgumentList() != null && aggregateFunctionCall.getArgumentList().size() >= 1) {
            ResolvedExpr resolvedExpr = aggregateFunctionCall.getArgumentList().get(0);
            for (int i = 0; i < aggregateFunctionCall.getArgumentList().size(); i++) {
                if (i == 0) {
                    // TODO: assume aggregate function's input is either a ColumnRef or a cast(ColumnRef).
                    // TODO: user might use multiple CAST so we need to handle this rare case.
                    projects.add(getExpressionConverter().convertRexNodeFromResolvedExpr(resolvedExpr, node.getInputScan().getColumnList(), input.getRowType().getFieldList(), ImmutableMap.of()));
                } else {
                    projects.add(getExpressionConverter().convertRexNodeFromResolvedExpr(aggregateFunctionCall.getArgumentList().get(i)));
                }
                fieldNames.add(getTrait().resolveAlias(resolvedComputedColumn.getColumn()));
            }
        }
    }
    return LogicalProject.create(input, ImmutableList.of(), projects, fieldNames);
}
Also used : ResolvedAggregateFunctionCall(com.google.zetasql.resolvedast.ResolvedNodes.ResolvedAggregateFunctionCall) ArrayList(java.util.ArrayList) ResolvedExpr(com.google.zetasql.resolvedast.ResolvedNodes.ResolvedExpr) ResolvedComputedColumn(com.google.zetasql.resolvedast.ResolvedNodes.ResolvedComputedColumn) RexNode(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rex.RexNode)

Example 19 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 LimitOffsetScanToLimitConverter method convert.

@Override
public RelNode convert(ResolvedLimitOffsetScan zetaNode, List<RelNode> inputs) {
    RelNode input = inputs.get(0);
    RelCollation relCollation = RelCollations.of(ImmutableList.of());
    RexNode offset = zetaNode.getOffset() == null ? null : getExpressionConverter().convertRexNodeFromResolvedExpr(zetaNode.getOffset());
    RexNode fetch = getExpressionConverter().convertRexNodeFromResolvedExpr(zetaNode.getLimit(), zetaNode.getColumnList(), input.getRowType().getFieldList(), ImmutableMap.of());
    // offset or fetch being RexDynamicParam means it is NULL (the only param supported currently)
    if (offset instanceof RexDynamicParam || RexLiteral.isNullLiteral(offset) || fetch instanceof RexDynamicParam || RexLiteral.isNullLiteral(fetch)) {
        throw new UnsupportedOperationException("Limit requires non-null count and offset");
    }
    return LogicalSort.create(input, relCollation, offset, fetch);
}
Also used : RelCollation(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rel.RelCollation) RelNode(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rel.RelNode) RexNode(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rex.RexNode) RexDynamicParam(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rex.RexDynamicParam)

Example 20 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 ZetaSqlUnnest method deriveUncollectRowType.

/**
 * Returns the row type returned by applying the 'UNNEST' operation to a relational expression.
 *
 * <p>Each column in the relational expression must be a multiset of structs or an array. The
 * return type is the type of that column, plus an ORDINALITY column if {@code withOrdinality}.
 */
public static RelDataType deriveUncollectRowType(RelNode rel, boolean withOrdinality) {
    RelDataType inputType = rel.getRowType();
    assert inputType.isStruct() : inputType + " is not a struct";
    final List<RelDataTypeField> fields = inputType.getFieldList();
    final RelDataTypeFactory typeFactory = rel.getCluster().getTypeFactory();
    final RelDataTypeFactory.Builder builder = typeFactory.builder();
    if (fields.size() == 1 && fields.get(0).getType().getSqlTypeName() == SqlTypeName.ANY) {
        // and Any type.
        return builder.add(fields.get(0).getName(), SqlTypeName.ANY).nullable(true).build();
    }
    for (RelDataTypeField field : fields) {
        if (field.getType() instanceof MapSqlType) {
            builder.add(SqlUnnestOperator.MAP_KEY_COLUMN_NAME, Preconditions.checkArgumentNotNull(field.getType().getKeyType(), "Encountered MAP type with null key type in field %s", field));
            builder.add(SqlUnnestOperator.MAP_VALUE_COLUMN_NAME, Preconditions.checkArgumentNotNull(field.getType().getValueType(), "Encountered MAP type with null value type in field %s", field));
        } else {
            assert field.getType() instanceof ArraySqlType;
            RelDataType ret = Preconditions.checkArgumentNotNull(field.getType().getComponentType(), "Encountered ARRAY type with null component type in field %s", field);
            // Only difference than Uncollect.java: treats record type and scalar type equally
            builder.add(SqlUtil.deriveAliasFromOrdinal(field.getIndex()), ret);
        }
    }
    if (withOrdinality) {
        builder.add(SqlUnnestOperator.ORDINALITY_COLUMN_NAME, SqlTypeName.INTEGER);
    }
    return builder.build();
}
Also used : RelDataTypeField(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rel.type.RelDataTypeField) MapSqlType(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.sql.type.MapSqlType) RelDataTypeFactory(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rel.type.RelDataTypeFactory) RelDataType(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rel.type.RelDataType) ArraySqlType(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.sql.type.ArraySqlType)

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