You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@flink.apache.org by di...@apache.org on 2021/07/02 12:55:35 UTC
[flink] branch master updated (488e037 -> 18dca62)
This is an automated email from the ASF dual-hosted git repository.
dianfu pushed a change to branch master
in repository https://gitbox.apache.org/repos/asf/flink.git.
from 488e037 [hotfix][docs] Fix variable name
new 4dec5eb [FLINK-23213][python] Move class SimpleStateRequestHandler into a separate file
new 68bc6c6 [FLINK-23213][python] Refactor BeamPythonFunctionRunner to make it more readable by grouping similar methods together
new 18dca62 [FLINK-23213][python] Remove BeamTableStatefulPythonFunctionRunner and BeamTableStatelessPythonFunctionRunner and merge them into BeamTablePythonFunctionRunner
The 3 revisions listed above as "new" are entirely new to this
repository and will be described in separate emails. The revisions
listed as "add" were already present in the repository and have only
been added to this reference.
Summary of changes:
.../java/org/apache/flink/python/Constants.java | 3 +
.../beam/BeamDataStreamPythonFunctionRunner.java | 2 +
.../python/beam/BeamPythonFunctionRunner.java | 888 +++------------------
.../python/beam/SimpleStateRequestHandler.java | 669 ++++++++++++++++
.../python/AbstractStatelessFunctionOperator.java | 7 +-
.../AbstractPythonStreamAggregateOperator.java | 4 +-
.../python/beam/BeamTablePythonFunctionRunner.java | 14 +-
.../BeamTableStatefulPythonFunctionRunner.java | 76 --
.../BeamTableStatelessPythonFunctionRunner.java | 73 --
.../PassThroughPythonAggregateFunctionRunner.java | 8 +-
.../PassThroughPythonScalarFunctionRunner.java | 7 +-
.../PassThroughPythonTableFunctionRunner.java | 11 +-
...ThroughStreamAggregatePythonFunctionRunner.java | 5 +-
...amGroupWindowAggregatePythonFunctionRunner.java | 4 +-
...ghStreamTableAggregatePythonFunctionRunner.java | 4 +-
15 files changed, 846 insertions(+), 929 deletions(-)
create mode 100644 flink-python/src/main/java/org/apache/flink/streaming/api/runners/python/beam/SimpleStateRequestHandler.java
delete mode 100644 flink-python/src/main/java/org/apache/flink/table/runtime/runners/python/beam/BeamTableStatefulPythonFunctionRunner.java
delete mode 100644 flink-python/src/main/java/org/apache/flink/table/runtime/runners/python/beam/BeamTableStatelessPythonFunctionRunner.java
[flink] 03/03: [FLINK-23213][python] Remove
BeamTableStatefulPythonFunctionRunner and
BeamTableStatelessPythonFunctionRunner and merge them into
BeamTablePythonFunctionRunner
Posted by di...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
dianfu pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/flink.git
commit 18dca620fa8a7df6d2b3778cf1a458d7edbd0183
Author: Dian Fu <di...@apache.org>
AuthorDate: Fri Jul 2 14:49:12 2021 +0800
[FLINK-23213][python] Remove BeamTableStatefulPythonFunctionRunner and BeamTableStatelessPythonFunctionRunner and merge them into BeamTablePythonFunctionRunner
---
.../python/AbstractStatelessFunctionOperator.java | 7 +-
.../AbstractPythonStreamAggregateOperator.java | 4 +-
.../python/beam/BeamTablePythonFunctionRunner.java | 13 +++-
.../BeamTableStatefulPythonFunctionRunner.java | 76 ----------------------
.../BeamTableStatelessPythonFunctionRunner.java | 73 ---------------------
.../PassThroughPythonAggregateFunctionRunner.java | 8 ++-
.../PassThroughPythonScalarFunctionRunner.java | 7 +-
.../PassThroughPythonTableFunctionRunner.java | 11 ++--
...ThroughStreamAggregatePythonFunctionRunner.java | 5 +-
...amGroupWindowAggregatePythonFunctionRunner.java | 4 +-
...ghStreamTableAggregatePythonFunctionRunner.java | 4 +-
11 files changed, 41 insertions(+), 171 deletions(-)
diff --git a/flink-python/src/main/java/org/apache/flink/table/runtime/operators/python/AbstractStatelessFunctionOperator.java b/flink-python/src/main/java/org/apache/flink/table/runtime/operators/python/AbstractStatelessFunctionOperator.java
index 04f38df..21dd935 100644
--- a/flink-python/src/main/java/org/apache/flink/table/runtime/operators/python/AbstractStatelessFunctionOperator.java
+++ b/flink-python/src/main/java/org/apache/flink/table/runtime/operators/python/AbstractStatelessFunctionOperator.java
@@ -29,7 +29,7 @@ import org.apache.flink.fnexecution.v1.FlinkFnApi;
import org.apache.flink.python.PythonFunctionRunner;
import org.apache.flink.streaming.api.operators.python.AbstractOneInputPythonFunctionOperator;
import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
-import org.apache.flink.table.runtime.runners.python.beam.BeamTableStatelessPythonFunctionRunner;
+import org.apache.flink.table.runtime.runners.python.beam.BeamTablePythonFunctionRunner;
import org.apache.flink.table.types.logical.RowType;
import org.apache.flink.util.Preconditions;
@@ -143,7 +143,7 @@ public abstract class AbstractStatelessFunctionOperator<IN, OUT, UDFIN>
@Override
public PythonFunctionRunner createPythonFunctionRunner() throws IOException {
- return new BeamTableStatelessPythonFunctionRunner(
+ return new BeamTablePythonFunctionRunner(
getRuntimeContext().getTaskName(),
createPythonEnvironmentManager(),
userDefinedFunctionInputType,
@@ -152,6 +152,9 @@ public abstract class AbstractStatelessFunctionOperator<IN, OUT, UDFIN>
getUserDefinedFunctionsProto(),
jobOptions,
getFlinkMetricContainer(),
+ null,
+ null,
+ null,
getContainingTask().getEnvironment().getMemoryManager(),
getOperatorConfig()
.getManagedMemoryFractionOperatorUseCaseOfSlot(
diff --git a/flink-python/src/main/java/org/apache/flink/table/runtime/operators/python/aggregate/AbstractPythonStreamAggregateOperator.java b/flink-python/src/main/java/org/apache/flink/table/runtime/operators/python/aggregate/AbstractPythonStreamAggregateOperator.java
index 0014bd6..ef8e3ea 100644
--- a/flink-python/src/main/java/org/apache/flink/table/runtime/operators/python/aggregate/AbstractPythonStreamAggregateOperator.java
+++ b/flink-python/src/main/java/org/apache/flink/table/runtime/operators/python/aggregate/AbstractPythonStreamAggregateOperator.java
@@ -40,7 +40,7 @@ import org.apache.flink.table.planner.plan.utils.KeySelectorUtil;
import org.apache.flink.table.planner.typeutils.DataViewUtils;
import org.apache.flink.table.runtime.keyselector.RowDataKeySelector;
import org.apache.flink.table.runtime.operators.python.utils.StreamRecordRowDataWrappingCollector;
-import org.apache.flink.table.runtime.runners.python.beam.BeamTableStatefulPythonFunctionRunner;
+import org.apache.flink.table.runtime.runners.python.beam.BeamTablePythonFunctionRunner;
import org.apache.flink.table.runtime.typeutils.InternalTypeInfo;
import org.apache.flink.table.runtime.typeutils.PythonTypeUtils;
import org.apache.flink.table.types.logical.RowType;
@@ -192,7 +192,7 @@ public abstract class AbstractPythonStreamAggregateOperator
@Override
public PythonFunctionRunner createPythonFunctionRunner() throws Exception {
- return new BeamTableStatefulPythonFunctionRunner(
+ return new BeamTablePythonFunctionRunner(
getRuntimeContext().getTaskName(),
createPythonEnvironmentManager(),
userDefinedFunctionInputType,
diff --git a/flink-python/src/main/java/org/apache/flink/table/runtime/runners/python/beam/BeamTablePythonFunctionRunner.java b/flink-python/src/main/java/org/apache/flink/table/runtime/runners/python/beam/BeamTablePythonFunctionRunner.java
index b59ef4c..1f5dde0 100644
--- a/flink-python/src/main/java/org/apache/flink/table/runtime/runners/python/beam/BeamTablePythonFunctionRunner.java
+++ b/flink-python/src/main/java/org/apache/flink/table/runtime/runners/python/beam/BeamTablePythonFunctionRunner.java
@@ -29,6 +29,7 @@ import org.apache.flink.streaming.api.runners.python.beam.BeamPythonFunctionRunn
import org.apache.flink.table.types.logical.RowType;
import org.apache.flink.util.Preconditions;
+import com.google.protobuf.GeneratedMessageV3;
import org.apache.beam.model.pipeline.v1.RunnerApi;
import java.util.Map;
@@ -38,17 +39,19 @@ import static org.apache.flink.table.runtime.typeutils.PythonTypeUtils.toProtoTy
/** A {@link BeamTablePythonFunctionRunner} used to execute Python functions in Table API. */
@Internal
-public abstract class BeamTablePythonFunctionRunner extends BeamPythonFunctionRunner {
+public class BeamTablePythonFunctionRunner extends BeamPythonFunctionRunner {
private final RowType inputType;
private final RowType outputType;
+ private final GeneratedMessageV3 userDefinedFunctionProto;
- BeamTablePythonFunctionRunner(
+ public BeamTablePythonFunctionRunner(
String taskName,
PythonEnvironmentManager environmentManager,
RowType inputType,
RowType outputType,
String functionUrn,
+ GeneratedMessageV3 userDefinedFunctionProto,
Map<String, String> jobOptions,
FlinkMetricContainer flinkMetricContainer,
KeyedStateBackend keyedStateBackend,
@@ -75,6 +78,12 @@ public abstract class BeamTablePythonFunctionRunner extends BeamPythonFunctionRu
outputMode);
this.inputType = Preconditions.checkNotNull(inputType);
this.outputType = Preconditions.checkNotNull(outputType);
+ this.userDefinedFunctionProto = Preconditions.checkNotNull(userDefinedFunctionProto);
+ }
+
+ @Override
+ protected byte[] getUserDefinedFunctionsProtoBytes() {
+ return userDefinedFunctionProto.toByteArray();
}
@Override
diff --git a/flink-python/src/main/java/org/apache/flink/table/runtime/runners/python/beam/BeamTableStatefulPythonFunctionRunner.java b/flink-python/src/main/java/org/apache/flink/table/runtime/runners/python/beam/BeamTableStatefulPythonFunctionRunner.java
deleted file mode 100644
index 7f1d8f7..0000000
--- a/flink-python/src/main/java/org/apache/flink/table/runtime/runners/python/beam/BeamTableStatefulPythonFunctionRunner.java
+++ /dev/null
@@ -1,76 +0,0 @@
-/*
- * 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.runners.python.beam;
-
-import org.apache.flink.api.common.typeutils.TypeSerializer;
-import org.apache.flink.fnexecution.v1.FlinkFnApi;
-import org.apache.flink.python.env.PythonEnvironmentManager;
-import org.apache.flink.python.metric.FlinkMetricContainer;
-import org.apache.flink.runtime.memory.MemoryManager;
-import org.apache.flink.runtime.state.KeyedStateBackend;
-import org.apache.flink.table.types.logical.RowType;
-
-import java.util.Map;
-
-/** A {@link BeamTableStatefulPythonFunctionRunner} used to execute Python stateful functions. */
-public class BeamTableStatefulPythonFunctionRunner extends BeamTablePythonFunctionRunner {
-
- private final FlinkFnApi.UserDefinedAggregateFunctions userDefinedAggregateFunctions;
-
- public BeamTableStatefulPythonFunctionRunner(
- String taskName,
- PythonEnvironmentManager environmentManager,
- RowType inputType,
- RowType outputType,
- String functionUrn,
- FlinkFnApi.UserDefinedAggregateFunctions userDefinedFunctions,
- Map<String, String> jobOptions,
- FlinkMetricContainer flinkMetricContainer,
- KeyedStateBackend keyedStateBackend,
- TypeSerializer keySerializer,
- TypeSerializer namespaceSerializer,
- MemoryManager memoryManager,
- double managedMemoryFraction,
- FlinkFnApi.CoderParam.DataType inputDataType,
- FlinkFnApi.CoderParam.DataType outputDataType,
- FlinkFnApi.CoderParam.OutputMode outputMode) {
- super(
- taskName,
- environmentManager,
- inputType,
- outputType,
- functionUrn,
- jobOptions,
- flinkMetricContainer,
- keyedStateBackend,
- keySerializer,
- namespaceSerializer,
- memoryManager,
- managedMemoryFraction,
- inputDataType,
- outputDataType,
- outputMode);
- this.userDefinedAggregateFunctions = userDefinedFunctions;
- }
-
- @Override
- protected byte[] getUserDefinedFunctionsProtoBytes() {
- return this.userDefinedAggregateFunctions.toByteArray();
- }
-}
diff --git a/flink-python/src/main/java/org/apache/flink/table/runtime/runners/python/beam/BeamTableStatelessPythonFunctionRunner.java b/flink-python/src/main/java/org/apache/flink/table/runtime/runners/python/beam/BeamTableStatelessPythonFunctionRunner.java
deleted file mode 100644
index 7f6b660..0000000
--- a/flink-python/src/main/java/org/apache/flink/table/runtime/runners/python/beam/BeamTableStatelessPythonFunctionRunner.java
+++ /dev/null
@@ -1,73 +0,0 @@
-/*
- * 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.runners.python.beam;
-
-import org.apache.flink.annotation.Internal;
-import org.apache.flink.fnexecution.v1.FlinkFnApi;
-import org.apache.flink.python.env.PythonEnvironmentManager;
-import org.apache.flink.python.metric.FlinkMetricContainer;
-import org.apache.flink.runtime.memory.MemoryManager;
-import org.apache.flink.table.types.logical.RowType;
-
-import java.util.Map;
-
-/** A {@link BeamTableStatelessPythonFunctionRunner} used to execute Python stateless functions. */
-@Internal
-public class BeamTableStatelessPythonFunctionRunner extends BeamTablePythonFunctionRunner {
-
- private final FlinkFnApi.UserDefinedFunctions userDefinedFunctions;
-
- public BeamTableStatelessPythonFunctionRunner(
- String taskName,
- PythonEnvironmentManager environmentManager,
- RowType inputType,
- RowType outputType,
- String functionUrn,
- FlinkFnApi.UserDefinedFunctions userDefinedFunctions,
- Map<String, String> jobOptions,
- FlinkMetricContainer flinkMetricContainer,
- MemoryManager memoryManager,
- double managedMemoryFraction,
- FlinkFnApi.CoderParam.DataType inputDataType,
- FlinkFnApi.CoderParam.DataType outputDataType,
- FlinkFnApi.CoderParam.OutputMode outputMode) {
- super(
- taskName,
- environmentManager,
- inputType,
- outputType,
- functionUrn,
- jobOptions,
- flinkMetricContainer,
- null,
- null,
- null,
- memoryManager,
- managedMemoryFraction,
- inputDataType,
- outputDataType,
- outputMode);
- this.userDefinedFunctions = userDefinedFunctions;
- }
-
- @Override
- protected byte[] getUserDefinedFunctionsProtoBytes() {
- return this.userDefinedFunctions.toByteArray();
- }
-}
diff --git a/flink-python/src/test/java/org/apache/flink/table/runtime/utils/PassThroughPythonAggregateFunctionRunner.java b/flink-python/src/test/java/org/apache/flink/table/runtime/utils/PassThroughPythonAggregateFunctionRunner.java
index f7ef576..76cb9ab 100644
--- a/flink-python/src/test/java/org/apache/flink/table/runtime/utils/PassThroughPythonAggregateFunctionRunner.java
+++ b/flink-python/src/test/java/org/apache/flink/table/runtime/utils/PassThroughPythonAggregateFunctionRunner.java
@@ -28,7 +28,7 @@ import org.apache.flink.python.env.PythonEnvironmentManager;
import org.apache.flink.python.metric.FlinkMetricContainer;
import org.apache.flink.table.data.RowData;
import org.apache.flink.table.runtime.arrow.serializers.RowDataArrowSerializer;
-import org.apache.flink.table.runtime.runners.python.beam.BeamTableStatelessPythonFunctionRunner;
+import org.apache.flink.table.runtime.runners.python.beam.BeamTablePythonFunctionRunner;
import org.apache.flink.table.types.logical.RowType;
import org.apache.beam.runners.fnexecution.control.JobBundleFactory;
@@ -43,8 +43,7 @@ import java.util.Map;
* A {@link PassThroughPythonAggregateFunctionRunner} runner that just return the first input
* element with the same key as the execution results.
*/
-public class PassThroughPythonAggregateFunctionRunner
- extends BeamTableStatelessPythonFunctionRunner {
+public class PassThroughPythonAggregateFunctionRunner extends BeamTablePythonFunctionRunner {
private static final IntSerializer windowBoundarySerializer = IntSerializer.INSTANCE;
@@ -84,6 +83,9 @@ public class PassThroughPythonAggregateFunctionRunner
jobOptions,
flinkMetricContainer,
null,
+ null,
+ null,
+ null,
0.0,
FlinkFnApi.CoderParam.DataType.ARROW,
FlinkFnApi.CoderParam.DataType.ARROW,
diff --git a/flink-python/src/test/java/org/apache/flink/table/runtime/utils/PassThroughPythonScalarFunctionRunner.java b/flink-python/src/test/java/org/apache/flink/table/runtime/utils/PassThroughPythonScalarFunctionRunner.java
index e353402..04e5c57 100644
--- a/flink-python/src/test/java/org/apache/flink/table/runtime/utils/PassThroughPythonScalarFunctionRunner.java
+++ b/flink-python/src/test/java/org/apache/flink/table/runtime/utils/PassThroughPythonScalarFunctionRunner.java
@@ -21,7 +21,7 @@ package org.apache.flink.table.runtime.utils;
import org.apache.flink.fnexecution.v1.FlinkFnApi;
import org.apache.flink.python.env.PythonEnvironmentManager;
import org.apache.flink.python.metric.FlinkMetricContainer;
-import org.apache.flink.table.runtime.runners.python.beam.BeamTableStatelessPythonFunctionRunner;
+import org.apache.flink.table.runtime.runners.python.beam.BeamTablePythonFunctionRunner;
import org.apache.flink.table.types.logical.RowType;
import org.apache.beam.runners.fnexecution.control.JobBundleFactory;
@@ -35,7 +35,7 @@ import java.util.Map;
* A {@link PassThroughPythonScalarFunctionRunner} runner that just return the input elements as the
* execution results.
*/
-public class PassThroughPythonScalarFunctionRunner extends BeamTableStatelessPythonFunctionRunner {
+public class PassThroughPythonScalarFunctionRunner extends BeamTablePythonFunctionRunner {
private final List<byte[]> buffer;
@@ -58,6 +58,9 @@ public class PassThroughPythonScalarFunctionRunner extends BeamTableStatelessPyt
jobOptions,
flinkMetricContainer,
null,
+ null,
+ null,
+ null,
0.0,
FlinkFnApi.CoderParam.DataType.FLATTEN_ROW,
FlinkFnApi.CoderParam.DataType.FLATTEN_ROW,
diff --git a/flink-python/src/test/java/org/apache/flink/table/runtime/utils/PassThroughPythonTableFunctionRunner.java b/flink-python/src/test/java/org/apache/flink/table/runtime/utils/PassThroughPythonTableFunctionRunner.java
index 279d2de..b2165f7 100644
--- a/flink-python/src/test/java/org/apache/flink/table/runtime/utils/PassThroughPythonTableFunctionRunner.java
+++ b/flink-python/src/test/java/org/apache/flink/table/runtime/utils/PassThroughPythonTableFunctionRunner.java
@@ -21,7 +21,7 @@ package org.apache.flink.table.runtime.utils;
import org.apache.flink.fnexecution.v1.FlinkFnApi;
import org.apache.flink.python.env.PythonEnvironmentManager;
import org.apache.flink.python.metric.FlinkMetricContainer;
-import org.apache.flink.table.runtime.runners.python.beam.BeamTableStatelessPythonFunctionRunner;
+import org.apache.flink.table.runtime.runners.python.beam.BeamTablePythonFunctionRunner;
import org.apache.flink.table.types.logical.RowType;
import org.apache.beam.runners.fnexecution.control.JobBundleFactory;
@@ -32,10 +32,10 @@ import java.util.List;
import java.util.Map;
/**
- * A {@link BeamTableStatelessPythonFunctionRunner} that emit each input element in inner join and
- * emit null in left join when certain test conditions are met.
+ * A {@link BeamTablePythonFunctionRunner} that emit each input element in inner join and emit null
+ * in left join when certain test conditions are met.
*/
-public class PassThroughPythonTableFunctionRunner extends BeamTableStatelessPythonFunctionRunner {
+public class PassThroughPythonTableFunctionRunner extends BeamTablePythonFunctionRunner {
private int num = 0;
@@ -60,6 +60,9 @@ public class PassThroughPythonTableFunctionRunner extends BeamTableStatelessPyth
jobOptions,
flinkMetricContainer,
null,
+ null,
+ null,
+ null,
0.0,
FlinkFnApi.CoderParam.DataType.FLATTEN_ROW,
FlinkFnApi.CoderParam.DataType.FLATTEN_ROW,
diff --git a/flink-python/src/test/java/org/apache/flink/table/runtime/utils/PassThroughStreamAggregatePythonFunctionRunner.java b/flink-python/src/test/java/org/apache/flink/table/runtime/utils/PassThroughStreamAggregatePythonFunctionRunner.java
index 7a35a9f..4fe2f0c 100644
--- a/flink-python/src/test/java/org/apache/flink/table/runtime/utils/PassThroughStreamAggregatePythonFunctionRunner.java
+++ b/flink-python/src/test/java/org/apache/flink/table/runtime/utils/PassThroughStreamAggregatePythonFunctionRunner.java
@@ -23,7 +23,7 @@ import org.apache.flink.fnexecution.v1.FlinkFnApi;
import org.apache.flink.python.env.PythonEnvironmentManager;
import org.apache.flink.python.metric.FlinkMetricContainer;
import org.apache.flink.runtime.state.KeyedStateBackend;
-import org.apache.flink.table.runtime.runners.python.beam.BeamTableStatefulPythonFunctionRunner;
+import org.apache.flink.table.runtime.runners.python.beam.BeamTablePythonFunctionRunner;
import org.apache.flink.table.types.logical.RowType;
import org.apache.beam.runners.fnexecution.control.JobBundleFactory;
@@ -39,8 +39,7 @@ import java.util.function.Function;
* stream group aggregate operators. It will process the input data with the provided
* `processFunction`.
*/
-public class PassThroughStreamAggregatePythonFunctionRunner
- extends BeamTableStatefulPythonFunctionRunner {
+public class PassThroughStreamAggregatePythonFunctionRunner extends BeamTablePythonFunctionRunner {
private final List<byte[]> buffer;
diff --git a/flink-python/src/test/java/org/apache/flink/table/runtime/utils/PassThroughStreamGroupWindowAggregatePythonFunctionRunner.java b/flink-python/src/test/java/org/apache/flink/table/runtime/utils/PassThroughStreamGroupWindowAggregatePythonFunctionRunner.java
index 8b1ed79..0126d4b 100644
--- a/flink-python/src/test/java/org/apache/flink/table/runtime/utils/PassThroughStreamGroupWindowAggregatePythonFunctionRunner.java
+++ b/flink-python/src/test/java/org/apache/flink/table/runtime/utils/PassThroughStreamGroupWindowAggregatePythonFunctionRunner.java
@@ -24,7 +24,7 @@ import org.apache.flink.python.env.PythonEnvironmentManager;
import org.apache.flink.python.metric.FlinkMetricContainer;
import org.apache.flink.runtime.state.KeyedStateBackend;
import org.apache.flink.table.runtime.operators.python.aggregate.PassThroughPythonStreamGroupWindowAggregateOperator;
-import org.apache.flink.table.runtime.runners.python.beam.BeamTableStatefulPythonFunctionRunner;
+import org.apache.flink.table.runtime.runners.python.beam.BeamTablePythonFunctionRunner;
import org.apache.flink.table.types.logical.RowType;
import org.apache.beam.runners.fnexecution.control.JobBundleFactory;
@@ -37,7 +37,7 @@ import java.util.Map;
* Python stream group window aggregate operators.
*/
public class PassThroughStreamGroupWindowAggregatePythonFunctionRunner
- extends BeamTableStatefulPythonFunctionRunner {
+ extends BeamTablePythonFunctionRunner {
private final PassThroughPythonStreamGroupWindowAggregateOperator operator;
diff --git a/flink-python/src/test/java/org/apache/flink/table/runtime/utils/PassThroughStreamTableAggregatePythonFunctionRunner.java b/flink-python/src/test/java/org/apache/flink/table/runtime/utils/PassThroughStreamTableAggregatePythonFunctionRunner.java
index 7339e97..77934cb 100644
--- a/flink-python/src/test/java/org/apache/flink/table/runtime/utils/PassThroughStreamTableAggregatePythonFunctionRunner.java
+++ b/flink-python/src/test/java/org/apache/flink/table/runtime/utils/PassThroughStreamTableAggregatePythonFunctionRunner.java
@@ -23,7 +23,7 @@ import org.apache.flink.fnexecution.v1.FlinkFnApi;
import org.apache.flink.python.env.PythonEnvironmentManager;
import org.apache.flink.python.metric.FlinkMetricContainer;
import org.apache.flink.runtime.state.KeyedStateBackend;
-import org.apache.flink.table.runtime.runners.python.beam.BeamTableStatefulPythonFunctionRunner;
+import org.apache.flink.table.runtime.runners.python.beam.BeamTablePythonFunctionRunner;
import org.apache.flink.table.types.logical.RowType;
import org.apache.beam.runners.fnexecution.control.JobBundleFactory;
@@ -41,7 +41,7 @@ import java.util.function.Function;
* `processFunction`.
*/
public class PassThroughStreamTableAggregatePythonFunctionRunner
- extends BeamTableStatefulPythonFunctionRunner {
+ extends BeamTablePythonFunctionRunner {
private final List<byte[]> buffer;
[flink] 02/03: [FLINK-23213][python] Refactor
BeamPythonFunctionRunner to make it more readable by grouping similar
methods together
Posted by di...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
dianfu pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/flink.git
commit 68bc6c61b41c7dfb77472035f15d535e2dd38f28
Author: Dian Fu <di...@apache.org>
AuthorDate: Fri Jul 2 14:00:01 2021 +0800
[FLINK-23213][python] Refactor BeamPythonFunctionRunner to make it more readable by grouping similar methods together
---
.../java/org/apache/flink/python/Constants.java | 3 +
.../beam/BeamDataStreamPythonFunctionRunner.java | 2 +
.../python/beam/BeamPythonFunctionRunner.java | 252 +++++++++++----------
.../python/beam/BeamTablePythonFunctionRunner.java | 1 +
4 files changed, 141 insertions(+), 117 deletions(-)
diff --git a/flink-python/src/main/java/org/apache/flink/python/Constants.java b/flink-python/src/main/java/org/apache/flink/python/Constants.java
index 5111cc6..54d1d06 100644
--- a/flink-python/src/main/java/org/apache/flink/python/Constants.java
+++ b/flink-python/src/main/java/org/apache/flink/python/Constants.java
@@ -24,4 +24,7 @@ public class Constants {
// function urns
public static final String STATELESS_FUNCTION_URN = "flink:transform:ds:stateless_function:v1";
public static final String STATEFUL_FUNCTION_URN = "flink:transform:ds:stateful_function:v1";
+
+ // coder urns
+ public static final String FLINK_CODER_URN = "flink:coder:v1";
}
diff --git a/flink-python/src/main/java/org/apache/flink/streaming/api/runners/python/beam/BeamDataStreamPythonFunctionRunner.java b/flink-python/src/main/java/org/apache/flink/streaming/api/runners/python/beam/BeamDataStreamPythonFunctionRunner.java
index 29c8f1d..fa138d6 100644
--- a/flink-python/src/main/java/org/apache/flink/streaming/api/runners/python/beam/BeamDataStreamPythonFunctionRunner.java
+++ b/flink-python/src/main/java/org/apache/flink/streaming/api/runners/python/beam/BeamDataStreamPythonFunctionRunner.java
@@ -34,6 +34,8 @@ import javax.annotation.Nullable;
import java.util.Map;
+import static org.apache.flink.python.Constants.FLINK_CODER_URN;
+
/**
* {@link BeamDataStreamPythonFunctionRunner} is responsible for starting a beam python harness to
* execute user defined python function.
diff --git a/flink-python/src/main/java/org/apache/flink/streaming/api/runners/python/beam/BeamPythonFunctionRunner.java b/flink-python/src/main/java/org/apache/flink/streaming/api/runners/python/beam/BeamPythonFunctionRunner.java
index 177e242..30992b6 100644
--- a/flink-python/src/main/java/org/apache/flink/streaming/api/runners/python/beam/BeamPythonFunctionRunner.java
+++ b/flink-python/src/main/java/org/apache/flink/streaming/api/runners/python/beam/BeamPythonFunctionRunner.java
@@ -19,6 +19,7 @@
package org.apache.flink.streaming.api.runners.python.beam;
import org.apache.flink.annotation.Internal;
+import org.apache.flink.annotation.VisibleForTesting;
import org.apache.flink.api.common.typeutils.TypeSerializer;
import org.apache.flink.api.java.tuple.Tuple2;
import org.apache.flink.fnexecution.v1.FlinkFnApi;
@@ -102,27 +103,33 @@ public abstract class BeamPythonFunctionRunner implements PythonFunctionRunner {
private static final String MANAGED_MEMORY_RESOURCE_ID = "python-process-managed-memory";
private static final String PYTHON_WORKER_MEMORY_LIMIT = "_PYTHON_WORKER_MEMORY_LIMIT";
- protected static final String FLINK_CODER_URN = "flink:coder:v1";
-
- protected final FlinkFnApi.CoderParam.OutputMode outputMode;
-
- protected final FlinkFnApi.CoderParam.DataType inputDataType;
- protected final FlinkFnApi.CoderParam.DataType outputDataType;
-
- private transient boolean bundleStarted;
-
private final String taskName;
/** The Python execution environment manager. */
private final PythonEnvironmentManager environmentManager;
+ /** The urn which represents the function kind to be executed. */
private final String functionUrn;
/** The options used to configure the Python worker process. */
private final Map<String, String> jobOptions;
- /** The Python function execution result tuple. */
- protected final Tuple2<byte[], Integer> resultTuple;
+ /** The flinkMetricContainer will be set to null if metric is configured to be turned off. */
+ @Nullable private FlinkMetricContainer flinkMetricContainer;
+
+ private final MemoryManager memoryManager;
+
+ /** The fraction of total managed memory in the slot that the Python worker could use. */
+ private final double managedMemoryFraction;
+
+ protected final FlinkFnApi.CoderParam.DataType inputDataType;
+ protected final FlinkFnApi.CoderParam.DataType outputDataType;
+
+ protected final FlinkFnApi.CoderParam.OutputMode outputMode;
+
+ // ------------------------------------------------------------------------
+
+ private transient boolean bundleStarted;
/**
* The bundle factory which has all job-scoped information and can be used to create a {@link
@@ -152,33 +159,28 @@ public abstract class BeamPythonFunctionRunner implements PythonFunctionRunner {
*/
private transient RemoteBundle remoteBundle;
- /** The Python function execution result receiver. */
- protected transient LinkedBlockingQueue<byte[]> resultBuffer;
-
- /** The receiver which forwards the input elements to a remote environment for processing. */
- protected transient FnDataReceiver<WindowedValue<byte[]>> mainInputReceiver;
-
- /** The flinkMetricContainer will be set to null if metric is configured to be turned off. */
- @Nullable private FlinkMetricContainer flinkMetricContainer;
+ /** The Python function execution result tuple: (raw bytes, length). */
+ private transient Tuple2<byte[], Integer> reusableResultTuple;
- @Nullable private final MemoryManager memoryManager;
+ /** Buffers the Python function execution result which has still not been processed. */
+ @VisibleForTesting protected transient LinkedBlockingQueue<byte[]> resultBuffer;
- /** The fraction of total managed memory in the slot that the Python worker could use. */
- private final double managedMemoryFraction;
+ /** The receiver which forwards the input elements to a remote environment for processing. */
+ @VisibleForTesting protected transient FnDataReceiver<WindowedValue<byte[]>> mainInputReceiver;
/** The shared resource among Python operators of the same slot. */
- private OpaqueMemoryResource<PythonSharedResources> sharedResources;
+ private transient OpaqueMemoryResource<PythonSharedResources> sharedResources;
public BeamPythonFunctionRunner(
String taskName,
PythonEnvironmentManager environmentManager,
String functionUrn,
Map<String, String> jobOptions,
- FlinkMetricContainer flinkMetricContainer,
+ @Nullable FlinkMetricContainer flinkMetricContainer,
@Nullable KeyedStateBackend keyedStateBackend,
@Nullable TypeSerializer keySerializer,
@Nullable TypeSerializer namespaceSerializer,
- @Nullable MemoryManager memoryManager,
+ MemoryManager memoryManager,
double managedMemoryFraction,
FlinkFnApi.CoderParam.DataType inputDataType,
FlinkFnApi.CoderParam.DataType outputDataType,
@@ -196,13 +198,15 @@ public abstract class BeamPythonFunctionRunner implements PythonFunctionRunner {
this.inputDataType = Preconditions.checkNotNull(inputDataType);
this.outputDataType = Preconditions.checkNotNull(outputDataType);
this.outputMode = Preconditions.checkNotNull(outputMode);
- this.resultTuple = new Tuple2<>();
}
+ // ------------------------------------------------------------------------
+
@Override
public void open(PythonConfig config) throws Exception {
this.bundleStarted = false;
this.resultBuffer = new LinkedBlockingQueue<>();
+ this.reusableResultTuple = new Tuple2<>();
// The creation of stageBundleFactory depends on the initialized environment manager.
environmentManager.open();
@@ -288,15 +292,38 @@ public abstract class BeamPythonFunctionRunner implements PythonFunctionRunner {
mainInputReceiver.accept(WindowedValue.valueInGlobalWindow(data));
}
+ /** Checks whether to invoke startBundle. */
+ private void checkInvokeStartBundle() {
+ if (!bundleStarted) {
+ startBundle();
+ bundleStarted = true;
+ }
+ }
+
+ @VisibleForTesting
+ protected void startBundle() {
+ try {
+ remoteBundle =
+ stageBundleFactory.getBundle(
+ createOutputReceiverFactory(), stateRequestHandler, progressHandler);
+ mainInputReceiver =
+ Preconditions.checkNotNull(
+ Iterables.getOnlyElement(remoteBundle.getInputReceivers().values()),
+ "Failed to retrieve main input receiver.");
+ } catch (Throwable t) {
+ throw new RuntimeException("Failed to start remote bundle", t);
+ }
+ }
+
@Override
public Tuple2<byte[], Integer> pollResult() throws Exception {
byte[] result = resultBuffer.poll();
if (result == null) {
return null;
} else {
- this.resultTuple.f0 = result;
- this.resultTuple.f1 = result.length;
- return this.resultTuple;
+ this.reusableResultTuple.f0 = result;
+ this.reusableResultTuple.f1 = result.length;
+ return this.reusableResultTuple;
}
}
@@ -311,15 +338,20 @@ public abstract class BeamPythonFunctionRunner implements PythonFunctionRunner {
}
}
- public JobBundleFactory createJobBundleFactory(Struct pipelineOptions) throws Exception {
- return DefaultJobBundleFactory.create(
- JobInfo.create(
- taskName,
- taskName,
- environmentManager.createRetrievalToken(),
- pipelineOptions));
+ private void finishBundle() {
+ try {
+ remoteBundle.close();
+ } catch (Throwable t) {
+ throw new RuntimeException("Failed to close remote bundle", t);
+ } finally {
+ remoteBundle = null;
+ }
}
+ // ------------------------------------------------------------------------
+ // Python Execution Environment Management
+ // ------------------------------------------------------------------------
+
/**
* Creates a specification which specifies the portability Python execution environment. It's
* used by Beam's portability framework to creates the actual Python execution environment.
@@ -338,81 +370,9 @@ public abstract class BeamPythonFunctionRunner implements PythonFunctionRunner {
throw new RuntimeException("Currently only ProcessPythonEnvironment is supported.");
}
- protected void startBundle() {
- try {
- remoteBundle =
- stageBundleFactory.getBundle(
- createOutputReceiverFactory(), stateRequestHandler, progressHandler);
- mainInputReceiver =
- Preconditions.checkNotNull(
- Iterables.getOnlyElement(remoteBundle.getInputReceivers().values()),
- "Failed to retrieve main input receiver.");
- } catch (Throwable t) {
- throw new RuntimeException("Failed to start remote bundle", t);
- }
- }
-
- private void finishBundle() {
- try {
- remoteBundle.close();
- } catch (Throwable t) {
- throw new RuntimeException("Failed to close remote bundle", t);
- } finally {
- remoteBundle = null;
- }
- }
-
- private OutputReceiverFactory createOutputReceiverFactory() {
- return new OutputReceiverFactory() {
-
- // the input value type is always byte array
- @SuppressWarnings("unchecked")
- @Override
- public FnDataReceiver<WindowedValue<byte[]>> create(String pCollectionId) {
- return input -> resultBuffer.add(input.getValue());
- }
- };
- }
-
- /**
- * Ignore bundle progress if flinkMetricContainer is null. The flinkMetricContainer will be set
- * to null if metric is configured to be turned off.
- */
- private BundleProgressHandler getProgressHandler(FlinkMetricContainer flinkMetricContainer) {
- if (flinkMetricContainer == null) {
- return BundleProgressHandler.ignored();
- } else {
- return new BundleProgressHandler() {
- @Override
- public void onProgress(BeamFnApi.ProcessBundleProgressResponse progress) {
- flinkMetricContainer.updateMetrics(taskName, progress.getMonitoringInfosList());
- }
-
- @Override
- public void onCompleted(BeamFnApi.ProcessBundleResponse response) {
- flinkMetricContainer.updateMetrics(taskName, response.getMonitoringInfosList());
- }
- };
- }
- }
-
- /** To make the error messages more user friendly, throws an exception with the boot logs. */
- private StageBundleFactory createStageBundleFactory(
- JobBundleFactory jobBundleFactory, RunnerApi.Environment environment) throws Exception {
- try {
- return jobBundleFactory.forStage(createExecutableStage(environment));
- } catch (Throwable e) {
- throw new RuntimeException(environmentManager.getBootLog(), e);
- }
- }
-
- /** Checks whether to invoke startBundle. */
- private void checkInvokeStartBundle() throws Exception {
- if (!bundleStarted) {
- startBundle();
- bundleStarted = true;
- }
- }
+ // ------------------------------------------------------------------------
+ // Construct ExecutableStage
+ // ------------------------------------------------------------------------
/**
* Creates a {@link ExecutableStage} which contains the Python user-defined functions to be
@@ -513,12 +473,6 @@ public abstract class BeamPythonFunctionRunner implements PythonFunctionRunner {
.build());
}
- protected abstract byte[] getUserDefinedFunctionsProtoBytes();
-
- protected abstract RunnerApi.Coder getInputCoderProto();
-
- protected abstract RunnerApi.Coder getOutputCoderProto();
-
/** Gets the proto representation of the window coder. */
private RunnerApi.Coder getWindowCoderProto() {
return RunnerApi.Coder.newBuilder()
@@ -529,6 +483,70 @@ public abstract class BeamPythonFunctionRunner implements PythonFunctionRunner {
.build();
}
+ protected abstract byte[] getUserDefinedFunctionsProtoBytes();
+
+ protected abstract RunnerApi.Coder getInputCoderProto();
+
+ protected abstract RunnerApi.Coder getOutputCoderProto();
+
+ // ------------------------------------------------------------------------
+ // Construct RemoteBundler
+ // ------------------------------------------------------------------------
+
+ private OutputReceiverFactory createOutputReceiverFactory() {
+ return new OutputReceiverFactory() {
+
+ // the input value type is always byte array
+ @SuppressWarnings("unchecked")
+ @Override
+ public FnDataReceiver<WindowedValue<byte[]>> create(String pCollectionId) {
+ return input -> resultBuffer.add(input.getValue());
+ }
+ };
+ }
+
+ @VisibleForTesting
+ public JobBundleFactory createJobBundleFactory(Struct pipelineOptions) throws Exception {
+ return DefaultJobBundleFactory.create(
+ JobInfo.create(
+ taskName,
+ taskName,
+ environmentManager.createRetrievalToken(),
+ pipelineOptions));
+ }
+
+ /** To make the error messages more user friendly, throws an exception with the boot logs. */
+ private StageBundleFactory createStageBundleFactory(
+ JobBundleFactory jobBundleFactory, RunnerApi.Environment environment) throws Exception {
+ try {
+ return jobBundleFactory.forStage(createExecutableStage(environment));
+ } catch (Throwable e) {
+ throw new RuntimeException(environmentManager.getBootLog(), e);
+ }
+ }
+
+ /**
+ * Ignore bundle progress if flinkMetricContainer is null. The flinkMetricContainer will be set
+ * to null if metric is configured to be turned off.
+ */
+ private BundleProgressHandler getProgressHandler(FlinkMetricContainer flinkMetricContainer) {
+ if (flinkMetricContainer == null) {
+ return BundleProgressHandler.ignored();
+ } else {
+ return new BundleProgressHandler() {
+ @Override
+ public void onProgress(BeamFnApi.ProcessBundleProgressResponse progress) {
+ flinkMetricContainer.updateMetrics(taskName, progress.getMonitoringInfosList());
+ }
+
+ @Override
+ public void onCompleted(BeamFnApi.ProcessBundleResponse response) {
+ flinkMetricContainer.updateMetrics(taskName, response.getMonitoringInfosList());
+ }
+ };
+ }
+ }
+
private static StateRequestHandler getStateRequestHandler(
KeyedStateBackend keyedStateBackend,
TypeSerializer keySerializer,
diff --git a/flink-python/src/main/java/org/apache/flink/table/runtime/runners/python/beam/BeamTablePythonFunctionRunner.java b/flink-python/src/main/java/org/apache/flink/table/runtime/runners/python/beam/BeamTablePythonFunctionRunner.java
index a32e8e0..b59ef4c 100644
--- a/flink-python/src/main/java/org/apache/flink/table/runtime/runners/python/beam/BeamTablePythonFunctionRunner.java
+++ b/flink-python/src/main/java/org/apache/flink/table/runtime/runners/python/beam/BeamTablePythonFunctionRunner.java
@@ -33,6 +33,7 @@ import org.apache.beam.model.pipeline.v1.RunnerApi;
import java.util.Map;
+import static org.apache.flink.python.Constants.FLINK_CODER_URN;
import static org.apache.flink.table.runtime.typeutils.PythonTypeUtils.toProtoType;
/** A {@link BeamTablePythonFunctionRunner} used to execute Python functions in Table API. */
[flink] 01/03: [FLINK-23213][python] Move class
SimpleStateRequestHandler into a separate file
Posted by di...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
dianfu pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/flink.git
commit 4dec5eba69efe853339fd4ead9f2c90fb17c6704
Author: Dian Fu <di...@apache.org>
AuthorDate: Fri Jul 2 12:16:25 2021 +0800
[FLINK-23213][python] Move class SimpleStateRequestHandler into a separate file
---
.../python/beam/BeamPythonFunctionRunner.java | 646 --------------------
.../python/beam/SimpleStateRequestHandler.java | 669 +++++++++++++++++++++
2 files changed, 669 insertions(+), 646 deletions(-)
diff --git a/flink-python/src/main/java/org/apache/flink/streaming/api/runners/python/beam/BeamPythonFunctionRunner.java b/flink-python/src/main/java/org/apache/flink/streaming/api/runners/python/beam/BeamPythonFunctionRunner.java
index 6dc9458..177e242 100644
--- a/flink-python/src/main/java/org/apache/flink/streaming/api/runners/python/beam/BeamPythonFunctionRunner.java
+++ b/flink-python/src/main/java/org/apache/flink/streaming/api/runners/python/beam/BeamPythonFunctionRunner.java
@@ -19,20 +19,8 @@
package org.apache.flink.streaming.api.runners.python.beam;
import org.apache.flink.annotation.Internal;
-import org.apache.flink.api.common.ExecutionConfig;
-import org.apache.flink.api.common.state.ListState;
-import org.apache.flink.api.common.state.ListStateDescriptor;
-import org.apache.flink.api.common.state.MapState;
-import org.apache.flink.api.common.state.MapStateDescriptor;
-import org.apache.flink.api.common.state.StateDescriptor;
-import org.apache.flink.api.common.typeinfo.PrimitiveArrayTypeInfo;
import org.apache.flink.api.common.typeutils.TypeSerializer;
import org.apache.flink.api.java.tuple.Tuple2;
-import org.apache.flink.api.java.typeutils.runtime.RowSerializer;
-import org.apache.flink.core.memory.ByteArrayInputStreamWithPos;
-import org.apache.flink.core.memory.ByteArrayOutputStreamWithPos;
-import org.apache.flink.core.memory.DataInputViewStreamWrapper;
-import org.apache.flink.core.memory.DataOutputViewStreamWrapper;
import org.apache.flink.fnexecution.v1.FlinkFnApi;
import org.apache.flink.python.PythonConfig;
import org.apache.flink.python.PythonFunctionRunner;
@@ -44,14 +32,6 @@ import org.apache.flink.python.metric.FlinkMetricContainer;
import org.apache.flink.runtime.memory.MemoryManager;
import org.apache.flink.runtime.memory.OpaqueMemoryResource;
import org.apache.flink.runtime.state.KeyedStateBackend;
-import org.apache.flink.runtime.state.VoidNamespace;
-import org.apache.flink.runtime.state.VoidNamespaceSerializer;
-import org.apache.flink.runtime.state.internal.InternalMergingState;
-import org.apache.flink.streaming.api.utils.ByteArrayWrapper;
-import org.apache.flink.streaming.api.utils.ByteArrayWrapperSerializer;
-import org.apache.flink.table.data.RowData;
-import org.apache.flink.table.runtime.typeutils.AbstractRowDataSerializer;
-import org.apache.flink.table.runtime.typeutils.RowDataSerializer;
import org.apache.flink.util.Preconditions;
import org.apache.flink.util.function.LongFunctionWithException;
@@ -83,8 +63,6 @@ import org.apache.beam.sdk.options.PortablePipelineOptions;
import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
import org.apache.beam.sdk.transforms.windowing.GlobalWindow;
import org.apache.beam.sdk.util.WindowedValue;
-import org.apache.beam.vendor.grpc.v1p26p0.com.google.common.base.Charsets;
-import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.ByteString;
import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.Struct;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Iterables;
import org.slf4j.Logger;
@@ -97,28 +75,17 @@ import java.io.IOException;
import java.util.Arrays;
import java.util.Collection;
import java.util.Collections;
-import java.util.HashMap;
-import java.util.HashSet;
-import java.util.Iterator;
-import java.util.LinkedList;
import java.util.List;
import java.util.Map;
-import java.util.Set;
-import java.util.UUID;
-import java.util.concurrent.CompletableFuture;
-import java.util.concurrent.CompletionStage;
import java.util.concurrent.LinkedBlockingQueue;
import static org.apache.beam.runners.core.construction.BeamUrns.getUrn;
-import static org.apache.flink.streaming.api.utils.PythonOperatorUtils.setCurrentKeyForStreaming;
/** A {@link BeamPythonFunctionRunner} used to execute Python functions. */
@Internal
public abstract class BeamPythonFunctionRunner implements PythonFunctionRunner {
protected static final Logger LOG = LoggerFactory.getLogger(BeamPythonFunctionRunner.class);
- private static final String PYTHON_STATE_PREFIX = "python-state-";
-
private static final String INPUT_ID = "input";
private static final String OUTPUT_ID = "output";
private static final String TRANSFORM_ID = "transform";
@@ -575,617 +542,4 @@ public abstract class BeamPythonFunctionRunner implements PythonFunctionRunner {
keyedStateBackend, keySerializer, namespaceSerializer, jobOptions);
}
}
-
- /** The type of the Python map state iterate request. */
- private enum IterateType {
-
- /** Equivalent to iterate {@link Map#entrySet() }. */
- ITEMS((byte) 0),
-
- /** Equivalent to iterate {@link Map#keySet() }. */
- KEYS((byte) 1),
-
- /** Equivalent to iterate {@link Map#values() }. */
- VALUES((byte) 2);
-
- private final byte ord;
-
- IterateType(byte ord) {
- this.ord = ord;
- }
-
- public byte getOrd() {
- return ord;
- }
-
- public static IterateType fromOrd(byte ord) {
- switch (ord) {
- case 0:
- return ITEMS;
- case 1:
- return KEYS;
- case 2:
- return VALUES;
- default:
- throw new RuntimeException("Unsupported ordinal: " + ord);
- }
- }
- }
-
- /** A state request handler which handles the state request from Python side. */
- private static class SimpleStateRequestHandler implements StateRequestHandler {
-
- private static final String CLEAR_CACHED_ITERATOR_MARK = "clear_iterators";
- private static final String MERGE_NAMESPACES_MARK = "merge_namespaces";
-
- // map state GET request flags
- private static final byte GET_FLAG = 0;
- private static final byte ITERATE_FLAG = 1;
- private static final byte CHECK_EMPTY_FLAG = 2;
-
- // map state GET response flags
- private static final byte EXIST_FLAG = 0;
- private static final byte IS_NONE_FLAG = 1;
- private static final byte NOT_EXIST_FLAG = 2;
- private static final byte IS_EMPTY_FLAG = 3;
- private static final byte NOT_EMPTY_FLAG = 4;
-
- // map state APPEND request flags
- private static final byte DELETE = 0;
- private static final byte SET_NONE = 1;
- private static final byte SET_VALUE = 2;
-
- private static final BeamFnApi.StateGetResponse.Builder NOT_EXIST_RESPONSE =
- BeamFnApi.StateGetResponse.newBuilder()
- .setData(ByteString.copyFrom(new byte[] {NOT_EXIST_FLAG}));
- private static final BeamFnApi.StateGetResponse.Builder IS_NONE_RESPONSE =
- BeamFnApi.StateGetResponse.newBuilder()
- .setData(ByteString.copyFrom(new byte[] {IS_NONE_FLAG}));
- private static final BeamFnApi.StateGetResponse.Builder IS_EMPTY_RESPONSE =
- BeamFnApi.StateGetResponse.newBuilder()
- .setData(ByteString.copyFrom(new byte[] {IS_EMPTY_FLAG}));
- private static final BeamFnApi.StateGetResponse.Builder NOT_EMPTY_RESPONSE =
- BeamFnApi.StateGetResponse.newBuilder()
- .setData(ByteString.copyFrom(new byte[] {NOT_EMPTY_FLAG}));
-
- private final TypeSerializer keySerializer;
- private final TypeSerializer namespaceSerializer;
- private final TypeSerializer<byte[]> valueSerializer;
- private final KeyedStateBackend keyedStateBackend;
-
- /** Reusable InputStream used to holding the elements to be deserialized. */
- private final ByteArrayInputStreamWithPos bais;
-
- /** InputStream Wrapper. */
- private final DataInputViewStreamWrapper baisWrapper;
-
- /** Reusable OutputStream used to holding the serialized input elements. */
- private final ByteArrayOutputStreamWithPos baos;
-
- /** OutputStream Wrapper. */
- private final DataOutputViewStreamWrapper baosWrapper;
-
- /** The cache of the stateDescriptors. */
- private final Map<String, StateDescriptor> stateDescriptorCache;
-
- /** The cache of the map state iterators. */
- private final Map<ByteArrayWrapper, Iterator> mapStateIteratorCache;
-
- private final int mapStateIterateResponseBatchSize;
-
- private final ByteArrayWrapper reuseByteArrayWrapper = new ByteArrayWrapper(new byte[0]);
-
- /** Let StateRequestHandler for user state only use a single cache token. */
- private final BeamFnApi.ProcessBundleRequest.CacheToken cacheToken;
-
- SimpleStateRequestHandler(
- KeyedStateBackend keyedStateBackend,
- TypeSerializer keySerializer,
- TypeSerializer namespaceSerializer,
- Map<String, String> config) {
- this.keyedStateBackend = keyedStateBackend;
- TypeSerializer frameworkKeySerializer = keyedStateBackend.getKeySerializer();
- if (!(frameworkKeySerializer instanceof AbstractRowDataSerializer
- || frameworkKeySerializer instanceof RowSerializer)) {
- throw new RuntimeException(
- "Currently SimpleStateRequestHandler only support row key!");
- }
- this.keySerializer = keySerializer;
- this.namespaceSerializer = namespaceSerializer;
- this.valueSerializer =
- PrimitiveArrayTypeInfo.BYTE_PRIMITIVE_ARRAY_TYPE_INFO.createSerializer(
- new ExecutionConfig());
- bais = new ByteArrayInputStreamWithPos();
- baisWrapper = new DataInputViewStreamWrapper(bais);
- baos = new ByteArrayOutputStreamWithPos();
- baosWrapper = new DataOutputViewStreamWrapper(baos);
- stateDescriptorCache = new HashMap<>();
- mapStateIteratorCache = new HashMap<>();
- mapStateIterateResponseBatchSize =
- Integer.valueOf(
- config.getOrDefault(
- PythonOptions.MAP_STATE_ITERATE_RESPONSE_BATCH_SIZE.key(),
- PythonOptions.MAP_STATE_ITERATE_RESPONSE_BATCH_SIZE
- .defaultValue()
- .toString()));
- if (mapStateIterateResponseBatchSize <= 0) {
- throw new RuntimeException(
- String.format(
- "The value of '%s' must be greater than 0!",
- PythonOptions.MAP_STATE_ITERATE_RESPONSE_BATCH_SIZE.key()));
- }
- cacheToken = createCacheToken();
- }
-
- @Override
- public Iterable<BeamFnApi.ProcessBundleRequest.CacheToken> getCacheTokens() {
- return Collections.singleton(cacheToken);
- }
-
- @Override
- public CompletionStage<BeamFnApi.StateResponse.Builder> handle(
- BeamFnApi.StateRequest request) throws Exception {
- BeamFnApi.StateKey.TypeCase typeCase = request.getStateKey().getTypeCase();
- synchronized (keyedStateBackend) {
- if (typeCase.equals(BeamFnApi.StateKey.TypeCase.BAG_USER_STATE)) {
- return handleBagState(request);
- } else if (typeCase.equals(BeamFnApi.StateKey.TypeCase.MULTIMAP_SIDE_INPUT)) {
- return handleMapState(request);
- } else {
- throw new RuntimeException("Unsupported state type: " + typeCase);
- }
- }
- }
-
- private CompletionStage<BeamFnApi.StateResponse.Builder> handleBagState(
- BeamFnApi.StateRequest request) throws Exception {
- if (request.getStateKey().hasBagUserState()) {
- BeamFnApi.StateKey.BagUserState bagUserState =
- request.getStateKey().getBagUserState();
- // get key
- byte[] keyBytes = bagUserState.getKey().toByteArray();
- bais.setBuffer(keyBytes, 0, keyBytes.length);
- Object key = keySerializer.deserialize(baisWrapper);
- if (keyedStateBackend.getKeySerializer() instanceof RowDataSerializer) {
- setCurrentKeyForStreaming(
- keyedStateBackend,
- ((RowDataSerializer) keyedStateBackend.getKeySerializer())
- .toBinaryRow((RowData) key));
- } else {
- setCurrentKeyForStreaming(keyedStateBackend, key);
- }
- } else {
- throw new RuntimeException("Unsupported bag state request: " + request);
- }
-
- switch (request.getRequestCase()) {
- case GET:
- return handleBagGetRequest(request);
- case APPEND:
- return handleBagAppendRequest(request);
- case CLEAR:
- return handleBagClearRequest(request);
- default:
- throw new RuntimeException(
- String.format(
- "Unsupported request type %s for user state.",
- request.getRequestCase()));
- }
- }
-
- private List<ByteString> convertToByteString(ListState<byte[]> listState) throws Exception {
- List<ByteString> ret = new LinkedList<>();
- if (listState.get() == null) {
- return ret;
- }
- for (byte[] v : listState.get()) {
- ret.add(ByteString.copyFrom(v));
- }
- return ret;
- }
-
- private CompletionStage<BeamFnApi.StateResponse.Builder> handleBagGetRequest(
- BeamFnApi.StateRequest request) throws Exception {
-
- ListState<byte[]> partitionedState = getListState(request);
- List<ByteString> byteStrings = convertToByteString(partitionedState);
-
- return CompletableFuture.completedFuture(
- BeamFnApi.StateResponse.newBuilder()
- .setId(request.getId())
- .setGet(
- BeamFnApi.StateGetResponse.newBuilder()
- .setData(ByteString.copyFrom(byteStrings))));
- }
-
- private CompletionStage<BeamFnApi.StateResponse.Builder> handleBagAppendRequest(
- BeamFnApi.StateRequest request) throws Exception {
-
- ListState<byte[]> partitionedState = getListState(request);
- if (request.getStateKey()
- .getBagUserState()
- .getTransformId()
- .equals(MERGE_NAMESPACES_MARK)) {
- // get namespaces to merge
- byte[] namespacesBytes = request.getAppend().getData().toByteArray();
- bais.setBuffer(namespacesBytes, 0, namespacesBytes.length);
- int namespaceCount = baisWrapper.readInt();
- Set<Object> namespaces = new HashSet<>();
- for (int i = 0; i < namespaceCount; i++) {
- namespaces.add(namespaceSerializer.deserialize(baisWrapper));
- }
- byte[] targetNamespaceByte =
- request.getStateKey().getBagUserState().getWindow().toByteArray();
- bais.setBuffer(targetNamespaceByte, 0, targetNamespaceByte.length);
- Object targetNamespace = namespaceSerializer.deserialize(baisWrapper);
- ((InternalMergingState) partitionedState)
- .mergeNamespaces(targetNamespace, namespaces);
- } else {
- // get values
- byte[] valueBytes = request.getAppend().getData().toByteArray();
- partitionedState.add(valueBytes);
- }
- return CompletableFuture.completedFuture(
- BeamFnApi.StateResponse.newBuilder()
- .setId(request.getId())
- .setAppend(BeamFnApi.StateAppendResponse.getDefaultInstance()));
- }
-
- private CompletionStage<BeamFnApi.StateResponse.Builder> handleBagClearRequest(
- BeamFnApi.StateRequest request) throws Exception {
-
- ListState<byte[]> partitionedState = getListState(request);
-
- partitionedState.clear();
- return CompletableFuture.completedFuture(
- BeamFnApi.StateResponse.newBuilder()
- .setId(request.getId())
- .setClear(BeamFnApi.StateClearResponse.getDefaultInstance()));
- }
-
- private ListState<byte[]> getListState(BeamFnApi.StateRequest request) throws Exception {
- BeamFnApi.StateKey.BagUserState bagUserState = request.getStateKey().getBagUserState();
- String stateName = PYTHON_STATE_PREFIX + bagUserState.getUserStateId();
- ListStateDescriptor<byte[]> listStateDescriptor;
- StateDescriptor cachedStateDescriptor = stateDescriptorCache.get(stateName);
- if (cachedStateDescriptor instanceof ListStateDescriptor) {
- listStateDescriptor = (ListStateDescriptor<byte[]>) cachedStateDescriptor;
- } else if (cachedStateDescriptor == null) {
- listStateDescriptor = new ListStateDescriptor<>(stateName, valueSerializer);
- stateDescriptorCache.put(stateName, listStateDescriptor);
- } else {
- throw new RuntimeException(
- String.format(
- "State name corrupt detected: "
- + "'%s' is used both as LIST state and '%s' state at the same time.",
- stateName, cachedStateDescriptor.getType()));
- }
- byte[] windowBytes = bagUserState.getWindow().toByteArray();
- if (windowBytes.length != 0) {
- bais.setBuffer(windowBytes, 0, windowBytes.length);
- Object namespace = namespaceSerializer.deserialize(baisWrapper);
- return (ListState<byte[]>)
- keyedStateBackend.getPartitionedState(
- namespace, namespaceSerializer, listStateDescriptor);
- } else {
- return (ListState<byte[]>)
- keyedStateBackend.getPartitionedState(
- VoidNamespace.INSTANCE,
- VoidNamespaceSerializer.INSTANCE,
- listStateDescriptor);
- }
- }
-
- private CompletionStage<BeamFnApi.StateResponse.Builder> handleMapState(
- BeamFnApi.StateRequest request) throws Exception {
- // Currently the `beam_fn_api.proto` does not support MapState, so we use the
- // the `MultimapSideInput` message to mark the state as a MapState for now.
- if (request.getStateKey().hasMultimapSideInput()) {
- BeamFnApi.StateKey.MultimapSideInput mapUserState =
- request.getStateKey().getMultimapSideInput();
- // get key
- byte[] keyBytes = mapUserState.getKey().toByteArray();
- bais.setBuffer(keyBytes, 0, keyBytes.length);
- Object key = keySerializer.deserialize(baisWrapper);
- if (keyedStateBackend.getKeySerializer() instanceof RowDataSerializer) {
- setCurrentKeyForStreaming(
- keyedStateBackend,
- ((RowDataSerializer) keyedStateBackend.getKeySerializer())
- .toBinaryRow((RowData) key));
- } else {
- setCurrentKeyForStreaming(keyedStateBackend, key);
- }
- } else {
- throw new RuntimeException("Unsupported bag state request: " + request);
- }
-
- switch (request.getRequestCase()) {
- case GET:
- return handleMapGetRequest(request);
- case APPEND:
- return handleMapAppendRequest(request);
- case CLEAR:
- return handleMapClearRequest(request);
- default:
- throw new RuntimeException(
- String.format(
- "Unsupported request type %s for user state.",
- request.getRequestCase()));
- }
- }
-
- private CompletionStage<BeamFnApi.StateResponse.Builder> handleMapGetRequest(
- BeamFnApi.StateRequest request) throws Exception {
- MapState<ByteArrayWrapper, byte[]> mapState = getMapState(request);
- // The continuation token structure of GET request is:
- // [flag (1 byte)][serialized map key]
- // The continuation token structure of CHECK_EMPTY request is:
- // [flag (1 byte)]
- // The continuation token structure of ITERATE request is:
- // [flag (1 byte)][iterate type (1 byte)][iterator token length (int32)][iterator token]
- byte[] getRequest = request.getGet().getContinuationToken().toByteArray();
- byte getFlag = getRequest[0];
- BeamFnApi.StateGetResponse.Builder response;
- switch (getFlag) {
- case GET_FLAG:
- reuseByteArrayWrapper.setData(getRequest);
- reuseByteArrayWrapper.setOffset(1);
- reuseByteArrayWrapper.setLimit(getRequest.length);
- response = handleMapGetValueRequest(reuseByteArrayWrapper, mapState);
- break;
- case CHECK_EMPTY_FLAG:
- response = handleMapCheckEmptyRequest(mapState);
- break;
- case ITERATE_FLAG:
- bais.setBuffer(getRequest, 1, getRequest.length - 1);
- IterateType iterateType = IterateType.fromOrd(baisWrapper.readByte());
- int iterateTokenLength = baisWrapper.readInt();
- ByteArrayWrapper iterateToken;
- if (iterateTokenLength > 0) {
- reuseByteArrayWrapper.setData(getRequest);
- reuseByteArrayWrapper.setOffset(bais.getPosition());
- reuseByteArrayWrapper.setLimit(bais.getPosition() + iterateTokenLength);
- iterateToken = reuseByteArrayWrapper;
- } else {
- iterateToken = null;
- }
- response = handleMapIterateRequest(mapState, iterateType, iterateToken);
- break;
- default:
- throw new RuntimeException(
- String.format(
- "Unsupported get request type: '%d' for map state.", getFlag));
- }
-
- return CompletableFuture.completedFuture(
- BeamFnApi.StateResponse.newBuilder().setId(request.getId()).setGet(response));
- }
-
- private CompletionStage<BeamFnApi.StateResponse.Builder> handleMapAppendRequest(
- BeamFnApi.StateRequest request) throws Exception {
- // The structure of append request bytes is:
- // [number of requests (int32)][append request flag (1 byte)][map key length
- // (int32)][map key]
- // [map value length (int32)][map value][append request flag (1 byte)][map key length
- // (int32)][map key]
- // ...
- byte[] appendBytes = request.getAppend().getData().toByteArray();
- bais.setBuffer(appendBytes, 0, appendBytes.length);
- MapState<ByteArrayWrapper, byte[]> mapState = getMapState(request);
- int subRequestNum = baisWrapper.readInt();
- for (int i = 0; i < subRequestNum; i++) {
- byte requestFlag = baisWrapper.readByte();
- int keyLength = baisWrapper.readInt();
- reuseByteArrayWrapper.setData(appendBytes);
- reuseByteArrayWrapper.setOffset(bais.getPosition());
- reuseByteArrayWrapper.setLimit(bais.getPosition() + keyLength);
- baisWrapper.skipBytesToRead(keyLength);
- switch (requestFlag) {
- case DELETE:
- mapState.remove(reuseByteArrayWrapper);
- break;
- case SET_NONE:
- mapState.put(reuseByteArrayWrapper.copy(), null);
- break;
- case SET_VALUE:
- int valueLength = baisWrapper.readInt();
- byte[] valueBytes = new byte[valueLength];
- int readLength = baisWrapper.read(valueBytes);
- assert valueLength == readLength;
- mapState.put(reuseByteArrayWrapper.copy(), valueBytes);
- break;
- default:
- throw new RuntimeException(
- String.format(
- "Unsupported append request type: '%d' for map state.",
- requestFlag));
- }
- }
- return CompletableFuture.completedFuture(
- BeamFnApi.StateResponse.newBuilder()
- .setId(request.getId())
- .setAppend(BeamFnApi.StateAppendResponse.getDefaultInstance()));
- }
-
- private CompletionStage<BeamFnApi.StateResponse.Builder> handleMapClearRequest(
- BeamFnApi.StateRequest request) throws Exception {
- if (request.getStateKey()
- .getMultimapSideInput()
- .getTransformId()
- .equals(CLEAR_CACHED_ITERATOR_MARK)) {
- mapStateIteratorCache.clear();
- } else {
- MapState<ByteArrayWrapper, byte[]> partitionedState = getMapState(request);
- partitionedState.clear();
- }
- return CompletableFuture.completedFuture(
- BeamFnApi.StateResponse.newBuilder()
- .setId(request.getId())
- .setClear(BeamFnApi.StateClearResponse.getDefaultInstance()));
- }
-
- private BeamFnApi.StateGetResponse.Builder handleMapGetValueRequest(
- ByteArrayWrapper key, MapState<ByteArrayWrapper, byte[]> mapState)
- throws Exception {
- if (mapState.contains(key)) {
- byte[] value = mapState.get(key);
- if (value == null) {
- return IS_NONE_RESPONSE;
- } else {
- baos.reset();
- baosWrapper.writeByte(EXIST_FLAG);
- baosWrapper.write(value);
- return BeamFnApi.StateGetResponse.newBuilder()
- .setData(ByteString.copyFrom(baos.toByteArray()));
- }
- } else {
- return NOT_EXIST_RESPONSE;
- }
- }
-
- private BeamFnApi.StateGetResponse.Builder handleMapCheckEmptyRequest(
- MapState<ByteArrayWrapper, byte[]> mapState) throws Exception {
- if (mapState.isEmpty()) {
- return IS_EMPTY_RESPONSE;
- } else {
- return NOT_EMPTY_RESPONSE;
- }
- }
-
- private BeamFnApi.StateGetResponse.Builder handleMapIterateRequest(
- MapState<ByteArrayWrapper, byte[]> mapState,
- IterateType iterateType,
- ByteArrayWrapper iteratorToken)
- throws Exception {
- final Iterator iterator;
- if (iteratorToken == null) {
- switch (iterateType) {
- case ITEMS:
- case VALUES:
- iterator = mapState.iterator();
- break;
- case KEYS:
- iterator = mapState.keys().iterator();
- break;
- default:
- throw new RuntimeException("Unsupported iterate type: " + iterateType);
- }
- } else {
- iterator = mapStateIteratorCache.get(iteratorToken);
- if (iterator == null) {
- throw new RuntimeException("The cached iterator does not exist!");
- }
- }
- baos.reset();
- switch (iterateType) {
- case ITEMS:
- case VALUES:
- Iterator<Map.Entry<ByteArrayWrapper, byte[]>> entryIterator = iterator;
- for (int i = 0; i < mapStateIterateResponseBatchSize; i++) {
- if (entryIterator.hasNext()) {
- Map.Entry<ByteArrayWrapper, byte[]> entry = entryIterator.next();
- ByteArrayWrapper key = entry.getKey();
- baosWrapper.write(
- key.getData(),
- key.getOffset(),
- key.getLimit() - key.getOffset());
- baosWrapper.writeBoolean(entry.getValue() != null);
- if (entry.getValue() != null) {
- baosWrapper.write(entry.getValue());
- }
- } else {
- break;
- }
- }
- break;
- case KEYS:
- Iterator<ByteArrayWrapper> keyIterator = iterator;
- for (int i = 0; i < mapStateIterateResponseBatchSize; i++) {
- if (keyIterator.hasNext()) {
- ByteArrayWrapper key = keyIterator.next();
- baosWrapper.write(
- key.getData(),
- key.getOffset(),
- key.getLimit() - key.getOffset());
- } else {
- break;
- }
- }
- break;
- default:
- throw new RuntimeException("Unsupported iterate type: " + iterateType);
- }
- if (!iterator.hasNext()) {
- if (iteratorToken != null) {
- mapStateIteratorCache.remove(iteratorToken);
- }
- iteratorToken = null;
- } else {
- if (iteratorToken == null) {
- iteratorToken = new ByteArrayWrapper(UUID.randomUUID().toString().getBytes());
- }
- mapStateIteratorCache.put(iteratorToken, iterator);
- }
- BeamFnApi.StateGetResponse.Builder responseBuilder =
- BeamFnApi.StateGetResponse.newBuilder()
- .setData(ByteString.copyFrom(baos.toByteArray()));
- if (iteratorToken != null) {
- responseBuilder.setContinuationToken(
- ByteString.copyFrom(
- iteratorToken.getData(),
- iteratorToken.getOffset(),
- iteratorToken.getLimit() - iteratorToken.getOffset()));
- }
- return responseBuilder;
- }
-
- private MapState<ByteArrayWrapper, byte[]> getMapState(BeamFnApi.StateRequest request)
- throws Exception {
- BeamFnApi.StateKey.MultimapSideInput mapUserState =
- request.getStateKey().getMultimapSideInput();
- String stateName = PYTHON_STATE_PREFIX + mapUserState.getSideInputId();
- StateDescriptor cachedStateDescriptor = stateDescriptorCache.get(stateName);
- MapStateDescriptor<ByteArrayWrapper, byte[]> mapStateDescriptor;
- if (cachedStateDescriptor instanceof MapStateDescriptor) {
- mapStateDescriptor =
- (MapStateDescriptor<ByteArrayWrapper, byte[]>) cachedStateDescriptor;
- } else if (cachedStateDescriptor == null) {
- mapStateDescriptor =
- new MapStateDescriptor<>(
- stateName, ByteArrayWrapperSerializer.INSTANCE, valueSerializer);
- stateDescriptorCache.put(stateName, mapStateDescriptor);
- } else {
- throw new RuntimeException(
- String.format(
- "State name corrupt detected: "
- + "'%s' is used both as MAP state and '%s' state at the same time.",
- stateName, cachedStateDescriptor.getType()));
- }
- byte[] windowBytes = mapUserState.getWindow().toByteArray();
- if (windowBytes.length != 0) {
- bais.setBuffer(windowBytes, 0, windowBytes.length);
- Object namespace = namespaceSerializer.deserialize(baisWrapper);
- return (MapState<ByteArrayWrapper, byte[]>)
- keyedStateBackend.getPartitionedState(
- namespace, namespaceSerializer, mapStateDescriptor);
- } else {
- return (MapState<ByteArrayWrapper, byte[]>)
- keyedStateBackend.getPartitionedState(
- VoidNamespace.INSTANCE,
- VoidNamespaceSerializer.INSTANCE,
- mapStateDescriptor);
- }
- }
-
- private BeamFnApi.ProcessBundleRequest.CacheToken createCacheToken() {
- ByteString token =
- ByteString.copyFrom(UUID.randomUUID().toString().getBytes(Charsets.UTF_8));
- return BeamFnApi.ProcessBundleRequest.CacheToken.newBuilder()
- .setUserState(
- BeamFnApi.ProcessBundleRequest.CacheToken.UserState
- .getDefaultInstance())
- .setToken(token)
- .build();
- }
- }
}
diff --git a/flink-python/src/main/java/org/apache/flink/streaming/api/runners/python/beam/SimpleStateRequestHandler.java b/flink-python/src/main/java/org/apache/flink/streaming/api/runners/python/beam/SimpleStateRequestHandler.java
new file mode 100644
index 0000000..95f571f
--- /dev/null
+++ b/flink-python/src/main/java/org/apache/flink/streaming/api/runners/python/beam/SimpleStateRequestHandler.java
@@ -0,0 +1,669 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.flink.streaming.api.runners.python.beam;
+
+import org.apache.flink.annotation.Internal;
+import org.apache.flink.api.common.ExecutionConfig;
+import org.apache.flink.api.common.state.ListState;
+import org.apache.flink.api.common.state.ListStateDescriptor;
+import org.apache.flink.api.common.state.MapState;
+import org.apache.flink.api.common.state.MapStateDescriptor;
+import org.apache.flink.api.common.state.StateDescriptor;
+import org.apache.flink.api.common.typeinfo.PrimitiveArrayTypeInfo;
+import org.apache.flink.api.common.typeutils.TypeSerializer;
+import org.apache.flink.api.java.typeutils.runtime.RowSerializer;
+import org.apache.flink.core.memory.ByteArrayInputStreamWithPos;
+import org.apache.flink.core.memory.ByteArrayOutputStreamWithPos;
+import org.apache.flink.core.memory.DataInputViewStreamWrapper;
+import org.apache.flink.core.memory.DataOutputViewStreamWrapper;
+import org.apache.flink.python.PythonOptions;
+import org.apache.flink.runtime.state.KeyedStateBackend;
+import org.apache.flink.runtime.state.VoidNamespace;
+import org.apache.flink.runtime.state.VoidNamespaceSerializer;
+import org.apache.flink.runtime.state.internal.InternalMergingState;
+import org.apache.flink.streaming.api.utils.ByteArrayWrapper;
+import org.apache.flink.streaming.api.utils.ByteArrayWrapperSerializer;
+import org.apache.flink.table.data.RowData;
+import org.apache.flink.table.runtime.typeutils.AbstractRowDataSerializer;
+import org.apache.flink.table.runtime.typeutils.RowDataSerializer;
+
+import org.apache.beam.model.fnexecution.v1.BeamFnApi;
+import org.apache.beam.runners.fnexecution.state.StateRequestHandler;
+import org.apache.beam.vendor.grpc.v1p26p0.com.google.common.base.Charsets;
+import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.ByteString;
+
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.Iterator;
+import java.util.LinkedList;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import java.util.UUID;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CompletionStage;
+
+import static org.apache.flink.streaming.api.utils.PythonOperatorUtils.setCurrentKeyForStreaming;
+
+/** A state request handler which handles the state request from Python side. */
+@Internal
+public class SimpleStateRequestHandler implements StateRequestHandler {
+
+ private static final String CLEAR_CACHED_ITERATOR_MARK = "clear_iterators";
+ private static final String MERGE_NAMESPACES_MARK = "merge_namespaces";
+ private static final String PYTHON_STATE_PREFIX = "python-state-";
+
+ // map state GET request flags
+ private static final byte GET_FLAG = 0;
+ private static final byte ITERATE_FLAG = 1;
+ private static final byte CHECK_EMPTY_FLAG = 2;
+
+ // map state GET response flags
+ private static final byte EXIST_FLAG = 0;
+ private static final byte IS_NONE_FLAG = 1;
+ private static final byte NOT_EXIST_FLAG = 2;
+ private static final byte IS_EMPTY_FLAG = 3;
+ private static final byte NOT_EMPTY_FLAG = 4;
+
+ // map state APPEND request flags
+ private static final byte DELETE = 0;
+ private static final byte SET_NONE = 1;
+ private static final byte SET_VALUE = 2;
+
+ private static final BeamFnApi.StateGetResponse.Builder NOT_EXIST_RESPONSE =
+ BeamFnApi.StateGetResponse.newBuilder()
+ .setData(ByteString.copyFrom(new byte[] {NOT_EXIST_FLAG}));
+ private static final BeamFnApi.StateGetResponse.Builder IS_NONE_RESPONSE =
+ BeamFnApi.StateGetResponse.newBuilder()
+ .setData(ByteString.copyFrom(new byte[] {IS_NONE_FLAG}));
+ private static final BeamFnApi.StateGetResponse.Builder IS_EMPTY_RESPONSE =
+ BeamFnApi.StateGetResponse.newBuilder()
+ .setData(ByteString.copyFrom(new byte[] {IS_EMPTY_FLAG}));
+ private static final BeamFnApi.StateGetResponse.Builder NOT_EMPTY_RESPONSE =
+ BeamFnApi.StateGetResponse.newBuilder()
+ .setData(ByteString.copyFrom(new byte[] {NOT_EMPTY_FLAG}));
+
+ private final TypeSerializer keySerializer;
+ private final TypeSerializer namespaceSerializer;
+ private final TypeSerializer<byte[]> valueSerializer;
+ private final KeyedStateBackend keyedStateBackend;
+
+ /** Reusable InputStream used to holding the elements to be deserialized. */
+ private final ByteArrayInputStreamWithPos bais;
+
+ /** InputStream Wrapper. */
+ private final DataInputViewStreamWrapper baisWrapper;
+
+ /** Reusable OutputStream used to holding the serialized input elements. */
+ private final ByteArrayOutputStreamWithPos baos;
+
+ /** OutputStream Wrapper. */
+ private final DataOutputViewStreamWrapper baosWrapper;
+
+ /** The cache of the stateDescriptors. */
+ private final Map<String, StateDescriptor> stateDescriptorCache;
+
+ /** The cache of the map state iterators. */
+ private final Map<ByteArrayWrapper, Iterator> mapStateIteratorCache;
+
+ private final int mapStateIterateResponseBatchSize;
+
+ private final ByteArrayWrapper reuseByteArrayWrapper = new ByteArrayWrapper(new byte[0]);
+
+ /** Let StateRequestHandler for user state only use a single cache token. */
+ private final BeamFnApi.ProcessBundleRequest.CacheToken cacheToken;
+
+ SimpleStateRequestHandler(
+ KeyedStateBackend keyedStateBackend,
+ TypeSerializer keySerializer,
+ TypeSerializer namespaceSerializer,
+ Map<String, String> config) {
+ this.keyedStateBackend = keyedStateBackend;
+ TypeSerializer frameworkKeySerializer = keyedStateBackend.getKeySerializer();
+ if (!(frameworkKeySerializer instanceof AbstractRowDataSerializer
+ || frameworkKeySerializer instanceof RowSerializer)) {
+ throw new RuntimeException("Currently SimpleStateRequestHandler only support row key!");
+ }
+ this.keySerializer = keySerializer;
+ this.namespaceSerializer = namespaceSerializer;
+ this.valueSerializer =
+ PrimitiveArrayTypeInfo.BYTE_PRIMITIVE_ARRAY_TYPE_INFO.createSerializer(
+ new ExecutionConfig());
+ bais = new ByteArrayInputStreamWithPos();
+ baisWrapper = new DataInputViewStreamWrapper(bais);
+ baos = new ByteArrayOutputStreamWithPos();
+ baosWrapper = new DataOutputViewStreamWrapper(baos);
+ stateDescriptorCache = new HashMap<>();
+ mapStateIteratorCache = new HashMap<>();
+ mapStateIterateResponseBatchSize =
+ Integer.valueOf(
+ config.getOrDefault(
+ PythonOptions.MAP_STATE_ITERATE_RESPONSE_BATCH_SIZE.key(),
+ PythonOptions.MAP_STATE_ITERATE_RESPONSE_BATCH_SIZE
+ .defaultValue()
+ .toString()));
+ if (mapStateIterateResponseBatchSize <= 0) {
+ throw new RuntimeException(
+ String.format(
+ "The value of '%s' must be greater than 0!",
+ PythonOptions.MAP_STATE_ITERATE_RESPONSE_BATCH_SIZE.key()));
+ }
+ cacheToken = createCacheToken();
+ }
+
+ @Override
+ public Iterable<BeamFnApi.ProcessBundleRequest.CacheToken> getCacheTokens() {
+ return Collections.singleton(cacheToken);
+ }
+
+ @Override
+ public CompletionStage<BeamFnApi.StateResponse.Builder> handle(BeamFnApi.StateRequest request)
+ throws Exception {
+ BeamFnApi.StateKey.TypeCase typeCase = request.getStateKey().getTypeCase();
+ synchronized (keyedStateBackend) {
+ if (typeCase.equals(BeamFnApi.StateKey.TypeCase.BAG_USER_STATE)) {
+ return handleBagState(request);
+ } else if (typeCase.equals(BeamFnApi.StateKey.TypeCase.MULTIMAP_SIDE_INPUT)) {
+ return handleMapState(request);
+ } else {
+ throw new RuntimeException("Unsupported state type: " + typeCase);
+ }
+ }
+ }
+
+ private CompletionStage<BeamFnApi.StateResponse.Builder> handleBagState(
+ BeamFnApi.StateRequest request) throws Exception {
+ if (request.getStateKey().hasBagUserState()) {
+ BeamFnApi.StateKey.BagUserState bagUserState = request.getStateKey().getBagUserState();
+ // get key
+ byte[] keyBytes = bagUserState.getKey().toByteArray();
+ bais.setBuffer(keyBytes, 0, keyBytes.length);
+ Object key = keySerializer.deserialize(baisWrapper);
+ if (keyedStateBackend.getKeySerializer() instanceof RowDataSerializer) {
+ setCurrentKeyForStreaming(
+ keyedStateBackend,
+ ((RowDataSerializer) keyedStateBackend.getKeySerializer())
+ .toBinaryRow((RowData) key));
+ } else {
+ setCurrentKeyForStreaming(keyedStateBackend, key);
+ }
+ } else {
+ throw new RuntimeException("Unsupported bag state request: " + request);
+ }
+
+ switch (request.getRequestCase()) {
+ case GET:
+ return handleBagGetRequest(request);
+ case APPEND:
+ return handleBagAppendRequest(request);
+ case CLEAR:
+ return handleBagClearRequest(request);
+ default:
+ throw new RuntimeException(
+ String.format(
+ "Unsupported request type %s for user state.",
+ request.getRequestCase()));
+ }
+ }
+
+ private List<ByteString> convertToByteString(ListState<byte[]> listState) throws Exception {
+ List<ByteString> ret = new LinkedList<>();
+ if (listState.get() == null) {
+ return ret;
+ }
+ for (byte[] v : listState.get()) {
+ ret.add(ByteString.copyFrom(v));
+ }
+ return ret;
+ }
+
+ private CompletionStage<BeamFnApi.StateResponse.Builder> handleBagGetRequest(
+ BeamFnApi.StateRequest request) throws Exception {
+
+ ListState<byte[]> partitionedState = getListState(request);
+ List<ByteString> byteStrings = convertToByteString(partitionedState);
+
+ return CompletableFuture.completedFuture(
+ BeamFnApi.StateResponse.newBuilder()
+ .setId(request.getId())
+ .setGet(
+ BeamFnApi.StateGetResponse.newBuilder()
+ .setData(ByteString.copyFrom(byteStrings))));
+ }
+
+ private CompletionStage<BeamFnApi.StateResponse.Builder> handleBagAppendRequest(
+ BeamFnApi.StateRequest request) throws Exception {
+
+ ListState<byte[]> partitionedState = getListState(request);
+ if (request.getStateKey()
+ .getBagUserState()
+ .getTransformId()
+ .equals(MERGE_NAMESPACES_MARK)) {
+ // get namespaces to merge
+ byte[] namespacesBytes = request.getAppend().getData().toByteArray();
+ bais.setBuffer(namespacesBytes, 0, namespacesBytes.length);
+ int namespaceCount = baisWrapper.readInt();
+ Set<Object> namespaces = new HashSet<>();
+ for (int i = 0; i < namespaceCount; i++) {
+ namespaces.add(namespaceSerializer.deserialize(baisWrapper));
+ }
+ byte[] targetNamespaceByte =
+ request.getStateKey().getBagUserState().getWindow().toByteArray();
+ bais.setBuffer(targetNamespaceByte, 0, targetNamespaceByte.length);
+ Object targetNamespace = namespaceSerializer.deserialize(baisWrapper);
+ ((InternalMergingState) partitionedState).mergeNamespaces(targetNamespace, namespaces);
+ } else {
+ // get values
+ byte[] valueBytes = request.getAppend().getData().toByteArray();
+ partitionedState.add(valueBytes);
+ }
+ return CompletableFuture.completedFuture(
+ BeamFnApi.StateResponse.newBuilder()
+ .setId(request.getId())
+ .setAppend(BeamFnApi.StateAppendResponse.getDefaultInstance()));
+ }
+
+ private CompletionStage<BeamFnApi.StateResponse.Builder> handleBagClearRequest(
+ BeamFnApi.StateRequest request) throws Exception {
+
+ ListState<byte[]> partitionedState = getListState(request);
+
+ partitionedState.clear();
+ return CompletableFuture.completedFuture(
+ BeamFnApi.StateResponse.newBuilder()
+ .setId(request.getId())
+ .setClear(BeamFnApi.StateClearResponse.getDefaultInstance()));
+ }
+
+ private ListState<byte[]> getListState(BeamFnApi.StateRequest request) throws Exception {
+ BeamFnApi.StateKey.BagUserState bagUserState = request.getStateKey().getBagUserState();
+ String stateName = PYTHON_STATE_PREFIX + bagUserState.getUserStateId();
+ ListStateDescriptor<byte[]> listStateDescriptor;
+ StateDescriptor cachedStateDescriptor = stateDescriptorCache.get(stateName);
+ if (cachedStateDescriptor instanceof ListStateDescriptor) {
+ listStateDescriptor = (ListStateDescriptor<byte[]>) cachedStateDescriptor;
+ } else if (cachedStateDescriptor == null) {
+ listStateDescriptor = new ListStateDescriptor<>(stateName, valueSerializer);
+ stateDescriptorCache.put(stateName, listStateDescriptor);
+ } else {
+ throw new RuntimeException(
+ String.format(
+ "State name corrupt detected: "
+ + "'%s' is used both as LIST state and '%s' state at the same time.",
+ stateName, cachedStateDescriptor.getType()));
+ }
+ byte[] windowBytes = bagUserState.getWindow().toByteArray();
+ if (windowBytes.length != 0) {
+ bais.setBuffer(windowBytes, 0, windowBytes.length);
+ Object namespace = namespaceSerializer.deserialize(baisWrapper);
+ return (ListState<byte[]>)
+ keyedStateBackend.getPartitionedState(
+ namespace, namespaceSerializer, listStateDescriptor);
+ } else {
+ return (ListState<byte[]>)
+ keyedStateBackend.getPartitionedState(
+ VoidNamespace.INSTANCE,
+ VoidNamespaceSerializer.INSTANCE,
+ listStateDescriptor);
+ }
+ }
+
+ private CompletionStage<BeamFnApi.StateResponse.Builder> handleMapState(
+ BeamFnApi.StateRequest request) throws Exception {
+ // Currently the `beam_fn_api.proto` does not support MapState, so we use the
+ // the `MultimapSideInput` message to mark the state as a MapState for now.
+ if (request.getStateKey().hasMultimapSideInput()) {
+ BeamFnApi.StateKey.MultimapSideInput mapUserState =
+ request.getStateKey().getMultimapSideInput();
+ // get key
+ byte[] keyBytes = mapUserState.getKey().toByteArray();
+ bais.setBuffer(keyBytes, 0, keyBytes.length);
+ Object key = keySerializer.deserialize(baisWrapper);
+ if (keyedStateBackend.getKeySerializer() instanceof RowDataSerializer) {
+ setCurrentKeyForStreaming(
+ keyedStateBackend,
+ ((RowDataSerializer) keyedStateBackend.getKeySerializer())
+ .toBinaryRow((RowData) key));
+ } else {
+ setCurrentKeyForStreaming(keyedStateBackend, key);
+ }
+ } else {
+ throw new RuntimeException("Unsupported bag state request: " + request);
+ }
+
+ switch (request.getRequestCase()) {
+ case GET:
+ return handleMapGetRequest(request);
+ case APPEND:
+ return handleMapAppendRequest(request);
+ case CLEAR:
+ return handleMapClearRequest(request);
+ default:
+ throw new RuntimeException(
+ String.format(
+ "Unsupported request type %s for user state.",
+ request.getRequestCase()));
+ }
+ }
+
+ private CompletionStage<BeamFnApi.StateResponse.Builder> handleMapGetRequest(
+ BeamFnApi.StateRequest request) throws Exception {
+ MapState<ByteArrayWrapper, byte[]> mapState = getMapState(request);
+ // The continuation token structure of GET request is:
+ // [flag (1 byte)][serialized map key]
+ // The continuation token structure of CHECK_EMPTY request is:
+ // [flag (1 byte)]
+ // The continuation token structure of ITERATE request is:
+ // [flag (1 byte)][iterate type (1 byte)][iterator token length (int32)][iterator token]
+ byte[] getRequest = request.getGet().getContinuationToken().toByteArray();
+ byte getFlag = getRequest[0];
+ BeamFnApi.StateGetResponse.Builder response;
+ switch (getFlag) {
+ case GET_FLAG:
+ reuseByteArrayWrapper.setData(getRequest);
+ reuseByteArrayWrapper.setOffset(1);
+ reuseByteArrayWrapper.setLimit(getRequest.length);
+ response = handleMapGetValueRequest(reuseByteArrayWrapper, mapState);
+ break;
+ case CHECK_EMPTY_FLAG:
+ response = handleMapCheckEmptyRequest(mapState);
+ break;
+ case ITERATE_FLAG:
+ bais.setBuffer(getRequest, 1, getRequest.length - 1);
+ IterateType iterateType = IterateType.fromOrd(baisWrapper.readByte());
+ int iterateTokenLength = baisWrapper.readInt();
+ ByteArrayWrapper iterateToken;
+ if (iterateTokenLength > 0) {
+ reuseByteArrayWrapper.setData(getRequest);
+ reuseByteArrayWrapper.setOffset(bais.getPosition());
+ reuseByteArrayWrapper.setLimit(bais.getPosition() + iterateTokenLength);
+ iterateToken = reuseByteArrayWrapper;
+ } else {
+ iterateToken = null;
+ }
+ response = handleMapIterateRequest(mapState, iterateType, iterateToken);
+ break;
+ default:
+ throw new RuntimeException(
+ String.format(
+ "Unsupported get request type: '%d' for map state.", getFlag));
+ }
+
+ return CompletableFuture.completedFuture(
+ BeamFnApi.StateResponse.newBuilder().setId(request.getId()).setGet(response));
+ }
+
+ private CompletionStage<BeamFnApi.StateResponse.Builder> handleMapAppendRequest(
+ BeamFnApi.StateRequest request) throws Exception {
+ // The structure of append request bytes is:
+ // [number of requests (int32)][append request flag (1 byte)][map key length
+ // (int32)][map key]
+ // [map value length (int32)][map value][append request flag (1 byte)][map key length
+ // (int32)][map key]
+ // ...
+ byte[] appendBytes = request.getAppend().getData().toByteArray();
+ bais.setBuffer(appendBytes, 0, appendBytes.length);
+ MapState<ByteArrayWrapper, byte[]> mapState = getMapState(request);
+ int subRequestNum = baisWrapper.readInt();
+ for (int i = 0; i < subRequestNum; i++) {
+ byte requestFlag = baisWrapper.readByte();
+ int keyLength = baisWrapper.readInt();
+ reuseByteArrayWrapper.setData(appendBytes);
+ reuseByteArrayWrapper.setOffset(bais.getPosition());
+ reuseByteArrayWrapper.setLimit(bais.getPosition() + keyLength);
+ baisWrapper.skipBytesToRead(keyLength);
+ switch (requestFlag) {
+ case DELETE:
+ mapState.remove(reuseByteArrayWrapper);
+ break;
+ case SET_NONE:
+ mapState.put(reuseByteArrayWrapper.copy(), null);
+ break;
+ case SET_VALUE:
+ int valueLength = baisWrapper.readInt();
+ byte[] valueBytes = new byte[valueLength];
+ int readLength = baisWrapper.read(valueBytes);
+ assert valueLength == readLength;
+ mapState.put(reuseByteArrayWrapper.copy(), valueBytes);
+ break;
+ default:
+ throw new RuntimeException(
+ String.format(
+ "Unsupported append request type: '%d' for map state.",
+ requestFlag));
+ }
+ }
+ return CompletableFuture.completedFuture(
+ BeamFnApi.StateResponse.newBuilder()
+ .setId(request.getId())
+ .setAppend(BeamFnApi.StateAppendResponse.getDefaultInstance()));
+ }
+
+ private CompletionStage<BeamFnApi.StateResponse.Builder> handleMapClearRequest(
+ BeamFnApi.StateRequest request) throws Exception {
+ if (request.getStateKey()
+ .getMultimapSideInput()
+ .getTransformId()
+ .equals(CLEAR_CACHED_ITERATOR_MARK)) {
+ mapStateIteratorCache.clear();
+ } else {
+ MapState<ByteArrayWrapper, byte[]> partitionedState = getMapState(request);
+ partitionedState.clear();
+ }
+ return CompletableFuture.completedFuture(
+ BeamFnApi.StateResponse.newBuilder()
+ .setId(request.getId())
+ .setClear(BeamFnApi.StateClearResponse.getDefaultInstance()));
+ }
+
+ private BeamFnApi.StateGetResponse.Builder handleMapGetValueRequest(
+ ByteArrayWrapper key, MapState<ByteArrayWrapper, byte[]> mapState) throws Exception {
+ if (mapState.contains(key)) {
+ byte[] value = mapState.get(key);
+ if (value == null) {
+ return IS_NONE_RESPONSE;
+ } else {
+ baos.reset();
+ baosWrapper.writeByte(EXIST_FLAG);
+ baosWrapper.write(value);
+ return BeamFnApi.StateGetResponse.newBuilder()
+ .setData(ByteString.copyFrom(baos.toByteArray()));
+ }
+ } else {
+ return NOT_EXIST_RESPONSE;
+ }
+ }
+
+ private BeamFnApi.StateGetResponse.Builder handleMapCheckEmptyRequest(
+ MapState<ByteArrayWrapper, byte[]> mapState) throws Exception {
+ if (mapState.isEmpty()) {
+ return IS_EMPTY_RESPONSE;
+ } else {
+ return NOT_EMPTY_RESPONSE;
+ }
+ }
+
+ private BeamFnApi.StateGetResponse.Builder handleMapIterateRequest(
+ MapState<ByteArrayWrapper, byte[]> mapState,
+ IterateType iterateType,
+ ByteArrayWrapper iteratorToken)
+ throws Exception {
+ final Iterator iterator;
+ if (iteratorToken == null) {
+ switch (iterateType) {
+ case ITEMS:
+ case VALUES:
+ iterator = mapState.iterator();
+ break;
+ case KEYS:
+ iterator = mapState.keys().iterator();
+ break;
+ default:
+ throw new RuntimeException("Unsupported iterate type: " + iterateType);
+ }
+ } else {
+ iterator = mapStateIteratorCache.get(iteratorToken);
+ if (iterator == null) {
+ throw new RuntimeException("The cached iterator does not exist!");
+ }
+ }
+ baos.reset();
+ switch (iterateType) {
+ case ITEMS:
+ case VALUES:
+ Iterator<Map.Entry<ByteArrayWrapper, byte[]>> entryIterator = iterator;
+ for (int i = 0; i < mapStateIterateResponseBatchSize; i++) {
+ if (entryIterator.hasNext()) {
+ Map.Entry<ByteArrayWrapper, byte[]> entry = entryIterator.next();
+ ByteArrayWrapper key = entry.getKey();
+ baosWrapper.write(
+ key.getData(), key.getOffset(), key.getLimit() - key.getOffset());
+ baosWrapper.writeBoolean(entry.getValue() != null);
+ if (entry.getValue() != null) {
+ baosWrapper.write(entry.getValue());
+ }
+ } else {
+ break;
+ }
+ }
+ break;
+ case KEYS:
+ Iterator<ByteArrayWrapper> keyIterator = iterator;
+ for (int i = 0; i < mapStateIterateResponseBatchSize; i++) {
+ if (keyIterator.hasNext()) {
+ ByteArrayWrapper key = keyIterator.next();
+ baosWrapper.write(
+ key.getData(), key.getOffset(), key.getLimit() - key.getOffset());
+ } else {
+ break;
+ }
+ }
+ break;
+ default:
+ throw new RuntimeException("Unsupported iterate type: " + iterateType);
+ }
+ if (!iterator.hasNext()) {
+ if (iteratorToken != null) {
+ mapStateIteratorCache.remove(iteratorToken);
+ }
+ iteratorToken = null;
+ } else {
+ if (iteratorToken == null) {
+ iteratorToken = new ByteArrayWrapper(UUID.randomUUID().toString().getBytes());
+ }
+ mapStateIteratorCache.put(iteratorToken, iterator);
+ }
+ BeamFnApi.StateGetResponse.Builder responseBuilder =
+ BeamFnApi.StateGetResponse.newBuilder()
+ .setData(ByteString.copyFrom(baos.toByteArray()));
+ if (iteratorToken != null) {
+ responseBuilder.setContinuationToken(
+ ByteString.copyFrom(
+ iteratorToken.getData(),
+ iteratorToken.getOffset(),
+ iteratorToken.getLimit() - iteratorToken.getOffset()));
+ }
+ return responseBuilder;
+ }
+
+ private MapState<ByteArrayWrapper, byte[]> getMapState(BeamFnApi.StateRequest request)
+ throws Exception {
+ BeamFnApi.StateKey.MultimapSideInput mapUserState =
+ request.getStateKey().getMultimapSideInput();
+ String stateName = PYTHON_STATE_PREFIX + mapUserState.getSideInputId();
+ StateDescriptor cachedStateDescriptor = stateDescriptorCache.get(stateName);
+ MapStateDescriptor<ByteArrayWrapper, byte[]> mapStateDescriptor;
+ if (cachedStateDescriptor instanceof MapStateDescriptor) {
+ mapStateDescriptor =
+ (MapStateDescriptor<ByteArrayWrapper, byte[]>) cachedStateDescriptor;
+ } else if (cachedStateDescriptor == null) {
+ mapStateDescriptor =
+ new MapStateDescriptor<>(
+ stateName, ByteArrayWrapperSerializer.INSTANCE, valueSerializer);
+ stateDescriptorCache.put(stateName, mapStateDescriptor);
+ } else {
+ throw new RuntimeException(
+ String.format(
+ "State name corrupt detected: "
+ + "'%s' is used both as MAP state and '%s' state at the same time.",
+ stateName, cachedStateDescriptor.getType()));
+ }
+ byte[] windowBytes = mapUserState.getWindow().toByteArray();
+ if (windowBytes.length != 0) {
+ bais.setBuffer(windowBytes, 0, windowBytes.length);
+ Object namespace = namespaceSerializer.deserialize(baisWrapper);
+ return (MapState<ByteArrayWrapper, byte[]>)
+ keyedStateBackend.getPartitionedState(
+ namespace, namespaceSerializer, mapStateDescriptor);
+ } else {
+ return (MapState<ByteArrayWrapper, byte[]>)
+ keyedStateBackend.getPartitionedState(
+ VoidNamespace.INSTANCE,
+ VoidNamespaceSerializer.INSTANCE,
+ mapStateDescriptor);
+ }
+ }
+
+ private BeamFnApi.ProcessBundleRequest.CacheToken createCacheToken() {
+ ByteString token =
+ ByteString.copyFrom(UUID.randomUUID().toString().getBytes(Charsets.UTF_8));
+ return BeamFnApi.ProcessBundleRequest.CacheToken.newBuilder()
+ .setUserState(
+ BeamFnApi.ProcessBundleRequest.CacheToken.UserState.getDefaultInstance())
+ .setToken(token)
+ .build();
+ }
+
+ /** The type of the Python map state iterate request. */
+ private enum IterateType {
+
+ /** Equivalent to iterate {@link Map#entrySet() }. */
+ ITEMS((byte) 0),
+
+ /** Equivalent to iterate {@link Map#keySet() }. */
+ KEYS((byte) 1),
+
+ /** Equivalent to iterate {@link Map#values() }. */
+ VALUES((byte) 2);
+
+ private final byte ord;
+
+ IterateType(byte ord) {
+ this.ord = ord;
+ }
+
+ public byte getOrd() {
+ return ord;
+ }
+
+ public static IterateType fromOrd(byte ord) {
+ switch (ord) {
+ case 0:
+ return ITEMS;
+ case 1:
+ return KEYS;
+ case 2:
+ return VALUES;
+ default:
+ throw new RuntimeException("Unsupported ordinal: " + ord);
+ }
+ }
+ }
+}