Search in sources :

Example 1 with AggregatePushDownSpec

use of org.apache.flink.table.planner.plan.abilities.source.AggregatePushDownSpec in project flink by apache.

the class PushLocalAggIntoScanRuleBase method pushLocalAggregateIntoScan.

protected void pushLocalAggregateIntoScan(RelOptRuleCall call, BatchPhysicalGroupAggregateBase localAgg, BatchPhysicalTableSourceScan oldScan, int[] calcRefFields) {
    RowType inputType = FlinkTypeFactory.toLogicalRowType(oldScan.getRowType());
    List<int[]> groupingSets = Collections.singletonList(ArrayUtils.addAll(localAgg.grouping(), localAgg.auxGrouping()));
    List<AggregateCall> aggCallList = JavaScalaConversionUtil.toJava(localAgg.getAggCallList());
    // map arg index in aggregate to field index in scan through referred fields by calc.
    if (calcRefFields != null) {
        groupingSets = translateGroupingArgIndex(groupingSets, calcRefFields);
        aggCallList = translateAggCallArgIndex(aggCallList, calcRefFields);
    }
    RowType producedType = FlinkTypeFactory.toLogicalRowType(localAgg.getRowType());
    TableSourceTable oldTableSourceTable = oldScan.tableSourceTable();
    DynamicTableSource newTableSource = oldScan.tableSource().copy();
    boolean isPushDownSuccess = AggregatePushDownSpec.apply(inputType, groupingSets, aggCallList, producedType, newTableSource, SourceAbilityContext.from(oldScan));
    if (!isPushDownSuccess) {
        // aggregate push down failed, just return without changing any nodes.
        return;
    }
    // create new source table with new spec and statistic.
    AggregatePushDownSpec aggregatePushDownSpec = new AggregatePushDownSpec(inputType, groupingSets, aggCallList, producedType);
    TableSourceTable newTableSourceTable = oldTableSourceTable.copy(newTableSource, localAgg.getRowType(), new SourceAbilitySpec[] { aggregatePushDownSpec }).copy(FlinkStatistic.UNKNOWN());
    // transform to new nodes.
    BatchPhysicalTableSourceScan newScan = oldScan.copy(oldScan.getTraitSet(), newTableSourceTable);
    BatchPhysicalExchange oldExchange = call.rel(0);
    BatchPhysicalExchange newExchange = oldExchange.copy(oldExchange.getTraitSet(), newScan, oldExchange.getDistribution());
    call.transformTo(newExchange);
}
Also used : AggregateCall(org.apache.calcite.rel.core.AggregateCall) SourceAbilitySpec(org.apache.flink.table.planner.plan.abilities.source.SourceAbilitySpec) AggregatePushDownSpec(org.apache.flink.table.planner.plan.abilities.source.AggregatePushDownSpec) RowType(org.apache.flink.table.types.logical.RowType) TableSourceTable(org.apache.flink.table.planner.plan.schema.TableSourceTable) BatchPhysicalExchange(org.apache.flink.table.planner.plan.nodes.physical.batch.BatchPhysicalExchange) DynamicTableSource(org.apache.flink.table.connector.source.DynamicTableSource) BatchPhysicalTableSourceScan(org.apache.flink.table.planner.plan.nodes.physical.batch.BatchPhysicalTableSourceScan)

Example 2 with AggregatePushDownSpec

use of org.apache.flink.table.planner.plan.abilities.source.AggregatePushDownSpec in project flink by apache.

the class PushLocalAggIntoScanRuleBase method canPushDown.

protected boolean canPushDown(RelOptRuleCall call, BatchPhysicalGroupAggregateBase aggregate, BatchPhysicalTableSourceScan tableSourceScan) {
    TableConfig tableConfig = ShortcutUtils.unwrapContext(call.getPlanner()).getTableConfig();
    if (!tableConfig.getConfiguration().getBoolean(OptimizerConfigOptions.TABLE_OPTIMIZER_SOURCE_AGGREGATE_PUSHDOWN_ENABLED)) {
        return false;
    }
    if (aggregate.isFinal() || aggregate.getAggCallList().isEmpty()) {
        return false;
    }
    List<AggregateCall> aggCallList = JavaScalaConversionUtil.toJava(aggregate.getAggCallList());
    for (AggregateCall aggCall : aggCallList) {
        if (aggCall.isDistinct() || aggCall.isApproximate() || aggCall.getArgList().size() > 1 || aggCall.hasFilter() || !aggCall.getCollation().getFieldCollations().isEmpty()) {
            return false;
        }
    }
    TableSourceTable tableSourceTable = tableSourceScan.tableSourceTable();
    // we can not push aggregates twice
    return tableSourceTable != null && tableSourceTable.tableSource() instanceof SupportsAggregatePushDown && Arrays.stream(tableSourceTable.abilitySpecs()).noneMatch(spec -> spec instanceof AggregatePushDownSpec);
}
Also used : AggregateCall(org.apache.calcite.rel.core.AggregateCall) SupportsAggregatePushDown(org.apache.flink.table.connector.source.abilities.SupportsAggregatePushDown) Arrays(java.util.Arrays) RexProgram(org.apache.calcite.rex.RexProgram) AggregatePushDownSpec(org.apache.flink.table.planner.plan.abilities.source.AggregatePushDownSpec) SourceAbilityContext(org.apache.flink.table.planner.plan.abilities.source.SourceAbilityContext) RexNodeExtractor(org.apache.flink.table.planner.plan.utils.RexNodeExtractor) ShortcutUtils(org.apache.flink.table.planner.utils.ShortcutUtils) ArrayUtils(org.apache.commons.lang3.ArrayUtils) FlinkTypeFactory(org.apache.flink.table.planner.calcite.FlinkTypeFactory) RowType(org.apache.flink.table.types.logical.RowType) ArrayList(java.util.ArrayList) OptimizerConfigOptions(org.apache.flink.table.api.config.OptimizerConfigOptions) BatchPhysicalCalc(org.apache.flink.table.planner.plan.nodes.physical.batch.BatchPhysicalCalc) RelOptRuleOperand(org.apache.calcite.plan.RelOptRuleOperand) RexNode(org.apache.calcite.rex.RexNode) DynamicTableSource(org.apache.flink.table.connector.source.DynamicTableSource) TableConfig(org.apache.flink.table.api.TableConfig) ProjectPushDownSpec(org.apache.flink.table.planner.plan.abilities.source.ProjectPushDownSpec) BatchPhysicalTableSourceScan(org.apache.flink.table.planner.plan.nodes.physical.batch.BatchPhysicalTableSourceScan) BatchPhysicalGroupAggregateBase(org.apache.flink.table.planner.plan.nodes.physical.batch.BatchPhysicalGroupAggregateBase) Collectors(java.util.stream.Collectors) SourceAbilitySpec(org.apache.flink.table.planner.plan.abilities.source.SourceAbilitySpec) TableSourceTable(org.apache.flink.table.planner.plan.schema.TableSourceTable) RelOptRuleCall(org.apache.calcite.plan.RelOptRuleCall) RexInputRef(org.apache.calcite.rex.RexInputRef) RelOptRule(org.apache.calcite.plan.RelOptRule) List(java.util.List) BatchPhysicalExchange(org.apache.flink.table.planner.plan.nodes.physical.batch.BatchPhysicalExchange) FlinkStatistic(org.apache.flink.table.planner.plan.stats.FlinkStatistic) JavaScalaConversionUtil(org.apache.flink.table.planner.utils.JavaScalaConversionUtil) AggregateCall(org.apache.calcite.rel.core.AggregateCall) Collections(java.util.Collections) AggregatePushDownSpec(org.apache.flink.table.planner.plan.abilities.source.AggregatePushDownSpec) TableConfig(org.apache.flink.table.api.TableConfig) TableSourceTable(org.apache.flink.table.planner.plan.schema.TableSourceTable) SupportsAggregatePushDown(org.apache.flink.table.connector.source.abilities.SupportsAggregatePushDown)

Aggregations

AggregateCall (org.apache.calcite.rel.core.AggregateCall)2 DynamicTableSource (org.apache.flink.table.connector.source.DynamicTableSource)2 AggregatePushDownSpec (org.apache.flink.table.planner.plan.abilities.source.AggregatePushDownSpec)2 SourceAbilitySpec (org.apache.flink.table.planner.plan.abilities.source.SourceAbilitySpec)2 BatchPhysicalExchange (org.apache.flink.table.planner.plan.nodes.physical.batch.BatchPhysicalExchange)2 BatchPhysicalTableSourceScan (org.apache.flink.table.planner.plan.nodes.physical.batch.BatchPhysicalTableSourceScan)2 TableSourceTable (org.apache.flink.table.planner.plan.schema.TableSourceTable)2 RowType (org.apache.flink.table.types.logical.RowType)2 ArrayList (java.util.ArrayList)1 Arrays (java.util.Arrays)1 Collections (java.util.Collections)1 List (java.util.List)1 Collectors (java.util.stream.Collectors)1 RelOptRule (org.apache.calcite.plan.RelOptRule)1 RelOptRuleCall (org.apache.calcite.plan.RelOptRuleCall)1 RelOptRuleOperand (org.apache.calcite.plan.RelOptRuleOperand)1 RexInputRef (org.apache.calcite.rex.RexInputRef)1 RexNode (org.apache.calcite.rex.RexNode)1 RexProgram (org.apache.calcite.rex.RexProgram)1 ArrayUtils (org.apache.commons.lang3.ArrayUtils)1