This is an automated email from the ASF dual-hosted git repository. guoweijie pushed a commit to branch master in repository https://gitbox.apache.org/repos/asf/flink.git
commit 4f71c5b4660be0da199123f085c8605f654803cf Author: Weijie Guo <[email protected]> AuthorDate: Fri Mar 22 17:03:32 2024 +0800 [FLINK-34548][API] Implement process function's underlying operators --- .../impl/common/KeyCheckedOutputCollector.java | 72 ++++++++++ .../datastream/impl/common/OutputCollector.java | 43 ++++++ .../datastream/impl/common/TimestampCollector.java | 43 ++++++ .../impl/context/DefaultNonPartitionedContext.java | 31 ++++ .../impl/context/DefaultRuntimeContext.java | 26 ++++ .../DefaultTwoOutputNonPartitionedContext.java | 32 +++++ .../impl/operators/KeyedProcessOperator.java | 53 +++++++ .../KeyedTwoInputBroadcastProcessOperator.java | 54 +++++++ .../KeyedTwoInputNonBroadcastProcessOperator.java | 56 ++++++++ .../operators/KeyedTwoOutputProcessOperator.java | 78 ++++++++++ .../datastream/impl/operators/ProcessOperator.java | 71 +++++++++ .../TwoInputBroadcastProcessOperator.java | 86 +++++++++++ .../TwoInputNonBroadcastProcessOperator.java | 86 +++++++++++ .../impl/operators/TwoOutputProcessOperator.java | 113 +++++++++++++++ .../impl/common/KeyCheckedOutputCollectorTest.java | 72 ++++++++++ .../impl/common/OutputCollectorTest.java | 53 +++++++ .../impl/common/TestingTimestampCollector.java | 72 ++++++++++ .../impl/operators/KeyedProcessOperatorTest.java | 123 ++++++++++++++++ .../KeyedTwoInputBroadcastProcessOperatorTest.java | 155 ++++++++++++++++++++ ...yedTwoInputNonBroadcastProcessOperatorTest.java | 156 ++++++++++++++++++++ .../KeyedTwoOutputProcessOperatorTest.java | 158 +++++++++++++++++++++ .../impl/operators/ProcessOperatorTest.java | 83 +++++++++++ .../TwoInputBroadcastProcessOperatorTest.java | 112 +++++++++++++++ .../TwoInputNonBroadcastProcessOperatorTest.java | 115 +++++++++++++++ .../operators/TwoOutputProcessOperatorTest.java | 108 ++++++++++++++ .../util/TwoInputStreamOperatorTestHarness.java | 13 ++ 26 files changed, 2064 insertions(+) diff --git a/flink-datastream/src/main/java/org/apache/flink/datastream/impl/common/KeyCheckedOutputCollector.java b/flink-datastream/src/main/java/org/apache/flink/datastream/impl/common/KeyCheckedOutputCollector.java new file mode 100644 index 00000000000..5902b0500d1 --- /dev/null +++ b/flink-datastream/src/main/java/org/apache/flink/datastream/impl/common/KeyCheckedOutputCollector.java @@ -0,0 +1,72 @@ +/* + * 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.datastream.impl.common; + +import org.apache.flink.api.java.functions.KeySelector; +import org.apache.flink.util.ExceptionUtils; + +import java.util.function.Supplier; + +import static org.apache.flink.util.Preconditions.checkNotNull; + +/** + * This output checks whether the current key of the output record and the key extracted with a + * specific key selector are exactly the same. + */ +public class KeyCheckedOutputCollector<KEY, OUT> extends TimestampCollector<OUT> { + private final TimestampCollector<OUT> collector; + + private final KeySelector<OUT, KEY> outKeySelector; + + private final Supplier<KEY> currentKeyGetter; + + public KeyCheckedOutputCollector( + TimestampCollector<OUT> collector, + KeySelector<OUT, KEY> outKeySelector, + Supplier<KEY> currentKeyGetter) { + this.collector = collector; + this.outKeySelector = outKeySelector; + this.currentKeyGetter = currentKeyGetter; + } + + @Override + public void collect(OUT outputRecord) { + checkOutputKey(outputRecord); + collector.collect(outputRecord); + } + + @Override + public void collectAndOverwriteTimestamp(OUT outputRecord, long timestamp) { + checkOutputKey(outputRecord); + collector.collectAndOverwriteTimestamp(outputRecord, timestamp); + } + + private void checkOutputKey(OUT outputRecord) { + try { + KEY currentKey = currentKeyGetter.get(); + KEY outputKey = checkNotNull(outKeySelector).getKey(outputRecord); + if (!outputKey.equals(currentKey)) { + throw new IllegalStateException( + "Output key must equals to input key if the output key selector is not null."); + } + } catch (Exception e) { + ExceptionUtils.rethrow(e); + } + } +} diff --git a/flink-datastream/src/main/java/org/apache/flink/datastream/impl/common/OutputCollector.java b/flink-datastream/src/main/java/org/apache/flink/datastream/impl/common/OutputCollector.java new file mode 100644 index 00000000000..1157749d81d --- /dev/null +++ b/flink-datastream/src/main/java/org/apache/flink/datastream/impl/common/OutputCollector.java @@ -0,0 +1,43 @@ +/* + * 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.datastream.impl.common; + +import org.apache.flink.datastream.api.common.Collector; +import org.apache.flink.streaming.api.operators.Output; +import org.apache.flink.streaming.runtime.streamrecord.StreamRecord; + +/** The default implementation of {@link Collector}. */ +public class OutputCollector<OUT> extends TimestampCollector<OUT> { + private final Output<StreamRecord<OUT>> output; + + public OutputCollector(Output<StreamRecord<OUT>> output) { + this.output = output; + } + + @Override + public void collect(OUT outputRecord) { + output.collect(reuse.replace(outputRecord)); + } + + @Override + public void collectAndOverwriteTimestamp(OUT record, long timestamp) { + setTimestamp(timestamp); + output.collect(reuse.replace(record)); + } +} diff --git a/flink-datastream/src/main/java/org/apache/flink/datastream/impl/common/TimestampCollector.java b/flink-datastream/src/main/java/org/apache/flink/datastream/impl/common/TimestampCollector.java new file mode 100644 index 00000000000..a980ce3918e --- /dev/null +++ b/flink-datastream/src/main/java/org/apache/flink/datastream/impl/common/TimestampCollector.java @@ -0,0 +1,43 @@ +/* + * 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.datastream.impl.common; + +import org.apache.flink.datastream.api.common.Collector; +import org.apache.flink.streaming.runtime.streamrecord.StreamRecord; + +/** The base {@link Collector} which take care of records timestamp. */ +public abstract class TimestampCollector<OUT> implements Collector<OUT> { + protected final StreamRecord<OUT> reuse = new StreamRecord<>(null); + + public void setTimestampFromStreamRecord(StreamRecord<?> timestampBase) { + if (timestampBase.hasTimestamp()) { + setTimestamp(timestampBase.getTimestamp()); + } else { + eraseTimestamp(); + } + } + + public void setTimestamp(long timestamp) { + reuse.setTimestamp(timestamp); + } + + public void eraseTimestamp() { + reuse.eraseTimestamp(); + } +} diff --git a/flink-datastream/src/main/java/org/apache/flink/datastream/impl/context/DefaultNonPartitionedContext.java b/flink-datastream/src/main/java/org/apache/flink/datastream/impl/context/DefaultNonPartitionedContext.java new file mode 100644 index 00000000000..017bdcb8e25 --- /dev/null +++ b/flink-datastream/src/main/java/org/apache/flink/datastream/impl/context/DefaultNonPartitionedContext.java @@ -0,0 +1,31 @@ +/* + * 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.datastream.impl.context; + +import org.apache.flink.datastream.api.context.NonPartitionedContext; +import org.apache.flink.datastream.api.function.ApplyPartitionFunction; + +/** The default implementation of {@link NonPartitionedContext}. */ +public class DefaultNonPartitionedContext<OUT> implements NonPartitionedContext<OUT> { + + @Override + public void applyToAllPartitions(ApplyPartitionFunction<OUT> applyPartitionFunction) { + // TODO implements this method. + } +} diff --git a/flink-datastream/src/main/java/org/apache/flink/datastream/impl/context/DefaultRuntimeContext.java b/flink-datastream/src/main/java/org/apache/flink/datastream/impl/context/DefaultRuntimeContext.java new file mode 100644 index 00000000000..22f61223ca1 --- /dev/null +++ b/flink-datastream/src/main/java/org/apache/flink/datastream/impl/context/DefaultRuntimeContext.java @@ -0,0 +1,26 @@ +/* + * 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.datastream.impl.context; + +import org.apache.flink.datastream.api.context.RuntimeContext; + +/** The default implementation of {@link RuntimeContext}. */ +public class DefaultRuntimeContext implements RuntimeContext { + // TODO implements this class +} diff --git a/flink-datastream/src/main/java/org/apache/flink/datastream/impl/context/DefaultTwoOutputNonPartitionedContext.java b/flink-datastream/src/main/java/org/apache/flink/datastream/impl/context/DefaultTwoOutputNonPartitionedContext.java new file mode 100644 index 00000000000..7a76e923191 --- /dev/null +++ b/flink-datastream/src/main/java/org/apache/flink/datastream/impl/context/DefaultTwoOutputNonPartitionedContext.java @@ -0,0 +1,32 @@ +/* + * 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.datastream.impl.context; + +import org.apache.flink.datastream.api.context.TwoOutputNonPartitionedContext; +import org.apache.flink.datastream.api.function.TwoOutputApplyPartitionFunction; + +/** The default implementation of {@link TwoOutputNonPartitionedContext}. */ +public class DefaultTwoOutputNonPartitionedContext<OUT1, OUT2> + implements TwoOutputNonPartitionedContext<OUT1, OUT2> { + @Override + public void applyToAllPartitions( + TwoOutputApplyPartitionFunction<OUT1, OUT2> applyPartitionFunction) { + // TODO implements this method. + } +} diff --git a/flink-datastream/src/main/java/org/apache/flink/datastream/impl/operators/KeyedProcessOperator.java b/flink-datastream/src/main/java/org/apache/flink/datastream/impl/operators/KeyedProcessOperator.java new file mode 100644 index 00000000000..92a07cf8623 --- /dev/null +++ b/flink-datastream/src/main/java/org/apache/flink/datastream/impl/operators/KeyedProcessOperator.java @@ -0,0 +1,53 @@ +/* + * 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.datastream.impl.operators; + +import org.apache.flink.api.java.functions.KeySelector; +import org.apache.flink.datastream.api.function.OneInputStreamProcessFunction; +import org.apache.flink.datastream.api.stream.KeyedPartitionStream; +import org.apache.flink.datastream.impl.common.KeyCheckedOutputCollector; +import org.apache.flink.datastream.impl.common.OutputCollector; +import org.apache.flink.datastream.impl.common.TimestampCollector; + +import javax.annotation.Nullable; + +/** Operator for {@link OneInputStreamProcessFunction} in {@link KeyedPartitionStream}. */ +public class KeyedProcessOperator<KEY, IN, OUT> extends ProcessOperator<IN, OUT> { + + @Nullable private final KeySelector<OUT, KEY> outKeySelector; + + public KeyedProcessOperator(OneInputStreamProcessFunction<IN, OUT> userFunction) { + this(userFunction, null); + } + + public KeyedProcessOperator( + OneInputStreamProcessFunction<IN, OUT> userFunction, + @Nullable KeySelector<OUT, KEY> outKeySelector) { + super(userFunction); + this.outKeySelector = outKeySelector; + } + + @Override + protected TimestampCollector<OUT> getOutputCollector() { + return outKeySelector != null + ? new KeyCheckedOutputCollector<>( + new OutputCollector<>(output), outKeySelector, () -> (KEY) getCurrentKey()) + : new OutputCollector<>(output); + } +} diff --git a/flink-datastream/src/main/java/org/apache/flink/datastream/impl/operators/KeyedTwoInputBroadcastProcessOperator.java b/flink-datastream/src/main/java/org/apache/flink/datastream/impl/operators/KeyedTwoInputBroadcastProcessOperator.java new file mode 100644 index 00000000000..c3a5664f287 --- /dev/null +++ b/flink-datastream/src/main/java/org/apache/flink/datastream/impl/operators/KeyedTwoInputBroadcastProcessOperator.java @@ -0,0 +1,54 @@ +/* + * 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.datastream.impl.operators; + +import org.apache.flink.api.java.functions.KeySelector; +import org.apache.flink.datastream.api.function.TwoInputBroadcastStreamProcessFunction; +import org.apache.flink.datastream.api.stream.KeyedPartitionStream; +import org.apache.flink.datastream.impl.common.KeyCheckedOutputCollector; +import org.apache.flink.datastream.impl.common.OutputCollector; +import org.apache.flink.datastream.impl.common.TimestampCollector; + +import javax.annotation.Nullable; + +/** Operator for {@link TwoInputBroadcastStreamProcessFunction} in {@link KeyedPartitionStream}. */ +public class KeyedTwoInputBroadcastProcessOperator<KEY, IN1, IN2, OUT> + extends TwoInputBroadcastProcessOperator<IN1, IN2, OUT> { + @Nullable private final KeySelector<OUT, KEY> outKeySelector; + + public KeyedTwoInputBroadcastProcessOperator( + TwoInputBroadcastStreamProcessFunction<IN1, IN2, OUT> userFunction) { + this(userFunction, null); + } + + public KeyedTwoInputBroadcastProcessOperator( + TwoInputBroadcastStreamProcessFunction<IN1, IN2, OUT> userFunction, + @Nullable KeySelector<OUT, KEY> outKeySelector) { + super(userFunction); + this.outKeySelector = outKeySelector; + } + + @Override + protected TimestampCollector<OUT> getOutputCollector() { + return outKeySelector == null + ? new OutputCollector<>(output) + : new KeyCheckedOutputCollector<>( + new OutputCollector<>(output), outKeySelector, () -> (KEY) getCurrentKey()); + } +} diff --git a/flink-datastream/src/main/java/org/apache/flink/datastream/impl/operators/KeyedTwoInputNonBroadcastProcessOperator.java b/flink-datastream/src/main/java/org/apache/flink/datastream/impl/operators/KeyedTwoInputNonBroadcastProcessOperator.java new file mode 100644 index 00000000000..5e4945cfe19 --- /dev/null +++ b/flink-datastream/src/main/java/org/apache/flink/datastream/impl/operators/KeyedTwoInputNonBroadcastProcessOperator.java @@ -0,0 +1,56 @@ +/* + * 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.datastream.impl.operators; + +import org.apache.flink.api.java.functions.KeySelector; +import org.apache.flink.datastream.api.function.TwoInputNonBroadcastStreamProcessFunction; +import org.apache.flink.datastream.api.stream.KeyedPartitionStream; +import org.apache.flink.datastream.impl.common.KeyCheckedOutputCollector; +import org.apache.flink.datastream.impl.common.OutputCollector; +import org.apache.flink.datastream.impl.common.TimestampCollector; + +import javax.annotation.Nullable; + +/** + * Operator for {@link TwoInputNonBroadcastStreamProcessFunction} in {@link KeyedPartitionStream}. + */ +public class KeyedTwoInputNonBroadcastProcessOperator<KEY, IN1, IN2, OUT> + extends TwoInputNonBroadcastProcessOperator<IN1, IN2, OUT> { + @Nullable private final KeySelector<OUT, KEY> outKeySelector; + + public KeyedTwoInputNonBroadcastProcessOperator( + TwoInputNonBroadcastStreamProcessFunction<IN1, IN2, OUT> userFunction) { + this(userFunction, null); + } + + public KeyedTwoInputNonBroadcastProcessOperator( + TwoInputNonBroadcastStreamProcessFunction<IN1, IN2, OUT> userFunction, + @Nullable KeySelector<OUT, KEY> outKeySelector) { + super(userFunction); + this.outKeySelector = outKeySelector; + } + + @Override + protected TimestampCollector<OUT> getOutputCollector() { + return outKeySelector == null + ? new OutputCollector<>(output) + : new KeyCheckedOutputCollector<>( + new OutputCollector<>(output), outKeySelector, () -> (KEY) getCurrentKey()); + } +} diff --git a/flink-datastream/src/main/java/org/apache/flink/datastream/impl/operators/KeyedTwoOutputProcessOperator.java b/flink-datastream/src/main/java/org/apache/flink/datastream/impl/operators/KeyedTwoOutputProcessOperator.java new file mode 100644 index 00000000000..76a2941f5a7 --- /dev/null +++ b/flink-datastream/src/main/java/org/apache/flink/datastream/impl/operators/KeyedTwoOutputProcessOperator.java @@ -0,0 +1,78 @@ +/* + * 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.datastream.impl.operators; + +import org.apache.flink.api.java.functions.KeySelector; +import org.apache.flink.datastream.api.function.TwoOutputStreamProcessFunction; +import org.apache.flink.datastream.impl.common.KeyCheckedOutputCollector; +import org.apache.flink.datastream.impl.common.OutputCollector; +import org.apache.flink.datastream.impl.common.TimestampCollector; +import org.apache.flink.util.OutputTag; +import org.apache.flink.util.Preconditions; + +import javax.annotation.Nullable; + +/** */ +public class KeyedTwoOutputProcessOperator<KEY, IN, OUT_MAIN, OUT_SIDE> + extends TwoOutputProcessOperator<IN, OUT_MAIN, OUT_SIDE> { + + @Nullable private final KeySelector<OUT_MAIN, KEY> mainOutKeySelector; + + @Nullable private final KeySelector<OUT_SIDE, KEY> sideOutKeySelector; + + public KeyedTwoOutputProcessOperator( + TwoOutputStreamProcessFunction<IN, OUT_MAIN, OUT_SIDE> userFunction, + OutputTag<OUT_SIDE> outputTag) { + this(userFunction, outputTag, null, null); + } + + public KeyedTwoOutputProcessOperator( + TwoOutputStreamProcessFunction<IN, OUT_MAIN, OUT_SIDE> userFunction, + OutputTag<OUT_SIDE> outputTag, + @Nullable KeySelector<OUT_MAIN, KEY> mainOutKeySelector, + @Nullable KeySelector<OUT_SIDE, KEY> sideOutKeySelector) { + super(userFunction, outputTag); + Preconditions.checkArgument( + (mainOutKeySelector == null && sideOutKeySelector == null) + || (mainOutKeySelector != null && sideOutKeySelector != null), + "Both mainOutKeySelector and sideOutKeySelector must be null or not null."); + this.mainOutKeySelector = mainOutKeySelector; + this.sideOutKeySelector = sideOutKeySelector; + } + + @Override + protected TimestampCollector<OUT_MAIN> getMainCollector() { + return mainOutKeySelector != null && sideOutKeySelector != null + ? new KeyCheckedOutputCollector<>( + new OutputCollector<>(output), + mainOutKeySelector, + () -> (KEY) getCurrentKey()) + : new OutputCollector<>(output); + } + + @Override + public TimestampCollector<OUT_SIDE> getSideCollector() { + return mainOutKeySelector != null && sideOutKeySelector != null + ? new KeyCheckedOutputCollector<>( + new SideOutputCollector(output), + sideOutKeySelector, + () -> (KEY) getCurrentKey()) + : new SideOutputCollector(output); + } +} diff --git a/flink-datastream/src/main/java/org/apache/flink/datastream/impl/operators/ProcessOperator.java b/flink-datastream/src/main/java/org/apache/flink/datastream/impl/operators/ProcessOperator.java new file mode 100644 index 00000000000..6346079cc21 --- /dev/null +++ b/flink-datastream/src/main/java/org/apache/flink/datastream/impl/operators/ProcessOperator.java @@ -0,0 +1,71 @@ +/* + * 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.datastream.impl.operators; + +import org.apache.flink.datastream.api.function.OneInputStreamProcessFunction; +import org.apache.flink.datastream.impl.common.OutputCollector; +import org.apache.flink.datastream.impl.common.TimestampCollector; +import org.apache.flink.datastream.impl.context.DefaultNonPartitionedContext; +import org.apache.flink.datastream.impl.context.DefaultRuntimeContext; +import org.apache.flink.streaming.api.operators.AbstractUdfStreamOperator; +import org.apache.flink.streaming.api.operators.BoundedOneInput; +import org.apache.flink.streaming.api.operators.ChainingStrategy; +import org.apache.flink.streaming.api.operators.OneInputStreamOperator; +import org.apache.flink.streaming.runtime.streamrecord.StreamRecord; + +/** Operator for {@link OneInputStreamProcessFunction}. */ +public class ProcessOperator<IN, OUT> + extends AbstractUdfStreamOperator<OUT, OneInputStreamProcessFunction<IN, OUT>> + implements OneInputStreamOperator<IN, OUT>, BoundedOneInput { + + protected transient DefaultRuntimeContext context; + + protected transient DefaultNonPartitionedContext<OUT> nonPartitionedContext; + + protected transient TimestampCollector<OUT> outputCollector; + + public ProcessOperator(OneInputStreamProcessFunction<IN, OUT> userFunction) { + super(userFunction); + + chainingStrategy = ChainingStrategy.ALWAYS; + } + + @Override + public void open() throws Exception { + super.open(); + context = new DefaultRuntimeContext(); + nonPartitionedContext = new DefaultNonPartitionedContext<>(); + outputCollector = getOutputCollector(); + } + + @Override + public void processElement(StreamRecord<IN> element) throws Exception { + outputCollector.setTimestampFromStreamRecord(element); + userFunction.processRecord(element.getValue(), outputCollector, context); + } + + protected TimestampCollector<OUT> getOutputCollector() { + return new OutputCollector<>(output); + } + + @Override + public void endInput() throws Exception { + userFunction.endInput(nonPartitionedContext); + } +} diff --git a/flink-datastream/src/main/java/org/apache/flink/datastream/impl/operators/TwoInputBroadcastProcessOperator.java b/flink-datastream/src/main/java/org/apache/flink/datastream/impl/operators/TwoInputBroadcastProcessOperator.java new file mode 100644 index 00000000000..5c0faf05bab --- /dev/null +++ b/flink-datastream/src/main/java/org/apache/flink/datastream/impl/operators/TwoInputBroadcastProcessOperator.java @@ -0,0 +1,86 @@ +/* + * 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.datastream.impl.operators; + +import org.apache.flink.datastream.api.function.TwoInputBroadcastStreamProcessFunction; +import org.apache.flink.datastream.impl.common.OutputCollector; +import org.apache.flink.datastream.impl.common.TimestampCollector; +import org.apache.flink.datastream.impl.context.DefaultNonPartitionedContext; +import org.apache.flink.datastream.impl.context.DefaultRuntimeContext; +import org.apache.flink.streaming.api.operators.AbstractUdfStreamOperator; +import org.apache.flink.streaming.api.operators.BoundedMultiInput; +import org.apache.flink.streaming.api.operators.ChainingStrategy; +import org.apache.flink.streaming.api.operators.TwoInputStreamOperator; +import org.apache.flink.streaming.runtime.streamrecord.StreamRecord; + +import static org.apache.flink.util.Preconditions.checkState; + +/** Operator for {@link TwoInputBroadcastStreamProcessFunction}. */ +public class TwoInputBroadcastProcessOperator<IN1, IN2, OUT> + extends AbstractUdfStreamOperator< + OUT, TwoInputBroadcastStreamProcessFunction<IN1, IN2, OUT>> + implements TwoInputStreamOperator<IN1, IN2, OUT>, BoundedMultiInput { + + protected transient TimestampCollector<OUT> collector; + + protected transient DefaultRuntimeContext context; + + protected transient DefaultNonPartitionedContext<OUT> nonPartitionedContext; + + public TwoInputBroadcastProcessOperator( + TwoInputBroadcastStreamProcessFunction<IN1, IN2, OUT> userFunction) { + super(userFunction); + chainingStrategy = ChainingStrategy.ALWAYS; + } + + @Override + public void open() throws Exception { + super.open(); + this.collector = getOutputCollector(); + this.context = new DefaultRuntimeContext(); + this.nonPartitionedContext = new DefaultNonPartitionedContext<>(); + } + + @Override + public void processElement1(StreamRecord<IN1> element) throws Exception { + collector.setTimestampFromStreamRecord(element); + userFunction.processRecordFromNonBroadcastInput(element.getValue(), collector, context); + } + + @Override + public void processElement2(StreamRecord<IN2> element) throws Exception { + collector.setTimestampFromStreamRecord(element); + userFunction.processRecordFromBroadcastInput(element.getValue(), nonPartitionedContext); + } + + protected TimestampCollector<OUT> getOutputCollector() { + return new OutputCollector<>(output); + } + + @Override + public void endInput(int inputId) throws Exception { + // sanity check. + checkState(inputId >= 1 && inputId <= 2); + if (inputId == 1) { + userFunction.endNonBroadcastInput(nonPartitionedContext); + } else { + userFunction.endBroadcastInput(nonPartitionedContext); + } + } +} diff --git a/flink-datastream/src/main/java/org/apache/flink/datastream/impl/operators/TwoInputNonBroadcastProcessOperator.java b/flink-datastream/src/main/java/org/apache/flink/datastream/impl/operators/TwoInputNonBroadcastProcessOperator.java new file mode 100644 index 00000000000..60d88682c2d --- /dev/null +++ b/flink-datastream/src/main/java/org/apache/flink/datastream/impl/operators/TwoInputNonBroadcastProcessOperator.java @@ -0,0 +1,86 @@ +/* + * 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.datastream.impl.operators; + +import org.apache.flink.datastream.api.function.TwoInputNonBroadcastStreamProcessFunction; +import org.apache.flink.datastream.impl.common.OutputCollector; +import org.apache.flink.datastream.impl.common.TimestampCollector; +import org.apache.flink.datastream.impl.context.DefaultNonPartitionedContext; +import org.apache.flink.datastream.impl.context.DefaultRuntimeContext; +import org.apache.flink.streaming.api.operators.AbstractUdfStreamOperator; +import org.apache.flink.streaming.api.operators.BoundedMultiInput; +import org.apache.flink.streaming.api.operators.ChainingStrategy; +import org.apache.flink.streaming.api.operators.TwoInputStreamOperator; +import org.apache.flink.streaming.runtime.streamrecord.StreamRecord; + +import static org.apache.flink.util.Preconditions.checkState; + +/** Operator for {@link TwoInputNonBroadcastStreamProcessFunction}. */ +public class TwoInputNonBroadcastProcessOperator<IN1, IN2, OUT> + extends AbstractUdfStreamOperator< + OUT, TwoInputNonBroadcastStreamProcessFunction<IN1, IN2, OUT>> + implements TwoInputStreamOperator<IN1, IN2, OUT>, BoundedMultiInput { + + protected transient TimestampCollector<OUT> collector; + + protected transient DefaultRuntimeContext context; + + protected transient DefaultNonPartitionedContext<OUT> nonPartitionedContext; + + public TwoInputNonBroadcastProcessOperator( + TwoInputNonBroadcastStreamProcessFunction<IN1, IN2, OUT> userFunction) { + super(userFunction); + chainingStrategy = ChainingStrategy.ALWAYS; + } + + @Override + public void open() throws Exception { + super.open(); + this.collector = getOutputCollector(); + this.context = new DefaultRuntimeContext(); + this.nonPartitionedContext = new DefaultNonPartitionedContext<>(); + } + + @Override + public void processElement1(StreamRecord<IN1> element) throws Exception { + collector.setTimestampFromStreamRecord(element); + userFunction.processRecordFromFirstInput(element.getValue(), collector, context); + } + + @Override + public void processElement2(StreamRecord<IN2> element) throws Exception { + collector.setTimestampFromStreamRecord(element); + userFunction.processRecordFromSecondInput(element.getValue(), collector, context); + } + + protected TimestampCollector<OUT> getOutputCollector() { + return new OutputCollector<>(output); + } + + @Override + public void endInput(int inputId) throws Exception { + // sanity check. + checkState(inputId >= 1 && inputId <= 2); + if (inputId == 1) { + userFunction.endFirstInput(nonPartitionedContext); + } else { + userFunction.endSecondInput(nonPartitionedContext); + } + } +} diff --git a/flink-datastream/src/main/java/org/apache/flink/datastream/impl/operators/TwoOutputProcessOperator.java b/flink-datastream/src/main/java/org/apache/flink/datastream/impl/operators/TwoOutputProcessOperator.java new file mode 100644 index 00000000000..2eb7dee587e --- /dev/null +++ b/flink-datastream/src/main/java/org/apache/flink/datastream/impl/operators/TwoOutputProcessOperator.java @@ -0,0 +1,113 @@ +/* + * 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.datastream.impl.operators; + +import org.apache.flink.datastream.api.context.TwoOutputNonPartitionedContext; +import org.apache.flink.datastream.api.function.TwoOutputStreamProcessFunction; +import org.apache.flink.datastream.impl.common.OutputCollector; +import org.apache.flink.datastream.impl.common.TimestampCollector; +import org.apache.flink.datastream.impl.context.DefaultRuntimeContext; +import org.apache.flink.datastream.impl.context.DefaultTwoOutputNonPartitionedContext; +import org.apache.flink.streaming.api.operators.AbstractUdfStreamOperator; +import org.apache.flink.streaming.api.operators.BoundedOneInput; +import org.apache.flink.streaming.api.operators.ChainingStrategy; +import org.apache.flink.streaming.api.operators.OneInputStreamOperator; +import org.apache.flink.streaming.api.operators.Output; +import org.apache.flink.streaming.runtime.streamrecord.StreamRecord; +import org.apache.flink.util.OutputTag; + +/** + * Operator for {@link TwoOutputStreamProcessFunction}. + * + * <p>We support the second output via flink side-output mechanism. + */ +public class TwoOutputProcessOperator<IN, OUT_MAIN, OUT_SIDE> + extends AbstractUdfStreamOperator< + OUT_MAIN, TwoOutputStreamProcessFunction<IN, OUT_MAIN, OUT_SIDE>> + implements OneInputStreamOperator<IN, OUT_MAIN>, BoundedOneInput { + protected transient TimestampCollector<OUT_MAIN> mainCollector; + + protected transient TimestampCollector<OUT_SIDE> sideCollector; + + protected transient DefaultRuntimeContext context; + + protected transient TwoOutputNonPartitionedContext<OUT_MAIN, OUT_SIDE> nonPartitionedContext; + + protected OutputTag<OUT_SIDE> outputTag; + + public TwoOutputProcessOperator( + TwoOutputStreamProcessFunction<IN, OUT_MAIN, OUT_SIDE> userFunction, + OutputTag<OUT_SIDE> outputTag) { + super(userFunction); + + this.outputTag = outputTag; + chainingStrategy = ChainingStrategy.ALWAYS; + } + + @Override + public void open() throws Exception { + this.mainCollector = getMainCollector(); + this.sideCollector = getSideCollector(); + this.context = new DefaultRuntimeContext(); + this.nonPartitionedContext = new DefaultTwoOutputNonPartitionedContext<>(); + } + + @Override + public void processElement(StreamRecord<IN> element) throws Exception { + mainCollector.setTimestampFromStreamRecord(element); + sideCollector.setTimestampFromStreamRecord(element); + userFunction.processRecord(element.getValue(), mainCollector, sideCollector, context); + } + + @Override + public void endInput() throws Exception { + userFunction.endInput(nonPartitionedContext); + } + + protected TimestampCollector<OUT_MAIN> getMainCollector() { + return new OutputCollector<>(output); + } + + public TimestampCollector<OUT_SIDE> getSideCollector() { + return new SideOutputCollector(output); + } + + /** + * This is a special implementation of {@link TimestampCollector} that using side-output + * mechanism to emit data. + */ + protected class SideOutputCollector extends TimestampCollector<OUT_SIDE> { + private final Output<StreamRecord<OUT_MAIN>> output; + + public SideOutputCollector(Output<StreamRecord<OUT_MAIN>> output) { + this.output = output; + } + + @Override + public void collect(OUT_SIDE outputRecord) { + output.collect(outputTag, reuse.replace(outputRecord)); + } + + @Override + public void collectAndOverwriteTimestamp(OUT_SIDE record, long timestamp) { + setTimestamp(timestamp); + output.collect(outputTag, reuse.replace(record)); + } + } +} diff --git a/flink-datastream/src/test/java/org/apache/flink/datastream/impl/common/KeyCheckedOutputCollectorTest.java b/flink-datastream/src/test/java/org/apache/flink/datastream/impl/common/KeyCheckedOutputCollectorTest.java new file mode 100644 index 00000000000..57b36ef5f41 --- /dev/null +++ b/flink-datastream/src/test/java/org/apache/flink/datastream/impl/common/KeyCheckedOutputCollectorTest.java @@ -0,0 +1,72 @@ +/* + * 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.datastream.impl.common; + +import org.junit.jupiter.api.Test; + +import java.util.concurrent.CompletableFuture; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; + +/** Tests for {@link KeyCheckedOutputCollector}. */ +class KeyCheckedOutputCollectorTest { + @Test + void testCollect() { + TestingTimestampCollector.Builder<Integer> builder = TestingTimestampCollector.builder(); + CompletableFuture<Integer> consumeRecord = new CompletableFuture<>(); + builder.setCollectConsumer(consumeRecord::complete); + KeyCheckedOutputCollector<Integer, Integer> collector = + new KeyCheckedOutputCollector<>(builder.build(), (ignore) -> 1, () -> 1); + + final int record = 1; + collector.collect(record); + assertThat(consumeRecord).isCompletedWithValue(record); + } + + @Test + void testCollectAndOverwriteTimestamp() { + TestingTimestampCollector.Builder<Integer> builder = TestingTimestampCollector.builder(); + CompletableFuture<Integer> consumeRecord = new CompletableFuture<>(); + CompletableFuture<Long> consumeTimeStamp = new CompletableFuture<>(); + builder.setCollectAndOverwriteTimestampConsumer( + (data, timeStamp) -> { + consumeRecord.complete(data); + consumeTimeStamp.complete(timeStamp); + }); + + final int record = 1; + final long timeStamp = 10L; + KeyCheckedOutputCollector<Integer, Integer> collector = + new KeyCheckedOutputCollector<>(builder.build(), (ignore) -> 1, () -> 1); + collector.collectAndOverwriteTimestamp(record, timeStamp); + assertThat(consumeRecord).isCompletedWithValue(record); + assertThat(consumeTimeStamp).isCompletedWithValue(timeStamp); + } + + @Test + void testNotEqualToCurrentKey() { + TestingTimestampCollector.Builder<Integer> builder = TestingTimestampCollector.builder(); + KeyCheckedOutputCollector<Integer, Integer> collector = + new KeyCheckedOutputCollector<>(builder.build(), (ignore) -> 1, () -> 2); + assertThatThrownBy(() -> collector.collect(1)).isInstanceOf(IllegalStateException.class); + assertThatThrownBy(() -> collector.collectAndOverwriteTimestamp(1, 10L)) + .isInstanceOf(IllegalStateException.class); + } +} diff --git a/flink-datastream/src/test/java/org/apache/flink/datastream/impl/common/OutputCollectorTest.java b/flink-datastream/src/test/java/org/apache/flink/datastream/impl/common/OutputCollectorTest.java new file mode 100644 index 00000000000..49a45f93489 --- /dev/null +++ b/flink-datastream/src/test/java/org/apache/flink/datastream/impl/common/OutputCollectorTest.java @@ -0,0 +1,53 @@ +/* + * 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.datastream.impl.common; + +import org.apache.flink.streaming.runtime.streamrecord.StreamElement; +import org.apache.flink.streaming.runtime.streamrecord.StreamRecord; +import org.apache.flink.streaming.util.CollectorOutput; + +import org.junit.jupiter.api.Test; + +import java.util.ArrayList; +import java.util.List; + +import static org.assertj.core.api.Assertions.assertThat; + +/** Tests for {@link OutputCollector}. */ +class OutputCollectorTest { + @Test + void testCollect() { + List<StreamElement> list = new ArrayList<>(); + CollectorOutput<Integer> collectorOutput = new CollectorOutput<>(list); + OutputCollector<Integer> collector = new OutputCollector<>(collectorOutput); + collector.collect(1); + collector.collect(2); + assertThat(list).containsExactly(new StreamRecord<>(1), new StreamRecord<>(2)); + } + + @Test + void testCollectAndOverwriteTimestamp() { + List<StreamElement> list = new ArrayList<>(); + CollectorOutput<Integer> collectorOutput = new CollectorOutput<>(list); + OutputCollector<Integer> collector = new OutputCollector<>(collectorOutput); + collector.collectAndOverwriteTimestamp(1, 10L); + collector.collectAndOverwriteTimestamp(2, 20L); + assertThat(list).containsExactly(new StreamRecord<>(1, 10L), new StreamRecord<>(2, 20L)); + } +} diff --git a/flink-datastream/src/test/java/org/apache/flink/datastream/impl/common/TestingTimestampCollector.java b/flink-datastream/src/test/java/org/apache/flink/datastream/impl/common/TestingTimestampCollector.java new file mode 100644 index 00000000000..d9d1f3abf5b --- /dev/null +++ b/flink-datastream/src/test/java/org/apache/flink/datastream/impl/common/TestingTimestampCollector.java @@ -0,0 +1,72 @@ +/* + * 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.datastream.impl.common; + +import java.util.function.BiConsumer; +import java.util.function.Consumer; + +/** Mock {@link TimestampCollector}'s subclass for testing. */ +public class TestingTimestampCollector<T> extends TimestampCollector<T> { + private final Consumer<T> collectConsumer; + + private final BiConsumer<T, Long> collectAndOverwriteTimestampConsumer; + + private TestingTimestampCollector( + Consumer<T> collectConsumer, BiConsumer<T, Long> collectAndOverwriteTimestampConsuemr) { + this.collectConsumer = collectConsumer; + this.collectAndOverwriteTimestampConsumer = collectAndOverwriteTimestampConsuemr; + } + + @Override + public void collect(T record) { + collectConsumer.accept(record); + } + + @Override + public void collectAndOverwriteTimestamp(T record, long timestamp) { + collectAndOverwriteTimestampConsumer.accept(record, timestamp); + } + + public static <T> Builder<T> builder() { + return new Builder<>(); + } + + /** Builder for {@link TestingTimestampCollector}. */ + public static class Builder<T> { + private Consumer<T> collectConsumer = (ignore) -> {}; + + private BiConsumer<T, Long> collectAndOverwriteTimestampConsumer = (ingore1, ignore2) -> {}; + + public Builder<T> setCollectConsumer(Consumer<T> collectConsumer) { + this.collectConsumer = collectConsumer; + return this; + } + + public Builder<T> setCollectAndOverwriteTimestampConsumer( + BiConsumer<T, Long> collectAndOverwriteTimestampConsumer) { + this.collectAndOverwriteTimestampConsumer = collectAndOverwriteTimestampConsumer; + return this; + } + + public TestingTimestampCollector<T> build() { + return new TestingTimestampCollector<>( + collectConsumer, collectAndOverwriteTimestampConsumer); + } + } +} diff --git a/flink-datastream/src/test/java/org/apache/flink/datastream/impl/operators/KeyedProcessOperatorTest.java b/flink-datastream/src/test/java/org/apache/flink/datastream/impl/operators/KeyedProcessOperatorTest.java new file mode 100644 index 00000000000..201645bd422 --- /dev/null +++ b/flink-datastream/src/test/java/org/apache/flink/datastream/impl/operators/KeyedProcessOperatorTest.java @@ -0,0 +1,123 @@ +/* + * 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.datastream.impl.operators; + +import org.apache.flink.api.common.typeinfo.Types; +import org.apache.flink.api.java.functions.KeySelector; +import org.apache.flink.datastream.api.common.Collector; +import org.apache.flink.datastream.api.context.NonPartitionedContext; +import org.apache.flink.datastream.api.context.RuntimeContext; +import org.apache.flink.datastream.api.function.OneInputStreamProcessFunction; +import org.apache.flink.streaming.runtime.streamrecord.StreamRecord; +import org.apache.flink.streaming.util.KeyedOneInputStreamOperatorTestHarness; + +import org.junit.jupiter.api.Test; + +import java.util.Collection; +import java.util.concurrent.CompletableFuture; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; + +/** Tests for {@link KeyedProcessOperator}. */ +class KeyedProcessOperatorTest { + @Test + void testProcessRecord() throws Exception { + KeyedProcessOperator<Integer, Integer, Integer> processOperator = + new KeyedProcessOperator<>( + new OneInputStreamProcessFunction<Integer, Integer>() { + @Override + public void processRecord( + Integer record, Collector<Integer> output, RuntimeContext ctx) { + output.collect(record + 1); + } + }); + + try (KeyedOneInputStreamOperatorTestHarness<Integer, Integer, Integer> testHarness = + new KeyedOneInputStreamOperatorTestHarness<>( + processOperator, + (KeySelector<Integer, Integer>) value -> value, + Types.INT)) { + testHarness.open(); + testHarness.processElement(new StreamRecord<>(1)); + testHarness.processElement(new StreamRecord<>(2)); + testHarness.processElement(new StreamRecord<>(3)); + + Collection<StreamRecord<Integer>> recordOutput = testHarness.getRecordOutput(); + assertThat(recordOutput) + .containsExactly( + new StreamRecord<>(2), new StreamRecord<>(3), new StreamRecord<>(4)); + } + } + + @Test + void testEndInput() throws Exception { + CompletableFuture<Void> future = new CompletableFuture<>(); + KeyedProcessOperator<Integer, Integer, Integer> processOperator = + new KeyedProcessOperator<>( + new OneInputStreamProcessFunction<Integer, Integer>() { + @Override + public void processRecord( + Integer record, Collector<Integer> output, RuntimeContext ctx) { + // do nothing. + } + + @Override + public void endInput(NonPartitionedContext<Integer> ctx) { + future.complete(null); + } + }); + + try (KeyedOneInputStreamOperatorTestHarness<Integer, Integer, Integer> testHarness = + new KeyedOneInputStreamOperatorTestHarness<>( + processOperator, + (KeySelector<Integer, Integer>) value -> value, + Types.INT)) { + testHarness.open(); + testHarness.endInput(); + assertThat(future).isCompleted(); + } + } + + @Test + void testCheckKey() throws Exception { + KeyedProcessOperator<Integer, Integer, Integer> processOperator = + new KeyedProcessOperator<>( + new OneInputStreamProcessFunction<Integer, Integer>() { + @Override + public void processRecord( + Integer record, Collector<Integer> output, RuntimeContext ctx) { + // forward the record to check input key. + output.collect(record); + } + }, + // -1 is an invalid key in this suite. + (ignore) -> -1); + + try (KeyedOneInputStreamOperatorTestHarness<Integer, Integer, Integer> testHarness = + new KeyedOneInputStreamOperatorTestHarness<>( + processOperator, + (KeySelector<Integer, Integer>) value -> value, + Types.INT)) { + testHarness.open(); + assertThatThrownBy(() -> testHarness.processElement(new StreamRecord<>(2))) + .isInstanceOf(IllegalStateException.class); + } + } +} diff --git a/flink-datastream/src/test/java/org/apache/flink/datastream/impl/operators/KeyedTwoInputBroadcastProcessOperatorTest.java b/flink-datastream/src/test/java/org/apache/flink/datastream/impl/operators/KeyedTwoInputBroadcastProcessOperatorTest.java new file mode 100644 index 00000000000..0c4fcf46c45 --- /dev/null +++ b/flink-datastream/src/test/java/org/apache/flink/datastream/impl/operators/KeyedTwoInputBroadcastProcessOperatorTest.java @@ -0,0 +1,155 @@ +/* + * 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.datastream.impl.operators; + +import org.apache.flink.api.common.typeinfo.Types; +import org.apache.flink.api.java.functions.KeySelector; +import org.apache.flink.datastream.api.common.Collector; +import org.apache.flink.datastream.api.context.NonPartitionedContext; +import org.apache.flink.datastream.api.context.RuntimeContext; +import org.apache.flink.datastream.api.function.TwoInputBroadcastStreamProcessFunction; +import org.apache.flink.streaming.runtime.streamrecord.StreamRecord; +import org.apache.flink.streaming.util.KeyedTwoInputStreamOperatorTestHarness; + +import org.junit.jupiter.api.Test; + +import java.util.ArrayList; +import java.util.List; +import java.util.concurrent.CompletableFuture; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; + +/** Tests for {@link KeyedTwoInputBroadcastProcessOperator}. */ +class KeyedTwoInputBroadcastProcessOperatorTest { + @Test + void testProcessRecord() throws Exception { + List<Integer> fromNonBroadcastInput = new ArrayList<>(); + List<Long> fromBroadcastInput = new ArrayList<>(); + KeyedTwoInputBroadcastProcessOperator<Long, Integer, Long, Long> processOperator = + new KeyedTwoInputBroadcastProcessOperator<>( + new TwoInputBroadcastStreamProcessFunction<Integer, Long, Long>() { + @Override + public void processRecordFromNonBroadcastInput( + Integer record, Collector<Long> output, RuntimeContext ctx) { + fromNonBroadcastInput.add(record); + } + + @Override + public void processRecordFromBroadcastInput( + Long record, NonPartitionedContext<Long> ctx) { + fromBroadcastInput.add(record); + } + }); + + try (KeyedTwoInputStreamOperatorTestHarness<Long, Integer, Long, Long> testHarness = + new KeyedTwoInputStreamOperatorTestHarness<>( + processOperator, + (KeySelector<Integer, Long>) (data) -> (long) (data + 1), + (KeySelector<Long, Long>) value -> value + 1, + Types.LONG)) { + testHarness.open(); + testHarness.processElement1(new StreamRecord<>(1)); + testHarness.processElement2(new StreamRecord<>(2L)); + testHarness.processElement2(new StreamRecord<>(4L)); + testHarness.processElement1(new StreamRecord<>(3)); + } + + assertThat(fromNonBroadcastInput).containsExactly(1, 3); + assertThat(fromBroadcastInput).containsExactly(2L, 4L); + } + + @Test + void testEndInput() throws Exception { + CompletableFuture<Void> nonBroadcastInputEnd = new CompletableFuture<>(); + CompletableFuture<Void> broadcastInputEnd = new CompletableFuture<>(); + KeyedTwoInputBroadcastProcessOperator<Long, Integer, Long, Long> processOperator = + new KeyedTwoInputBroadcastProcessOperator<>( + new TwoInputBroadcastStreamProcessFunction<Integer, Long, Long>() { + @Override + public void processRecordFromNonBroadcastInput( + Integer record, Collector<Long> output, RuntimeContext ctx) { + // do nothing. + } + + @Override + public void processRecordFromBroadcastInput( + Long record, NonPartitionedContext<Long> ctx) { + // do nothing. + } + + @Override + public void endNonBroadcastInput(NonPartitionedContext<Long> ctx) { + nonBroadcastInputEnd.complete(null); + } + + @Override + public void endBroadcastInput(NonPartitionedContext<Long> ctx) { + broadcastInputEnd.complete(null); + } + }); + + try (KeyedTwoInputStreamOperatorTestHarness<Long, Integer, Long, Long> testHarness = + new KeyedTwoInputStreamOperatorTestHarness<>( + processOperator, + (KeySelector<Integer, Long>) Long::valueOf, + (KeySelector<Long, Long>) value -> value, + Types.LONG)) { + testHarness.open(); + testHarness.endInput1(); + assertThat(nonBroadcastInputEnd).isCompleted(); + testHarness.endInput2(); + assertThat(broadcastInputEnd).isCompleted(); + } + } + + @Test + void testCheckKey() throws Exception { + KeyedTwoInputBroadcastProcessOperator<Long, Integer, Long, Long> processOperator = + new KeyedTwoInputBroadcastProcessOperator<>( + new TwoInputBroadcastStreamProcessFunction<Integer, Long, Long>() { + @Override + public void processRecordFromNonBroadcastInput( + Integer record, Collector<Long> output, RuntimeContext ctx) { + output.collect(Long.valueOf(record)); + } + + @Override + public void processRecordFromBroadcastInput( + Long record, NonPartitionedContext<Long> ctx) throws Exception { + ctx.applyToAllPartitions( + (collect, context) -> collect.collect(record)); + } + }, + // -1 is an invalid key in this suite. + (out) -> -1L); + + try (KeyedTwoInputStreamOperatorTestHarness<Long, Integer, Long, Long> testHarness = + new KeyedTwoInputStreamOperatorTestHarness<>( + processOperator, + (KeySelector<Integer, Long>) Long::valueOf, + (KeySelector<Long, Long>) value -> value, + Types.LONG)) { + testHarness.open(); + assertThatThrownBy(() -> testHarness.processElement1(new StreamRecord<>(1))) + .isInstanceOf(IllegalStateException.class); + // TODO also test processElement2 after non-partitioned context fully implemented. + } + } +} diff --git a/flink-datastream/src/test/java/org/apache/flink/datastream/impl/operators/KeyedTwoInputNonBroadcastProcessOperatorTest.java b/flink-datastream/src/test/java/org/apache/flink/datastream/impl/operators/KeyedTwoInputNonBroadcastProcessOperatorTest.java new file mode 100644 index 00000000000..e45e482ea38 --- /dev/null +++ b/flink-datastream/src/test/java/org/apache/flink/datastream/impl/operators/KeyedTwoInputNonBroadcastProcessOperatorTest.java @@ -0,0 +1,156 @@ +/* + * 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.datastream.impl.operators; + +import org.apache.flink.api.common.typeinfo.Types; +import org.apache.flink.api.java.functions.KeySelector; +import org.apache.flink.datastream.api.common.Collector; +import org.apache.flink.datastream.api.context.NonPartitionedContext; +import org.apache.flink.datastream.api.context.RuntimeContext; +import org.apache.flink.datastream.api.function.TwoInputNonBroadcastStreamProcessFunction; +import org.apache.flink.streaming.runtime.streamrecord.StreamRecord; +import org.apache.flink.streaming.util.KeyedTwoInputStreamOperatorTestHarness; + +import org.junit.jupiter.api.Test; + +import java.util.Collection; +import java.util.concurrent.CompletableFuture; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; + +/** Tests for {@link KeyedTwoInputNonBroadcastProcessOperator}. */ +class KeyedTwoInputNonBroadcastProcessOperatorTest { + @Test + void testProcessRecord() throws Exception { + KeyedTwoInputNonBroadcastProcessOperator<Long, Integer, Long, Long> processOperator = + new KeyedTwoInputNonBroadcastProcessOperator<>( + new TwoInputNonBroadcastStreamProcessFunction<Integer, Long, Long>() { + @Override + public void processRecordFromFirstInput( + Integer record, Collector<Long> output, RuntimeContext ctx) { + output.collect(Long.valueOf(record)); + } + + @Override + public void processRecordFromSecondInput( + Long record, Collector<Long> output, RuntimeContext ctx) { + output.collect(record); + } + }); + + try (KeyedTwoInputStreamOperatorTestHarness<Long, Integer, Long, Long> testHarness = + new KeyedTwoInputStreamOperatorTestHarness<>( + processOperator, + (KeySelector<Integer, Long>) (data) -> (long) (data + 1), + (KeySelector<Long, Long>) value -> value + 1, + Types.LONG)) { + testHarness.open(); + testHarness.processElement1(new StreamRecord<>(1)); + testHarness.processElement2(new StreamRecord<>(2L)); + testHarness.processElement2(new StreamRecord<>(4L)); + testHarness.processElement1(new StreamRecord<>(3)); + Collection<StreamRecord<Long>> recordOutput = testHarness.getRecordOutput(); + assertThat(recordOutput) + .containsExactly( + new StreamRecord<>(1L), + new StreamRecord<>(2L), + new StreamRecord<>(4L), + new StreamRecord<>(3L)); + } + } + + @Test + void testEndInput() throws Exception { + CompletableFuture<Void> firstFuture = new CompletableFuture<>(); + CompletableFuture<Void> secondFuture = new CompletableFuture<>(); + KeyedTwoInputNonBroadcastProcessOperator<Long, Integer, Long, Long> processOperator = + new KeyedTwoInputNonBroadcastProcessOperator<>( + new TwoInputNonBroadcastStreamProcessFunction<Integer, Long, Long>() { + @Override + public void processRecordFromFirstInput( + Integer record, Collector<Long> output, RuntimeContext ctx) { + // do nothing. + } + + @Override + public void processRecordFromSecondInput( + Long record, Collector<Long> output, RuntimeContext ctx) { + // do nothing. + } + + @Override + public void endFirstInput(NonPartitionedContext<Long> ctx) { + firstFuture.complete(null); + } + + @Override + public void endSecondInput(NonPartitionedContext<Long> ctx) { + secondFuture.complete(null); + } + }); + + try (KeyedTwoInputStreamOperatorTestHarness<Long, Integer, Long, Long> testHarness = + new KeyedTwoInputStreamOperatorTestHarness<>( + processOperator, + (KeySelector<Integer, Long>) Long::valueOf, + (KeySelector<Long, Long>) value -> value, + Types.LONG)) { + testHarness.open(); + testHarness.endInput1(); + assertThat(firstFuture).isCompleted(); + testHarness.endInput2(); + assertThat(secondFuture).isCompleted(); + } + } + + @Test + void testCheckKey() throws Exception { + KeyedTwoInputNonBroadcastProcessOperator<Long, Integer, Long, Long> processOperator = + new KeyedTwoInputNonBroadcastProcessOperator<>( + new TwoInputNonBroadcastStreamProcessFunction<Integer, Long, Long>() { + @Override + public void processRecordFromFirstInput( + Integer record, Collector<Long> output, RuntimeContext ctx) { + output.collect(Long.valueOf(record)); + } + + @Override + public void processRecordFromSecondInput( + Long record, Collector<Long> output, RuntimeContext ctx) { + output.collect(record); + } + }, + // -1 is an invalid key in this suite. + (out) -> -1L); + + try (KeyedTwoInputStreamOperatorTestHarness<Long, Integer, Long, Long> testHarness = + new KeyedTwoInputStreamOperatorTestHarness<>( + processOperator, + (KeySelector<Integer, Long>) Long::valueOf, + (KeySelector<Long, Long>) value -> value, + Types.LONG)) { + testHarness.open(); + assertThatThrownBy(() -> testHarness.processElement1(new StreamRecord<>(1))) + .isInstanceOf(IllegalStateException.class); + assertThatThrownBy(() -> testHarness.processElement2(new StreamRecord<>(1L))) + .isInstanceOf(IllegalStateException.class); + } + } +} diff --git a/flink-datastream/src/test/java/org/apache/flink/datastream/impl/operators/KeyedTwoOutputProcessOperatorTest.java b/flink-datastream/src/test/java/org/apache/flink/datastream/impl/operators/KeyedTwoOutputProcessOperatorTest.java new file mode 100644 index 00000000000..b39b6cede96 --- /dev/null +++ b/flink-datastream/src/test/java/org/apache/flink/datastream/impl/operators/KeyedTwoOutputProcessOperatorTest.java @@ -0,0 +1,158 @@ +/* + * 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.datastream.impl.operators; + +import org.apache.flink.api.common.typeinfo.Types; +import org.apache.flink.api.java.functions.KeySelector; +import org.apache.flink.datastream.api.common.Collector; +import org.apache.flink.datastream.api.context.RuntimeContext; +import org.apache.flink.datastream.api.context.TwoOutputNonPartitionedContext; +import org.apache.flink.datastream.api.function.TwoOutputStreamProcessFunction; +import org.apache.flink.streaming.runtime.streamrecord.StreamRecord; +import org.apache.flink.streaming.util.KeyedOneInputStreamOperatorTestHarness; +import org.apache.flink.util.OutputTag; + +import org.junit.jupiter.api.Test; + +import java.util.Collection; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.ConcurrentLinkedQueue; +import java.util.concurrent.atomic.AtomicBoolean; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; + +/** Tests for {@link KeyedTwoOutputProcessOperator}. */ +class KeyedTwoOutputProcessOperatorTest { + @Test + void testProcessRecord() throws Exception { + OutputTag<Long> sideOutputTag = new OutputTag<Long>("side-output") {}; + + KeyedTwoOutputProcessOperator<Integer, Integer, Integer, Long> processOperator = + new KeyedTwoOutputProcessOperator<>( + new TwoOutputStreamProcessFunction<Integer, Integer, Long>() { + @Override + public void processRecord( + Integer record, + Collector<Integer> output1, + Collector<Long> output2, + RuntimeContext ctx) { + output1.collect(record); + output2.collect((long) (record * 2)); + } + }, + sideOutputTag); + + try (KeyedOneInputStreamOperatorTestHarness<Integer, Integer, Integer> testHarness = + new KeyedOneInputStreamOperatorTestHarness<>( + processOperator, + (KeySelector<Integer, Integer>) value -> value, + Types.INT)) { + testHarness.open(); + testHarness.processElement(new StreamRecord<>(1)); + testHarness.processElement(new StreamRecord<>(2)); + testHarness.processElement(new StreamRecord<>(3)); + Collection<StreamRecord<Integer>> firstOutput = testHarness.getRecordOutput(); + ConcurrentLinkedQueue<StreamRecord<Long>> secondOutput = + testHarness.getSideOutput(sideOutputTag); + assertThat(firstOutput) + .containsExactly( + new StreamRecord<>(1), new StreamRecord<>(2), new StreamRecord<>(3)); + assertThat(secondOutput) + .containsExactly( + new StreamRecord<>(2L), new StreamRecord<>(4L), new StreamRecord<>(6L)); + } + } + + @Test + void testEndInput() throws Exception { + CompletableFuture<Void> future = new CompletableFuture<>(); + OutputTag<Long> sideOutputTag = new OutputTag<Long>("side-output") {}; + + KeyedTwoOutputProcessOperator<Integer, Integer, Integer, Long> processOperator = + new KeyedTwoOutputProcessOperator<>( + new TwoOutputStreamProcessFunction<Integer, Integer, Long>() { + @Override + public void processRecord( + Integer record, + Collector<Integer> output1, + Collector<Long> output2, + RuntimeContext ctx) { + // do nothing. + } + + @Override + public void endInput( + TwoOutputNonPartitionedContext<Integer, Long> ctx) { + future.complete(null); + } + }, + sideOutputTag); + + try (KeyedOneInputStreamOperatorTestHarness<Integer, Integer, Integer> testHarness = + new KeyedOneInputStreamOperatorTestHarness<>( + processOperator, + (KeySelector<Integer, Integer>) value -> value, + Types.INT)) { + testHarness.open(); + testHarness.endInput(); + assertThat(future).isCompleted(); + } + } + + @Test + void testKeyCheck() throws Exception { + OutputTag<Long> sideOutputTag = new OutputTag<Long>("side-output") {}; + AtomicBoolean emitToFirstOutput = new AtomicBoolean(true); + KeyedTwoOutputProcessOperator<Integer, Integer, Integer, Long> processOperator = + new KeyedTwoOutputProcessOperator<>( + new TwoOutputStreamProcessFunction<Integer, Integer, Long>() { + @Override + public void processRecord( + Integer record, + Collector<Integer> output1, + Collector<Long> output2, + RuntimeContext ctx) { + if (emitToFirstOutput.get()) { + output1.collect(record); + } else { + output2.collect((long) (record)); + } + } + }, + sideOutputTag, + // -1 is an invalid key in this suite. + (KeySelector<Integer, Integer>) value -> -1, + // -1 is an invalid key in this suite. + (KeySelector<Long, Integer>) value -> -1); + + try (KeyedOneInputStreamOperatorTestHarness<Integer, Integer, Integer> testHarness = + new KeyedOneInputStreamOperatorTestHarness<>( + processOperator, + (KeySelector<Integer, Integer>) value -> value, + Types.INT)) { + testHarness.open(); + assertThatThrownBy(() -> testHarness.processElement(new StreamRecord<>(1))) + .isInstanceOf(IllegalStateException.class); + emitToFirstOutput.set(false); + assertThatThrownBy(() -> testHarness.processElement(new StreamRecord<>(1))) + .isInstanceOf(IllegalStateException.class); + } + } +} diff --git a/flink-datastream/src/test/java/org/apache/flink/datastream/impl/operators/ProcessOperatorTest.java b/flink-datastream/src/test/java/org/apache/flink/datastream/impl/operators/ProcessOperatorTest.java new file mode 100644 index 00000000000..c28802dd267 --- /dev/null +++ b/flink-datastream/src/test/java/org/apache/flink/datastream/impl/operators/ProcessOperatorTest.java @@ -0,0 +1,83 @@ +/* + * 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.datastream.impl.operators; + +import org.apache.flink.datastream.api.common.Collector; +import org.apache.flink.datastream.api.context.NonPartitionedContext; +import org.apache.flink.datastream.api.context.RuntimeContext; +import org.apache.flink.datastream.api.function.OneInputStreamProcessFunction; +import org.apache.flink.streaming.runtime.streamrecord.StreamRecord; +import org.apache.flink.streaming.util.OneInputStreamOperatorTestHarness; + +import org.junit.jupiter.api.Test; + +import java.util.Collection; +import java.util.concurrent.CompletableFuture; + +import static org.assertj.core.api.Assertions.assertThat; + +/** Tests for {@link ProcessOperator}. */ +class ProcessOperatorTest { + @Test + void testProcessRecord() throws Exception { + ProcessOperator<Integer, String> processOperator = + new ProcessOperator<>((record, output, ctx) -> output.collect(record + "-")); + + try (OneInputStreamOperatorTestHarness<Integer, String> testHarness = + new OneInputStreamOperatorTestHarness<>(processOperator)) { + testHarness.open(); + testHarness.processElement(new StreamRecord<>(1)); + testHarness.processElement(new StreamRecord<>(2)); + testHarness.processElement(new StreamRecord<>(3)); + + Collection<StreamRecord<String>> recordOutput = testHarness.getRecordOutput(); + assertThat(recordOutput) + .containsExactly( + new StreamRecord<>("1-"), + new StreamRecord<>("2-"), + new StreamRecord<>("3-")); + } + } + + @Test + void testEndInput() throws Exception { + CompletableFuture<Void> future = new CompletableFuture<>(); + ProcessOperator<Integer, String> processOperator = + new ProcessOperator<>( + new OneInputStreamProcessFunction<Integer, String>() { + @Override + public void processRecord( + Integer record, Collector<String> output, RuntimeContext ctx) { + // do nothing. + } + + @Override + public void endInput(NonPartitionedContext<String> ctx) { + future.complete(null); + } + }); + + try (OneInputStreamOperatorTestHarness<Integer, String> testHarness = + new OneInputStreamOperatorTestHarness<>(processOperator)) { + testHarness.open(); + testHarness.endInput(); + assertThat(future).isCompleted(); + } + } +} diff --git a/flink-datastream/src/test/java/org/apache/flink/datastream/impl/operators/TwoInputBroadcastProcessOperatorTest.java b/flink-datastream/src/test/java/org/apache/flink/datastream/impl/operators/TwoInputBroadcastProcessOperatorTest.java new file mode 100644 index 00000000000..d92bb531faa --- /dev/null +++ b/flink-datastream/src/test/java/org/apache/flink/datastream/impl/operators/TwoInputBroadcastProcessOperatorTest.java @@ -0,0 +1,112 @@ +/* + * 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.datastream.impl.operators; + +import org.apache.flink.datastream.api.common.Collector; +import org.apache.flink.datastream.api.context.NonPartitionedContext; +import org.apache.flink.datastream.api.context.RuntimeContext; +import org.apache.flink.datastream.api.function.TwoInputBroadcastStreamProcessFunction; +import org.apache.flink.streaming.runtime.streamrecord.StreamRecord; +import org.apache.flink.streaming.util.TwoInputStreamOperatorTestHarness; + +import org.junit.jupiter.api.Test; + +import java.util.ArrayList; +import java.util.List; +import java.util.concurrent.CompletableFuture; + +import static org.assertj.core.api.Assertions.assertThat; + +/** Tests for {@link TwoInputBroadcastProcessOperator}. */ +class TwoInputBroadcastProcessOperatorTest { + @Test + void testProcessRecord() throws Exception { + List<Long> fromNonBroadcastInput = new ArrayList<>(); + List<Long> fromBroadcastInput = new ArrayList<>(); + TwoInputBroadcastProcessOperator<Integer, Long, Long> processOperator = + new TwoInputBroadcastProcessOperator<>( + new TwoInputBroadcastStreamProcessFunction<Integer, Long, Long>() { + + @Override + public void processRecordFromNonBroadcastInput( + Integer record, Collector<Long> output, RuntimeContext ctx) { + fromNonBroadcastInput.add(Long.valueOf(record)); + } + + @Override + public void processRecordFromBroadcastInput( + Long record, NonPartitionedContext<Long> ctx) { + fromBroadcastInput.add(record); + } + }); + + try (TwoInputStreamOperatorTestHarness<Integer, Long, Long> testHarness = + new TwoInputStreamOperatorTestHarness<>(processOperator)) { + testHarness.open(); + testHarness.processElement1(new StreamRecord<>(1)); + testHarness.processElement2(new StreamRecord<>(2L)); + testHarness.processElement1(new StreamRecord<>(3)); + testHarness.processElement1(new StreamRecord<>(5)); + testHarness.processElement2(new StreamRecord<>(4L)); + assertThat(fromNonBroadcastInput).containsExactly(1L, 3L, 5L); + assertThat(fromBroadcastInput).containsExactly(2L, 4L); + } + } + + @Test + void testEndInput() throws Exception { + CompletableFuture<Void> nonBroadcastInputEnd = new CompletableFuture<>(); + CompletableFuture<Void> broadcastInputEnd = new CompletableFuture<>(); + TwoInputBroadcastProcessOperator<Integer, Long, Long> processOperator = + new TwoInputBroadcastProcessOperator<>( + new TwoInputBroadcastStreamProcessFunction<Integer, Long, Long>() { + + @Override + public void processRecordFromNonBroadcastInput( + Integer record, Collector<Long> output, RuntimeContext ctx) { + // do nothing. + } + + @Override + public void processRecordFromBroadcastInput( + Long record, NonPartitionedContext<Long> ctx) { + // do nothing. + } + + @Override + public void endNonBroadcastInput(NonPartitionedContext<Long> ctx) { + nonBroadcastInputEnd.complete(null); + } + + @Override + public void endBroadcastInput(NonPartitionedContext<Long> ctx) { + broadcastInputEnd.complete(null); + } + }); + + try (TwoInputStreamOperatorTestHarness<Integer, Long, Long> testHarness = + new TwoInputStreamOperatorTestHarness<>(processOperator)) { + testHarness.open(); + testHarness.endInput1(); + assertThat(nonBroadcastInputEnd).isCompleted(); + testHarness.endInput2(); + assertThat(broadcastInputEnd).isCompleted(); + } + } +} diff --git a/flink-datastream/src/test/java/org/apache/flink/datastream/impl/operators/TwoInputNonBroadcastProcessOperatorTest.java b/flink-datastream/src/test/java/org/apache/flink/datastream/impl/operators/TwoInputNonBroadcastProcessOperatorTest.java new file mode 100644 index 00000000000..6fce6bc0b20 --- /dev/null +++ b/flink-datastream/src/test/java/org/apache/flink/datastream/impl/operators/TwoInputNonBroadcastProcessOperatorTest.java @@ -0,0 +1,115 @@ +/* + * 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.datastream.impl.operators; + +import org.apache.flink.datastream.api.common.Collector; +import org.apache.flink.datastream.api.context.NonPartitionedContext; +import org.apache.flink.datastream.api.context.RuntimeContext; +import org.apache.flink.datastream.api.function.TwoInputNonBroadcastStreamProcessFunction; +import org.apache.flink.streaming.runtime.streamrecord.StreamRecord; +import org.apache.flink.streaming.util.TwoInputStreamOperatorTestHarness; + +import org.junit.jupiter.api.Test; + +import java.util.Collection; +import java.util.concurrent.CompletableFuture; + +import static org.assertj.core.api.Assertions.assertThat; + +/** Tests for {@link TwoInputNonBroadcastProcessOperator}. */ +class TwoInputNonBroadcastProcessOperatorTest { + @Test + void testProcessRecord() throws Exception { + TwoInputNonBroadcastProcessOperator<Integer, Long, Long> processOperator = + new TwoInputNonBroadcastProcessOperator<>( + new TwoInputNonBroadcastStreamProcessFunction<Integer, Long, Long>() { + @Override + public void processRecordFromFirstInput( + Integer record, Collector<Long> output, RuntimeContext ctx) { + output.collect(Long.valueOf(record)); + } + + @Override + public void processRecordFromSecondInput( + Long record, Collector<Long> output, RuntimeContext ctx) { + output.collect(record); + } + }); + + try (TwoInputStreamOperatorTestHarness<Integer, Long, Long> testHarness = + new TwoInputStreamOperatorTestHarness<>(processOperator)) { + testHarness.open(); + testHarness.processElement1(new StreamRecord<>(1)); + testHarness.processElement2(new StreamRecord<>(2L)); + testHarness.processElement1(new StreamRecord<>(3)); + testHarness.processElement1(new StreamRecord<>(5)); + testHarness.processElement2(new StreamRecord<>(4L)); + Collection<StreamRecord<Long>> recordOutput = testHarness.getRecordOutput(); + assertThat(recordOutput) + .containsExactly( + new StreamRecord<>(1L), + new StreamRecord<>(2L), + new StreamRecord<>(3L), + new StreamRecord<>(5L), + new StreamRecord<>(4L)); + } + } + + @Test + void testEndInput() throws Exception { + CompletableFuture<Void> firstFuture = new CompletableFuture<>(); + CompletableFuture<Void> secondFuture = new CompletableFuture<>(); + TwoInputNonBroadcastProcessOperator<Integer, Long, Long> processOperator = + new TwoInputNonBroadcastProcessOperator<>( + new TwoInputNonBroadcastStreamProcessFunction<Integer, Long, Long>() { + @Override + public void processRecordFromFirstInput( + Integer record, Collector<Long> output, RuntimeContext ctx) + throws Exception { + // do nothing. + } + + @Override + public void processRecordFromSecondInput( + Long record, Collector<Long> output, RuntimeContext ctx) + throws Exception { + // do nothing. + } + + @Override + public void endFirstInput(NonPartitionedContext<Long> ctx) { + firstFuture.complete(null); + } + + @Override + public void endSecondInput(NonPartitionedContext<Long> ctx) { + secondFuture.complete(null); + } + }); + + try (TwoInputStreamOperatorTestHarness<Integer, Long, Long> testHarness = + new TwoInputStreamOperatorTestHarness<>(processOperator)) { + testHarness.open(); + testHarness.endInput1(); + assertThat(firstFuture).isCompleted(); + testHarness.endInput2(); + assertThat(secondFuture).isCompleted(); + } + } +} diff --git a/flink-datastream/src/test/java/org/apache/flink/datastream/impl/operators/TwoOutputProcessOperatorTest.java b/flink-datastream/src/test/java/org/apache/flink/datastream/impl/operators/TwoOutputProcessOperatorTest.java new file mode 100644 index 00000000000..52842429597 --- /dev/null +++ b/flink-datastream/src/test/java/org/apache/flink/datastream/impl/operators/TwoOutputProcessOperatorTest.java @@ -0,0 +1,108 @@ +/* + * 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.datastream.impl.operators; + +import org.apache.flink.datastream.api.common.Collector; +import org.apache.flink.datastream.api.context.RuntimeContext; +import org.apache.flink.datastream.api.context.TwoOutputNonPartitionedContext; +import org.apache.flink.datastream.api.function.TwoOutputStreamProcessFunction; +import org.apache.flink.streaming.runtime.streamrecord.StreamRecord; +import org.apache.flink.streaming.util.OneInputStreamOperatorTestHarness; +import org.apache.flink.util.OutputTag; + +import org.junit.jupiter.api.Test; + +import java.util.Collection; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.ConcurrentLinkedQueue; + +import static org.assertj.core.api.Assertions.assertThat; + +/** Tests for {@link TwoOutputProcessOperator}. */ +class TwoOutputProcessOperatorTest { + @Test + void testProcessRecord() throws Exception { + OutputTag<Long> sideOutputTag = new OutputTag<Long>("side-output") {}; + + TwoOutputProcessOperator<Integer, Integer, Long> processOperator = + new TwoOutputProcessOperator<>( + new TwoOutputStreamProcessFunction<Integer, Integer, Long>() { + @Override + public void processRecord( + Integer record, + Collector<Integer> output1, + Collector<Long> output2, + RuntimeContext ctx) { + output1.collect(record); + output2.collect((long) (record * 2)); + } + }, + sideOutputTag); + + try (OneInputStreamOperatorTestHarness<Integer, Integer> testHarness = + new OneInputStreamOperatorTestHarness<>(processOperator)) { + testHarness.open(); + testHarness.processElement(new StreamRecord<>(1)); + testHarness.processElement(new StreamRecord<>(2)); + testHarness.processElement(new StreamRecord<>(3)); + Collection<StreamRecord<Integer>> firstOutput = testHarness.getRecordOutput(); + ConcurrentLinkedQueue<StreamRecord<Long>> secondOutput = + testHarness.getSideOutput(sideOutputTag); + assertThat(firstOutput) + .containsExactly( + new StreamRecord<>(1), new StreamRecord<>(2), new StreamRecord<>(3)); + assertThat(secondOutput) + .containsExactly( + new StreamRecord<>(2L), new StreamRecord<>(4L), new StreamRecord<>(6L)); + } + } + + @Test + void testEndInput() throws Exception { + CompletableFuture<Void> future = new CompletableFuture<>(); + OutputTag<Long> sideOutputTag = new OutputTag<Long>("side-output") {}; + + TwoOutputProcessOperator<Integer, Integer, Long> processOperator = + new TwoOutputProcessOperator<>( + new TwoOutputStreamProcessFunction<Integer, Integer, Long>() { + @Override + public void processRecord( + Integer record, + Collector<Integer> output1, + Collector<Long> output2, + RuntimeContext ctx) { + // do nothing. + } + + @Override + public void endInput( + TwoOutputNonPartitionedContext<Integer, Long> ctx) { + future.complete(null); + } + }, + sideOutputTag); + + try (OneInputStreamOperatorTestHarness<Integer, Integer> testHarness = + new OneInputStreamOperatorTestHarness<>(processOperator)) { + testHarness.open(); + testHarness.endInput(); + assertThat(future).isCompleted(); + } + } +} diff --git a/flink-streaming-java/src/test/java/org/apache/flink/streaming/util/TwoInputStreamOperatorTestHarness.java b/flink-streaming-java/src/test/java/org/apache/flink/streaming/util/TwoInputStreamOperatorTestHarness.java index 457a774211d..3bc0cdeab71 100644 --- a/flink-streaming-java/src/test/java/org/apache/flink/streaming/util/TwoInputStreamOperatorTestHarness.java +++ b/flink-streaming-java/src/test/java/org/apache/flink/streaming/util/TwoInputStreamOperatorTestHarness.java @@ -18,6 +18,7 @@ package org.apache.flink.streaming.util; +import org.apache.flink.streaming.api.operators.BoundedMultiInput; import org.apache.flink.streaming.api.operators.TwoInputStreamOperator; import org.apache.flink.streaming.api.watermark.Watermark; import org.apache.flink.streaming.runtime.streamrecord.RecordAttributes; @@ -98,4 +99,16 @@ public class TwoInputStreamOperatorTestHarness<IN1, IN2, OUT> public void processRecordAttributes2(RecordAttributes recordAttributes) throws Exception { twoInputOperator.processRecordAttributes2(recordAttributes); } + + public void endInput1() throws Exception { + if (operator instanceof BoundedMultiInput) { + ((BoundedMultiInput) operator).endInput(1); + } + } + + public void endInput2() throws Exception { + if (operator instanceof BoundedMultiInput) { + ((BoundedMultiInput) operator).endInput(2); + } + } }
