autophagy commented on code in PR #29345:
URL: https://github.com/apache/flink/pull/29345#discussion_r4246518565
##########
flink-table/flink-table-test-utils/src/main/java/org/apache/flink/table/runtime/functions/ProcessTableFunctionTestHarness.java:
##########
@@ -234,6 +239,58 @@ public static <OUT> Builder<OUT> ofClass(
return new Builder<>(functionClass);
}
+ //
-------------------------------------------------------------------------
+ // Test State Snapshotting
+ //
-------------------------------------------------------------------------
+
+ /**
+ * Takes a snapshot of this harness: the configuration it was built with,
the state of every
+ * partition, the pending and the fired timers, the watermark of every
table argument, and the
+ * collected output.
+ *
+ * <p>The snapshot is isolated from this harness, so processing further
elements, advancing a
+ * watermark or mutating state afterwards does not change what the
snapshot holds. State is
+ * copied in its internal representation, which means no user-facing state
object is touched.
+ * The one part that is not copied is the values collected by the PTF
itself (see {@link
+ * #getFunctionOutput()}), since those are the instances the function
emitted.
Review Comment:
I wonder if we should also copy the values collected by the PTF itself,
since I think this means that output from a pojo-typed ptf would leak across
restore/snapshot.
##########
flink-table/flink-table-test-utils/src/main/java/org/apache/flink/table/runtime/functions/ProcessTableFunctionTestHarness.java:
##########
@@ -234,6 +239,58 @@ public static <OUT> Builder<OUT> ofClass(
return new Builder<>(functionClass);
}
+ //
-------------------------------------------------------------------------
+ // Test State Snapshotting
+ //
-------------------------------------------------------------------------
+
+ /**
+ * Takes a snapshot of this harness: the configuration it was built with,
the state of every
+ * partition, the pending and the fired timers, the watermark of every
table argument, and the
+ * collected output.
+ *
+ * <p>The snapshot is isolated from this harness, so processing further
elements, advancing a
+ * watermark or mutating state afterwards does not change what the
snapshot holds. State is
+ * copied in its internal representation, which means no user-facing state
object is touched.
+ * The one part that is not copied is the values collected by the PTF
itself (see {@link
+ * #getFunctionOutput()}), since those are the instances the function
emitted.
+ *
+ * @see #restoreFromSnapshot(TestSnapshot)
+ */
+ public TestSnapshot<OUT> snapshot() {
+ return new TestSnapshot<>(
+ builder,
+ stateManager.snapshotState(),
+ timerManager.snapshot(),
+ new ArrayList<>(functionOutput),
+ copyRows(output));
+ }
+
+ /**
+ * Creates a new harness that continues from the given snapshot.
+ *
+ * <p>The returned harness runs a fresh instance of the PTF, opened as
{@link Builder#build()}
+ * would, and starts out with the state, timers, watermarks and output
that the snapshot was
+ * taken with. A snapshot can be restored more than once, and each
restored harness is
+ * independent of the others.
+ *
+ * @see #snapshot()
+ */
+ public static <OUT> ProcessTableFunctionTestHarness<OUT>
restoreFromSnapshot(
+ TestSnapshot<OUT> snapshot) throws Exception {
+ checkNotNull(snapshot, "snapshot must not be null");
+
+ final ProcessTableFunctionTestHarness<OUT> harness =
snapshot.builder.build();
+ harness.stateManager.restoreState(snapshot.state);
+ harness.timerManager.restore(snapshot.timers);
+ harness.functionOutput.addAll(snapshot.functionOutput);
+ harness.output.addAll(copyRows(snapshot.output));
+ return harness;
+ }
+
+ private static List<Row> copyRows(List<Row> rows) {
+ return rows.stream().map(Row::copy).collect(Collectors.toList());
Review Comment:
According to the doc for `Row.copy`, it doesn't perform a deep copy. In a
situation with a PTF with a nested output type, like `"ROW<arr ARRAY<INT>>"`,
this means you could modify an output row from the harness and it would modify
the output row in the snapshot (or vice verse), like:
```
final TestSnapshot<Row> snapshot = harness.snapshot();
final Integer[] liveArray = (Integer[])
harness.getOutput().get(0).getField(0);
liveArray[0] = -1;
```
If you then did
`ProcessTableFunctionTestHarness.restoreFromSnapshot(snapshot)`, the restored
harness' `getOutput().get(0).getField(0)[0]` would contain `-1`, and not
whatever the value was when the snapshot was taken.
It's a bit contrived but worth keeping in mind, I'm not sure of the best way
of doing a deep copy of a row is (other than roundtripping it through something
like `ExternalSerializer`).
##########
flink-table/flink-table-test-utils/src/main/java/org/apache/flink/table/runtime/functions/ListViewStateConverter.java:
##########
@@ -76,4 +80,15 @@ public Object toExternal(Object internal) {
public Object createNewInternalState() {
return new GenericArrayData(new Object[0]);
}
+
+ @Override
+ public Object copyInternal(Object internal) {
+ final ArrayData arrayData = (ArrayData) internal;
Review Comment:
I think if you replaced the `elementSerializer` with a `ArrayDataSerializer`
of `arrayType.getElementType()`, this function could just use its existing copy
rather than re-implementing it:
```
public Object copyInternal(Object internal) {
return arrayDataSerializer.copy((ArrayData) internal);
}
```
##########
flink-table/flink-table-test-utils/src/main/java/org/apache/flink/table/runtime/functions/ProcessTableFunctionTestHarness.java:
##########
@@ -234,6 +239,58 @@ public static <OUT> Builder<OUT> ofClass(
return new Builder<>(functionClass);
}
+ //
-------------------------------------------------------------------------
+ // Test State Snapshotting
+ //
-------------------------------------------------------------------------
+
+ /**
+ * Takes a snapshot of this harness: the configuration it was built with,
the state of every
+ * partition, the pending and the fired timers, the watermark of every
table argument, and the
+ * collected output.
+ *
+ * <p>The snapshot is isolated from this harness, so processing further
elements, advancing a
+ * watermark or mutating state afterwards does not change what the
snapshot holds. State is
+ * copied in its internal representation, which means no user-facing state
object is touched.
Review Comment:
I'm not sure stating how the snapshotting is done is very relevant to the
user, I think detail like this description of how state is snapshotted can be
omitted.
##########
flink-table/flink-table-test-utils/src/main/java/org/apache/flink/table/runtime/functions/ProcessTableFunctionTestHarness.java:
##########
@@ -1124,6 +1181,20 @@ private Builder(Class<? extends
ProcessTableFunction<OUT>> functionClass) {
this.functionClass = checkNotNull(functionClass, "functionClass
must not be null");
}
+ private Builder(Builder<OUT> other) {
+ this.functionClass = other.functionClass;
+ this.scalarArgs.putAll(other.scalarArgs);
+ this.tableArgs.putAll(other.tableArgs);
+ this.partitionConfigs.putAll(other.partitionConfigs);
+ other.stateArgs.forEach((name, config) -> this.stateArgs.put(name,
config.copy()));
+ this.onTimeColumnName = other.onTimeColumnName;
+ }
+
+ /** Creates a copy of this builder, detached from any further
configuration of it. */
+ private Builder<OUT> copy() {
+ return new Builder<>(this);
+ }
+
Review Comment:
As stated above, if we instead have the builder construct and pass in a
frozen immutable representation of its build state/initial conditions to the
harness, we wouldn't need this constructor or the copy.
##########
flink-table/flink-table-test-utils/src/main/java/org/apache/flink/table/runtime/functions/Timer.java:
##########
@@ -81,6 +81,15 @@ void markFired() {
this.fired = true;
}
+ /** Creates an isolated copy of this timer, including whether it has
already fired. */
+ Timer copy() {
+ final Timer copy = new Timer(timestamp, name, Row.copy(partitionKey));
Review Comment:
Might be more simpler to have a `Timer` constructor that also takes in the
fired state too, rather than marking it after construction.
##########
flink-table/flink-table-test-utils/src/main/java/org/apache/flink/table/runtime/functions/TestHarnessStateManager.java:
##########
@@ -198,6 +198,43 @@ <T> Map<Row, T> getStateForAllKeys(String stateName) {
return result;
}
+ /**
+ * Returns a deep copy of all state, keyed by partition key. Later
mutations of the live state
+ * are not visible through the returned map.
+ */
+ Map<Row, Map<String, Object>> snapshotState() {
+ return copyStateByKey(stateByKey);
+ }
+
+ /** Replaces all state with a deep copy of the given snapshot. */
+ void restoreState(Map<Row, Map<String, Object>> snapshot) {
+ stateByKey.clear();
+ stateByKey.putAll(copyStateByKey(snapshot));
+ }
Review Comment:
I think if we tackle the restoration of state/timers as commented in
https://github.com/apache/flink/pull/29345/changes#r4246665648 then I think
this pattern of clearing and then `putAll` would not be necessary, as we'd be
able to initialise the state manager with the correct state.
##########
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:
I think this test could be folded into `testSnapshotAndRestoreOutput`
##########
flink-table/flink-table-test-utils/src/main/java/org/apache/flink/table/runtime/functions/ProcessTableFunctionTestHarness.java:
##########
@@ -234,6 +239,58 @@ public static <OUT> Builder<OUT> ofClass(
return new Builder<>(functionClass);
}
+ //
-------------------------------------------------------------------------
+ // Test State Snapshotting
+ //
-------------------------------------------------------------------------
+
+ /**
+ * Takes a snapshot of this harness: the configuration it was built with,
the state of every
+ * partition, the pending and the fired timers, the watermark of every
table argument, and the
+ * collected output.
+ *
+ * <p>The snapshot is isolated from this harness, so processing further
elements, advancing a
+ * watermark or mutating state afterwards does not change what the
snapshot holds. State is
+ * copied in its internal representation, which means no user-facing state
object is touched.
+ * The one part that is not copied is the values collected by the PTF
itself (see {@link
+ * #getFunctionOutput()}), since those are the instances the function
emitted.
+ *
+ * @see #restoreFromSnapshot(TestSnapshot)
+ */
+ public TestSnapshot<OUT> snapshot() {
+ return new TestSnapshot<>(
+ builder,
+ stateManager.snapshotState(),
+ timerManager.snapshot(),
+ new ArrayList<>(functionOutput),
+ copyRows(output));
+ }
+
+ /**
+ * Creates a new harness that continues from the given snapshot.
+ *
+ * <p>The returned harness runs a fresh instance of the PTF, opened as
{@link Builder#build()}
+ * would, and starts out with the state, timers, watermarks and output
that the snapshot was
+ * taken with. A snapshot can be restored more than once, and each
restored harness is
+ * independent of the others.
+ *
+ * @see #snapshot()
+ */
+ public static <OUT> ProcessTableFunctionTestHarness<OUT>
restoreFromSnapshot(
+ TestSnapshot<OUT> snapshot) throws Exception {
+ checkNotNull(snapshot, "snapshot must not be null");
+
+ final ProcessTableFunctionTestHarness<OUT> harness =
snapshot.builder.build();
+ harness.stateManager.restoreState(snapshot.state);
+ harness.timerManager.restore(snapshot.timers);
+ harness.functionOutput.addAll(snapshot.functionOutput);
+ harness.output.addAll(copyRows(snapshot.output));
Review Comment:
I guess here what we are doing is building the harness from its initial
builder again (with initial conditions like initial state), creating the
state/timer manager from scratch and the updating all of the harness properties
with the restoration state. I suppose this is why we also test for
`testRestoreDropsStateThatWasClearedBeforeTheSnapshot`.
I wonder if it would be easier & clearer to have a private
`buildFromSnapshot(snapshot)` method on the builder itself that handle this. By
this i mean that instead of setting up the initial state and then overwriting
it on the harness' internal state/timer managers, the builder itself could
construct the right initial state from the snapshot and initialise the harness
with it directly?
It would also make it easier to handle `close()`ing the opened harness too
on restoration failure (since in this current construction I dont think this
happens if `restoreFromSnapshot` fails.
##########
flink-table/flink-table-test-utils/src/main/java/org/apache/flink/table/runtime/functions/ProcessTableFunctionTestHarness.java:
##########
@@ -155,7 +155,11 @@ private static class TableArgumentConverters {
@Nullable private InvocationContext currentInvocation;
+ /** The configuration this harness was built with, kept for {@link
#snapshot()}. */
+ private final Builder<OUT> builder;
Review Comment:
I'm not sure passing in the builder to the harness itself (then having to
copy it) and having a bit of a circular dependency is right to do here. Like
your comment says we don't care about this as a `Builder` but as the initial
configuration that the harness was built with. Typing it as a `Builder` means
it still allows adding more initial args, state etc even though we know its
meant to be a fixed representation.
I think instead having some sort of representation of the builder's build
state would make more sense, as an immutable value class that the builder
constructs at `build()` time. This way it'd both be clear what this is from the
type (not a mutable builder but as a frozen initial configuration) and would
prevent mutating the initial state.
##########
flink-table/flink-table-test-utils/src/main/java/org/apache/flink/table/runtime/functions/TestHarnessStateManager.java:
##########
@@ -198,6 +198,43 @@ <T> Map<Row, T> getStateForAllKeys(String stateName) {
return result;
}
+ /**
+ * Returns a deep copy of all state, keyed by partition key. Later
mutations of the live state
+ * are not visible through the returned map.
+ */
+ Map<Row, Map<String, Object>> snapshotState() {
Review Comment:
Small nit, but I think just `snapshot` and `restore` would be fine to use
here. its clear from them being in the state manager what we are snapshotting
and restoring.
##########
docs/content/docs/dev/table/functions/ptfs.md:
##########
@@ -2892,6 +2892,65 @@ void testAtomicOutputFunctionOutput() throws Exception {
{{< top >}}
+### Snapshotting and Restoring a Test
+
+`snapshot()` captures a test at a point in time: the state of every partition,
the pending and the
+fired timers, the watermark of every table argument, and the collected output.
Passing the snapshot
+to `ProcessTableFunctionTestHarness.restoreFromSnapshot()` creates a new
harness that continues from
+there, which lets several test cases share one setup instead of repeating it.
+
+A snapshot is isolated from the harness it was taken from, so processing
further elements, advancing
+a watermark or changing state afterwards does not change what it holds. It can
also be restored more
+than once, and the restored harnesses are independent of each other. Every
restore runs a fresh
Review Comment:
I guess this is not _entirely_ independent since I think the
`getFunctionOutput` might leak since we dont copy it.
--
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]
