This is an automated email from the ASF dual-hosted git repository.
wenjin272 pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/flink-agents.git
The following commit(s) were added to refs/heads/main by this push:
new 169de4a4 [runtime][java] Reject raw bytes from PyFlink inputs (#1075)
169de4a4 is described below
commit 169de4a44d72bb19d03134fc813038867a84d413
Author: Wenjin Xie <[email protected]>
AuthorDate: Mon Aug 31 15:15:22 2026 +0800
[runtime][java] Reject raw bytes from PyFlink inputs (#1075)
Validate PyFlink value type information before building the agent operator
so only the supported pickled representation reaches Python deserialization.
Add regression coverage for pickled, explicit, and malformed input types.
Generated-by: OpenAI Codex 0.144.5 (GPT-5)
Co-authored-by: Codex <[email protected]>
---
.../apache/flink/agents/runtime/CompileUtils.java | 39 +++++++----
.../flink/agents/runtime/CompileUtilsTest.java | 81 ++++++++++++++++++++--
2 files changed, 103 insertions(+), 17 deletions(-)
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 37351f2c..c1bf0805 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
@@ -35,10 +35,22 @@ 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_KEY_FIELD_INDEX = 0;
+ 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(
@@ -46,7 +58,7 @@ public class CompileUtils {
agentPlan,
TypeInformation.of(byte[].class),
false,
- isPickledPythonKeyType(inputDataStream.getKeyType()));
+ isPickledPythonFieldType(inputDataStream.getKeyType(),
PYTHON_KEY_FIELD_INDEX));
}
// ============================ invoke by java
====================================
@@ -95,20 +107,21 @@ public class CompileUtils {
.setParallelism(keyedInputStream.getParallelism());
}
- /**
- * Returns whether PyFlink's single logical key field uses its default
pickle representation.
- */
- static boolean isPickledPythonKeyType(TypeInformation<?> keyType) {
+ /** 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(
- keyType instanceof RowTypeInfo,
- "Expected PyFlink key type to be a single-field RowTypeInfo,
but got %s",
- keyType);
- RowTypeInfo rowType = (RowTypeInfo) keyType;
+ 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() == 1,
- "Expected PyFlink key type to contain one logical field, but
got arity %s",
+ rowType.getArity() == expectedArity,
+ "Expected PyFlink type to contain %s fields, but got arity %s",
+ expectedArity,
rowType.getArity());
- TypeInformation<?> logicalKeyType = rowType.getTypeAt(0);
- return logicalKeyType instanceof PickledByteArrayTypeInfo;
+ 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 5d53c187..96a0ce8f 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,8 +17,10 @@
*/
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;
@@ -26,15 +28,18 @@ 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 {
@@ -101,15 +106,83 @@ public class CompileUtilsTest {
@Test
void detectsPickledAndExplicitByteArrayPythonKeyTypes() {
assertThat(
- CompileUtils.isPickledPythonKeyType(
-
Types.ROW(PickledByteArrayTypeInfo.PICKLED_BYTE_ARRAY_TYPE_INFO)))
+ CompileUtils.isPickledPythonFieldType(
+
Types.ROW(PickledByteArrayTypeInfo.PICKLED_BYTE_ARRAY_TYPE_INFO),
+ 0))
.isTrue();
assertThat(
- CompileUtils.isPickledPythonKeyType(
- Types.ROW(Types.PRIMITIVE_ARRAY(Types.BYTE))))
+ CompileUtils.isPickledPythonFieldType(
+ Types.ROW(Types.PRIMITIVE_ARRAY(Types.BYTE)),
0))
.isFalse();
}
+ @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++) {
