java/feathub-udf/flink-udf/src/main/java/com/alibaba/feathub/flink/udf/processfunction/SlidingWindowKeyedProcessFunction.java
<filename>java/feathub-udf/flink-udf/src/main/java/com/alibaba/feathub/flink/udf/processfunction/SlidingWindowKeyedProcessFunction.java/*
 * Copyright 2022 The FeatHub Authors
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     https://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package com.alibaba.feathub.flink.udf.processfunction;

import org.apache.flink.api.common.functions.RuntimeContext;
import org.apache.flink.api.common.state.ListState;
import org.apache.flink.api.common.state.ListStateDescriptor;
import org.apache.flink.api.common.state.MapState;
import org.apache.flink.api.common.state.MapStateDescriptor;
import org.apache.flink.api.common.state.ValueState;
import org.apache.flink.api.common.state.ValueStateDescriptor;
import org.apache.flink.api.common.typeinfo.TypeInformation;
import org.apache.flink.api.common.typeinfo.Types;
import org.apache.flink.api.common.typeutils.TypeSerializer;
import org.apache.flink.api.common.typeutils.base.ListSerializer;
import org.apache.flink.api.common.typeutils.base.LongSerializer;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.streaming.api.functions.KeyedProcessFunction;
import org.apache.flink.types.Row;
import org.apache.flink.util.Collector;
import org.apache.flink.util.Preconditions;

import com.alibaba.feathub.flink.udf.AggregationFieldsDescriptor;

import java.io.IOException;
import java.time.Instant;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
import java.util.HashSet;
import java.util.Iterator;
import java.util.List;
import java.util.Objects;
import java.util.Set;

/**
 * A KeyedProcessFunction that aggregate sliding windows with different sizes. The ProcessFunction
 * should be applied after the input has been aggregated with a tumble window with size equals to
 * the step size of sliding window.
 *
 * <p>One native Flink sliding window operator can only aggregate under one sliding window size. If
 * we need aggregation under different window sizes, we need to apply multiple native Flink window
 * operators with different sizes and join the results together. Each window operator keeps the rows
 * for its window size. This will result in a lots of duplicates row in the state backend. For
 * example, if we want to do aggregation under 1 hour and 2 hours, both sliding window will keep the
 * rows in the last one hour, which are duplicated. With this process function, we only keep the
 * rows for the maximum window size so that we can avoid duplicated rows in state backend.
 */
public class SlidingWindowKeyedProcessFunction extends KeyedProcessFunction<Row, Row, Row> {
    private final AggregationFieldsDescriptor aggregationFieldsDescriptor;
    private final TypeSerializer<Row> inputRowTypeSerializer;
    private final TypeSerializer<Row> outputRowTypeSerializer;
    private final String rowTimeFieldName;
    private final String[] keyFieldNames;
    private final long stepSizeMs;
    private SlidingWindowState state;
    private final SlidingWindowExpiredRowHandler expiredRowHandler;
    private final boolean skipSameWindowOutput;

    /**
     * The grace period specifies the maximum duration a row can stay in the state after it is
     * retracted from all windows. The grace period can reduce the frequency of pruning row and
     * reduce state accessing.
     */
    private final long pruneRowGracePeriod;

    public SlidingWindowKeyedProcessFunction(
            AggregationFieldsDescriptor aggregationFieldsDescriptor,
            TypeSerializer<Row> inputRowTypeSerializer,
            TypeSerializer<Row> outputRowTypeSerializer,
            String[] keyFieldNames,
            String rowTimeFieldName,
            long stepSizeMs,
            SlidingWindowExpiredRowHandler expiredRowHandler,
            boolean skipSameWindowOutput) {
        this.aggregationFieldsDescriptor = aggregationFieldsDescriptor;
        this.inputRowTypeSerializer = inputRowTypeSerializer;
        this.outputRowTypeSerializer = outputRowTypeSerializer;
        this.rowTimeFieldName = rowTimeFieldName;
        this.keyFieldNames = keyFieldNames;
        this.stepSizeMs = stepSizeMs;
        this.expiredRowHandler = expiredRowHandler;
        this.skipSameWindowOutput = skipSameWindowOutput;

        // We set the grace period to be 1/10 of the maximum window time, so that the state size is
        // at most 1/10 larger. And we only need to call prune row once every 1/10 of the maximum
        // window time.
        this.pruneRowGracePeriod =
                Math.floorDiv(aggregationFieldsDescriptor.getMaxWindowSizeMs(), 10);
    }

    @Override
    public void open(Configuration parameters) {
        state =
                SlidingWindowState.create(
                        getRuntimeContext(),
                        aggregationFieldsDescriptor,
                        inputRowTypeSerializer,
                        outputRowTypeSerializer);
    }

    @Override
    public void processElement(
            Row row, KeyedProcessFunction<Row, Row, Row>.Context ctx, Collector<Row> out)
            throws Exception {
        final long rowTime = ((Instant) row.getFieldAs(rowTimeFieldName)).toEpochMilli();

        long maxTriggeredTimer =
                state.maxTriggeredTimer.value() == null
                        ? Long.MIN_VALUE
                        : state.maxTriggeredTimer.value();

        if (rowTime + aggregationFieldsDescriptor.getMaxWindowSizeMs()
                <= ctx.timerService().currentWatermark()) {
            // TODO: Add metrics to record such data.
            return;
        }

        if (rowTime > maxTriggeredTimer) {
            state.addRow(rowTime, row, null);

            if (skipSameWindowOutput) {
                ctx.timerService().registerEventTimeTimer(rowTime);

                for (AggregationFieldsDescriptor.AggregationFieldDescriptor descriptor :
                        aggregationFieldsDescriptor.getAggFieldDescriptors()) {
                    ctx.timerService().registerEventTimeTimer(rowTime + descriptor.windowSizeMs);
                }
            }
        } else {
            Row accumulatorStates = state.getAccumulatorStates();
            Set<Integer> affectedFieldIndices = new HashSet<>();

            for (AggregationFieldsDescriptor.AggregationFieldDescriptor descriptor :
                    aggregationFieldsDescriptor.getAggFieldDescriptors()) {
                final int aggFieldIdx = aggregationFieldsDescriptor.getAggFieldIdx(descriptor);
                final long timeToRetract = rowTime + descriptor.windowSizeMs;

                if (timeToRetract <= maxTriggeredTimer) {
                    continue;
                }

                affectedFieldIndices.add(aggFieldIdx);

                Object accumulatorState = accumulatorStates.getField(aggFieldIdx);
                Object preAggResult = row.getField(descriptor.fieldName);
                descriptor.aggFunc.merge(accumulatorState, preAggResult);

                if (skipSameWindowOutput) {
                    ctx.timerService().registerEventTimeTimer(timeToRetract);
                }
            }

            state.addRow(rowTime, row, affectedFieldIndices);
            state.updateAccumulatorStates(accumulatorStates);

            if (skipSameWindowOutput) {
                // Triggers the output of the latest aggregation result instantly at the next window
                // time after late data is added to the aggregation field accumulators.
                // Note that processElement() should not be used to collect outputs in order to
                // avoid outputting aggregation results with duplicated window time.
                ctx.timerService().registerEventTimeTimer(maxTriggeredTimer + stepSizeMs);
            }
        }

        if (!skipSameWindowOutput) {
            Long triggerTime = state.maxRegisteredTimer.value();
            if (triggerTime == null) {
                triggerTime = rowTime;
            } else {
                triggerTime = Math.max(triggerTime + stepSizeMs, rowTime);
            }
            long rowExpireTime = rowTime + aggregationFieldsDescriptor.getMaxWindowSizeMs();

            // TODO: Register at most one timer at a timestamp regardless of the number of the keys
            //  processed by the operator.
            while (triggerTime <= rowExpireTime) {
                ctx.timerService().registerEventTimeTimer(triggerTime);
                triggerTime += stepSizeMs;
            }
            state.maxRegisteredTimer.update(triggerTime - stepSizeMs);
        }
    }

    @Override
    public void onTimer(
            long timestamp,
            KeyedProcessFunction<Row, Row, Row>.OnTimerContext ctx,
            Collector<Row> out)
            throws Exception {
        state.maxTriggeredTimer.update(timestamp);
        boolean hasRow = false;

        final List<Long> timestampList = state.getTimestampList();
        final int timestampListSize = timestampList.size();

        Row outputRow = Row.withNames();
        for (int i = 0; i < keyFieldNames.length; i++) {
            outputRow.setField(keyFieldNames[i], ctx.getCurrentKey().getField(i));
        }

        final Row accumulatorStates = state.getAccumulatorStates();
        final List<Integer> leftIdxList = state.getLeftTimestampIdx();
        final List<Row> rowsToAdd = state.timestampToRows.get(timestamp);

        for (AggregationFieldsDescriptor.AggregationFieldDescriptor descriptor :
                aggregationFieldsDescriptor.getAggFieldDescriptors()) {
            final int aggFieldIdx = aggregationFieldsDescriptor.getAggFieldIdx(descriptor);
            Object accumulatorState = accumulatorStates.getField(aggFieldIdx);
            if (rowsToAdd != null) {
                for (Row row : rowsToAdd) {
                    descriptor.aggFunc.merge(accumulatorState, row.getField(descriptor.fieldName));
                }
            }

            // Advance left idx and retract rows whose rowTime is out of the time window
            // (timestamp - descriptor.windowSizeMs, timestamp]
            int leftIdx = leftIdxList.get(aggFieldIdx);
            for (; leftIdx < timestampListSize; ++leftIdx) {
                long rowTime = timestampList.get(leftIdx);
                if (timestamp - descriptor.windowSizeMs < rowTime) {
                    break;
                }
                for (Row row : state.timestampToRows.get(rowTime)) {
                    descriptor.aggFunc.retractAccumulator(
                            accumulatorState, row.getField(descriptor.fieldName));
                }
            }
            if (leftIdx < timestampList.size() && timestampList.get(leftIdx) <= timestamp) {
                // If the row time of the earliest row that is not retracted is less than or
                // equal to the current time, we know there is at least one row in the
                // current window.
                hasRow = true;
            }
            leftIdxList.set(aggFieldIdx, leftIdx);

            outputRow.setField(
                    descriptor.fieldName, descriptor.aggFunc.getResult(accumulatorState));
        }
        state.updateLeftTimestampIdx(leftIdxList);
        state.updateAccumulatorStates(accumulatorStates);
        outputRow.setField(rowTimeFieldName, Instant.ofEpochMilli(timestamp));

        if (!hasRow) {
            if (expiredRowHandler != null) {
                final Row lastRow =
                        Preconditions.checkNotNull(
                                state.lastOutputRow.value(), "The last row should not be null.");
                expiredRowHandler.handleExpiredRow(out, lastRow, timestamp);
            }

            state.pruneRow(timestamp, aggregationFieldsDescriptor);
            state.lastOutputRow.clear();
            return;
        }

        if (!skipSameWindowOutput
                || !hasEqualAggregationResult(state.lastOutputRow.value(), outputRow)) {
            state.lastOutputRow.update(outputRow);
            out.collect(outputRow);
        }

        final Long earliestRowTime = timestampList.get(0);
        if (earliestRowTime
                <= timestamp
                        - aggregationFieldsDescriptor.getMaxWindowSizeMs()
                        - pruneRowGracePeriod) {
            state.pruneRow(timestamp, aggregationFieldsDescriptor);
        }
    }

    private boolean hasEqualAggregationResult(Row row0, Row row1) {
        if (row0 == null || row1 == null) {
            return row0 == null && row1 == null;
        }

        for (AggregationFieldsDescriptor.AggregationFieldDescriptor descriptor :
                aggregationFieldsDescriptor.getAggFieldDescriptors()) {
            if (!Objects.equals(
                    row0.getField(descriptor.fieldName), row1.getField(descriptor.fieldName))) {
                return false;
            }
        }

        return true;
    }

    /** The state of {@link SlidingWindowKeyedProcessFunction}. */
    private static class SlidingWindowState {

        private final AggregationFieldsDescriptor aggregationFieldsDescriptor;

        /** This MapState maps from row timestamp to the rows. */
        private final MapState<Long, List<Row>> timestampToRows;

        /**
         * This ListState keeps all the row timestamp that has been added to the timestampToRows
         * state. Since the operator assume rows are ordered by time, the timestamp list should be
         * ordered as well.
         */
        private final ListState<Long> timestampList;

        /**
         * This ValueState keeps a list of index of the entry in the timestampList for each
         * aggregation field such that all entries before it must have been outside the time window
         * for this field. And it always points to the timestamp of the earliest row that has not
         * been retracted from the aggregation function of that field. This list has the same order
         * as the aggregation fields.
         */
        private final ValueState<List<Integer>> leftTimestampIdxList;

        /**
         * This ValueState keeps the maximum registered timer so that we don't try to register the
         * same timer twice. Only used when skipSameWindowOutput is disabled.
         */
        private final ValueState<Long> maxRegisteredTimer;

        /**
         * This ValueState keeps the maximum triggered timer so that we only try to register timers
         * whose timestamps are larger than this one.
         */
        private final ValueState<Long> maxTriggeredTimer;

        /**
         * This ValueState keeps the last output row so that we can handle last output row when it
         * is expired.
         */
        private final ValueState<Row> lastOutputRow;

        /**
         * This ValueState keeps the accumulator state for the given aggregationFieldsDescriptor.
         */
        private final ValueState<Row> accumulatorStates;

        private SlidingWindowState(
                AggregationFieldsDescriptor aggregationFieldsDescriptor,
                MapState<Long, List<Row>> timestampToRows,
                ListState<Long> timestampList,
                ValueState<List<Integer>> leftTimestampIdxList,
                ValueState<Long> maxRegisteredTimer,
                ValueState<Long> maxTriggeredTimer,
                ValueState<Row> lastOutputRow,
                ValueState<Row> accumulatorStates) {
            this.aggregationFieldsDescriptor = aggregationFieldsDescriptor;
            this.timestampToRows = timestampToRows;
            this.timestampList = timestampList;
            this.leftTimestampIdxList = leftTimestampIdxList;
            this.maxRegisteredTimer = maxRegisteredTimer;
            this.maxTriggeredTimer = maxTriggeredTimer;
            this.lastOutputRow = lastOutputRow;
            this.accumulatorStates = accumulatorStates;
        }

        @SuppressWarnings({"rawtypes"})
        public static SlidingWindowState create(
                RuntimeContext context,
                AggregationFieldsDescriptor aggregationFieldsDescriptor,
                TypeSerializer<Row> inputRowTypeSerializer,
                TypeSerializer<Row> outputRowTypeSerializer) {
            final MapState<Long, List<Row>> timestampToRows =
                    context.getMapState(
                            new MapStateDescriptor<>(
                                    "TimestampToRows",
                                    LongSerializer.INSTANCE,
                                    new ListSerializer<>(inputRowTypeSerializer)));

            final ListState<Long> listState =
                    context.getListState(
                            new ListStateDescriptor<>(
                                    "TimestampListState", LongSerializer.INSTANCE));

            final ValueState<List<Integer>> outFieldNameToLeftTimestampIdx =
                    context.getState(
                            new ValueStateDescriptor<>(
                                    "OutFieldNameToLeftTimestampIdx", Types.LIST(Types.INT)));

            final ValueState<Long> maxRegisteredTimer =
                    context.getState(
                            new ValueStateDescriptor<>(
                                    "MaxRegisteredTimer", LongSerializer.INSTANCE));

            final ValueState<Long> maxTriggeredTimer =
                    context.getState(
                            new ValueStateDescriptor<>(
                                    "maxTriggeredTimer", LongSerializer.INSTANCE));

            final ValueState<Row> lastOutputRow =
                    context.getState(
                            new ValueStateDescriptor<>("LastOutputRow", outputRowTypeSerializer));

            final TypeInformation[] accumulatorTypeInformation =
                    aggregationFieldsDescriptor.getAggFieldDescriptors().stream()
                            .map(descriptor -> descriptor.aggFunc.getAccumulatorTypeInformation())
                            .toArray(TypeInformation[]::new);
            ValueState<Row> accumulatorStates =
                    context.getState(
                            new ValueStateDescriptor<>(
                                    "AccumulatorStates", Types.ROW(accumulatorTypeInformation)));

            return new SlidingWindowState(
                    aggregationFieldsDescriptor,
                    timestampToRows,
                    listState,
                    outFieldNameToLeftTimestampIdx,
                    maxRegisteredTimer,
                    maxTriggeredTimer,
                    lastOutputRow,
                    accumulatorStates);
        }

        /**
         * Add the row to the state and adjust left idx accordingly.
         *
         * @param timestamp The row time of the row.
         * @param row The row to be added.
         * @param affectedAggFieldIndices The indices of the agg fields that would add or have added
         *     this row to their accumulator. If this value is null, it means that all agg fields
         *     add this row.
         */
        public void addRow(long timestamp, Row row, Set<Integer> affectedAggFieldIndices)
                throws Exception {
            if (timestampToRows.contains(timestamp)) {
                List<Row> rows = timestampToRows.get(timestamp);
                rows.add(row);
                timestampToRows.put(timestamp, rows);
            } else {
                timestampToRows.put(timestamp, new ArrayList<>(Collections.singletonList(row)));

                List<Long> timestamps = getTimestampList();

                if (timestamps.isEmpty() || timestamps.get(timestamps.size() - 1) < timestamp) {
                    timestampList.add(timestamp);
                } else {
                    // TODO: Use skip list to optimize the performance of inserting late data.
                    for (int insertionIdx = 0; insertionIdx < timestamps.size(); insertionIdx++) {
                        if (timestamps.get(insertionIdx) < timestamp) {
                            continue;
                        }
                        timestamps.add(insertionIdx, timestamp);
                        break;
                    }
                    timestampList.update(timestamps);
                }

                // Because inserting an element in the middle of the timestamp list breaks the
                // mapping from indices to timestamp, we need to update the leftIdxList to
                // correctly reflect the timestamps of the rows to retract. How to update a leftIdx
                // depends on the comparison result between `timestamp` and `maxTriggeredTimer -
                // windowSizeMs`, which is reflected in `affectedAggFieldIndices` and may not be
                // equal to the comparison result between `insertionIdx` and `leftIdx`.
                List<Integer> leftIdxList = getLeftTimestampIdx();
                for (AggregationFieldsDescriptor.AggregationFieldDescriptor aggFieldDescriptor :
                        aggregationFieldsDescriptor.getAggFieldDescriptors()) {
                    int aggFieldIdx =
                            aggregationFieldsDescriptor.getAggFieldIdx(aggFieldDescriptor);
                    if (affectedAggFieldIndices != null
                            && !affectedAggFieldIndices.contains(aggFieldIdx)) {
                        leftIdxList.set(aggFieldIdx, leftIdxList.get(aggFieldIdx) + 1);
                    }
                }
                leftTimestampIdxList.update(leftIdxList);
            }
        }

        /**
         * Prune all the row with timestamp that are less than or equals to currentTimestamp -
         * maxWindowSizeMs. It also cleans up the state if the window is empty after prune.
         *
         * @param currentTimestamp The current timestamp.
         * @param aggFieldsDescriptor The aggFieldsDescriptor.
         */
        public void pruneRow(long currentTimestamp, AggregationFieldsDescriptor aggFieldsDescriptor)
                throws Exception {
            long lowerBound = currentTimestamp - aggFieldsDescriptor.getMaxWindowSizeMs();
            List<Long> timestamps = getTimestampList();

            final int originalSize = timestamps.size();
            final Iterator<Long> iterator = timestamps.iterator();

            while (iterator.hasNext()) {
                final long cur = iterator.next();
                if (cur > lowerBound) {
                    break;
                }
                timestampToRows.remove(cur);
                iterator.remove();
            }

            final List<Integer> leftIdxList = getLeftTimestampIdx();
            for (AggregationFieldsDescriptor.AggregationFieldDescriptor descriptor :
                    aggFieldsDescriptor.getAggFieldDescriptors()) {
                final int removedRowCnt = originalSize - timestamps.size();
                final int aggFieldIdx = aggFieldsDescriptor.getAggFieldIdx(descriptor);
                final int leftIdx = leftIdxList.get(aggFieldIdx) - removedRowCnt;
                leftIdxList.set(aggFieldIdx, Math.max(leftIdx, 0));
            }
            updateLeftTimestampIdx(leftIdxList);

            if (timestamps.isEmpty()) {
                timestampList.clear();
                leftTimestampIdxList.clear();
                maxRegisteredTimer.clear();
                lastOutputRow.clear();
                accumulatorStates.clear();
            } else {
                timestampList.update(timestamps);
            }
        }

        public List<Integer> getLeftTimestampIdx() throws Exception {
            List<Integer> leftTimestampIdx = leftTimestampIdxList.value();
            if (leftTimestampIdx == null) {
                leftTimestampIdx =
                        new ArrayList<>(
                                Arrays.asList(
                                        new Integer
                                                [aggregationFieldsDescriptor
                                                        .getAggFieldDescriptors()
                                                        .size()]));
                Collections.fill(leftTimestampIdx, 0);
            }
            return leftTimestampIdx;
        }

        public void updateLeftTimestampIdx(List<Integer> leftTimestampIdx) throws Exception {
            leftTimestampIdxList.update(leftTimestampIdx);
        }

        public List<Long> getTimestampList() throws Exception {
            List<Long> timestamps = new ArrayList<>();
            final Iterable<Long> iter = timestampList.get();
            if (iter != null) {
                iter.forEach(timestamps::add);
            }
            return timestamps;
        }

        public Row getAccumulatorStates() throws IOException {
            Row acc = accumulatorStates.value();
            if (acc == null) {
                final Object[] accumulators =
                        aggregationFieldsDescriptor.getAggFieldDescriptors().stream()
                                .map(descriptor -> descriptor.aggFunc.createAccumulator())
                                .toArray();
                acc = Row.of(accumulators);
            }
            return acc;
        }

        public void updateAccumulatorStates(Row accumulators) throws IOException {
            accumulatorStates.update(accumulators);
        }
    }
}
