Search in sources :

Example 1 with BatchPhysicalExchange

use of org.apache.flink.table.planner.plan.nodes.physical.batch.BatchPhysicalExchange 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)

Aggregations

AggregateCall (org.apache.calcite.rel.core.AggregateCall)1 DynamicTableSource (org.apache.flink.table.connector.source.DynamicTableSource)1 AggregatePushDownSpec (org.apache.flink.table.planner.plan.abilities.source.AggregatePushDownSpec)1 SourceAbilitySpec (org.apache.flink.table.planner.plan.abilities.source.SourceAbilitySpec)1 BatchPhysicalExchange (org.apache.flink.table.planner.plan.nodes.physical.batch.BatchPhysicalExchange)1 BatchPhysicalTableSourceScan (org.apache.flink.table.planner.plan.nodes.physical.batch.BatchPhysicalTableSourceScan)1 TableSourceTable (org.apache.flink.table.planner.plan.schema.TableSourceTable)1 RowType (org.apache.flink.table.types.logical.RowType)1