SEPURI-SAI-KRISHNA commented on code in PR #29345: URL: https://github.com/apache/flink/pull/29345#discussion_r4284245266
########## flink-table/flink-table-test-utils/src/test/java/org/apache/flink/table/runtime/functions/ProcessTableFunctionTestHarnessSnapshotTest.java: ########## @@ -0,0 +1,586 @@ +/* + * 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.table.runtime.functions; + +import org.apache.flink.table.annotation.ArgumentHint; +import org.apache.flink.table.annotation.ArgumentTrait; +import org.apache.flink.table.annotation.DataTypeHint; +import org.apache.flink.table.annotation.StateHint; +import org.apache.flink.table.api.DataTypes; +import org.apache.flink.table.api.dataview.ListView; +import org.apache.flink.table.api.dataview.MapView; +import org.apache.flink.table.api.dataview.ValueView; +import org.apache.flink.table.functions.ProcessTableFunction; +import org.apache.flink.table.runtime.functions.ProcessTableFunctionTestHarness.TableArgument; +import org.apache.flink.table.runtime.functions.ProcessTableFunctionTestHarness.TestSnapshot; +import org.apache.flink.types.Row; + +import org.junit.jupiter.api.Test; + +import java.time.Duration; +import java.time.Instant; +import java.time.LocalDateTime; +import java.util.ArrayList; +import java.util.List; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.junit.jupiter.api.Assertions.assertThrows; + +/** Tests for {@link ProcessTableFunctionTestHarness#snapshot()} and its restoration. */ +class ProcessTableFunctionTestHarnessSnapshotTest { + + /** Emits the input row if its value reaches the given threshold. */ + @DataTypeHint("ROW<value INT>") + public static class ThresholdFilterPTF extends ProcessTableFunction<Row> { + public void eval( + @ArgumentHint(ArgumentTrait.ROW_SEMANTIC_TABLE) Row input, + @ArgumentHint(ArgumentTrait.SCALAR) Integer threshold) { + int value = input.getFieldAs("value"); + if (value >= threshold) { + collect(input); + } + } + } + + /** Counts rows per partition in POJO state. */ + @DataTypeHint("ROW<count BIGINT>") + public static class CountingPTF extends ProcessTableFunction<Row> { + /** Counter held as structured-type state. */ + public static class CounterState { + public long counter = 0L; + } + + public void eval( + @StateHint CounterState state, + @ArgumentHint(ArgumentTrait.SET_SEMANTIC_TABLE) Row input) { + state.counter++; + collect(Row.of(state.counter)); + } + } + + /** Accumulates all seen values per partition in {@link ListView} state. */ + @DataTypeHint("ROW<values ARRAY<INT>>") + public static class ListStatePTF extends ProcessTableFunction<Row> { + public void eval( + @StateHint(type = @DataTypeHint("ARRAY<INT>")) ListView<Integer> listState, + @ArgumentHint(ArgumentTrait.SET_SEMANTIC_TABLE) Row input) + throws Exception { + listState.add(input.<Integer>getFieldAs("value")); + + final List<Integer> values = new ArrayList<>(); + for (Integer value : listState.get()) { + values.add(value); + } + collect(Row.of((Object) values.toArray(new Integer[0]))); + } + } + + /** Counts occurrences of every key per partition in {@link MapView} state. */ + @DataTypeHint("ROW<key STRING, count INT>") + public static class MapStatePTF extends ProcessTableFunction<Row> { + public void eval( + @StateHint MapView<String, Integer> mapState, + @ArgumentHint(ArgumentTrait.SET_SEMANTIC_TABLE) Row input) + throws Exception { + final String key = input.getFieldAs("key"); + final Integer count = mapState.get(key); + mapState.put(key, count == null ? 1 : count + 1); + collect(Row.of(key, mapState.get(key))); + } + } + + /** Counts rows per partition in {@link Row} state. */ + @DataTypeHint("ROW<count BIGINT>") + public static class RowStatePTF extends ProcessTableFunction<Row> { + public void eval( + @StateHint(type = @DataTypeHint("ROW<count BIGINT>")) Row memory, + @ArgumentHint(ArgumentTrait.SET_SEMANTIC_TABLE) Row input) { + long count = 1L; + if (memory.getField("count") != null) { + count += memory.<Long>getFieldAs("count"); + } + memory.setField("count", count); + collect(Row.of(count)); + } + } + + /** Accumulates a running total per partition in {@link ValueView} state. */ + @DataTypeHint("ROW<total INT>") + public static class ValueViewStatePTF extends ProcessTableFunction<Row> { + public void eval( + @StateHint ValueView<Integer> total, + @ArgumentHint(ArgumentTrait.SET_SEMANTIC_TABLE) Row input) { + Integer current = total.getValue(); + if (current == null) { + current = 0; + } + current += input.<Integer>getFieldAs("value"); + total.setValue(current); + collect(Row.of(current)); + } + } + + /** Registers a named timer five seconds after the row's event time. */ + @DataTypeHint("ROW<message STRING>") + public static class NamedTimerPTF extends ProcessTableFunction<Row> { + public void eval( + Context ctx, + @ArgumentHint({ArgumentTrait.SET_SEMANTIC_TABLE, ArgumentTrait.REQUIRE_ON_TIME}) + Row input) { + final TimeContext<LocalDateTime> timeCtx = ctx.timeContext(LocalDateTime.class); + timeCtx.registerOnTime("myTimer", timeCtx.time().plus(Duration.ofSeconds(5))); + } + + public void onTimer(OnTimerContext ctx) { + collect(Row.of("fired-" + ctx.currentTimer())); + } + } + + /** Registers a timer one millisecond after the current watermark. */ + @DataTypeHint("ROW<message STRING>") + public static class WatermarkTimerPTF extends ProcessTableFunction<Row> { + public void eval(Context ctx, @ArgumentHint(ArgumentTrait.SET_SEMANTIC_TABLE) Row input) { + final TimeContext<Long> timeCtx = ctx.timeContext(Long.class); + final Long watermark = timeCtx.currentWatermark(); + timeCtx.registerOnTime("t", watermark == null ? 1L : watermark + 1); + collect(Row.of("registered")); + } + + public void onTimer(OnTimerContext ctx) { + collect(Row.of("fired-" + ctx.currentTimer())); + } + } + + /** Counts its invocations on the function instance instead of in state. */ + @DataTypeHint("ROW<invocation INT>") + public static class InstanceCountingPTF extends ProcessTableFunction<Row> { + private int invocations = 0; + + public void eval(@ArgumentHint(ArgumentTrait.ROW_SEMANTIC_TABLE) Row input) { + invocations++; + collect(Row.of(invocations)); + } + } + + // ------------------------------------------------------------------------- + // Output + // ------------------------------------------------------------------------- + + @Test + void testSnapshotAndRestoreOutput() throws Exception { + final TestSnapshot<Row> snapshot; + + try (ProcessTableFunctionTestHarness<Row> harness = filterHarness()) { + harness.processElement(Row.of(100)); + harness.processElement(Row.of(30)); + assertThat(harness.getOutput()).containsExactly(Row.of(100)); + + snapshot = harness.snapshot(); + + harness.processElement(Row.of(75)); + assertThat(harness.getOutput()).containsExactly(Row.of(100), Row.of(75)); + } + + try (ProcessTableFunctionTestHarness<Row> restored = + ProcessTableFunctionTestHarness.restoreFromSnapshot(snapshot)) { + assertThat(restored.getOutput()).containsExactly(Row.of(100)); + assertThat(restored.getFunctionOutput()).containsExactly(Row.of(100)); + + restored.processElement(Row.of(75)); + assertThat(restored.getOutput()).containsExactly(Row.of(100), Row.of(75)); + } + } + + @Test + void testSnapshotIsNotAffectedByLaterOutput() throws Exception { Review Comment: Folded, keeping the `clearOutput()` assertion. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
