You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@flink.apache.org by we...@apache.org on 2026/08/31 08:46:40 UTC
(flink-agents) 01/02: [runtime][java] Reject raw bytes from PyFlink inputs
This is an automated email from the ASF dual-hosted git repository.
wenjin272 pushed a commit to branch release-0.3
in repository https://gitbox.apache.org/repos/asf/flink-agents.git
commit 81dd1ef6e0c5c73a9f02b7c56fe277f6ef40df67
Author: WenjinXie <we...@gmail.com>
AuthorDate: Mon Aug 31 15:32:26 2026 +0800
[runtime][java] Reject raw bytes from PyFlink inputs
Backport the PyFlink input value type validation to the 0.3 release line.
Generated-by: OpenAI Codex 0.144.5 (GPT-5)
Co-authored-by: Codex <co...@openai.com>
---
.../apache/flink/agents/runtime/CompileUtils.java | 33 ++++++++++
.../flink/agents/runtime/CompileUtilsTest.java | 74 ++++++++++++++++++++++
2 files changed, 107 insertions(+)
diff --git a/runtime/src/main/java/org/apache/flink/agents/runtime/CompileUtils.java b/runtime/src/main/java/org/apache/flink/agents/runtime/CompileUtils.java
index 753fd7ea..68afdc33 100644
--- a/runtime/src/main/java/org/apache/flink/agents/runtime/CompileUtils.java
+++ b/runtime/src/main/java/org/apache/flink/agents/runtime/CompileUtils.java
@@ -24,17 +24,32 @@ import org.apache.flink.agents.plan.AgentPlan;
import org.apache.flink.agents.runtime.operator.ActionExecutionOperatorFactory;
import org.apache.flink.api.common.typeinfo.TypeInformation;
import org.apache.flink.api.java.functions.KeySelector;
+import org.apache.flink.api.java.typeutils.RowTypeInfo;
import org.apache.flink.streaming.api.datastream.DataStream;
import org.apache.flink.streaming.api.datastream.KeyedStream;
+import org.apache.flink.streaming.api.typeinfo.python.PickledByteArrayTypeInfo;
import org.apache.flink.types.Row;
+import static org.apache.flink.util.Preconditions.checkArgument;
+
/** A utility class that bridges Flink DataStream/SQL with the Flink Agents agent. */
public class CompileUtils {
+ private static final int PYTHON_VALUE_FIELD_INDEX = 1;
+
// ============================ invoke by python ====================================
public static DataStream<byte[]> connectToAgent(
KeyedStream<Row, Row> inputDataStream, String agentPlanJson)
throws JsonProcessingException {
+ TypeInformation<?> inputType = inputDataStream.getType();
+ checkArgument(
+ isPickledPythonFieldType(inputType, PYTHON_VALUE_FIELD_INDEX),
+ "Flink Agents only supports PyFlink input values serialized with "
+ + "PickledByteArrayTypeInfo. Convert raw byte-array inputs with a Python "
+ + "operator using the default pickle output type before connecting them "
+ + "to Flink Agents, but got %s",
+ inputType);
+
// deserialize agent plan json.
AgentPlan agentPlan = new ObjectMapper().readValue(agentPlanJson, AgentPlan.class);
return connectToAgent(inputDataStream, agentPlan, TypeInformation.of(byte[].class), false);
@@ -82,4 +97,22 @@ public class CompileUtils {
new ActionExecutionOperatorFactory(agentPlan, inputIsJava))
.setParallelism(keyedInputStream.getParallelism());
}
+
+ /** Returns whether a PyFlink Row field uses its default pickle representation. */
+ static boolean isPickledPythonFieldType(TypeInformation<?> typeInformation, int fieldIndex) {
+ checkArgument(fieldIndex >= 0, "Field index must not be negative, but got %s", fieldIndex);
+ checkArgument(
+ typeInformation instanceof RowTypeInfo,
+ "Expected PyFlink type to be a RowTypeInfo, but got %s",
+ typeInformation);
+ RowTypeInfo rowType = (RowTypeInfo) typeInformation;
+ int expectedArity = fieldIndex + 1;
+ checkArgument(
+ rowType.getArity() == expectedArity,
+ "Expected PyFlink type to contain %s fields, but got arity %s",
+ expectedArity,
+ rowType.getArity());
+ TypeInformation<?> fieldType = rowType.getTypeAt(fieldIndex);
+ return fieldType instanceof PickledByteArrayTypeInfo;
+ }
}
diff --git a/runtime/src/test/java/org/apache/flink/agents/runtime/CompileUtilsTest.java b/runtime/src/test/java/org/apache/flink/agents/runtime/CompileUtilsTest.java
index ebfc15f4..a578abdd 100644
--- a/runtime/src/test/java/org/apache/flink/agents/runtime/CompileUtilsTest.java
+++ b/runtime/src/test/java/org/apache/flink/agents/runtime/CompileUtilsTest.java
@@ -17,22 +17,29 @@
*/
package org.apache.flink.agents.runtime;
+import com.fasterxml.jackson.core.JsonProcessingException;
import org.apache.flink.agents.plan.AgentPlan;
import org.apache.flink.agents.runtime.operator.ActionExecutionOperatorTest;
+import org.apache.flink.api.common.typeinfo.TypeInformation;
+import org.apache.flink.api.common.typeinfo.Types;
import org.apache.flink.api.java.functions.KeySelector;
import org.apache.flink.streaming.api.datastream.DataStream;
import org.apache.flink.streaming.api.datastream.DataStreamSource;
import org.apache.flink.streaming.api.datastream.KeyedStream;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
+import org.apache.flink.streaming.api.typeinfo.python.PickledByteArrayTypeInfo;
+import org.apache.flink.types.Row;
import org.apache.flink.util.CloseableIterator;
import org.junit.jupiter.api.BeforeAll;
import org.junit.jupiter.api.Test;
import java.util.ArrayList;
+import java.util.Collections;
import java.util.List;
import java.util.stream.Collectors;
import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
/** Tests for {@link CompileUtils}. */
public class CompileUtilsTest {
@@ -96,6 +103,73 @@ public class CompileUtilsTest {
checkResult(resultList);
}
+ @Test
+ void detectsPickledAndNonPickledPythonValueTypes() {
+ assertThat(
+ CompileUtils.isPickledPythonFieldType(
+ Types.ROW(
+ PickledByteArrayTypeInfo.PICKLED_BYTE_ARRAY_TYPE_INFO,
+ PickledByteArrayTypeInfo.PICKLED_BYTE_ARRAY_TYPE_INFO),
+ 1))
+ .isTrue();
+ assertThat(
+ CompileUtils.isPickledPythonFieldType(
+ Types.ROW(
+ PickledByteArrayTypeInfo.PICKLED_BYTE_ARRAY_TYPE_INFO,
+ Types.PRIMITIVE_ARRAY(Types.BYTE)),
+ 1))
+ .isFalse();
+ assertThat(
+ CompileUtils.isPickledPythonFieldType(
+ Types.ROW(
+ PickledByteArrayTypeInfo.PICKLED_BYTE_ARRAY_TYPE_INFO,
+ Types.STRING),
+ 1))
+ .isFalse();
+ }
+
+ @Test
+ void rejectsRawByteArrayPythonInputBeforeDeserializingTheAgentPlan() {
+ KeyedStream<Row, Row> inputDataStream =
+ createPythonInputStream(Types.PRIMITIVE_ARRAY(Types.BYTE));
+
+ assertThatThrownBy(() -> CompileUtils.connectToAgent(inputDataStream, "not-json"))
+ .isInstanceOf(IllegalArgumentException.class)
+ .hasMessageContaining("only supports PyFlink input values")
+ .hasMessageContaining("raw byte-array");
+ }
+
+ @Test
+ void acceptsPickledPythonInputBeforeDeserializingTheAgentPlan() {
+ KeyedStream<Row, Row> inputDataStream =
+ createPythonInputStream(PickledByteArrayTypeInfo.PICKLED_BYTE_ARRAY_TYPE_INFO);
+
+ assertThatThrownBy(() -> CompileUtils.connectToAgent(inputDataStream, "not-json"))
+ .isInstanceOf(JsonProcessingException.class);
+ }
+
+ @Test
+ void rejectsMalformedPythonInputType() {
+ assertThatThrownBy(
+ () ->
+ CompileUtils.isPickledPythonFieldType(
+ Types.ROW(
+ PickledByteArrayTypeInfo
+ .PICKLED_BYTE_ARRAY_TYPE_INFO),
+ 1))
+ .isInstanceOf(IllegalArgumentException.class)
+ .hasMessageContaining("contain 2 fields");
+ }
+
+ private static KeyedStream<Row, Row> createPythonInputStream(TypeInformation<?> valueType) {
+ StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
+ TypeInformation<Row> inputType =
+ Types.ROW(PickledByteArrayTypeInfo.PICKLED_BYTE_ARRAY_TYPE_INFO, valueType);
+ Row input = Row.of(new byte[0], new byte[0]);
+ return env.fromData(Collections.singletonList(input), inputType)
+ .keyBy(value -> Row.of(value.getField(0)));
+ }
+
private static List<Long> getTestSequence() {
List<Long> testSequence = new ArrayList<>();
for (int i = 0; i < TEST_SEQUENCE_REPEAT; i++) {