All Downloads are FREE. Search and download functionalities are using the official Maven repository.

org.apache.flink.streaming.api.operators.co.CoBroadcastWithNonKeyedOperator Maven / Gradle / Ivy

There is a newer version: 1.14.6
Show newest version
/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you 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
 *
 *     http://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 org.apache.flink.streaming.api.operators.co;

import org.apache.flink.annotation.Internal;
import org.apache.flink.api.common.ExecutionConfig;
import org.apache.flink.api.common.state.BroadcastState;
import org.apache.flink.api.common.state.MapStateDescriptor;
import org.apache.flink.api.common.state.ReadOnlyBroadcastState;
import org.apache.flink.streaming.api.functions.co.BroadcastProcessFunction;
import org.apache.flink.streaming.api.functions.co.BroadcastProcessFunction.Context;
import org.apache.flink.streaming.api.operators.AbstractUdfStreamOperator;
import org.apache.flink.streaming.api.operators.InternalTimerService;
import org.apache.flink.streaming.api.operators.TimestampedCollector;
import org.apache.flink.streaming.api.operators.TwoInputStreamOperator;
import org.apache.flink.streaming.api.watermark.Watermark;
import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
import org.apache.flink.streaming.runtime.tasks.ProcessingTimeService;
import org.apache.flink.util.OutputTag;
import org.apache.flink.util.Preconditions;

import java.util.HashMap;
import java.util.List;
import java.util.Map;

import static org.apache.flink.util.Preconditions.checkArgument;
import static org.apache.flink.util.Preconditions.checkState;

/**
 * A {@link TwoInputStreamOperator} for executing {@link BroadcastProcessFunction
 * BroadcastProcessFunctions}.
 *
 * @param  The input type of the keyed (non-broadcast) side.
 * @param  The input type of the broadcast side.
 * @param  The output type of the operator.
 */
@Internal
public class CoBroadcastWithNonKeyedOperator
        extends AbstractUdfStreamOperator>
        implements TwoInputStreamOperator {

    private static final long serialVersionUID = -1869740381935471752L;

    /** We listen to this ourselves because we don't have an {@link InternalTimerService}. */
    private long currentWatermark = Long.MIN_VALUE;

    private final List> broadcastStateDescriptors;

    private transient TimestampedCollector collector;

    private transient Map, BroadcastState> broadcastStates;

    private transient ReadWriteContextImpl rwContext;

    private transient ReadOnlyContextImpl rContext;

    public CoBroadcastWithNonKeyedOperator(
            final BroadcastProcessFunction function,
            final List> broadcastStateDescriptors) {
        super(function);
        this.broadcastStateDescriptors = Preconditions.checkNotNull(broadcastStateDescriptors);
    }

    @Override
    public void open() throws Exception {
        super.open();

        collector = new TimestampedCollector<>(output);

        this.broadcastStates = new HashMap<>(broadcastStateDescriptors.size());
        for (MapStateDescriptor descriptor : broadcastStateDescriptors) {
            broadcastStates.put(
                    descriptor, getOperatorStateBackend().getBroadcastState(descriptor));
        }

        rwContext =
                new ReadWriteContextImpl(
                        getExecutionConfig(),
                        userFunction,
                        broadcastStates,
                        getProcessingTimeService());
        rContext =
                new ReadOnlyContextImpl(
                        getExecutionConfig(),
                        userFunction,
                        broadcastStates,
                        getProcessingTimeService());
    }

    @Override
    public void processElement1(StreamRecord element) throws Exception {
        collector.setTimestamp(element);
        rContext.setElement(element);
        userFunction.processElement(element.getValue(), rContext, collector);
        rContext.setElement(null);
    }

    @Override
    public void processElement2(StreamRecord element) throws Exception {
        collector.setTimestamp(element);
        rwContext.setElement(element);
        userFunction.processBroadcastElement(element.getValue(), rwContext, collector);
        rwContext.setElement(null);
    }

    @Override
    public void processWatermark(Watermark mark) throws Exception {
        super.processWatermark(mark);
        currentWatermark = mark.getTimestamp();
    }

    private class ReadWriteContextImpl extends Context {

        private final ExecutionConfig config;

        private final Map, BroadcastState> states;

        private final ProcessingTimeService timerService;

        private StreamRecord element;

        ReadWriteContextImpl(
                final ExecutionConfig executionConfig,
                final BroadcastProcessFunction function,
                final Map, BroadcastState> broadcastStates,
                final ProcessingTimeService timerService) {

            function.super();
            this.config = Preconditions.checkNotNull(executionConfig);
            this.states = Preconditions.checkNotNull(broadcastStates);
            this.timerService = Preconditions.checkNotNull(timerService);
        }

        void setElement(StreamRecord e) {
            this.element = e;
        }

        @Override
        public Long timestamp() {
            checkState(element != null);
            return element.getTimestamp();
        }

        @Override
        public  BroadcastState getBroadcastState(
                MapStateDescriptor stateDescriptor) {
            Preconditions.checkNotNull(stateDescriptor);

            stateDescriptor.initializeSerializerUnlessSet(config);
            BroadcastState state = (BroadcastState) states.get(stateDescriptor);
            if (state == null) {
                throw new IllegalArgumentException(
                        "The requested state does not exist. "
                                + "Check for typos in your state descriptor, or specify the state descriptor "
                                + "in the datastream.broadcast(...) call if you forgot to register it.");
            }
            return state;
        }

        @Override
        public  void output(OutputTag outputTag, X value) {
            checkArgument(outputTag != null, "OutputTag must not be null.");
            output.collect(outputTag, new StreamRecord<>(value, element.getTimestamp()));
        }

        @Override
        public long currentProcessingTime() {
            return timerService.getCurrentProcessingTime();
        }

        @Override
        public long currentWatermark() {
            return currentWatermark;
        }
    }

    private class ReadOnlyContextImpl
            extends BroadcastProcessFunction.ReadOnlyContext {

        private final ExecutionConfig config;

        private final Map, BroadcastState> states;

        private final ProcessingTimeService timerService;

        private StreamRecord element;

        ReadOnlyContextImpl(
                final ExecutionConfig executionConfig,
                final BroadcastProcessFunction function,
                final Map, BroadcastState> broadcastStates,
                final ProcessingTimeService timerService) {

            function.super();
            this.config = Preconditions.checkNotNull(executionConfig);
            this.states = Preconditions.checkNotNull(broadcastStates);
            this.timerService = Preconditions.checkNotNull(timerService);
        }

        void setElement(StreamRecord e) {
            this.element = e;
        }

        @Override
        public Long timestamp() {
            checkState(element != null);
            return element.hasTimestamp() ? element.getTimestamp() : null;
        }

        @Override
        public  void output(OutputTag outputTag, X value) {
            checkArgument(outputTag != null, "OutputTag must not be null.");
            output.collect(outputTag, new StreamRecord<>(value, element.getTimestamp()));
        }

        @Override
        public long currentProcessingTime() {
            return timerService.getCurrentProcessingTime();
        }

        @Override
        public long currentWatermark() {
            return currentWatermark;
        }

        @Override
        public  ReadOnlyBroadcastState getBroadcastState(
                MapStateDescriptor stateDescriptor) {
            Preconditions.checkNotNull(stateDescriptor);

            stateDescriptor.initializeSerializerUnlessSet(config);
            ReadOnlyBroadcastState state =
                    (ReadOnlyBroadcastState) states.get(stateDescriptor);
            if (state == null) {
                throw new IllegalArgumentException(
                        "The requested state does not exist. "
                                + "Check for typos in your state descriptor, or specify the state descriptor "
                                + "in the datastream.broadcast(...) call if you forgot to register it.");
            }
            return state;
        }
    }
}




© 2015 - 2024 Weber Informatics LLC | Privacy Policy