Search in sources :

Example 1 with AbstractTopNFunction

use of org.apache.flink.table.runtime.operators.rank.AbstractTopNFunction in project flink by apache.

the class StreamExecRank method translateToPlanInternal.

@SuppressWarnings("unchecked")
@Override
protected Transformation<RowData> translateToPlanInternal(PlannerBase planner, ExecNodeConfig config) {
    switch(rankType) {
        case ROW_NUMBER:
            break;
        case RANK:
            throw new TableException("RANK() on streaming table is not supported currently");
        case DENSE_RANK:
            throw new TableException("DENSE_RANK() on streaming table is not supported currently");
        default:
            throw new TableException(String.format("Streaming tables do not support %s rank function.", rankType));
    }
    ExecEdge inputEdge = getInputEdges().get(0);
    Transformation<RowData> inputTransform = (Transformation<RowData>) inputEdge.translateToPlan(planner);
    RowType inputType = (RowType) inputEdge.getOutputType();
    InternalTypeInfo<RowData> inputRowTypeInfo = InternalTypeInfo.of(inputType);
    int[] sortFields = sortSpec.getFieldIndices();
    RowDataKeySelector sortKeySelector = KeySelectorUtil.getRowDataSelector(sortFields, inputRowTypeInfo);
    // create a sort spec on sort keys.
    int[] sortKeyPositions = IntStream.range(0, sortFields.length).toArray();
    SortSpec.SortSpecBuilder builder = SortSpec.builder();
    IntStream.range(0, sortFields.length).forEach(idx -> builder.addField(idx, sortSpec.getFieldSpec(idx).getIsAscendingOrder(), sortSpec.getFieldSpec(idx).getNullIsLast()));
    SortSpec sortSpecInSortKey = builder.build();
    GeneratedRecordComparator sortKeyComparator = ComparatorCodeGenerator.gen(config.getTableConfig(), "StreamExecSortComparator", RowType.of(sortSpec.getFieldTypes(inputType)), sortSpecInSortKey);
    long cacheSize = config.get(TABLE_EXEC_RANK_TOPN_CACHE_SIZE);
    StateTtlConfig ttlConfig = StateConfigUtil.createTtlConfig(config.getStateRetentionTime());
    AbstractTopNFunction processFunction;
    if (rankStrategy instanceof RankProcessStrategy.AppendFastStrategy) {
        if (sortFields.length == 1 && TypeCheckUtils.isProcTime(inputType.getChildren().get(sortFields[0])) && sortSpec.getFieldSpec(0).getIsAscendingOrder()) {
            processFunction = new AppendOnlyFirstNFunction(ttlConfig, inputRowTypeInfo, sortKeyComparator, sortKeySelector, rankType, rankRange, generateUpdateBefore, outputRankNumber);
        } else if (RankUtil.isTop1(rankRange)) {
            processFunction = new FastTop1Function(ttlConfig, inputRowTypeInfo, sortKeyComparator, sortKeySelector, rankType, rankRange, generateUpdateBefore, outputRankNumber, cacheSize);
        } else {
            processFunction = new AppendOnlyTopNFunction(ttlConfig, inputRowTypeInfo, sortKeyComparator, sortKeySelector, rankType, rankRange, generateUpdateBefore, outputRankNumber, cacheSize);
        }
    } else if (rankStrategy instanceof RankProcessStrategy.UpdateFastStrategy) {
        if (RankUtil.isTop1(rankRange)) {
            processFunction = new FastTop1Function(ttlConfig, inputRowTypeInfo, sortKeyComparator, sortKeySelector, rankType, rankRange, generateUpdateBefore, outputRankNumber, cacheSize);
        } else {
            RankProcessStrategy.UpdateFastStrategy updateFastStrategy = (RankProcessStrategy.UpdateFastStrategy) rankStrategy;
            int[] primaryKeys = updateFastStrategy.getPrimaryKeys();
            RowDataKeySelector rowKeySelector = KeySelectorUtil.getRowDataSelector(primaryKeys, inputRowTypeInfo);
            processFunction = new UpdatableTopNFunction(ttlConfig, inputRowTypeInfo, rowKeySelector, sortKeyComparator, sortKeySelector, rankType, rankRange, generateUpdateBefore, outputRankNumber, cacheSize);
        }
    // TODO Use UnaryUpdateTopNFunction after SortedMapState is merged
    } else if (rankStrategy instanceof RankProcessStrategy.RetractStrategy) {
        EqualiserCodeGenerator equaliserCodeGen = new EqualiserCodeGenerator(inputType.getFields().stream().map(RowType.RowField::getType).toArray(LogicalType[]::new));
        GeneratedRecordEqualiser generatedEqualiser = equaliserCodeGen.generateRecordEqualiser("RankValueEqualiser");
        ComparableRecordComparator comparator = new ComparableRecordComparator(sortKeyComparator, sortKeyPositions, sortSpec.getFieldTypes(inputType), sortSpec.getAscendingOrders(), sortSpec.getNullsIsLast());
        processFunction = new RetractableTopNFunction(ttlConfig, inputRowTypeInfo, comparator, sortKeySelector, rankType, rankRange, generatedEqualiser, generateUpdateBefore, outputRankNumber);
    } else {
        throw new TableException(String.format("rank strategy:%s is not supported.", rankStrategy));
    }
    KeyedProcessOperator<RowData, RowData, RowData> operator = new KeyedProcessOperator<>(processFunction);
    processFunction.setKeyContext(operator);
    OneInputTransformation<RowData, RowData> transform = ExecNodeUtil.createOneInputTransformation(inputTransform, createTransformationMeta(RANK_TRANSFORMATION, config), operator, InternalTypeInfo.of((RowType) getOutputType()), inputTransform.getParallelism());
    // set KeyType and Selector for state
    RowDataKeySelector selector = KeySelectorUtil.getRowDataSelector(partitionSpec.getFieldIndices(), inputRowTypeInfo);
    transform.setStateKeySelector(selector);
    transform.setStateKeyType(selector.getProducedType());
    return transform;
}
Also used : OneInputTransformation(org.apache.flink.streaming.api.transformations.OneInputTransformation) Transformation(org.apache.flink.api.dag.Transformation) UpdatableTopNFunction(org.apache.flink.table.runtime.operators.rank.UpdatableTopNFunction) ExecEdge(org.apache.flink.table.planner.plan.nodes.exec.ExecEdge) RowType(org.apache.flink.table.types.logical.RowType) LogicalType(org.apache.flink.table.types.logical.LogicalType) FastTop1Function(org.apache.flink.table.runtime.operators.rank.FastTop1Function) GeneratedRecordEqualiser(org.apache.flink.table.runtime.generated.GeneratedRecordEqualiser) RankProcessStrategy(org.apache.flink.table.planner.plan.utils.RankProcessStrategy) RowData(org.apache.flink.table.data.RowData) AbstractTopNFunction(org.apache.flink.table.runtime.operators.rank.AbstractTopNFunction) RowDataKeySelector(org.apache.flink.table.runtime.keyselector.RowDataKeySelector) KeyedProcessOperator(org.apache.flink.streaming.api.operators.KeyedProcessOperator) TableException(org.apache.flink.table.api.TableException) AppendOnlyTopNFunction(org.apache.flink.table.runtime.operators.rank.AppendOnlyTopNFunction) AppendOnlyFirstNFunction(org.apache.flink.table.runtime.operators.rank.AppendOnlyFirstNFunction) EqualiserCodeGenerator(org.apache.flink.table.planner.codegen.EqualiserCodeGenerator) StateTtlConfig(org.apache.flink.api.common.state.StateTtlConfig) RetractableTopNFunction(org.apache.flink.table.runtime.operators.rank.RetractableTopNFunction) ComparableRecordComparator(org.apache.flink.table.runtime.operators.rank.ComparableRecordComparator) GeneratedRecordComparator(org.apache.flink.table.runtime.generated.GeneratedRecordComparator) SortSpec(org.apache.flink.table.planner.plan.nodes.exec.spec.SortSpec)

Aggregations

StateTtlConfig (org.apache.flink.api.common.state.StateTtlConfig)1 Transformation (org.apache.flink.api.dag.Transformation)1 KeyedProcessOperator (org.apache.flink.streaming.api.operators.KeyedProcessOperator)1 OneInputTransformation (org.apache.flink.streaming.api.transformations.OneInputTransformation)1 TableException (org.apache.flink.table.api.TableException)1 RowData (org.apache.flink.table.data.RowData)1 EqualiserCodeGenerator (org.apache.flink.table.planner.codegen.EqualiserCodeGenerator)1 ExecEdge (org.apache.flink.table.planner.plan.nodes.exec.ExecEdge)1 SortSpec (org.apache.flink.table.planner.plan.nodes.exec.spec.SortSpec)1 RankProcessStrategy (org.apache.flink.table.planner.plan.utils.RankProcessStrategy)1 GeneratedRecordComparator (org.apache.flink.table.runtime.generated.GeneratedRecordComparator)1 GeneratedRecordEqualiser (org.apache.flink.table.runtime.generated.GeneratedRecordEqualiser)1 RowDataKeySelector (org.apache.flink.table.runtime.keyselector.RowDataKeySelector)1 AbstractTopNFunction (org.apache.flink.table.runtime.operators.rank.AbstractTopNFunction)1 AppendOnlyFirstNFunction (org.apache.flink.table.runtime.operators.rank.AppendOnlyFirstNFunction)1 AppendOnlyTopNFunction (org.apache.flink.table.runtime.operators.rank.AppendOnlyTopNFunction)1 ComparableRecordComparator (org.apache.flink.table.runtime.operators.rank.ComparableRecordComparator)1 FastTop1Function (org.apache.flink.table.runtime.operators.rank.FastTop1Function)1 RetractableTopNFunction (org.apache.flink.table.runtime.operators.rank.RetractableTopNFunction)1 UpdatableTopNFunction (org.apache.flink.table.runtime.operators.rank.UpdatableTopNFunction)1