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 c22e93b83 [api][runtime][java][python] Simplify agent output
materialization with typed toDataStream/toTable (#1183)
c22e93b83 is described below
commit c22e93b835c4d251898a4a7ea71981215483f828
Author: yunfengzhou-hub <[email protected]>
AuthorDate: Fri Oct 9 11:16:46 2026 +0800
[api][runtime][java][python] Simplify agent output materialization with
typed toDataStream/toTable (#1183)
Generated-by: Qoder 1.32.1 (Qwen3.8-Max)
---
.../org/apache/flink/agents/api/AgentBuilder.java | 74 ++++-
.../agents/api/AgentBuilderApplyByNameTest.java | 21 ++
.../docs/development/integrate_with_flink.md | 78 +++--
.../docs/get-started/quickstart/workflow_agent.md | 6 +-
.../integration/test/FlinkIntegrationAgent.java | 34 +++
.../integration/test/FlinkIntegrationTest.java | 103 +++++++
.../agents/integration/test/ReActAgentTest.java | 11 +-
.../integration/test/SkillsIntegrationTest.java | 11 +-
.../examples/WorkflowMultipleAgentExample.java | 10 +-
python/flink_agents/api/execution_environment.py | 39 ++-
.../quickstart/workflow_multiple_agent_example.py | 6 +-
python/flink_agents/runtime/output_type_utils.py | 331 +++++++++++++++++++++
.../runtime/remote_execution_environment.py | 161 ++++++++--
.../runtime/tests/output_type_fixtures.py | 52 ++++
.../runtime/tests/test_output_type_utils.py | 297 ++++++++++++++++++
.../tests/test_remote_execution_environment.py | 214 +++++++++++++
runtime/pom.xml | 9 +
.../flink/agents/runtime/OutputTypeUtils.java | 317 ++++++++++++++++++++
.../runtime/env/RemoteExecutionEnvironment.java | 51 +++-
.../flink/agents/runtime/OutputTypeUtilsTest.java | 303 +++++++++++++++++++
.../runtime/env/TypedOutputTerminalTest.java | 176 +++++++++++
21 files changed, 2216 insertions(+), 88 deletions(-)
diff --git a/api/src/main/java/org/apache/flink/agents/api/AgentBuilder.java
b/api/src/main/java/org/apache/flink/agents/api/AgentBuilder.java
index 4aac8b5ee..362e88b1f 100644
--- a/api/src/main/java/org/apache/flink/agents/api/AgentBuilder.java
+++ b/api/src/main/java/org/apache/flink/agents/api/AgentBuilder.java
@@ -19,6 +19,7 @@
package org.apache.flink.agents.api;
import org.apache.flink.agents.api.agents.Agent;
+import org.apache.flink.api.common.typeinfo.TypeInformation;
import org.apache.flink.streaming.api.datastream.DataStream;
import org.apache.flink.table.api.Schema;
import org.apache.flink.table.api.Table;
@@ -72,20 +73,85 @@ public interface AgentBuilder {
* Get output DataStream of agent execution.
*
* <p>This method converts the agent's output events into a Flink
DataStream that can be further
- * processed in the Flink pipeline.
+ * processed in the Flink pipeline. The returned view is unrestricted: its
elements are whatever
+ * the agent emitted ({@code Object}). To obtain a typed view, declare the
output type with
+ * {@link #toDataStream(TypeInformation)} or {@link #toDataStream(Class)}.
*
* @return DataStream containing outputs from agent execution.
*/
DataStream<Object> toDataStream();
/**
- * Get output Table of agent execution.
+ * Get output DataStream of agent execution, typed as {@code T}.
*
- * <p>This method converts the agent's output events into a Flink Table
using the provided
- * schema for structure definition.
+ * <p>The agent operator keeps emitting {@code Object} on the shared raw
stream and the declared
+ * type is applied by a downstream conversion operator, so the
unrestricted {@link
+ * #toDataStream()} view is unaffected and heterogeneous output stays
available through it. Each
+ * call is independent, so several typed views of different types can
coexist on one execution.
+ *
+ * @param typeInformation the declared type of the agent's output elements.
+ * @param <T> the declared output element type.
+ * @return DataStream whose elements are typed as {@code T}.
+ */
+ <T> DataStream<T> toDataStream(TypeInformation<T> typeInformation);
+
+ /**
+ * Get output DataStream of agent execution, typed as {@code T}, deriving
the {@link
+ * TypeInformation} from the given output class.
+ *
+ * <p>Convenience overload for the common case. Prefer {@link
#toDataStream(TypeInformation)}
+ * for generic or custom types that a {@link Class} cannot express.
+ *
+ * @param outputType the declared class of the agent's output elements.
+ * @param <T> the declared output element type.
+ * @return DataStream whose elements are typed as {@code T}.
+ */
+ default <T> DataStream<T> toDataStream(Class<T> outputType) {
+ return toDataStream(TypeInformation.of(outputType));
+ }
+
+ /**
+ * Get output Table of agent execution, materializing the unrestricted
view with an explicit
+ * physical schema.
+ *
+ * <p>The table's row type is derived from the schema's physical columns,
and a downstream
+ * conversion operator adapts each agent output element into a row
matching those columns by
+ * name. An element may be a row, a map keyed by column name, an object
exposing each column
+ * through a getter or field, or a scalar for a single-column schema.
Computed and metadata
+ * columns declared in the schema are derived by the planner rather than
read from the agent
+ * output. To declare the output with a type instead of a physical schema,
use {@link
+ * #toTable(TypeInformation)} or {@link #toTable(Class)}.
*
* @param schema Schema indicating the structure of the output table.
* @return Table containing outputs from agent execution.
*/
Table toTable(Schema schema);
+
+ /**
+ * Get output Table of agent execution, typed as {@code T}, deriving the
physical schema from
+ * the declared type (for example, a POJO's fields become columns).
+ *
+ * <p>The declared type is applied by a downstream conversion operator on
the shared raw stream,
+ * so the unrestricted {@link #toDataStream()} view is unaffected.
+ *
+ * @param typeInformation the declared type of the agent's output elements.
+ * @param <T> the declared output element type.
+ * @return Table containing the agent's output, typed as {@code T}.
+ */
+ <T> Table toTable(TypeInformation<T> typeInformation);
+
+ /**
+ * Get output Table of agent execution, typed as {@code T}, deriving the
{@link TypeInformation}
+ * from the given output class.
+ *
+ * <p>Convenience overload for the common case. Prefer {@link
#toTable(TypeInformation)} for
+ * generic or custom types that a {@link Class} cannot express.
+ *
+ * @param outputType the declared class of the agent's output elements.
+ * @param <T> the declared output element type.
+ * @return Table containing the agent's output, typed as {@code T}.
+ */
+ default <T> Table toTable(Class<T> outputType) {
+ return toTable(TypeInformation.of(outputType));
+ }
}
diff --git
a/api/src/test/java/org/apache/flink/agents/api/AgentBuilderApplyByNameTest.java
b/api/src/test/java/org/apache/flink/agents/api/AgentBuilderApplyByNameTest.java
index 8e3d6d3fd..e4f739912 100644
---
a/api/src/test/java/org/apache/flink/agents/api/AgentBuilderApplyByNameTest.java
+++
b/api/src/test/java/org/apache/flink/agents/api/AgentBuilderApplyByNameTest.java
@@ -19,6 +19,7 @@
package org.apache.flink.agents.api;
import org.apache.flink.agents.api.agents.Agent;
+import org.apache.flink.api.common.typeinfo.TypeInformation;
import org.apache.flink.streaming.api.datastream.DataStream;
import org.apache.flink.table.api.Schema;
import org.apache.flink.table.api.Table;
@@ -62,10 +63,20 @@ class AgentBuilderApplyByNameTest {
return null;
}
+ @Override
+ public <T> DataStream<T> toDataStream(TypeInformation<T>
typeInformation) {
+ return null;
+ }
+
@Override
public Table toTable(Schema schema) {
return null;
}
+
+ @Override
+ public <T> Table toTable(TypeInformation<T> typeInformation) {
+ return null;
+ }
}
/** Stub builder that overrides apply(String) to resolve from a one-agent
registry. */
@@ -104,9 +115,19 @@ class AgentBuilderApplyByNameTest {
return null;
}
+ @Override
+ public <T> DataStream<T> toDataStream(TypeInformation<T>
typeInformation) {
+ return null;
+ }
+
@Override
public Table toTable(Schema schema) {
return null;
}
+
+ @Override
+ public <T> Table toTable(TypeInformation<T> typeInformation) {
+ return null;
+ }
}
}
diff --git a/docs/content/docs/development/integrate_with_flink.md
b/docs/content/docs/development/integrate_with_flink.md
index dc93bfda4..c264baebd 100644
--- a/docs/content/docs/development/integrate_with_flink.md
+++ b/docs/content/docs/development/integrate_with_flink.md
@@ -161,7 +161,6 @@ Integrate the agent with input `Table`, and return the
output `Table` can be con
{{< tab "Python" >}}
```python
-from pyflink.common.typeinfo import BasicTypeInfo, ExternalTypeInfo,
RowTypeInfo
from pyflink.datastream import KeySelector
from pyflink.table import DataTypes, Schema
@@ -179,19 +178,14 @@ input_table = t_env.from_elements(
["id", "input"],
)
-# The output TypeInformation and Schema must be mutually consistent: both
-# describe a single "result" INT column here.
-output_type = ExternalTypeInfo(RowTypeInfo(
- [BasicTypeInfo.INT_TYPE_INFO()],
- ["result"],
-))
-
-schema = (Schema.new_builder().column("result", DataTypes.INT())).build()
+# A single output-type declaration: the Schema's physical columns give the
+# output row type, and each agent output element is adapted into a matching
row.
+schema = Schema.new_builder().column("result", DataTypes.INT()).build()
output_table = (
agents_env.from_table(input=input_table, key_selector=MyKeySelector())
.apply(your_agent)
- .to_table(schema=schema, output_type=output_type)
+ .to_table(schema=schema)
)
```
{{< /tab >}}
@@ -225,11 +219,14 @@ Table inputTable =
Row.of(2, "Bob", 92.0),
Row.of(3, "Charlie", 78.3));
-// The agent output is exposed as a single anonymous column named "f0".
-// Declare "f0" with the agent's OUTPUT type: a scalar type (e.g.
DataTypes.STRING())
-// when the agent emits a scalar value, or a nested DataTypes.ROW(...) only
when the
-// agent emits a composite row.
-Schema outputSchema = Schema.newBuilder().column("f0",
DataTypes.STRING()).build();
+// Declare the output columns. Each agent output element is adapted into a row
by
+// matching columns by name (from a POJO's getters/public fields, a Row, or a
Map);
+// a scalar output maps to a single-column schema.
+Schema outputSchema =
+ Schema.newBuilder()
+ .column("name", DataTypes.STRING())
+ .column("score", DataTypes.DOUBLE())
+ .build();
Table outputTable =
agentsEnv
@@ -244,18 +241,49 @@ Table outputTable =
User should provide `KeySelector` in `from_table()` to tell how to convert the
input `Table` to `KeyedStream` internally.
-The arguments required by `to_table()` differ by language:
+## Typed outputs
+
+`to_datastream()` / `toDataStream()` return an unrestricted stream whose
elements are whatever the agent emitted (`Object` in Java). To give the output
a concrete type, pass the type straight to the terminal:
`to_datastream(output_type)` / `toDataStream(TypeInformation)` materialize a
typed stream, and `to_table(output_type=...)` / `toTable(TypeInformation)`
materialize a typed table. The unrestricted call stays available and unchanged:
+
+{{< tabs "Typed outputs" >}}
+
+{{< tab "Python" >}}
+```python
+builder = agents_env.from_datastream(input_stream,
key_selector).apply(your_agent)
+
+raw = builder.to_datastream() # unrestricted: whatever the agent emitted
+
+# Pass the output type to the terminal to materialize a typed stream or table.
+typed = builder.to_datastream(ReviewOutput) # stream of validated
ReviewOutput
+table = builder.to_table(output_type=ReviewOutput) # schema derived from the
type
+# Or pass both: the Schema drives the physical table and is cross-checked
+# against the declared type.
+both = builder.to_table(review_schema, ReviewOutput)
+```
+{{< /tab >}}
+
+{{< tab "Java" >}}
+```java
+AgentBuilder builder =
+ agentsEnv.fromDataStream(inputStream, keySelector).apply(yourAgent);
+
+DataStream<Object> raw = builder.toDataStream(); // unrestricted
+
+// Pass the output type to the terminal to materialize a typed stream or table.
+DataStream<ReviewOutput> typed = builder.toDataStream(ReviewOutput.class);
+Table fromType = builder.toTable(ReviewOutput.class); // schema derived from
the type
+```
+{{< /tab >}}
+
+{{< /tabs >}}
-- **Python**: provide both `Schema` and `TypeInformation` to define the output
`Table` schema.
-- **Java**: provide only `Schema` (`toTable(Schema)`); `TypeInformation` is
not required.
+For the `Table` terminal:
-The two languages also name the output columns differently:
+- Calling `to_table(output_type=...)` / `toTable(TypeInformation)` with no
schema derives the physical columns from the declared type, so a POJO's fields
(or a structured Python type's fields) become columns.
+- In Python, `to_table(schema, output_type)` accepts both: the `Schema` drives
the physical table, preserving Table-domain information such as a primary key,
computed columns, metadata columns, or a watermark, and is cross-checked
against the declared type, raising if the two describe different row types. At
least one of the two is required. In Java the schema-only `toTable(Schema)` and
the type-only `toTable(TypeInformation)` are separate overloads.
-- **Python**: the `TypeInformation` passed to `to_table()` is a `RowTypeInfo`
whose field names become the output columns, so you name them directly (the
`"result"` column above matches `RowTypeInfo([...], ["result"])`).
-- **Java**: `toTable(Schema)` exposes the agent output as a single anonymous
column named `f0` (internally it calls
`StreamTableEnvironment.fromDataStream(DataStream<Object>, schema)`), so the
`Schema` must reference `f0` — wrap it in a `ROW(...)` only when the agent
emits a composite row.
+In Java, `toDataStream(Class)` / `toTable(Class)` are conveniences that derive
the `TypeInformation` from the class; the `TypeInformation` overloads remain
for generic or custom types.
-{{< hint info >}}
-In Python, `to_table()` currently requires both `Schema` and
`TypeInformation`; we plan to support providing only one of them in the future.
-{{< /hint >}}
+The schema-only `to_table(schema=...)` / `toTable(Schema)` stays available
when you have a physical schema but no output type to declare.
-For complete, runnable examples, see
[`WorkflowMultipleAgentExample.java`](https://github.com/apache/flink-agents/blob/main/examples/src/main/java/org/apache/flink/agents/examples/WorkflowMultipleAgentExample.java)
(Java) and
[`workflow_multiple_agent_example.py`](https://github.com/apache/flink-agents/blob/main/python/flink_agents/examples/quickstart/workflow_multiple_agent_example.py)
(Python).
+Each agent output element is adapted to the declared row type by matching
columns by name — from a POJO's getters or public fields, a `Row`, or a `Map`.
A scalar output maps to a single-column row.
diff --git a/docs/content/docs/get-started/quickstart/workflow_agent.md
b/docs/content/docs/get-started/quickstart/workflow_agent.md
index 4e844fa88..d71ca7ff9 100644
--- a/docs/content/docs/get-started/quickstart/workflow_agent.md
+++ b/docs/content/docs/get-started/quickstart/workflow_agent.md
@@ -414,7 +414,7 @@ input_table = t_env.from_path("product_reviews")
review_analysis_res_stream = (
agents_env.from_table(input=input_table, key_selector=TableKeySelector())
.apply(TableReviewAnalysisAgent())
- .to_datastream()
+ .to_datastream(ProductReviewAnalysisRes)
)
```
{{< /tab >}}
@@ -423,11 +423,11 @@ review_analysis_res_stream = (
```java
Table inputTable = tableEnv.from("product_reviews");
-DataStream<Object> reviewAnalysisResStream =
+DataStream<ProductReviewAnalysisRes> reviewAnalysisResStream =
agentsEnv
.fromTable(inputTable, new
TableReviewAnalysisAgent.RowKeySelector())
.apply(new TableReviewAnalysisAgent())
- .toDataStream();
+ .toDataStream(ProductReviewAnalysisRes.class);
```
{{< /tab >}}
diff --git
a/e2e-test/flink-agents-end-to-end-tests-integration/src/test/java/org/apache/flink/agents/integration/test/FlinkIntegrationAgent.java
b/e2e-test/flink-agents-end-to-end-tests-integration/src/test/java/org/apache/flink/agents/integration/test/FlinkIntegrationAgent.java
index bacb4aa0c..e0c934df3 100644
---
a/e2e-test/flink-agents-end-to-end-tests-integration/src/test/java/org/apache/flink/agents/integration/test/FlinkIntegrationAgent.java
+++
b/e2e-test/flink-agents-end-to-end-tests-integration/src/test/java/org/apache/flink/agents/integration/test/FlinkIntegrationAgent.java
@@ -187,4 +187,38 @@ public class FlinkIntegrationAgent {
ctx.sendEvent(new OutputEvent(output));
}
}
+
+ /** A structured output type used to exercise the typed output terminals.
*/
+ public static class TypedOutput {
+ public int value;
+ public String label;
+
+ public TypedOutput() {}
+
+ public TypedOutput(int value, String label) {
+ this.value = value;
+ this.label = label;
+ }
+
+ @Override
+ public String toString() {
+ return "TypedOutput{value=" + value + ", label='" + label + "'}";
+ }
+ }
+
+ /** Agent emitting a structured {@link TypedOutput} derived from each
integer input. */
+ public static class TypedOutputAgent extends Agent {
+ /**
+ * Turns each integer input into a deterministic structured output.
+ *
+ * @param event the input event carrying an integer id.
+ * @param ctx the runner context for sending the output event.
+ */
+ @Action(EventType.InputEvent)
+ public static void processInput(Event event, RunnerContext ctx) throws
Exception {
+ InputEvent inputEvent = InputEvent.fromEvent(event);
+ int id = (Integer) inputEvent.getInput();
+ ctx.sendEvent(new OutputEvent(new TypedOutput(id * 10, "item-" +
id)));
+ }
+ }
}
diff --git
a/e2e-test/flink-agents-end-to-end-tests-integration/src/test/java/org/apache/flink/agents/integration/test/FlinkIntegrationTest.java
b/e2e-test/flink-agents-end-to-end-tests-integration/src/test/java/org/apache/flink/agents/integration/test/FlinkIntegrationTest.java
index 11c677ac2..d7b34970a 100644
---
a/e2e-test/flink-agents-end-to-end-tests-integration/src/test/java/org/apache/flink/agents/integration/test/FlinkIntegrationTest.java
+++
b/e2e-test/flink-agents-end-to-end-tests-integration/src/test/java/org/apache/flink/agents/integration/test/FlinkIntegrationTest.java
@@ -19,6 +19,7 @@
package org.apache.flink.agents.integration.test;
import org.apache.flink.agents.api.AgentsExecutionEnvironment;
+import org.apache.flink.api.common.typeinfo.TypeInformation;
import org.apache.flink.api.java.functions.KeySelector;
import org.apache.flink.streaming.api.datastream.DataStream;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
@@ -181,6 +182,108 @@ public class FlinkIntegrationTest {
checkResult(results, "test_from_datastream_to_table.txt");
}
+ @Test
+ public void testToDataStreamWithTypeInformation() throws Exception {
+ StreamExecutionEnvironment env =
StreamExecutionEnvironment.getExecutionEnvironment();
+ env.setParallelism(1);
+
+ DataStream<Integer> inputStream = env.fromData(1, 2, 3);
+ AgentsExecutionEnvironment agentsEnv =
+ AgentsExecutionEnvironment.getExecutionEnvironment(env);
+
+ DataStream<FlinkIntegrationAgent.TypedOutput> outputStream =
+ agentsEnv
+ .fromDataStream(inputStream)
+ .apply(new FlinkIntegrationAgent.TypedOutputAgent())
+
.toDataStream(TypeInformation.of(FlinkIntegrationAgent.TypedOutput.class));
+
+ CloseableIterator<FlinkIntegrationAgent.TypedOutput> results =
outputStream.collectAsync();
+ agentsEnv.execute();
+
+ List<String> actual = new ArrayList<>();
+ while (results.hasNext()) {
+ actual.add(results.next().toString());
+ }
+ actual.sort(Comparator.naturalOrder());
+
+ List<String> expected =
+ new ArrayList<>(
+ List.of(
+ "TypedOutput{value=10, label='item-1'}",
+ "TypedOutput{value=20, label='item-2'}",
+ "TypedOutput{value=30, label='item-3'}"));
+ expected.sort(Comparator.naturalOrder());
+
+ Assertions.assertEquals(expected, actual);
+ }
+
+ @Test
+ public void testToTableWithTypeInformation() throws Exception {
+ StreamExecutionEnvironment env =
StreamExecutionEnvironment.getExecutionEnvironment();
+ env.setParallelism(1);
+
+ DataStream<Integer> inputStream = env.fromData(1, 2, 3);
+ AgentsExecutionEnvironment agentsEnv =
+ AgentsExecutionEnvironment.getExecutionEnvironment(env);
+
+ Table outputTable =
+ agentsEnv
+ .fromDataStream(inputStream)
+ .apply(new FlinkIntegrationAgent.TypedOutputAgent())
+
.toTable(TypeInformation.of(FlinkIntegrationAgent.TypedOutput.class));
+
+ // The schema is derived from the POJO type; locate columns by name.
+ List<String> columns =
outputTable.getResolvedSchema().getColumnNames();
+ int valueIdx = columns.indexOf("value");
+ int labelIdx = columns.indexOf("label");
+ Assertions.assertTrue(valueIdx >= 0 && labelIdx >= 0, "expected value
and label columns");
+
+ List<String> actual = collectValueLabel(outputTable, valueIdx,
labelIdx);
+ Assertions.assertEquals(List.of("10:item-1", "20:item-2",
"30:item-3"), actual);
+ }
+
+ @Test
+ public void testToTableWithMultiColumnSchema() throws Exception {
+ StreamExecutionEnvironment env =
StreamExecutionEnvironment.getExecutionEnvironment();
+ env.setParallelism(1);
+
+ DataStream<Integer> inputStream = env.fromData(1, 2, 3);
+ AgentsExecutionEnvironment agentsEnv =
+ AgentsExecutionEnvironment.getExecutionEnvironment(env);
+
+ // A POJO output materialized into the named columns declared by the
schema.
+ Schema outputSchema =
+ Schema.newBuilder()
+ .column("value", DataTypes.INT())
+ .column("label", DataTypes.STRING())
+ .build();
+
+ Table outputTable =
+ agentsEnv
+ .fromDataStream(inputStream)
+ .apply(new FlinkIntegrationAgent.TypedOutputAgent())
+ .toTable(outputSchema);
+
+ Assertions.assertEquals(
+ List.of("value", "label"),
outputTable.getResolvedSchema().getColumnNames());
+
+ List<String> actual = collectValueLabel(outputTable, 0, 1);
+ Assertions.assertEquals(List.of("10:item-1", "20:item-2",
"30:item-3"), actual);
+ }
+
+ private List<String> collectValueLabel(Table table, int valueIdx, int
labelIdx)
+ throws Exception {
+ List<String> actual = new ArrayList<>();
+ try (CloseableIterator<Row> results = table.execute().collect()) {
+ while (results.hasNext()) {
+ Row row = results.next();
+ actual.add(row.getField(valueIdx) + ":" +
row.getField(labelIdx));
+ }
+ }
+ actual.sort(Comparator.naturalOrder());
+ return actual;
+ }
+
private void checkResult(CloseableIterator<?> results, String fileName)
throws IOException {
String path =
Objects.requireNonNull(
diff --git
a/e2e-test/flink-agents-end-to-end-tests-integration/src/test/java/org/apache/flink/agents/integration/test/ReActAgentTest.java
b/e2e-test/flink-agents-end-to-end-tests-integration/src/test/java/org/apache/flink/agents/integration/test/ReActAgentTest.java
index e17f970bf..459ad37a1 100644
---
a/e2e-test/flink-agents-end-to-end-tests-integration/src/test/java/org/apache/flink/agents/integration/test/ReActAgentTest.java
+++
b/e2e-test/flink-agents-end-to-end-tests-integration/src/test/java/org/apache/flink/agents/integration/test/ReActAgentTest.java
@@ -29,7 +29,6 @@ import
org.apache.flink.agents.api.resource.ResourceDescriptor;
import org.apache.flink.agents.api.resource.ResourceName;
import org.apache.flink.agents.api.resource.ResourceType;
import org.apache.flink.agents.api.tools.Tool;
-import org.apache.flink.api.common.functions.MapFunction;
import org.apache.flink.api.common.typeinfo.BasicTypeInfo;
import org.apache.flink.api.common.typeinfo.TypeInformation;
import org.apache.flink.api.java.functions.KeySelector;
@@ -126,10 +125,7 @@ public class ReActAgentTest {
Row.of(2131, 29847, 3));
// Define output schema
- Schema outputSchema =
- Schema.newBuilder()
- .column("f0", DataTypes.ROW(DataTypes.FIELD("result",
DataTypes.DOUBLE())))
- .build();
+ Schema outputSchema = Schema.newBuilder().column("result",
DataTypes.DOUBLE()).build();
// Apply agent to the Table
Table outputTable =
@@ -142,10 +138,7 @@ public class ReActAgentTest {
.toTable(outputSchema);
// Collect the results to fully display the data
- CloseableIterator<Row> results =
- tableEnv.toDataStream(outputTable)
- .map((MapFunction<Row, Row>) x -> (Row)
x.getField("f0"))
- .collectAsync();
+ CloseableIterator<Row> results =
tableEnv.toDataStream(outputTable).collectAsync();
env.execute();
diff --git
a/e2e-test/flink-agents-end-to-end-tests-integration/src/test/java/org/apache/flink/agents/integration/test/SkillsIntegrationTest.java
b/e2e-test/flink-agents-end-to-end-tests-integration/src/test/java/org/apache/flink/agents/integration/test/SkillsIntegrationTest.java
index db5fbdd8e..b67e7aeed 100644
---
a/e2e-test/flink-agents-end-to-end-tests-integration/src/test/java/org/apache/flink/agents/integration/test/SkillsIntegrationTest.java
+++
b/e2e-test/flink-agents-end-to-end-tests-integration/src/test/java/org/apache/flink/agents/integration/test/SkillsIntegrationTest.java
@@ -28,7 +28,6 @@ import
org.apache.flink.agents.api.resource.ResourceDescriptor;
import org.apache.flink.agents.api.resource.ResourceName;
import org.apache.flink.agents.api.resource.ResourceType;
import org.apache.flink.agents.api.skills.Skills;
-import org.apache.flink.api.common.functions.MapFunction;
import org.apache.flink.api.common.typeinfo.BasicTypeInfo;
import org.apache.flink.api.common.typeinfo.TypeInformation;
import org.apache.flink.api.java.functions.KeySelector;
@@ -198,10 +197,7 @@ public class SkillsIntegrationTest {
DataTypes.FIELD("b", DataTypes.INT())),
Row.of(2, 3));
- Schema outputSchema =
- Schema.newBuilder()
- .column("f0", DataTypes.ROW(DataTypes.FIELD("result",
DataTypes.INT())))
- .build();
+ Schema outputSchema = Schema.newBuilder().column("result",
DataTypes.INT()).build();
Table outputTable =
agentsEnv
@@ -212,10 +208,7 @@ public class SkillsIntegrationTest {
.apply(agent)
.toTable(outputSchema);
- CloseableIterator<Row> results =
- tableEnv.toDataStream(outputTable)
- .map((MapFunction<Row, Row>) x -> (Row)
x.getField("f0"))
- .collectAsync();
+ CloseableIterator<Row> results =
tableEnv.toDataStream(outputTable).collectAsync();
env.execute();
diff --git
a/examples/src/main/java/org/apache/flink/agents/examples/WorkflowMultipleAgentExample.java
b/examples/src/main/java/org/apache/flink/agents/examples/WorkflowMultipleAgentExample.java
index 17ee75892..7218faa2d 100644
---
a/examples/src/main/java/org/apache/flink/agents/examples/WorkflowMultipleAgentExample.java
+++
b/examples/src/main/java/org/apache/flink/agents/examples/WorkflowMultipleAgentExample.java
@@ -166,19 +166,21 @@ public class WorkflowMultipleAgentExample {
Table inputTable = tableEnv.from("product_reviews");
// Use the TableReviewAnalysisAgent (LLM) to analyze each review.
- // The agent extracts the review score and unsatisfied reasons.
- DataStream<Object> reviewAnalysisResStream =
+ // The agent extracts the review score and unsatisfied reasons.
Declaring the
+ // output type via toDataStream(ProductReviewAnalysisRes.class) yields
a typed
+ // stream directly, so the result is already a
+ // DataStream<ProductReviewAnalysisRes> and needs no manual downcast
below.
+ DataStream<ProductReviewAnalysisRes> reviewAnalysisResStream =
agentsEnv
.fromTable(inputTable, new
TableReviewAnalysisAgent.RowKeySelector())
.apply(new TableReviewAnalysisAgent())
- .toDataStream();
+ .toDataStream(ProductReviewAnalysisRes.class);
// Aggregate the analysis results in 1-minute tumbling windows.
// This produces a score distribution and collects all unsatisfied
reasons for
// each product.
DataStream<String> aggregatedAnalysisResStream =
reviewAnalysisResStream
- .map(element -> (ProductReviewAnalysisRes) element)
.keyBy(ProductReviewAnalysisRes::getId)
.window(TumblingProcessingTimeWindows.of(Duration.ofMinutes(1)))
.process(new
AggregateScoreDistributionAndDislikeReasons());
diff --git a/python/flink_agents/api/execution_environment.py
b/python/flink_agents/api/execution_environment.py
index 02ba083db..d368fd9a6 100644
--- a/python/flink_agents/api/execution_environment.py
+++ b/python/flink_agents/api/execution_environment.py
@@ -23,7 +23,6 @@ if TYPE_CHECKING:
from pathlib import Path
from importlib_resources import files
-from pyflink.common import TypeInformation
from pyflink.datastream import DataStream, KeySelector,
StreamExecutionEnvironment
from pyflink.table import Schema, StreamTableEnvironment, Table
@@ -53,8 +52,22 @@ class AgentBuilder(ABC):
"""
@abstractmethod
- def to_datastream(self) -> DataStream:
- """Get output datastream of agent execution.
+ def to_datastream(self, output_type: Any = None) -> DataStream:
+ """Get the output datastream of agent execution.
+
+ Without ``output_type`` the returned view is unrestricted: it carries
+ whatever the agent emitted, with no output-type declaration. Passing
+ ``output_type`` layers a downstream conversion operator that
materializes
+ each element as that declared type, leaving the unrestricted view
+ unaffected, so several typed calls of different types can coexist on
one
+ execution.
+
+ Parameters
+ ----------
+ output_type : Any
+ Optional output-type declaration -- a Pydantic model / dataclass /
+ named tuple / ``TypedDict`` / ``RowTypeInfo``, or an explicit
+ ``TypeInformation``. When omitted, the unrestricted view is
returned.
Returns:
-------
@@ -62,17 +75,25 @@ class AgentBuilder(ABC):
Output datastream of agent execution.
"""
- # TODO: auto generate output_type.
@abstractmethod
- def to_table(self, schema: Schema, output_type: TypeInformation) -> Table:
+ def to_table(self, schema: Schema | None = None, output_type: Any = None)
-> Table:
"""Get output table of agent execution.
+ At least one of ``schema`` or ``output_type`` must be given. A
``schema``
+ fixes the physical columns and may carry Table-domain information such
as
+ a primary key, computed columns, or a watermark; an ``output_type``
+ derives the physical schema from the declared type (for example a
POJO's
+ fields become columns). Passing both cross-checks that they describe
the
+ same row type.
+
Parameters
----------
- schema : Schema
- Indicate schema of the output table.
- output_type : TypeInformation
- Indicate schema corresponding type information.
+ schema : Schema | None
+ A Table ``Schema`` whose physical columns give the row type.
+ output_type : Any
+ The output-type declaration -- a Pydantic model / dataclass / named
+ tuple / ``TypedDict`` / ``RowTypeInfo``, or an explicit
+ ``TypeInformation``.
Returns:
-------
diff --git
a/python/flink_agents/examples/quickstart/workflow_multiple_agent_example.py
b/python/flink_agents/examples/quickstart/workflow_multiple_agent_example.py
index fafa5c30c..b9bc3e6d5 100644
--- a/python/flink_agents/examples/quickstart/workflow_multiple_agent_example.py
+++ b/python/flink_agents/examples/quickstart/workflow_multiple_agent_example.py
@@ -142,11 +142,13 @@ def main() -> None:
)
# Use the TableReviewAnalysisAgent (LLM) to analyze each review.
- # The agent extracts the review score and unsatisfied reasons.
+ # The agent extracts the review score and unsatisfied reasons. Declaring
the
+ # output type with to_datastream(ProductReviewAnalysisRes) yields a typed
+ # stream, so downstream operators receive validated instances.
review_analysis_res_stream = (
agents_env.from_table(input=input_table,
key_selector=TableKeySelector())
.apply(TableReviewAnalysisAgent())
- .to_datastream()
+ .to_datastream(ProductReviewAnalysisRes)
)
# Aggregate the analysis results in 1-minute tumbling windows.
diff --git a/python/flink_agents/runtime/output_type_utils.py
b/python/flink_agents/runtime/output_type_utils.py
new file mode 100644
index 000000000..0e08ddbe0
--- /dev/null
+++ b/python/flink_agents/runtime/output_type_utils.py
@@ -0,0 +1,331 @@
+################################################################################
+# 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.
+#################################################################################
+"""Output-type inference and conversion helpers for typed agent terminals.
+
+An agent output terminal (``to_datastream`` / ``to_table``) accepts at most one
+type declaration. This module normalizes the accepted declarations -- a
Pydantic
+model, dataclass, named tuple, ``TypedDict``, an explicit ``RowTypeInfo``, or a
+Table ``Schema`` -- into the Flink type metadata each terminal needs, converts
an
+emitted value to a physical ``Row`` for the Table/Row serializer, and
+reconstructs declared instances for the typed ``to_datastream`` operator.
+
+``Schema`` and ``RowTypeInfo`` are converted into each other through Flink's
own
+``TypeConversions`` over the py4j gateway, so the derivation covers the full
type
+system (nested ``ROW``, arrays, temporal types) instead of a fixed scalar
subset.
+Only inference from a bare Python annotation (``infer_row_type_info``) stays
+limited to common scalars, because a Python type carries no Flink type on its
+own.
+"""
+
+from __future__ import annotations
+
+import dataclasses
+from typing import Any, get_type_hints, is_typeddict
+
+from pydantic import BaseModel
+from pyflink.common import Row
+from pyflink.common.typeinfo import (
+ ExternalTypeInfo,
+ RowTypeInfo,
+ TypeInformation,
+ Types,
+ _from_java_type,
+)
+from pyflink.java_gateway import get_gateway
+from pyflink.table import Schema
+from pyflink.table.types import _from_java_data_type
+
+__all__ = [
+ "infer_row_type_info",
+ "is_structured_declaration",
+ "reconstruct_instance",
+ "resolve_row_type_info",
+ "row_shape",
+ "row_type_info_to_schema",
+ "schema_to_row_type_info",
+ "to_row",
+]
+
+# python scalar type -> DataStream/Table TypeInformation. A bare Python
+# annotation carries no Flink type, so inference from a structured declaration
+# is limited to these common scalars; the py4j-backed Schema <-> RowTypeInfo
+# derivation below covers the full type system.
+_PY_TO_TYPEINFO: dict[type, TypeInformation] = {
+ str: Types.STRING(),
+ int: Types.LONG(),
+ float: Types.DOUBLE(),
+ bool: Types.BOOLEAN(),
+}
+
+
+def _is_named_tuple(cls: Any) -> bool:
+ """Whether ``cls`` is a ``NamedTuple`` class (a tuple subclass with
fields)."""
+ return isinstance(cls, type) and issubclass(cls, tuple) and hasattr(cls,
"_fields")
+
+
+def _field_types(cls: Any) -> list[tuple[str, Any]]:
+ """Return ordered ``(field_name, python_type)`` pairs for a structured
type.
+
+ Supports Pydantic models, dataclasses, named tuples and ``TypedDict``.
+ """
+ if isinstance(cls, type) and issubclass(cls, BaseModel):
+ return [(name, field.annotation) for name, field in
cls.model_fields.items()]
+ if isinstance(cls, type) and dataclasses.is_dataclass(cls):
+ hints = get_type_hints(cls)
+ return [(f.name, hints[f.name]) for f in dataclasses.fields(cls)]
+ if _is_named_tuple(cls):
+ hints = get_type_hints(cls)
+ return [(name, hints[name]) for name in cls._fields]
+ if is_typeddict(cls):
+ return list(get_type_hints(cls).items())
+ msg = (
+ f"cannot infer output fields from {cls!r}; expected a Pydantic model, "
+ "dataclass, named tuple, or TypedDict"
+ )
+ raise TypeError(msg)
+
+
+def _typeinfo_of(py_type: Any) -> TypeInformation:
+ if py_type not in _PY_TO_TYPEINFO:
+ supported = ", ".join(sorted(t.__name__ for t in _PY_TO_TYPEINFO))
+ msg = f"unsupported output field type {py_type!r}; supported types:
{supported}"
+ raise TypeError(msg)
+ return _PY_TO_TYPEINFO[py_type]
+
+
+def is_structured_declaration(decl: Any) -> bool:
+ """Whether ``decl`` is a python structured type whose fields can be
inferred."""
+ if isinstance(decl, TypeInformation | Schema):
+ return False
+ if not isinstance(decl, type):
+ return False
+ try:
+ _field_types(decl)
+ except TypeError:
+ return False
+ return True
+
+
+def infer_row_type_info(cls: Any) -> RowTypeInfo:
+ """Infer a ``RowTypeInfo`` from a Pydantic/dataclass/named
tuple/TypedDict."""
+ fields = _field_types(cls)
+ return RowTypeInfo(
+ [_typeinfo_of(py_type) for _, py_type in fields],
+ [name for name, _ in fields],
+ )
+
+
+def _unwrap_type_info(decl: Any) -> Any:
+ """Unwrap an ``ExternalTypeInfo`` to its inner ``TypeInformation``.
+
+ A Table terminal accepts an ``ExternalTypeInfo``-wrapped ``RowTypeInfo``
(the
+ form its Row serializer needs); resolution and canonicalization compare the
+ inner row type, so the wrapper is transparent to declaration equality.
+ """
+ if isinstance(decl, ExternalTypeInfo):
+ return decl._type_info
+ return decl
+
+
+def resolve_row_type_info(decl: Any) -> RowTypeInfo:
+ """Normalize a Table output declaration to a ``RowTypeInfo``.
+
+ Accepts an existing ``RowTypeInfo`` (optionally
``ExternalTypeInfo``-wrapped,
+ passed through) or a structured python type (inferred). A non-row
+ ``TypeInformation`` is rejected because a Table terminal needs a row-shaped
+ output.
+ """
+ decl = _unwrap_type_info(decl)
+ if isinstance(decl, RowTypeInfo):
+ return decl
+ if isinstance(decl, TypeInformation):
+ msg = (
+ "to_table requires a row-shaped output type, but got a non-row "
+ f"TypeInformation {decl!r}; declare a Pydantic model / dataclass /
"
+ "named tuple / TypedDict, or a RowTypeInfo"
+ )
+ raise TypeError(msg)
+ if is_structured_declaration(decl):
+ return infer_row_type_info(decl)
+ msg = f"cannot resolve {decl!r} to a row type for Table output"
+ raise TypeError(msg)
+
+
+def row_type_info_to_schema(typeinfo: RowTypeInfo) -> Schema:
+ """Derive a physical Table ``Schema`` from a ``RowTypeInfo``.
+
+ Delegates to Flink's ``TypeConversions.fromLegacyInfoToDataType`` over the
+ py4j gateway, so nested ``ROW`` and non-scalar fields round-trip instead of
+ being rejected.
+ """
+ gateway = get_gateway()
+ type_conversions =
gateway.jvm.org.apache.flink.table.types.utils.TypeConversions
+ j_data_type = type_conversions.fromLegacyInfoToDataType(
+ typeinfo.get_java_type_info()
+ )
+ row_data_type = _from_java_data_type(j_data_type)
+ return Schema.new_builder().from_row_data_type(row_data_type).build()
+
+
+def _resolve_data_type(abstract_type: Any) -> Any:
+ """Resolve a schema column's ``AbstractDataType`` to a concrete
``DataType``.
+
+ A column declared with a ``DataTypes`` object is already a resolved
+ ``DataType`` and is returned as-is. A column declared with a SQL type
string
+ (for example ``column("id", "BIGINT")``) is an ``UnresolvedDataType`` that
+ ``TypeConversions.fromDataTypeToLegacyInfo`` rejects, so its SQL text is
+ parsed through Flink's ``LogicalTypeParser`` first -- the same resolution
+ PyFlink applies when the schema is handed to ``from_data_stream``.
+ """
+ gateway = get_gateway()
+ data_type_class = abstract_type.getClass().forName(
+ "org.apache.flink.table.types.DataType"
+ )
+ if data_type_class.isInstance(abstract_type):
+ return abstract_type
+ sql_type = abstract_type.toString()
+ if sql_type.startswith("[") and sql_type.endswith("]"):
+ sql_type = sql_type[1:-1]
+ type_conversions =
gateway.jvm.org.apache.flink.table.types.utils.TypeConversions
+ logical_type_parser = (
+
gateway.jvm.org.apache.flink.table.types.logical.utils.LogicalTypeParser
+ )
+ class_loader =
gateway.jvm.java.lang.Thread.currentThread().getContextClassLoader()
+ logical_type = logical_type_parser.parse(sql_type, class_loader)
+ return type_conversions.fromLogicalToDataType(logical_type)
+
+
+def schema_to_row_type_info(schema: Schema) -> RowTypeInfo:
+ """Derive a ``RowTypeInfo`` from the physical columns of a Table
``Schema``.
+
+ Each column type is resolved to a concrete ``DataType`` -- a SQL type
string
+ such as ``"BIGINT"`` is parsed first -- then converted to a legacy
+ ``TypeInformation`` through Flink's ``fromDataTypeToLegacyInfo``. The
+ assembled Java ``RowTypeInfo`` is mapped back by PyFlink's
+ ``_from_java_type``, so nested ``ROW``, array, and temporal columns are
+ preserved.
+ """
+ gateway = get_gateway()
+ type_conversions =
gateway.jvm.org.apache.flink.table.types.utils.TypeConversions
+ j_row_type_info =
gateway.jvm.org.apache.flink.api.java.typeutils.RowTypeInfo
+ j_type_info =
gateway.jvm.org.apache.flink.api.common.typeinfo.TypeInformation
+ j_string = gateway.jvm.java.lang.String
+
+ # Only physical columns feed the DataStream Row; computed and metadata
+ # columns are derived by the Table planner rather than read from the agent
+ # output, mirroring the Java toTable(Schema) conversion.
+ columns = [
+ column
+ for column in schema._j_schema.getColumns()
+ if column.getClass().getSimpleName() == "UnresolvedPhysicalColumn"
+ ]
+ j_types = gateway.new_array(j_type_info, len(columns))
+ j_names = gateway.new_array(j_string, len(columns))
+ for i, column in enumerate(columns):
+ data_type = _resolve_data_type(column.getDataType())
+ j_types[i] = type_conversions.fromDataTypeToLegacyInfo(data_type)
+ j_names[i] = column.getName()
+ return _from_java_type(j_row_type_info(j_types, j_names))
+
+
+def row_shape(row_type_info: RowTypeInfo) -> tuple:
+ """Picklable nested field-name shape of a ``RowTypeInfo``.
+
+ The result is ``(names, children)`` where ``children[i]`` is the shape of a
+ nested ``RowTypeInfo`` field and ``None`` otherwise. It carries only the
row
+ nesting that value conversion needs, so it can be captured in a worker
+ closure; a ``RowTypeInfo`` itself holds py4j objects once its Java type has
+ been materialized and cannot be pickled into the closure.
+ """
+ names = list(row_type_info.get_field_names())
+ children = [
+ row_shape(field_type) if isinstance(field_type, RowTypeInfo) else None
+ for field_type in row_type_info.get_field_types()
+ ]
+ return (names, children)
+
+
+def to_row(value: Any, shape: tuple) -> Any:
+ """Adapt an emitted value to a physical ``Row`` using a :func:`row_shape`.
+
+ The raw agent output crosses a JSON boundary, so a structured value arrives
+ as a dict (a Pydantic model / dataclass degraded to its fields) or, for a
+ ``NamedTuple``, as a positional array; a value the agent emitted as a
``Row``
+ is reconstructed as one. Only fields the shape marks as nested rows are
+ converted recursively; every other value -- scalar, map, or list -- is left
+ untouched for its own coder, so no column kind is silently reshaped. This
is
+ the payload of the Table/Row conversion operator.
+ """
+ if isinstance(value, Row):
+ return value
+ names, children = shape
+ if isinstance(value, dict):
+ fields = [
+ to_row(value.get(name), child) if child is not None else
value.get(name)
+ for name, child in zip(names, children, strict=False)
+ ]
+ return Row(*fields)
+ if isinstance(value, list | tuple) and len(value) == len(names):
+ # A NamedTuple degrades to a positional array across the OutputEvent
JSON
+ # boundary (mirroring reconstruct_instance), so rebuild the Row by
field
+ # order. A length mismatch means the sequence is one column's own value
+ # (for example an ARRAY), which is left untouched for its coder.
+ fields = [
+ to_row(item, child) if child is not None else item
+ for item, child in zip(value, children, strict=False)
+ ]
+ return Row(*fields)
+ return value
+
+
+def reconstruct_instance(cls: Any, data: Any) -> Any:
+ """Rebuild an instance of a structured declaration from its emitted form.
+
+ This is the payload of the typed ``to_datastream`` conversion operator. A
+ Pydantic model / dataclass degrades to a dict across the JSON output
boundary
+ (only ``Row`` is reconstructed there) while a ``NamedTuple`` degrades to a
+ positional array, so the declared type is restored here; a ``Row`` is
adapted
+ through its field mapping, a ``TypedDict`` declaration is already satisfied
+ by a dict, and a value that is already an instance is passed through.
+ """
+ if isinstance(data, Row):
+ data = data.as_dict()
+ if is_typeddict(cls):
+ if not isinstance(data, dict):
+ msg = f"cannot reconstruct TypedDict {cls!r} from
{type(data).__name__}"
+ raise TypeError(msg)
+ return data
+ if isinstance(cls, type) and isinstance(data, cls):
+ return data
+ if _is_named_tuple(cls) and isinstance(data, list | tuple):
+ # A NamedTuple crosses the OutputEvent JSON boundary as a positional
+ # array, so rebuild it by field order instead of by name.
+ return cls(*data)
+ if not isinstance(data, dict):
+ msg = (
+ f"cannot reconstruct {cls!r} from {type(data).__name__}; expected
a "
+ "mapping of field names to values"
+ )
+ raise TypeError(msg)
+ if isinstance(cls, type) and issubclass(cls, BaseModel):
+ return cls.model_validate(data)
+ try:
+ return cls(**data)
+ except TypeError as error:
+ msg = f"cannot reconstruct {cls!r} from fields {sorted(data)}: {error}"
+ raise TypeError(msg) from error
diff --git a/python/flink_agents/runtime/remote_execution_environment.py
b/python/flink_agents/runtime/remote_execution_environment.py
index a5d1515b2..d26c14e5f 100644
--- a/python/flink_agents/runtime/remote_execution_environment.py
+++ b/python/flink_agents/runtime/remote_execution_environment.py
@@ -21,9 +21,10 @@ from pathlib import Path
from typing import Any, Callable, Dict
import cloudpickle
-from pyflink.common import TypeInformation
from pyflink.common.typeinfo import (
+ ExternalTypeInfo,
PickledBytesTypeInfo,
+ RowTypeInfo,
)
from pyflink.datastream import (
DataStream,
@@ -42,6 +43,16 @@ from flink_agents.api.execution_environment import (
from flink_agents.api.resource import ResourceType
from flink_agents.plan.agent_plan import AgentPlan
from flink_agents.plan.configuration import AgentConfiguration
+from flink_agents.runtime.output_type_utils import (
+ _unwrap_type_info,
+ is_structured_declaration,
+ reconstruct_instance,
+ resolve_row_type_info,
+ row_shape,
+ row_type_info_to_schema,
+ schema_to_row_type_info,
+ to_row,
+)
_CONFIG_FILE_NAME = "config.yaml"
_LEGACY_CONFIG_FILE_NAME = "flink-conf.yaml"
@@ -49,12 +60,23 @@ _AGENT_PLAN_JSON_VALIDATOR_CLASS =
"org.apache.flink.agents.plan.AgentPlanJsonVa
_AGENT_PLAN_JSON_VALIDATOR_METHOD = "validateAgentPlan"
+def _row_types_equal(left: RowTypeInfo, right: RowTypeInfo) -> bool:
+ """Whether two row types describe the same fields (names and types).
+
+ Java ``RowTypeInfo.equals`` ignores field names, so names are compared
+ separately while the Java check covers the (possibly nested) field types.
+ """
+ return list(left.get_field_names()) == list(right.get_field_names()) and
bool(
+ left.get_java_type_info().equals(right.get_java_type_info())
+ )
+
+
class RemoteAgentBuilder(AgentBuilder):
"""RemoteAgentBuilder for integrating datastream/table and agent."""
__input: DataStream
__agent_plan_json: str | None = None
- __output: DataStream = None
+ __raw_output: DataStream = None
__t_env: StreamTableEnvironment
__config: AgentConfiguration
__resources: Dict[ResourceType, Dict[str, Any]] = None
@@ -136,20 +158,19 @@ class RemoteAgentBuilder(AgentBuilder):
if error_message is not None:
raise ValueError(error_message)
- def to_datastream(self, output_type: TypeInformation | None = None) ->
DataStream:
- """Get output datastream of agent execution.
+ def _raw_output_stream(self) -> DataStream:
+ """Return the shared, untyped agent output stream (pickled bytes).
- Returns:
- -------
- DataStream
- Output datastream of agent execution.
+ Every terminal layers its own conversion operator on this single cached
+ stream, so requesting a typed view never re-runs the agent operator and
+ never changes the element type of the unrestricted view. Caching here,
at
+ the untyped boundary, is what removes the previous bug where the first
+ requested ``output_type`` was silently reused by every later terminal.
"""
if self.__agent_plan_json is None:
err_msg = "Must apply agent before call to_datastream/to_table."
raise RuntimeError(err_msg)
-
- # return the same output datastream when call to_datastream multiple.
- if self.__output is None:
+ if self.__raw_output is None:
j_data_stream_output = invoke_method(
None,
"org.apache.flink.agents.runtime.CompileUtils",
@@ -163,28 +184,126 @@ class RemoteAgentBuilder(AgentBuilder):
"java.lang.String",
],
)
- output_stream = DataStream(j_data_stream_output)
- self.__output = output_stream.map(
- lambda x: cloudpickle.loads(x), output_type=output_type
+ self.__raw_output = DataStream(j_data_stream_output)
+ return self.__raw_output
+
+ def _row_stream(self, row_type_info: RowTypeInfo) -> DataStream:
+ """Conversion operator adapting the raw output to physical ``Row``s.
+
+ Only the picklable nested field-name shape is captured in the closure;
+ the py4j-backed ``RowTypeInfo`` stays on the driver and is handed to
the
+ operator as ``ExternalTypeInfo`` for the Table/Row serializer.
+ """
+ raw = self._raw_output_stream()
+ shape = row_shape(row_type_info)
+ return raw.map(
+ lambda b: to_row(cloudpickle.loads(b), shape),
+ output_type=ExternalTypeInfo(row_type_info),
+ )
+
+ def to_datastream(self, output_type: Any = None) -> DataStream:
+ """Get the output datastream of agent execution.
+
+ The typed view is a downstream conversion operator on the shared,
cached
+ raw stream, so it never re-runs the agent operator nor changes the
element
+ type of the unrestricted view, and several typed views of different
types
+ can coexist on one execution.
+
+ Parameters
+ ----------
+ output_type : Any
+ Optional output-type declaration. When omitted, the unrestricted
view
+ carries whatever the agent emitted; when given, a downstream
conversion
+ operator materializes each element as that type.
+
+ Returns:
+ -------
+ DataStream
+ Output datastream of agent execution.
+ """
+ raw = self._raw_output_stream()
+ if output_type is None:
+ # Unrestricted view: keep the agent's own output elements.
+ return raw.map(lambda b: cloudpickle.loads(b))
+ if is_structured_declaration(output_type):
+ # Typed view of a python structured type: rebuild declared
instances.
+ # A Pydantic model has no native Flink TypeInformation, so the
stream
+ # stays pickle-typed and the increment is the validated instances.
+ model_cls = output_type
+ return raw.map(
+ lambda b: reconstruct_instance(model_cls, cloudpickle.loads(b))
)
- return self.__output
+ declared = _unwrap_type_info(output_type)
+ if isinstance(declared, RowTypeInfo):
+ return self._row_stream(declared)
+ # Scalar / other explicit TypeInformation: elements already match the
type.
+ return raw.map(lambda b: cloudpickle.loads(b), output_type=output_type)
- def to_table(self, schema: Schema, output_type: TypeInformation) -> Table:
+ def to_table(self, schema: Schema | None = None, output_type: Any = None)
-> Table:
"""Get output Table of agent execution.
+ At least one of ``schema`` or ``output_type`` must be given; passing
both
+ cross-checks that they describe the same row type.
+
Parameters
----------
- schema : Schema
- Indicate schema of the output table.
- output_type : TypeInformation
- Indicate schema corresponding type information.
+ schema : Schema | None
+ A Table ``Schema`` whose physical columns give the row type. It may
+ also carry Table-domain information such as a primary key, computed
+ columns, or a watermark.
+ output_type : Any
+ The output-type declaration -- a Pydantic model / dataclass / named
+ tuple / ``TypedDict`` / ``RowTypeInfo``, or an explicit
+ ``TypeInformation``. The physical schema is derived from it.
Returns:
-------
Table
Output Table of agent execution.
"""
- return self.t_env.from_data_stream(self.to_datastream(output_type),
schema)
+ if schema is not None and not isinstance(schema, Schema):
+ msg = (
+ "to_table 'schema' accepts a Table Schema only; to declare the
"
+ "output with a type (Pydantic model / dataclass / named tuple
/ "
+ "TypedDict / RowTypeInfo), pass it as output_type=... instead."
+ )
+ raise TypeError(msg)
+ if isinstance(output_type, Schema):
+ msg = (
+ "to_table 'output_type' accepts a type declaration (Pydantic
model "
+ "/ dataclass / named tuple / TypedDict / RowTypeInfo / "
+ "TypeInformation); pass a Table Schema as schema=... instead."
+ )
+ raise TypeError(msg)
+ if schema is None and output_type is None:
+ msg = (
+ "to_table requires at least one of 'schema' or 'output_type'; "
+ "neither was given."
+ )
+ raise ValueError(msg)
+ schema_row_type = (
+ schema_to_row_type_info(schema) if schema is not None else None
+ )
+ output_row_type = (
+ None if output_type is None else resolve_row_type_info(output_type)
+ )
+ if schema_row_type is not None and output_row_type is not None:
+ if not _row_types_equal(schema_row_type, output_row_type):
+ msg = (
+ "to_table got a schema and an output type describing
different "
+ f"row types: {schema_row_type!r} vs {output_row_type!r}"
+ )
+ raise ValueError(msg)
+ row_type_info = schema_row_type or output_row_type
+
+ # Preserve the caller's Schema when present so Table-domain information
+ # (primary key, computed columns, watermark) survives; otherwise
derive it.
+ table_schema = (
+ schema if schema is not None else
row_type_info_to_schema(row_type_info)
+ )
+ return self.t_env.from_data_stream(
+ self._row_stream(row_type_info), table_schema
+ )
class RemoteExecutionEnvironment(AgentsExecutionEnvironment):
diff --git a/python/flink_agents/runtime/tests/output_type_fixtures.py
b/python/flink_agents/runtime/tests/output_type_fixtures.py
new file mode 100644
index 000000000..e2155f2ea
--- /dev/null
+++ b/python/flink_agents/runtime/tests/output_type_fixtures.py
@@ -0,0 +1,52 @@
+################################################################################
+# 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.
+#################################################################################
+"""Shared structured output types for the typed-terminal tests.
+
+The output-type helpers and the typed ``to_datastream`` / ``to_table``
terminals
+are exercised with the same four kinds of python structured declaration: a
+Pydantic model, a dataclass, a named tuple and a ``TypedDict``. They live in
one
+module so the helper tests and the environment tests share a single definition
+instead of each redeclaring equivalent types.
+"""
+
+import dataclasses
+from typing import NamedTuple, TypedDict
+
+from pydantic import BaseModel
+
+
+class ModelOutput(BaseModel):
+ id: int
+ label: str
+ score: float
+
+
[email protected]
+class DcOutput:
+ id: int
+ label: str
+
+
+class NtOutput(NamedTuple):
+ id: int
+ label: str
+
+
+class TdOutput(TypedDict):
+ id: int
+ label: str
diff --git a/python/flink_agents/runtime/tests/test_output_type_utils.py
b/python/flink_agents/runtime/tests/test_output_type_utils.py
new file mode 100644
index 000000000..67d5d137e
--- /dev/null
+++ b/python/flink_agents/runtime/tests/test_output_type_utils.py
@@ -0,0 +1,297 @@
+################################################################################
+# 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.
+#################################################################################
+"""Unit tests for output type inference and conversion helpers."""
+
+import pytest
+from pydantic import BaseModel
+from pyflink.common import Row
+from pyflink.common.typeinfo import ExternalTypeInfo, RowTypeInfo, Types
+from pyflink.table import DataTypes, Schema
+
+from flink_agents.api.events.event import Event, OutputEvent
+from flink_agents.runtime.output_type_utils import (
+ infer_row_type_info,
+ is_structured_declaration,
+ reconstruct_instance,
+ resolve_row_type_info,
+ row_shape,
+ row_type_info_to_schema,
+ schema_to_row_type_info,
+ to_row,
+)
+from flink_agents.runtime.tests.output_type_fixtures import (
+ DcOutput,
+ ModelOutput,
+ NtOutput,
+ TdOutput,
+)
+
+
+def _row_spec(typeinfo: RowTypeInfo) -> tuple:
+ """Canonical comparable form of a flat row type: names plus type
strings."""
+ return (
+ list(typeinfo.get_field_names()),
+ [str(t) for t in typeinfo.get_field_types()],
+ )
+
+
+def test_infer_agrees_across_structured_types():
+ expected = _row_spec(RowTypeInfo([Types.LONG(), Types.STRING()], ["id",
"label"]))
+ for cls in (DcOutput, NtOutput, TdOutput):
+ assert _row_spec(infer_row_type_info(cls)) == expected
+
+
+def test_infer_pydantic_fields():
+ typeinfo = infer_row_type_info(ModelOutput)
+ assert typeinfo.get_field_names() == ["id", "label", "score"]
+ assert [str(t) for t in typeinfo.get_field_types()] == ["Long", "String",
"Double"]
+
+
+def test_is_structured_declaration():
+ assert is_structured_declaration(ModelOutput)
+ assert is_structured_declaration(DcOutput)
+ assert not is_structured_declaration(Types.LONG())
+ assert not is_structured_declaration(42)
+
+
+def test_row_type_info_to_schema_round_trip():
+ typeinfo = infer_row_type_info(ModelOutput)
+ schema = row_type_info_to_schema(typeinfo)
+ assert _row_spec(schema_to_row_type_info(schema)) == _row_spec(typeinfo)
+
+
+def test_schema_to_row_type_info():
+ schema = (
+ Schema.new_builder()
+ .column("id", DataTypes.BIGINT())
+ .column("label", DataTypes.STRING())
+ .build()
+ )
+ typeinfo = schema_to_row_type_info(schema)
+ assert typeinfo.get_field_names() == ["id", "label"]
+ assert [str(t) for t in typeinfo.get_field_types()] == ["Long", "String"]
+
+
+def test_schema_to_row_type_info_sql_type_strings():
+ """A column declared with a SQL type string resolves like a ``DataTypes``
one.
+
+ ``column("id", "BIGINT")`` stores an ``UnresolvedDataType``; resolving it
+ before the legacy conversion yields the same row type as the equivalent
+ ``DataTypes`` declaration, for scalars and a nested ``ROW`` alike.
+ """
+ schema = (
+ Schema.new_builder()
+ .column("id", "BIGINT")
+ .column("label", "STRING")
+ .column("nested", "ROW<x INT, y STRING>")
+ .build()
+ )
+ typeinfo = schema_to_row_type_info(schema)
+ assert typeinfo.get_field_names() == ["id", "label", "nested"]
+ assert [str(t) for t in typeinfo.get_field_types()[:2]] == ["Long",
"String"]
+ nested = typeinfo.get_field_types()[2]
+ assert isinstance(nested, RowTypeInfo)
+ assert nested.get_field_names() == ["x", "y"]
+
+ # The SQL strings agree with the equivalent ``DataTypes`` object columns.
+ object_schema = (
+ Schema.new_builder()
+ .column("id", DataTypes.BIGINT())
+ .column("label", DataTypes.STRING())
+ .build()
+ )
+ sql_schema = (
+ Schema.new_builder().column("id", "BIGINT").column("label",
"STRING").build()
+ )
+ assert _row_spec(schema_to_row_type_info(sql_schema)) == _row_spec(
+ schema_to_row_type_info(object_schema)
+ )
+
+
+def test_schema_to_row_type_info_skips_computed_and_metadata_columns():
+ """Only physical columns feed the derived row type.
+
+ Computed and metadata columns are planner-derived rather than read from the
+ agent output, so ``schema_to_row_type_info`` skips them; this mirrors the
+ Java ``toTable(Schema)`` conversion.
+ """
+ schema = (
+ Schema.new_builder()
+ .column("id", DataTypes.BIGINT())
+ .column("label", DataTypes.STRING())
+ .column_by_expression("id_plus_one", "id + 1")
+ .column_by_metadata("rowtime", "TIMESTAMP_LTZ(3)")
+ .build()
+ )
+ typeinfo = schema_to_row_type_info(schema)
+ assert typeinfo.get_field_names() == ["id", "label"]
+ assert [str(t) for t in typeinfo.get_field_types()] == ["Long", "String"]
+
+
+def test_nested_row_schema_derives_and_converts():
+ """A nested ROW column survives py4j derivation in both directions.
+
+ This is the full-type coverage the py4j-backed conversion adds over the
+ previous scalar-subset mapping, which rejected a nested ROW column
outright.
+ """
+ schema = (
+ Schema.new_builder()
+ .column("id", DataTypes.BIGINT())
+ .column(
+ "nested",
+ DataTypes.ROW(
+ [
+ DataTypes.FIELD("x", DataTypes.INT()),
+ DataTypes.FIELD("y", DataTypes.STRING()),
+ ]
+ ),
+ )
+ .build()
+ )
+ typeinfo = schema_to_row_type_info(schema)
+ assert typeinfo.get_field_names() == ["id", "nested"]
+ nested = typeinfo.get_field_types()[1]
+ assert isinstance(nested, RowTypeInfo)
+ assert nested.get_field_names() == ["x", "y"]
+
+ # A nested dict is converted to a nested Row along the picklable shape.
+ shape = row_shape(typeinfo)
+ row = to_row({"id": 1, "nested": {"x": 2, "y": "z"}}, shape)
+ assert row == Row(1, Row(2, "z"))
+
+ # The nested row type derives an equivalent schema and maps back.
+ round_tripped = schema_to_row_type_info(row_type_info_to_schema(typeinfo))
+ assert round_tripped.get_field_names() == ["id", "nested"]
+ rt_nested = round_tripped.get_field_types()[1]
+ assert isinstance(rt_nested, RowTypeInfo)
+ assert rt_nested.get_field_names() == ["x", "y"]
+
+
+def test_resolve_passes_through_and_infers():
+ typeinfo = infer_row_type_info(ModelOutput)
+ assert _row_spec(resolve_row_type_info(typeinfo)) == _row_spec(typeinfo)
+ assert _row_spec(resolve_row_type_info(DcOutput)) == _row_spec(
+ infer_row_type_info(DcOutput)
+ )
+
+
+def test_resolve_rejects_scalar_typeinfo():
+ with pytest.raises(TypeError, match="row-shaped"):
+ resolve_row_type_info(Types.LONG())
+
+
+def test_resolve_unwraps_external_type_info():
+ typeinfo = infer_row_type_info(ModelOutput)
+ wrapped = ExternalTypeInfo(typeinfo)
+ assert _row_spec(resolve_row_type_info(wrapped)) == _row_spec(typeinfo)
+
+
+def test_row_shape_marks_only_nested_rows():
+ names, children = row_shape(infer_row_type_info(ModelOutput))
+ assert names == ["id", "label", "score"]
+ assert children == [None, None, None]
+
+
+def test_to_row_from_dict():
+ shape = row_shape(infer_row_type_info(ModelOutput))
+ assert list(to_row({"id": 2, "label": "x", "score": 2.5}, shape)) == [2,
"x", 2.5]
+
+
+def test_to_row_passes_through_row():
+ shape = row_shape(infer_row_type_info(ModelOutput))
+ existing = Row(3, "y", 3.5)
+ assert to_row(existing, shape) is existing
+
+
+def test_to_row_passes_through_non_dict():
+ # A non-mapping value is left untouched for its own coder.
+ shape = row_shape(infer_row_type_info(ModelOutput))
+ assert to_row(42, shape) == 42
+
+
+def test_to_row_named_tuple_through_output_event_boundary():
+ """A ``NamedTuple`` degrades to a positional array across the JSON
boundary.
+
+ The Table path runs ``to_row`` (not ``reconstruct_instance``), so the
+ positional array the ``OutputEvent`` boundary produces must be rebuilt
into a
+ ``Row`` by field order here; otherwise the declared Row coder receives a
bare
+ list and fails with 'list' object has no attribute 'get_fields_by_names'.
+ """
+ event = OutputEvent(output=NtOutput(1, "good"))
+ degraded = Event.from_json(event.model_dump_json()).attributes["output"]
+ assert degraded == [1, "good"]
+
+ shape = row_shape(infer_row_type_info(NtOutput))
+ assert to_row(degraded, shape) == Row(1, "good")
+
+
+def test_to_row_passes_through_sequence_of_mismatched_length():
+ """The positional rebuild fires only when the sequence length matches the
row width.
+
+ A sequence whose length differs from the number of columns is a single
+ column's own value (an ARRAY, for instance) rather than a degraded
NamedTuple,
+ so it is left untouched for its coder instead of being split across
columns.
+ """
+ shape = row_shape(RowTypeInfo([Types.STRING()], ["arr"]))
+ assert to_row([1, 2, 3], shape) == [1, 2, 3]
+
+
+def test_infer_rejects_unsupported_field_type():
+ class Bad(BaseModel):
+ payload: bytes
+
+ with pytest.raises(TypeError, match="unsupported output field type"):
+ infer_row_type_info(Bad)
+
+
+def test_reconstruct_instance_from_dict():
+ data = {"id": 1, "label": "a"}
+ assert reconstruct_instance(DcOutput, data) == DcOutput(id=1, label="a")
+ assert reconstruct_instance(NtOutput, data) == NtOutput(id=1, label="a")
+ assert reconstruct_instance(TdOutput, data) == data
+ instance = reconstruct_instance(ModelOutput, {"id": 1, "label": "a",
"score": 1.5})
+ assert isinstance(instance, ModelOutput)
+
+
+def test_reconstruct_named_tuple_through_output_event_boundary():
+ """A ``NamedTuple`` degrades to a positional array across the JSON
boundary.
+
+ The agent emits ``NtOutput``; ``OutputEvent.model_dump_json`` serializes
the
+ tuple as a JSON array, so the value ``reconstruct_instance`` receives is
+ ``[1, "good"]`` rather than a mapping and must be rebuilt by field order.
+ """
+ event = OutputEvent(output=NtOutput(1, "good"))
+ degraded = Event.from_json(event.model_dump_json()).attributes["output"]
+ assert degraded == [1, "good"]
+ assert reconstruct_instance(NtOutput, degraded) == NtOutput(1, "good")
+
+
+def test_reconstruct_instance_from_row():
+ assert reconstruct_instance(DcOutput, Row(id=1, label="a")) == DcOutput(
+ id=1, label="a"
+ )
+
+
+def test_reconstruct_instance_passes_through_instance():
+ existing = ModelOutput(id=1, label="a", score=1.5)
+ assert reconstruct_instance(ModelOutput, existing) is existing
+
+
+def test_reconstruct_instance_rejects_non_mapping():
+ with pytest.raises(TypeError, match="expected a mapping"):
+ reconstruct_instance(DcOutput, 42)
diff --git
a/python/flink_agents/runtime/tests/test_remote_execution_environment.py
b/python/flink_agents/runtime/tests/test_remote_execution_environment.py
index ed8dcbdc8..7705c7629 100644
--- a/python/flink_agents/runtime/tests/test_remote_execution_environment.py
+++ b/python/flink_agents/runtime/tests/test_remote_execution_environment.py
@@ -18,20 +18,32 @@
import json
import os
import tempfile
+from collections.abc import Iterator
+from contextlib import contextmanager
from pathlib import Path
from unittest.mock import MagicMock, patch
+import cloudpickle
import pytest
import yaml
+from pyflink.common import Row
+from pyflink.common.typeinfo import ExternalTypeInfo, Types
+from pyflink.table import DataTypes, Schema
from flink_agents.api.agents.agent import Agent
from flink_agents.api.events.event import Event
from flink_agents.api.runner_context import RunnerContext
from flink_agents.plan.configuration import AgentConfiguration
+from flink_agents.runtime.output_type_utils import (
+ infer_row_type_info,
+ schema_to_row_type_info,
+)
from flink_agents.runtime.remote_execution_environment import (
RemoteAgentBuilder,
RemoteExecutionEnvironment,
+ _row_types_equal,
)
+from flink_agents.runtime.tests.output_type_fixtures import DcOutput,
ModelOutput
test_data = {
"agent": {
@@ -358,3 +370,205 @@ def test_to_datastream_submits_java_validated_plan_json()
-> None:
assert
json.loads(submitted_plan_json)["config"]["conf_data"]["plan.version"] == (
"validated"
)
+
+
+def _output_schema() -> Schema:
+ return (
+ Schema.new_builder()
+ .column("id", DataTypes.BIGINT())
+ .column("label", DataTypes.STRING())
+ .column("score", DataTypes.DOUBLE())
+ .build()
+ )
+
+
+@contextmanager
+def _applied_builder(t_env=None) -> Iterator[tuple]:
+ """Yield an applied builder with a capturable raw output stream.
+
+ ``invoke_method`` is patched so plan validation passes and
``connectToAgent``
+ returns a stub; ``DataStream`` is patched so the cached raw stream is
+ ``raw_mock``. A test then inspects each terminal's conversion operator
through
+ ``raw_mock.map.call_args_list`` without a Flink runtime.
+ """
+ if t_env is None:
+ t_env = MagicMock()
+ builder = RemoteAgentBuilder(
+ input=MagicMock(),
+ config=AgentConfiguration(),
+ t_env=t_env,
+ resources={},
+ agents={},
+ )
+ agent = _agent_with_conditions(["_input_event"])
+ raw_mock = MagicMock(name="raw_output")
+ with (
+ patch(
+ "flink_agents.runtime.remote_execution_environment.invoke_method",
+ side_effect=[None, MagicMock()],
+ ) as invoke_mock,
+ patch(
+ "flink_agents.runtime.remote_execution_environment.DataStream",
+ return_value=raw_mock,
+ ),
+ ):
+ builder.apply(agent)
+ yield builder, raw_mock, t_env, invoke_mock
+
+
+def _last_map(raw_mock) -> tuple:
+ """Return ``(func, output_type)`` of the most recent ``raw.map`` call."""
+ call = raw_mock.map.call_args_list[-1]
+ return call.args[0], call.kwargs.get("output_type")
+
+
+def test_to_datastream_raw_view_is_untyped():
+ with _applied_builder() as (builder, raw_mock, _, _):
+ builder.to_datastream()
+ func, output_type = _last_map(raw_mock)
+ assert output_type is None
+ assert func(cloudpickle.dumps({"id": 1})) == {"id": 1}
+
+
+def test_to_datastream_scalar_type():
+ long_ti = Types.LONG()
+ with _applied_builder() as (builder, raw_mock, _, _):
+ builder.to_datastream(long_ti)
+ func, output_type = _last_map(raw_mock)
+ assert output_type is long_ti
+ assert func(cloudpickle.dumps(42)) == 42
+
+
+def test_to_datastream_structured_model_yields_instances():
+ with _applied_builder() as (builder, raw_mock, _, _):
+ builder.to_datastream(ModelOutput)
+ func, output_type = _last_map(raw_mock)
+ # A Pydantic model has no native Flink type: the stream stays pickle-typed.
+ assert output_type is None
+ payload = cloudpickle.dumps({"id": 1, "label": "a", "score": 1.5})
+ assert func(payload) == ModelOutput(id=1, label="a", score=1.5)
+ # The closure captures no py4j object, so it survives a pickle round-trip.
+ restored = cloudpickle.loads(cloudpickle.dumps(func))
+ assert restored(payload) == ModelOutput(id=1, label="a", score=1.5)
+
+
+def test_to_datastream_row_type_info_yields_rows():
+ row_ti = infer_row_type_info(ModelOutput)
+ with _applied_builder() as (builder, raw_mock, _, _):
+ builder.to_datastream(row_ti)
+ func, output_type = _last_map(raw_mock)
+ assert isinstance(output_type, ExternalTypeInfo)
+ payload = cloudpickle.dumps({"id": 1, "label": "a", "score": 1.5})
+ assert list(func(payload)) == [1, "a", 1.5]
+ # The closure captures only the picklable shape, not the py4j RowTypeInfo.
+ restored = cloudpickle.loads(cloudpickle.dumps(func))
+ assert list(restored(payload)) == [1, "a", 1.5]
+
+
+def test_to_datastream_external_row_type_info_yields_rows():
+ # An ExternalTypeInfo-wrapped RowTypeInfo (the form a Table terminal hands
+ # back) is unwrapped to its inner row type before the row conversion.
+ wrapped = ExternalTypeInfo(infer_row_type_info(ModelOutput))
+ with _applied_builder() as (builder, raw_mock, _, _):
+ builder.to_datastream(wrapped)
+ func, output_type = _last_map(raw_mock)
+ assert isinstance(output_type, ExternalTypeInfo)
+ payload = cloudpickle.dumps({"id": 1, "label": "a", "score": 1.5})
+ assert list(func(payload)) == [1, "a", 1.5]
+
+
+def test_to_datastream_caches_raw_stream_without_welding_type():
+ long_ti = Types.LONG()
+ with _applied_builder() as (builder, raw_mock, _, invoke_mock):
+ builder.to_datastream()
+ builder.to_datastream(long_ti)
+ # The agent operator is connected once; each view maps its own conversion.
+ assert invoke_mock.call_count == 2
+ assert raw_mock.map.call_count == 2
+ first_type = raw_mock.map.call_args_list[0].kwargs.get("output_type")
+ second_type = raw_mock.map.call_args_list[1].kwargs.get("output_type")
+ assert first_type is None
+ assert second_type is long_ti
+
+
+def test_multiple_typed_datastreams_of_different_types_coexist():
+ # With no single-declaration guard, one execution exposes several
independent
+ # typed streams. Each typed call binds its own type through a separate
+ # conversion operator on the shared raw stream, and the unrestricted view
+ # stays untyped.
+ with _applied_builder() as (builder, raw_mock, _, _):
+ builder.to_datastream()
+ builder.to_datastream(ModelOutput)
+ builder.to_datastream(DcOutput)
+ # Raw view + two typed streams = three independent operators; no conflict.
+ assert raw_mock.map.call_count == 3
+ assert raw_mock.map.call_args_list[0].kwargs.get("output_type") is None
+
+
+def test_to_table_rejects_type_declaration_in_schema_slot():
+ # The builder's 'schema' slot is Schema-only; a type declaration must be
+ # passed as output_type=... instead.
+ with (
+ _applied_builder() as (builder, _, _, _),
+ pytest.raises(TypeError, match="accepts a Table Schema only"),
+ ):
+ builder.to_table(ModelOutput)
+
+
+def test_to_table_schema_only_builds_row_stream():
+ schema = _output_schema()
+ with _applied_builder() as (builder, raw_mock, t_env, _):
+ builder.to_table(schema)
+ func, output_type = _last_map(raw_mock)
+ assert isinstance(output_type, ExternalTypeInfo)
+ payload = cloudpickle.dumps({"id": 1, "label": "a", "score": 1.5})
+ assert isinstance(func(payload), Row)
+ # The caller's Schema is passed through to preserve Table-domain
information.
+ assert t_env.from_data_stream.call_args.args[1] is schema
+
+
+def test_to_table_output_type_derives_schema():
+ with _applied_builder() as (builder, _, t_env, _):
+ builder.to_table(output_type=ModelOutput)
+ derived = t_env.from_data_stream.call_args.args[1]
+ assert isinstance(derived, Schema)
+ assert _row_types_equal(
+ schema_to_row_type_info(derived), infer_row_type_info(ModelOutput)
+ )
+
+
+def test_to_table_cross_check_matching():
+ schema = _output_schema()
+ output_type = ExternalTypeInfo(infer_row_type_info(ModelOutput))
+ with _applied_builder() as (builder, _, t_env, _):
+ builder.to_table(schema=schema, output_type=output_type)
+ assert t_env.from_data_stream.call_args.args[1] is schema
+
+
+def test_to_table_cross_check_mismatched():
+ schema = _output_schema()
+ output_type = ExternalTypeInfo(infer_row_type_info(DcOutput))
+ with (
+ _applied_builder() as (builder, _, _, _),
+ pytest.raises(ValueError, match="different row types"),
+ ):
+ builder.to_table(schema=schema, output_type=output_type)
+
+
+def test_to_table_requires_schema_or_output_type():
+ # Neither a schema nor an output type: the terminal cannot derive a row
type.
+ with (
+ _applied_builder() as (builder, _, _, _),
+ pytest.raises(ValueError, match="at least one of 'schema' or
'output_type'"),
+ ):
+ builder.to_table()
+
+
+def test_to_table_rejects_schema_in_output_type_slot():
+ # The 'output_type' slot takes a type declaration; a Table Schema belongs
in
+ # the schema= slot, mirroring the schema-slot rejection above.
+ with (
+ _applied_builder() as (builder, _, _, _),
+ pytest.raises(TypeError, match="pass a Table Schema as schema="),
+ ):
+ builder.to_table(output_type=_output_schema())
diff --git a/runtime/pom.xml b/runtime/pom.xml
index e9e3521f2..4e96d98e4 100644
--- a/runtime/pom.xml
+++ b/runtime/pom.xml
@@ -98,6 +98,15 @@ under the License.
<version>${flink.version}</version>
<scope>provided</scope>
</dependency>
+ <!-- Test-scoped planner so the typed toTable terminals can be
unit-tested headlessly
+ (StreamTableEnvironment.create needs an ExecutorFactory);
production gets the planner
+ from the cluster classpath. -->
+ <dependency>
+ <groupId>org.apache.flink</groupId>
+ <artifactId>flink-table-planner_2.12</artifactId>
+ <version>${flink.version}</version>
+ <scope>test</scope>
+ </dependency>
<dependency>
<groupId>org.apache.flink</groupId>
<artifactId>flink-test-utils-junit</artifactId>
diff --git
a/runtime/src/main/java/org/apache/flink/agents/runtime/OutputTypeUtils.java
b/runtime/src/main/java/org/apache/flink/agents/runtime/OutputTypeUtils.java
new file mode 100644
index 000000000..983eec500
--- /dev/null
+++ b/runtime/src/main/java/org/apache/flink/agents/runtime/OutputTypeUtils.java
@@ -0,0 +1,317 @@
+/*
+ * 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.agents.runtime;
+
+import org.apache.flink.api.common.typeinfo.TypeInformation;
+import org.apache.flink.api.java.typeutils.RowTypeInfo;
+import org.apache.flink.table.api.Schema;
+import org.apache.flink.table.types.AbstractDataType;
+import org.apache.flink.table.types.DataType;
+import org.apache.flink.table.types.logical.LogicalType;
+import org.apache.flink.table.types.logical.utils.LogicalTypeParser;
+import org.apache.flink.table.types.utils.TypeConversions;
+import org.apache.flink.types.Row;
+
+import java.lang.reflect.Field;
+import java.lang.reflect.Method;
+import java.time.temporal.Temporal;
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Date;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import java.util.UUID;
+
+/**
+ * Output-type conversion helpers for the typed {@code toTable(Schema)}
terminal.
+ *
+ * <p>An agent emits arbitrary output objects through a {@code
DataStream<Object>}. Turning that
+ * untyped stream into a {@link org.apache.flink.table.api.Table} with a
user-declared {@link
+ * Schema} happens in a downstream conversion operator, so the raw stream
keeps its unrestricted
+ * {@code Object} element type and heterogeneous output stays available
through the raw view.
+ *
+ * <p>This class supplies the two pieces that operator needs:
+ *
+ * <ul>
+ * <li>{@link #schemaToRowTypeInfo(Schema)} derives the {@link RowTypeInfo}
of the physical
+ * columns a {@link Schema} declares. Computed and metadata columns are
skipped because the
+ * planner derives them, and only physical columns are read from the
agent output.
+ * <li>{@link #adaptToRow(Object, RowTypeInfo)} adapts one emitted value
into a {@link Row} that
+ * matches that row type, accepting a {@link Row}, a {@link Map}, a POJO
(read by field name),
+ * or a scalar (mapped into a single-column row type).
+ * </ul>
+ */
+public final class OutputTypeUtils {
+
+ private OutputTypeUtils() {}
+
+ /**
+ * Derives a {@link RowTypeInfo} from the physical columns of a Table
{@link Schema}.
+ *
+ * <p>Each physical column's type is resolved to a concrete {@link
DataType}: a {@code
+ * DataTypes} object is used as-is, while a SQL type string such as {@code
column("id",
+ * "BIGINT")} is parsed first. Computed and metadata columns are ignored
because the planner
+ * derives them rather than reading them from the agent output.
+ *
+ * @param schema the output schema declared by the user.
+ * @return a row type whose field names and types follow the schema's
physical columns, in
+ * order.
+ * @throws IllegalArgumentException if the schema has no physical columns,
or a physical
+ * column's type is neither a resolved {@link DataType} nor a
parseable SQL type string.
+ */
+ public static RowTypeInfo schemaToRowTypeInfo(Schema schema) {
+ List<String> names = new ArrayList<>();
+ List<TypeInformation<?>> types = new ArrayList<>();
+ for (Schema.UnresolvedColumn column : schema.getColumns()) {
+ if (!(column instanceof Schema.UnresolvedPhysicalColumn)) {
+ continue;
+ }
+ AbstractDataType<?> abstractType =
+ ((Schema.UnresolvedPhysicalColumn) column).getDataType();
+ names.add(column.getName());
+ types.add(
+ TypeConversions.fromDataTypeToLegacyInfo(
+ resolveDataType(abstractType, column.getName())));
+ }
+ if (names.isEmpty()) {
+ throw new IllegalArgumentException(
+ "Cannot derive an output row type: the schema declares no
physical columns.");
+ }
+ return new RowTypeInfo(types.toArray(new TypeInformation[0]),
names.toArray(new String[0]));
+ }
+
+ /**
+ * Resolves a schema column's {@link AbstractDataType} to a concrete
{@link DataType}.
+ *
+ * <p>A column declared with a {@code DataTypes} object is already
resolved and is returned
+ * as-is. A column declared with a SQL type string such as {@code
column("id", "BIGINT")}
+ * carries an {@link org.apache.flink.table.types.UnresolvedDataType} that
{@link
+ * TypeConversions#fromDataTypeToLegacyInfo} rejects, so its SQL text is
parsed through Flink's
+ * {@link LogicalTypeParser} first -- the same resolution the planner
applies once the schema
+ * reaches it, and the one the Python {@code schema_to_row_type_info}
performs.
+ */
+ private static DataType resolveDataType(AbstractDataType<?> abstractType,
String columnName) {
+ if (abstractType instanceof DataType) {
+ return (DataType) abstractType;
+ }
+ String sqlType = abstractType.toString();
+ if (sqlType.startsWith("[") && sqlType.endsWith("]")) {
+ sqlType = sqlType.substring(1, sqlType.length() - 1);
+ }
+ try {
+ LogicalType logicalType =
+ LogicalTypeParser.parse(
+ sqlType,
Thread.currentThread().getContextClassLoader());
+ return TypeConversions.fromLogicalToDataType(logicalType);
+ } catch (RuntimeException e) {
+ throw new IllegalArgumentException(
+ "Cannot derive the output row type from column '"
+ + columnName
+ + "': its type "
+ + abstractType
+ + " is neither a resolved DataType nor a parseable
SQL type string.",
+ e);
+ }
+ }
+
+ /**
+ * Adapts a single agent output value into a {@link Row} matching {@code
rowType}.
+ *
+ * <p>This is the payload of the {@code toTable(Schema)} conversion
operator. The accepted
+ * shapes are:
+ *
+ * <ul>
+ * <li>a {@link Row}: projected by name when named, otherwise passed
through positionally
+ * after an arity check;
+ * <li>a {@link Map}: read by the row type's field names;
+ * <li>a scalar value with a single-column row type: wrapped into a
one-column row, checked
+ * before the POJO lookup so a column name matching an internal
member of the scalar's own
+ * type does not extract that member;
+ * <li>an object exposing every column name (a POJO): read by name via
getters or fields.
+ * </ul>
+ *
+ * @param value the emitted output object.
+ * @param rowType the target row type derived from the output schema.
+ * @return a positional {@link Row} whose fields follow {@code rowType}'s
order.
+ * @throws IllegalArgumentException if the value cannot supply the
declared fields.
+ */
+ public static Row adaptToRow(Object value, RowTypeInfo rowType) {
+ String[] names = rowType.getFieldNames();
+ int arity = names.length;
+ if (value instanceof Row) {
+ return fromRow((Row) value, names, arity);
+ }
+ if (value instanceof Map) {
+ Map<?, ?> map = (Map<?, ?>) value;
+ Object[] fields = new Object[arity];
+ for (int i = 0; i < arity; i++) {
+ fields[i] = map.get(names[i]);
+ }
+ return Row.of(fields);
+ }
+ if (arity == 1 && isScalar(value)) {
+ // A scalar output mapped into a single-column table. Recognized
before the POJO
+ // property lookup so a column name that happens to match an
internal member of the
+ // scalar's own type -- String.value, String.getBytes() -- does
not extract that
+ // member instead of wrapping the value itself.
+ return Row.of(value);
+ }
+ if (canReadAll(value, names)) {
+ Object[] fields = new Object[arity];
+ for (int i = 0; i < arity; i++) {
+ fields[i] = readProperty(value, names[i], rowType);
+ }
+ return Row.of(fields);
+ }
+ if (arity == 1) {
+ // A single-column value that is not a recognized scalar (e.g. a
nested structure the
+ // schema maps into one column); wrap it whole.
+ return Row.of(value);
+ }
+ throw new IllegalArgumentException(
+ "Cannot adapt an output of type "
+ + value.getClass().getName()
+ + " to the declared schema "
+ + rowType
+ + ": it is neither a Row/Map nor an object exposing
the columns "
+ + Arrays.toString(names)
+ + ".");
+ }
+
+ private static Row fromRow(Row value, String[] names, int arity) {
+ Set<String> rowNames = tryFieldNames(value);
+ if (rowNames != null && rowNames.containsAll(Arrays.asList(names))) {
+ Object[] fields = new Object[arity];
+ for (int i = 0; i < arity; i++) {
+ fields[i] = value.getField(names[i]);
+ }
+ return Row.of(fields);
+ }
+ if (value.getArity() != arity) {
+ throw new IllegalArgumentException(
+ "Agent emitted a Row with "
+ + value.getArity()
+ + " fields, but the output schema declares "
+ + arity
+ + " columns "
+ + Arrays.toString(names)
+ + ".");
+ }
+ return value;
+ }
+
+ private static Set<String> tryFieldNames(Row row) {
+ try {
+ return row.getFieldNames(false);
+ } catch (RuntimeException e) {
+ // A positional row has no field names; fall back to positional
adaptation.
+ return null;
+ }
+ }
+
+ private static Object readProperty(Object bean, String name, RowTypeInfo
rowType) {
+ Class<?> type = bean.getClass();
+ Method getter = findGetter(type, name);
+ if (getter != null) {
+ try {
+ return getter.invoke(bean);
+ } catch (ReflectiveOperationException | RuntimeException e) {
+ throw new IllegalArgumentException(
+ "Failed to read output field '" + name + "' from " +
type.getName() + ".",
+ e);
+ }
+ }
+ Field field = findField(type, name);
+ if (field == null) {
+ throw new IllegalArgumentException(
+ "Cannot read output field '"
+ + name
+ + "' from "
+ + type.getName()
+ + " to match the declared schema "
+ + rowType
+ + ". Provide a matching getter or field.");
+ }
+ try {
+ field.setAccessible(true);
+ return field.get(bean);
+ } catch (ReflectiveOperationException | RuntimeException e) {
+ throw new IllegalArgumentException(
+ "Failed to read output field '" + name + "' from " +
type.getName() + ".", e);
+ }
+ }
+
+ /**
+ * Returns whether {@code value} is a scalar (non-decomposable) output
rather than a POJO whose
+ * fields should be projected into columns.
+ *
+ * <p>JDK value types carry internal members -- {@code String} has a
private {@code value} field
+ * and a {@code getBytes()} accessor -- that a by-name property lookup
would otherwise mistake
+ * for schema columns. Treating them as scalars keeps a single-column
schema wrapping the value
+ * itself, while a user POJO still projects its fields by name.
+ */
+ private static boolean isScalar(Object value) {
+ return value instanceof CharSequence
+ || value instanceof Number
+ || value instanceof Boolean
+ || value instanceof Character
+ || value instanceof Enum<?>
+ || value instanceof Temporal
+ || value instanceof Date
+ || value instanceof UUID
+ || value instanceof byte[];
+ }
+
+ private static boolean canReadAll(Object bean, String[] names) {
+ Class<?> type = bean.getClass();
+ for (String name : names) {
+ if (findGetter(type, name) == null && findField(type, name) ==
null) {
+ return false;
+ }
+ }
+ return true;
+ }
+
+ private static Method findGetter(Class<?> type, String name) {
+ String suffix = Character.toUpperCase(name.charAt(0)) +
name.substring(1);
+ for (String methodName : new String[] {"get" + suffix, "is" + suffix})
{
+ try {
+ Method method = type.getMethod(methodName);
+ if (method.getParameterCount() == 0) {
+ return method;
+ }
+ } catch (NoSuchMethodException e) {
+ // try the next candidate
+ }
+ }
+ return null;
+ }
+
+ private static Field findField(Class<?> type, String name) {
+ for (Class<?> current = type; current != null; current =
current.getSuperclass()) {
+ try {
+ return current.getDeclaredField(name);
+ } catch (NoSuchFieldException e) {
+ // keep walking up the hierarchy
+ }
+ }
+ return null;
+ }
+}
diff --git
a/runtime/src/main/java/org/apache/flink/agents/runtime/env/RemoteExecutionEnvironment.java
b/runtime/src/main/java/org/apache/flink/agents/runtime/env/RemoteExecutionEnvironment.java
index c89e4c556..8cb9e599d 100644
---
a/runtime/src/main/java/org/apache/flink/agents/runtime/env/RemoteExecutionEnvironment.java
+++
b/runtime/src/main/java/org/apache/flink/agents/runtime/env/RemoteExecutionEnvironment.java
@@ -25,7 +25,11 @@ import org.apache.flink.agents.api.resource.ResourceType;
import org.apache.flink.agents.plan.AgentConfiguration;
import org.apache.flink.agents.plan.AgentPlan;
import org.apache.flink.agents.runtime.CompileUtils;
+import org.apache.flink.agents.runtime.OutputTypeUtils;
+import org.apache.flink.api.common.functions.MapFunction;
+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.configuration.ConfigConstants;
import org.apache.flink.configuration.YamlParserUtils;
import org.apache.flink.streaming.api.datastream.DataStream;
@@ -33,6 +37,7 @@ import
org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.table.api.Schema;
import org.apache.flink.table.api.Table;
import org.apache.flink.table.api.bridge.java.StreamTableEnvironment;
+import org.apache.flink.types.Row;
import javax.annotation.Nullable;
@@ -234,10 +239,52 @@ public class RemoteExecutionEnvironment extends
AgentsExecutionEnvironment {
return outputDataStream;
}
+ @Override
+ public <V> DataStream<V> toDataStream(TypeInformation<V>
typeInformation) {
+ // Layer a downstream conversion operator on the shared raw stream
instead of
+ // retyping it, so the raw Object view stays unrestricted and
heterogeneous output
+ // remains available through toDataStream(). Each call is
independent, so several
+ // typed views of different types can coexist on one execution.
+ return toDataStream().map(new
OutputCaster<V>()).returns(typeInformation);
+ }
+
@Override
public Table toTable(Schema schema) {
- DataStream<Object> dataStream = toDataStream();
- return getTableEnvironment().fromDataStream(dataStream, schema);
+ // Convert each agent output into a Row matching the schema's
physical columns before
+ // handing the stream to the planner, so named columns declared in
the schema resolve
+ // against the emitted values. The conversion is a downstream
operator on the shared
+ // raw stream, so the unrestricted toDataStream() view stays
unaffected.
+ RowTypeInfo rowType = OutputTypeUtils.schemaToRowTypeInfo(schema);
+ DataStream<Row> rowStream =
+ toDataStream()
+ .map(value -> OutputTypeUtils.adaptToRow(value,
rowType))
+ .returns(rowType);
+ return getTableEnvironment().fromDataStream(rowStream, schema);
+ }
+
+ @Override
+ public <V> Table toTable(TypeInformation<V> typeInformation) {
+ // Derive the physical schema from the declared type:
fromDataStream reads the
+ // table structure off the typed stream's TypeInformation (a
POJO's fields become
+ // columns). The type is applied by a downstream operator on the
shared raw stream,
+ // so the unrestricted toDataStream() view is unaffected.
+ return
getTableEnvironment().fromDataStream(toDataStream(typeInformation));
+ }
+ }
+
+ /**
+ * Casts the unrestricted agent output objects to the type declared by a
typed terminal. The
+ * elements are already instances of the declared type; the operator
exists so the declared type
+ * is attached downstream of the shared raw stream rather than welded onto
it.
+ */
+ private static final class OutputCaster<T> implements MapFunction<Object,
T> {
+
+ private static final long serialVersionUID = 1L;
+
+ @Override
+ @SuppressWarnings("unchecked")
+ public T map(Object value) {
+ return (T) value;
}
}
}
diff --git
a/runtime/src/test/java/org/apache/flink/agents/runtime/OutputTypeUtilsTest.java
b/runtime/src/test/java/org/apache/flink/agents/runtime/OutputTypeUtilsTest.java
new file mode 100644
index 000000000..69a87b4c1
--- /dev/null
+++
b/runtime/src/test/java/org/apache/flink/agents/runtime/OutputTypeUtilsTest.java
@@ -0,0 +1,303 @@
+/*
+ * 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.agents.runtime;
+
+import org.apache.flink.api.java.typeutils.RowTypeInfo;
+import org.apache.flink.table.api.DataTypes;
+import org.apache.flink.table.api.Schema;
+import org.apache.flink.types.Row;
+import org.junit.jupiter.api.Test;
+
+import java.util.LinkedHashMap;
+import java.util.Map;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
+
+/** Tests for {@link OutputTypeUtils}. */
+public class OutputTypeUtilsTest {
+
+ /** A POJO exposing its values through getters only. */
+ public static class ReviewOutput {
+ private final int score;
+ private final String comment;
+
+ public ReviewOutput(int score, String comment) {
+ this.score = score;
+ this.comment = comment;
+ }
+
+ public int getScore() {
+ return score;
+ }
+
+ public String getComment() {
+ return comment;
+ }
+ }
+
+ /** A POJO exposing its values through public fields only. */
+ public static class PublicFields {
+ public String name;
+ public double value;
+
+ public PublicFields(String name, double value) {
+ this.name = name;
+ this.value = value;
+ }
+ }
+
+ private static RowTypeInfo rowType(Schema schema) {
+ return OutputTypeUtils.schemaToRowTypeInfo(schema);
+ }
+
+ @Test
+ void schemaToRowTypeInfoDerivesPhysicalColumnsInOrder() {
+ RowTypeInfo rowType =
+ rowType(
+ Schema.newBuilder()
+ .column("name", DataTypes.STRING())
+ .column("score", DataTypes.INT())
+ .column("ratio", DataTypes.DOUBLE())
+ .column("flag", DataTypes.BOOLEAN())
+ .build());
+
+ assertThat(rowType.getFieldNames()).containsExactly("name", "score",
"ratio", "flag");
+
assertThat(rowType.getTypeAt("name").getTypeClass()).isEqualTo(String.class);
+
assertThat(rowType.getTypeAt("score").getTypeClass()).isEqualTo(Integer.class);
+
assertThat(rowType.getTypeAt("ratio").getTypeClass()).isEqualTo(Double.class);
+
assertThat(rowType.getTypeAt("flag").getTypeClass()).isEqualTo(Boolean.class);
+ }
+
+ @Test
+ void schemaToRowTypeInfoSkipsComputedColumns() {
+ RowTypeInfo rowType =
+ rowType(
+ Schema.newBuilder()
+ .column("id", DataTypes.BIGINT())
+ .columnByExpression("doubled", "id * 2")
+ .build());
+
+ assertThat(rowType.getFieldNames()).containsExactly("id");
+
assertThat(rowType.getTypeAt("id").getTypeClass()).isEqualTo(Long.class);
+ }
+
+ @Test
+ void schemaToRowTypeInfoRejectsSchemaWithoutPhysicalColumns() {
+ Schema schema = Schema.newBuilder().columnByExpression("c",
"1").build();
+
+ assertThatThrownBy(() -> OutputTypeUtils.schemaToRowTypeInfo(schema))
+ .isInstanceOf(IllegalArgumentException.class)
+ .hasMessageContaining("no physical columns");
+ }
+
+ @Test
+ void schemaToRowTypeInfoResolvesSqlStringColumns() {
+ // A column declared with a SQL type string is an UnresolvedDataType;
it resolves to the
+ // same row type as the equivalent DataTypes declaration, for scalars
and a nested ROW.
+ RowTypeInfo fromSql =
+ rowType(
+ Schema.newBuilder()
+ .column("id", "BIGINT")
+ .column("label", "STRING")
+ .column("nested", "ROW<x INT, y STRING>")
+ .build());
+
+ assertThat(fromSql.getFieldNames()).containsExactly("id", "label",
"nested");
+
assertThat(fromSql.getTypeAt("id").getTypeClass()).isEqualTo(Long.class);
+
assertThat(fromSql.getTypeAt("label").getTypeClass()).isEqualTo(String.class);
+ assertThat(((RowTypeInfo)
fromSql.<Row>getTypeAt("nested")).getFieldNames())
+ .containsExactly("x", "y");
+
+ // The SQL strings agree with the equivalent DataTypes object columns.
+ RowTypeInfo fromDataTypes =
+ rowType(
+ Schema.newBuilder()
+ .column("id", DataTypes.BIGINT())
+ .column("label", DataTypes.STRING())
+ .build());
+ RowTypeInfo sqlScalars =
+ rowType(
+ Schema.newBuilder()
+ .column("id", "BIGINT")
+ .column("label", "STRING")
+ .build());
+
assertThat(sqlScalars.getFieldNames()).containsExactly(fromDataTypes.getFieldNames());
+ assertThat(sqlScalars.getTypeAt("id").getTypeClass())
+ .isEqualTo(fromDataTypes.getTypeAt("id").getTypeClass());
+ assertThat(sqlScalars.getTypeAt("label").getTypeClass())
+ .isEqualTo(fromDataTypes.getTypeAt("label").getTypeClass());
+ }
+
+ @Test
+ void adaptToRowWrapsScalarIntoSingleColumn() {
+ RowTypeInfo rowType = rowType(Schema.newBuilder().column("f0",
DataTypes.STRING()).build());
+
+ Row row = OutputTypeUtils.adaptToRow("hello", rowType);
+
+ assertThat(row.getArity()).isEqualTo(1);
+ assertThat(row.getField(0)).isEqualTo("hello");
+ }
+
+ @Test
+ void adaptToRowWrapsScalarWhenColumnNameCollidesWithStringInternals() {
+ // "value" is String's private byte[] field and "bytes" matches
String.getBytes(); a scalar
+ // String must be wrapped whole rather than have that internal member
extracted, which
+ // would otherwise surface as ClassCastException: byte[] cannot be
cast to String.
+ RowTypeInfo valueColumn =
+ rowType(Schema.newBuilder().column("value",
DataTypes.STRING()).build());
+ assertThat(OutputTypeUtils.adaptToRow("hello",
valueColumn).getField(0)).isEqualTo("hello");
+
+ RowTypeInfo bytesColumn =
+ rowType(Schema.newBuilder().column("bytes",
DataTypes.STRING()).build());
+ assertThat(OutputTypeUtils.adaptToRow("hello",
bytesColumn).getField(0)).isEqualTo("hello");
+ }
+
+ @Test
+ void adaptToRowReadsPojoGettersBySchemaName() {
+ // Schema column order intentionally differs from the POJO declaration
order.
+ RowTypeInfo rowType =
+ rowType(
+ Schema.newBuilder()
+ .column("comment", DataTypes.STRING())
+ .column("score", DataTypes.INT())
+ .build());
+
+ Row row = OutputTypeUtils.adaptToRow(new ReviewOutput(5, "great"),
rowType);
+
+ assertThat(row.getField(0)).isEqualTo("great");
+ assertThat(row.getField(1)).isEqualTo(5);
+ }
+
+ @Test
+ void adaptToRowReadsPublicFieldsWhenNoGetter() {
+ RowTypeInfo rowType =
+ rowType(
+ Schema.newBuilder()
+ .column("name", DataTypes.STRING())
+ .column("value", DataTypes.DOUBLE())
+ .build());
+
+ Row row = OutputTypeUtils.adaptToRow(new PublicFields("x", 1.5),
rowType);
+
+ assertThat(row.getField(0)).isEqualTo("x");
+ assertThat(row.getField(1)).isEqualTo(1.5);
+ }
+
+ @Test
+ void adaptToRowPrefersPojoPropertyOverScalarForSingleColumn() {
+ RowTypeInfo rowType = rowType(Schema.newBuilder().column("score",
DataTypes.INT()).build());
+
+ Row row = OutputTypeUtils.adaptToRow(new ReviewOutput(7, "c"),
rowType);
+
+ assertThat(row.getField(0)).isEqualTo(7);
+ }
+
+ @Test
+ void adaptToRowReadsMapBySchemaName() {
+ RowTypeInfo rowType =
+ rowType(
+ Schema.newBuilder()
+ .column("a", DataTypes.STRING())
+ .column("b", DataTypes.INT())
+ .build());
+ Map<String, Object> map = new LinkedHashMap<>();
+ map.put("b", 7);
+ map.put("a", "v");
+
+ Row row = OutputTypeUtils.adaptToRow(map, rowType);
+
+ assertThat(row.getField(0)).isEqualTo("v");
+ assertThat(row.getField(1)).isEqualTo(7);
+ }
+
+ @Test
+ void adaptToRowPassesThroughPositionalRow() {
+ RowTypeInfo rowType =
+ rowType(
+ Schema.newBuilder()
+ .column("a", DataTypes.STRING())
+ .column("b", DataTypes.INT())
+ .build());
+
+ Row row = OutputTypeUtils.adaptToRow(Row.of("v", 7), rowType);
+
+ assertThat(row.getField(0)).isEqualTo("v");
+ assertThat(row.getField(1)).isEqualTo(7);
+ }
+
+ @Test
+ void adaptToRowProjectsNamedRowBySchemaName() {
+ RowTypeInfo rowType =
+ rowType(
+ Schema.newBuilder()
+ .column("a", DataTypes.STRING())
+ .column("b", DataTypes.INT())
+ .build());
+ Row named = Row.withNames();
+ named.setField("b", 7);
+ named.setField("a", "v");
+
+ Row row = OutputTypeUtils.adaptToRow(named, rowType);
+
+ assertThat(row.getField(0)).isEqualTo("v");
+ assertThat(row.getField(1)).isEqualTo(7);
+ }
+
+ @Test
+ void adaptToRowRejectsRowWithWrongArity() {
+ RowTypeInfo rowType =
+ rowType(
+ Schema.newBuilder()
+ .column("a", DataTypes.STRING())
+ .column("b", DataTypes.INT())
+ .build());
+
+ assertThatThrownBy(() ->
OutputTypeUtils.adaptToRow(Row.of("only-one"), rowType))
+ .isInstanceOf(IllegalArgumentException.class)
+ .hasMessageContaining("declares 2 columns");
+ }
+
+ @Test
+ void adaptToRowRejectsPojoMissingColumn() {
+ RowTypeInfo rowType =
+ rowType(
+ Schema.newBuilder()
+ .column("score", DataTypes.INT())
+ .column("missing", DataTypes.STRING())
+ .build());
+
+ assertThatThrownBy(() -> OutputTypeUtils.adaptToRow(new
ReviewOutput(1, "c"), rowType))
+ .isInstanceOf(IllegalArgumentException.class)
+ .hasMessageContaining("missing");
+ }
+
+ @Test
+ void adaptToRowRejectsScalarForMultiColumnSchema() {
+ RowTypeInfo rowType =
+ rowType(
+ Schema.newBuilder()
+ .column("a", DataTypes.STRING())
+ .column("b", DataTypes.STRING())
+ .build());
+
+ assertThatThrownBy(() -> OutputTypeUtils.adaptToRow("scalar", rowType))
+ .isInstanceOf(IllegalArgumentException.class)
+ .hasMessageContaining("Cannot adapt an output of type");
+ }
+}
diff --git
a/runtime/src/test/java/org/apache/flink/agents/runtime/env/TypedOutputTerminalTest.java
b/runtime/src/test/java/org/apache/flink/agents/runtime/env/TypedOutputTerminalTest.java
new file mode 100644
index 000000000..e99c808a5
--- /dev/null
+++
b/runtime/src/test/java/org/apache/flink/agents/runtime/env/TypedOutputTerminalTest.java
@@ -0,0 +1,176 @@
+/*
+ * 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.agents.runtime.env;
+
+import org.apache.flink.agents.api.AgentBuilder;
+import org.apache.flink.agents.api.AgentsExecutionEnvironment;
+import org.apache.flink.agents.api.Event;
+import org.apache.flink.agents.api.EventType;
+import org.apache.flink.agents.api.agents.Agent;
+import org.apache.flink.agents.api.annotation.Action;
+import org.apache.flink.agents.api.context.RunnerContext;
+import org.apache.flink.api.common.typeinfo.TypeInformation;
+import org.apache.flink.streaming.api.datastream.DataStream;
+import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
+import org.apache.flink.table.api.Table;
+import org.junit.jupiter.api.Test;
+
+import java.util.List;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+/**
+ * Tests for the typed output terminals of the remote {@link AgentBuilder}:
the {@code
+ * toDataStream(TypeInformation)} / {@code toDataStream(Class)} and {@code
toTable(TypeInformation)}
+ * / {@code toTable(Class)} overloads.
+ *
+ * <p>A typed terminal must attach its declared type through a downstream
conversion operator rather
+ * than retyping the shared, cached raw stream. Because the type is applied
downstream, several
+ * typed terminals of different types can coexist on one execution while the
unrestricted {@code
+ * toDataStream()} view stays available and unchanged. These tests inspect the
constructed graph and
+ * never execute a job.
+ */
+class TypedOutputTerminalTest {
+
+ /** Minimal agent so a plan can be built; the graph is inspected without
executing it. */
+ public static class EchoAgent extends Agent {
+ @Action(EventType.InputEvent)
+ public static void handle(Event event, RunnerContext ctx) {
+ // Never runs: the terminal wiring is asserted at
graph-construction time.
+ }
+ }
+
+ /** Minimal POJO so a typed table terminal derives one physical column per
field. */
+ public static class Score {
+ public int value;
+ public String label;
+
+ public Score() {}
+
+ public Score(int value, String label) {
+ this.value = value;
+ this.label = label;
+ }
+ }
+
+ private static AgentBuilder appliedBuilder() {
+ StreamExecutionEnvironment env =
StreamExecutionEnvironment.getExecutionEnvironment();
+ env.setParallelism(1);
+ AgentsExecutionEnvironment agentsEnv =
+ AgentsExecutionEnvironment.getExecutionEnvironment(env);
+ DataStream<String> inputStream = env.fromData(List.of("a", "b", "c"));
+ return agentsEnv.fromDataStream(inputStream).apply(new EchoAgent());
+ }
+
+ @Test
+ void rawDataStreamIsCachedAcrossCalls() {
+ AgentBuilder builder = appliedBuilder();
+
+ DataStream<Object> raw = builder.toDataStream();
+
+ // The agent operator is connected once; repeated calls reuse the same
cached stream.
+ assertThat(builder.toDataStream()).isSameAs(raw);
+ }
+
+ @Test
+ void typedTerminalDoesNotWeldTypeOntoSharedRawStream() {
+ AgentBuilder builder = appliedBuilder();
+
+ DataStream<Object> raw = builder.toDataStream();
+ TypeInformation<?> rawTypeBefore = raw.getType();
+
+ DataStream<Integer> typed =
builder.toDataStream(TypeInformation.of(Integer.class));
+
+ // The declared type is attached to the typed stream...
+
assertThat(typed.getType()).isEqualTo(TypeInformation.of(Integer.class));
+ // ...and is not welded onto the shared raw stream, which keeps its
original type. Retyping
+ // the cached stream in place (for example via returns() on it) would
change raw.getType().
+ assertThat(raw.getType()).isEqualTo(rawTypeBefore);
+
assertThat(raw.getType()).isNotEqualTo(TypeInformation.of(Integer.class));
+ // The unrestricted view is still the same cached stream after the
typed terminal was built.
+ assertThat(builder.toDataStream()).isSameAs(raw);
+ }
+
+ @Test
+ void classOverloadDerivesTheSameTypeAsTypeInformation() {
+ AgentBuilder builder = appliedBuilder();
+
+ DataStream<Object> raw = builder.toDataStream();
+ TypeInformation<?> rawTypeBefore = raw.getType();
+
+ // The Class overload is a convenience that derives the
TypeInformation internally.
+ DataStream<Integer> typed = builder.toDataStream(Integer.class);
+
+
assertThat(typed.getType()).isEqualTo(TypeInformation.of(Integer.class));
+ // It carries the same anti-welding guarantee as the explicit
TypeInformation overload.
+ assertThat(raw.getType()).isEqualTo(rawTypeBefore);
+ }
+
+ @Test
+ void multipleTypedTerminalsOfDifferentTypesCoexist() {
+ AgentBuilder builder = appliedBuilder();
+
+ DataStream<Object> raw = builder.toDataStream();
+ TypeInformation<?> rawTypeBefore = raw.getType();
+
+ // Two independent typed terminals of different types on the same
execution. Each layers its
+ // own conversion operator on the shared raw stream, so they do not
conflict with one
+ // another.
+ DataStream<Integer> numbers =
builder.toDataStream(TypeInformation.of(Integer.class));
+ DataStream<String> names =
builder.toDataStream(TypeInformation.of(String.class));
+
+
assertThat(numbers.getType()).isEqualTo(TypeInformation.of(Integer.class));
+
assertThat(names.getType()).isEqualTo(TypeInformation.of(String.class));
+ // Neither terminal welds its type onto the shared raw stream, which
stays unrestricted.
+ assertThat(raw.getType()).isEqualTo(rawTypeBefore);
+ assertThat(builder.toDataStream()).isSameAs(raw);
+ }
+
+ @Test
+ void typedTableDerivesColumnsFromTypeInformation() {
+ AgentBuilder builder = appliedBuilder();
+
+ DataStream<Object> raw = builder.toDataStream();
+ TypeInformation<?> rawTypeBefore = raw.getType();
+
+ Table table = builder.toTable(TypeInformation.of(Score.class));
+
+ // The physical schema is derived from the declared type: each POJO
field becomes a column.
+ assertThat(table.getResolvedSchema().getColumnNames())
+ .containsExactlyInAnyOrder("value", "label");
+ // The type is attached by a downstream operator, so the shared raw
stream stays
+ // unrestricted.
+ assertThat(raw.getType()).isEqualTo(rawTypeBefore);
+ }
+
+ @Test
+ void typedTableClassOverloadDerivesTheSameColumns() {
+ AgentBuilder builder = appliedBuilder();
+
+ DataStream<Object> raw = builder.toDataStream();
+ TypeInformation<?> rawTypeBefore = raw.getType();
+
+ // The Class overload is a convenience that derives the
TypeInformation internally.
+ Table table = builder.toTable(Score.class);
+
+ assertThat(table.getResolvedSchema().getColumnNames())
+ .containsExactlyInAnyOrder("value", "label");
+ // It carries the same anti-welding guarantee as the explicit
TypeInformation overload.
+ assertThat(raw.getType()).isEqualTo(rawTypeBefore);
+ }
+}
