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 <[email protected]> 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 <[email protected]> --- .../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++) {
