From ccebaea33022bb1b3dcdbd0ee41583e0de379c7c Mon Sep 17 00:00:00 2001
From: Mika Naylor <mnaylor@confluent.io>
Date: Mon, 24 Aug 2026 17:36:36 +0200
Subject: [PATCH] [FLINK-40466][table] Refactor defining table arguments in PTF
 Test Harness for more complex table arguments

---
 .../docs/dev/table/functions/ptfs.md          | 111 ++-
 docs/content/docs/dev/table/functions/ptfs.md | 111 ++-
 .../ProcessTableFunctionTestHarness.java      | 145 ++--
 .../ProcessTableFunctionTestHarnessTest.java  | 659 ++++++++++++------
 4 files changed, 708 insertions(+), 318 deletions(-)

diff --git a/docs/content.zh/docs/dev/table/functions/ptfs.md b/docs/content.zh/docs/dev/table/functions/ptfs.md
index 96bd2dca3534a..067573e5979d2 100644
--- a/docs/content.zh/docs/dev/table/functions/ptfs.md
+++ b/docs/content.zh/docs/dev/table/functions/ptfs.md
@@ -2056,6 +2056,7 @@ import org.apache.flink.table.annotation.*;
 import org.apache.flink.table.api.DataTypes;
 import org.apache.flink.table.functions.ProcessTableFunction;
 import org.apache.flink.table.runtime.functions.ProcessTableFunctionTestHarness;
+import org.apache.flink.table.runtime.functions.ProcessTableFunctionTestHarness.TableArgument;
 import org.apache.flink.types.Row;
 import org.junit.jupiter.api.Test;
 
@@ -2077,7 +2078,10 @@ public class DoublePTF extends ProcessTableFunction<Row> {
 void testDoublePTF() throws Exception {
   try (ProcessTableFunctionTestHarness<Row> harness =
     ProcessTableFunctionTestHarness.ofClass(DoublePTF.class)
-    .withTableArgument("input", DataTypes.of("ROW<value INT>"))
+    .withTableArgument(
+        TableArgument.forArgument("input")
+            .type(DataTypes.of("ROW<value INT>"))
+            .build())
     .build()) {
 
     harness.processElement(Row.of(5));
@@ -2099,7 +2103,7 @@ void testDoublePTF() throws Exception {
 
 #### Testing Row-Semantic Tables
 
-Use `.withTableArgument()` to configure the input table schema:
+Use `TableArgument.forArgument(name).type(...)` to configure the input table schema:
 
 {{< tabs "row-semantic" >}}
 {{< tab "Java" >}}
@@ -2114,7 +2118,10 @@ public class PassthroughPTF extends ProcessTableFunction<Integer> {
 void testPassthrough() throws Exception {
   try (ProcessTableFunctionTestHarness<Integer> harness =
     ProcessTableFunctionTestHarness.ofClass(PassthroughPTF.class)
-    .withTableArgument("input", DataTypes.of("ROW<value INT>"))
+    .withTableArgument(
+        TableArgument.forArgument("input")
+            .type(DataTypes.of("ROW<value INT>"))
+            .build())
     .build()) {
 
     harness.processElement(Row.of(42));
@@ -2130,7 +2137,7 @@ void testPassthrough() throws Exception {
 
 #### Testing Set-Semantic Tables with Partitioning
 
-For `SET_SEMANTIC_TABLE`, use `.withPartitionBy()` to configure partition columns:
+For `SET_SEMANTIC_TABLE`, use `TableArgument.Builder#partitionBy()` to configure partition columns:
 
 {{< tabs "set-semantic" >}}
 {{< tab "Java" >}}
@@ -2147,8 +2154,11 @@ public class PartitionedPTF extends ProcessTableFunction<Row> {
 void testPartitionedPTF() throws Exception {
   try (ProcessTableFunctionTestHarness<Row> harness =
     ProcessTableFunctionTestHarness.ofClass(PartitionedPTF.class)
-    .withTableArgument("input", DataTypes.of("ROW<key STRING, value INT>"))
-    .withPartitionBy("input", "key")
+    .withTableArgument(
+        TableArgument.forArgument("input")
+            .type(DataTypes.of("ROW<key STRING, value INT>"))
+            .partitionBy("key")
+            .build())
     .build()) {
 
     harness.processElement(Row.of("A", 10));
@@ -2188,10 +2198,16 @@ public class JoinPTF extends ProcessTableFunction<Row> {
 void testMultiTable() throws Exception {
   try (ProcessTableFunctionTestHarness<Row> harness =
     ProcessTableFunctionTestHarness.ofClass(JoinPTF.class)
-    .withTableArgument("left", DataTypes.of("ROW<id INT, name STRING>"))
-    .withPartitionBy("left", "id")
-    .withTableArgument("right", DataTypes.of("ROW<id INT, city STRING>"))
-    .withPartitionBy("right", "id")
+    .withTableArgument(
+        TableArgument.forArgument("left")
+            .type(DataTypes.of("ROW<id INT, name STRING>"))
+            .partitionBy("id")
+            .build())
+    .withTableArgument(
+        TableArgument.forArgument("right")
+            .type(DataTypes.of("ROW<id INT, city STRING>"))
+            .partitionBy("id")
+            .build())
     .build()) {
 
     // Use processElementForTable() to target specific tables
@@ -2230,7 +2246,10 @@ public class FilterPTF extends ProcessTableFunction<Row> {
 void testFilter() throws Exception {
   try (ProcessTableFunctionTestHarness<Row> harness =
     ProcessTableFunctionTestHarness.ofClass(FilterPTF.class)
-    .withTableArgument("input", DataTypes.of("ROW<value INT>"))
+    .withTableArgument(
+        TableArgument.forArgument("input")
+            .type(DataTypes.of("ROW<value INT>"))
+            .build())
     .withScalarArgument("threshold", 50) // Configure scalar value
     .build()) {
 
@@ -2319,8 +2338,11 @@ public class StatefulPTF extends ProcessTableFunction<Row> {
 void testWithState() throws Exception {
   try (ProcessTableFunctionTestHarness<Row> harness =
     ProcessTableFunctionTestHarness.ofClass(StatefulPTF.class)
-    .withTableArgument("input", DataTypes.of("ROW<name STRING, value INT>"))
-    .withPartitionBy("input", "name")
+    .withTableArgument(
+        TableArgument.forArgument("input")
+            .type(DataTypes.of("ROW<name STRING, value INT>"))
+            .partitionBy("name")
+            .build())
     .build()) {
 
     harness.processElement(Row.of("Alice", 10));
@@ -2362,8 +2384,11 @@ void testWithInitialState() throws Exception {
 
   try (ProcessTableFunctionTestHarness<Row> harness =
     ProcessTableFunctionTestHarness.ofClass(StatefulPTF.class)
-    .withTableArgument("input", DataTypes.of("ROW<name STRING, value INT>"))
-    .withPartitionBy("input", "name")
+    .withTableArgument(
+        TableArgument.forArgument("input")
+            .type(DataTypes.of("ROW<name STRING, value INT>"))
+            .partitionBy("name")
+            .build())
     // Initial state is set per partition key
     .withInitialStateForKey("valueState", Row.of("Alice"), initialValue)
     .withInitialStateForKey("rowState", Row.of("Alice"), initialRow)
@@ -2391,8 +2416,11 @@ inspect state during tests:
 void testStateIntrospection() throws Exception {
   try (ProcessTableFunctionTestHarness<Row> harness =
     ProcessTableFunctionTestHarness.ofClass(StatefulPTF.class)
-    .withTableArgument("input", DataTypes.of("ROW<name STRING, value INT>"))
-    .withPartitionBy("input", "name")
+    .withTableArgument(
+        TableArgument.forArgument("input")
+            .type(DataTypes.of("ROW<name STRING, value INT>"))
+            .partitionBy("name")
+            .build())
     .build()) {
 
     harness.processElement(Row.of("Alice", 10));
@@ -2439,8 +2467,11 @@ modify state during tests:
 void testStateMutation() throws Exception {
   try (ProcessTableFunctionTestHarness<Row> harness =
     ProcessTableFunctionTestHarness.ofClass(StatefulPTF.class)
-    .withTableArgument("input", DataTypes.of("ROW<name STRING, value INT>"))
-    .withPartitionBy("input", "name")
+    .withTableArgument(
+        TableArgument.forArgument("input")
+            .type(DataTypes.of("ROW<name STRING, value INT>"))
+            .partitionBy("name")
+            .build())
     .build()) {
 
     harness.processElement(Row.of("Alice", 10));
@@ -2496,9 +2527,11 @@ public class TimerPTF extends ProcessTableFunction<Row> {
 void testTimerRegistrationAndFiring() throws Exception {
   try (ProcessTableFunctionTestHarness<Row> harness =
       ProcessTableFunctionTestHarness.ofClass(TimerPTF.class)
-          .withTableArgument("input",
-              DataTypes.of("ROW<partition STRING, name STRING, ts TIMESTAMP(3)>"))
-          .withPartitionBy("input", "partition")
+          .withTableArgument(
+              TableArgument.forArgument("input")
+                  .type(DataTypes.of("ROW<partition STRING, name STRING, ts TIMESTAMP(3)>"))
+                  .partitionBy("partition")
+                  .build())
           .withOnTimeColumn("ts")
           .build()) {
 
@@ -2562,9 +2595,11 @@ public class TimerWithStatePTF extends ProcessTableFunction<Row> {
 void testTimerWithState() throws Exception {
   try (ProcessTableFunctionTestHarness<Row> harness =
       ProcessTableFunctionTestHarness.ofClass(TimerWithStatePTF.class)
-          .withTableArgument("input",
-              DataTypes.of("ROW<partition STRING, ts TIMESTAMP(3)>"))
-          .withPartitionBy("input", "partition")
+          .withTableArgument(
+              TableArgument.forArgument("input")
+                  .type(DataTypes.of("ROW<partition STRING, ts TIMESTAMP(3)>"))
+                  .partitionBy("partition")
+                  .build())
           .withOnTimeColumn("ts")
           .build()) {
 
@@ -2587,7 +2622,8 @@ void testTimerWithState() throws Exception {
 
 #### Optional Partitioning
 
-For PTFs with `OPTIONAL_PARTITION_BY`, you can omit `withPartitionBy()` during harness setup. The
+For PTFs with `OPTIONAL_PARTITION_BY`, you can omit `.partitionBy(...)` on `TableArgument.Builder`
+during harness setup. The
 harness executes the function as if it had a parallelism of 1, with the default `Row.of()` key,
 so all data is routed through the same function instance. Use `Row.of()` to access state:
 
@@ -2614,7 +2650,10 @@ public class GlobalCountPTF extends ProcessTableFunction<Row> {
 void testOptionalPartitionWithoutPartitionBy() throws Exception {
   try (ProcessTableFunctionTestHarness<Row> harness =
       ProcessTableFunctionTestHarness.ofClass(GlobalCountPTF.class)
-      .withTableArgument("input", DataTypes.of("ROW<key STRING, value INT>"))
+      .withTableArgument(
+          TableArgument.forArgument("input")
+              .type(DataTypes.of("ROW<key STRING, value INT>"))
+              .build())
       .build()) {
 
       harness.processElement(Row.of("A", 10));
@@ -2652,7 +2691,10 @@ public static class DoublePTF extends ProcessTableFunction<Row> {
 void testBuilderType() throws Exception {
   try (ProcessTableFunctionTestHarness<Row> harness =
     ProcessTableFunctionTestHarness.ofClass(DoublePTF.class)
-    .withTableArgument("input", DataTypes.of("ROW<value INT>"))
+    .withTableArgument(
+        TableArgument.forArgument("input")
+            .type(DataTypes.of("ROW<value INT>"))
+            .build())
     .build()) {
 
     harness.processElement(Row.of(5));
@@ -2684,7 +2726,10 @@ public class CustomerPTF extends ProcessTableFunction<Customer> {
 void testPOJO() throws Exception {
   try (ProcessTableFunctionTestHarness<Customer> harness =
       ProcessTableFunctionTestHarness.ofClass(CustomerPTF.class)
-          .withTableArgument("c", DataTypes.of(Customer.class))
+          .withTableArgument(
+              TableArgument.forArgument("c")
+                  .type(DataTypes.of(Customer.class))
+                  .build())
           .build()) {
 
     harness.processElement(Row.of(30, "Alice"));
@@ -2725,9 +2770,11 @@ public class PartitionedAtomicPTF extends ProcessTableFunction<Integer> {
 void testAtomicOutputFunctionOutput() throws Exception {
   try (ProcessTableFunctionTestHarness<Integer> harness =
       ProcessTableFunctionTestHarness.ofClass(PartitionedAtomicPTF.class)
-          .withTableArgument("input",
-              DataTypes.of("ROW<key STRING, value INT, ts TIMESTAMP(3)>"))
-          .withPartitionBy("input", "key")
+          .withTableArgument(
+              TableArgument.forArgument("input")
+                  .type(DataTypes.of("ROW<key STRING, value INT, ts TIMESTAMP(3)>"))
+                  .partitionBy("key")
+                  .build())
           .withOnTimeColumn("ts")
           .build()) {
 
diff --git a/docs/content/docs/dev/table/functions/ptfs.md b/docs/content/docs/dev/table/functions/ptfs.md
index 6fd3efd8fa957..65e16e8b01a35 100644
--- a/docs/content/docs/dev/table/functions/ptfs.md
+++ b/docs/content/docs/dev/table/functions/ptfs.md
@@ -2059,6 +2059,7 @@ import org.apache.flink.table.annotation.*;
 import org.apache.flink.table.api.DataTypes;
 import org.apache.flink.table.functions.ProcessTableFunction;
 import org.apache.flink.table.runtime.functions.ProcessTableFunctionTestHarness;
+import org.apache.flink.table.runtime.functions.ProcessTableFunctionTestHarness.TableArgument;
 import org.apache.flink.types.Row;
 import org.junit.jupiter.api.Test;
 
@@ -2080,7 +2081,10 @@ public class DoublePTF extends ProcessTableFunction<Row> {
 void testDoublePTF() throws Exception {
   try (ProcessTableFunctionTestHarness<Row> harness =
     ProcessTableFunctionTestHarness.ofClass(DoublePTF.class)
-    .withTableArgument("input", DataTypes.of("ROW<value INT>"))
+    .withTableArgument(
+        TableArgument.forArgument("input")
+            .type(DataTypes.of("ROW<value INT>"))
+            .build())
     .build()) {
 
     harness.processElement(Row.of(5));
@@ -2102,7 +2106,7 @@ void testDoublePTF() throws Exception {
 
 #### Testing Row-Semantic Tables
 
-Use `.withTableArgument()` to configure the input table schema:
+Use `TableArgument.forArgument(name).type(...)` to configure the input table schema:
 
 {{< tabs "row-semantic" >}}
 {{< tab "Java" >}}
@@ -2117,7 +2121,10 @@ public class PassthroughPTF extends ProcessTableFunction<Integer> {
 void testPassthrough() throws Exception {
   try (ProcessTableFunctionTestHarness<Integer> harness =
     ProcessTableFunctionTestHarness.ofClass(PassthroughPTF.class)
-    .withTableArgument("input", DataTypes.of("ROW<value INT>"))
+    .withTableArgument(
+        TableArgument.forArgument("input")
+            .type(DataTypes.of("ROW<value INT>"))
+            .build())
     .build()) {
 
     harness.processElement(Row.of(42));
@@ -2133,7 +2140,7 @@ void testPassthrough() throws Exception {
 
 #### Testing Set-Semantic Tables with Partitioning
 
-For `SET_SEMANTIC_TABLE`, use `.withPartitionBy()` to configure partition columns:
+For `SET_SEMANTIC_TABLE`, use `TableArgument.Builder#partitionBy()` to configure partition columns:
 
 {{< tabs "set-semantic" >}}
 {{< tab "Java" >}}
@@ -2150,8 +2157,11 @@ public class PartitionedPTF extends ProcessTableFunction<Row> {
 void testPartitionedPTF() throws Exception {
   try (ProcessTableFunctionTestHarness<Row> harness =
     ProcessTableFunctionTestHarness.ofClass(PartitionedPTF.class)
-    .withTableArgument("input", DataTypes.of("ROW<key STRING, value INT>"))
-    .withPartitionBy("input", "key")
+    .withTableArgument(
+        TableArgument.forArgument("input")
+            .type(DataTypes.of("ROW<key STRING, value INT>"))
+            .partitionBy("key")
+            .build())
     .build()) {
 
     harness.processElement(Row.of("A", 10));
@@ -2191,10 +2201,16 @@ public class JoinPTF extends ProcessTableFunction<Row> {
 void testMultiTable() throws Exception {
   try (ProcessTableFunctionTestHarness<Row> harness =
     ProcessTableFunctionTestHarness.ofClass(JoinPTF.class)
-    .withTableArgument("left", DataTypes.of("ROW<id INT, name STRING>"))
-    .withPartitionBy("left", "id")
-    .withTableArgument("right", DataTypes.of("ROW<id INT, city STRING>"))
-    .withPartitionBy("right", "id")
+    .withTableArgument(
+        TableArgument.forArgument("left")
+            .type(DataTypes.of("ROW<id INT, name STRING>"))
+            .partitionBy("id")
+            .build())
+    .withTableArgument(
+        TableArgument.forArgument("right")
+            .type(DataTypes.of("ROW<id INT, city STRING>"))
+            .partitionBy("id")
+            .build())
     .build()) {
 
     // Use processElementForTable() to target specific tables
@@ -2233,7 +2249,10 @@ public class FilterPTF extends ProcessTableFunction<Row> {
 void testFilter() throws Exception {
   try (ProcessTableFunctionTestHarness<Row> harness =
     ProcessTableFunctionTestHarness.ofClass(FilterPTF.class)
-    .withTableArgument("input", DataTypes.of("ROW<value INT>"))
+    .withTableArgument(
+        TableArgument.forArgument("input")
+            .type(DataTypes.of("ROW<value INT>"))
+            .build())
     .withScalarArgument("threshold", 50) // Configure scalar value
     .build()) {
 
@@ -2322,8 +2341,11 @@ public class StatefulPTF extends ProcessTableFunction<Row> {
 void testWithState() throws Exception {
   try (ProcessTableFunctionTestHarness<Row> harness =
     ProcessTableFunctionTestHarness.ofClass(StatefulPTF.class)
-    .withTableArgument("input", DataTypes.of("ROW<name STRING, value INT>"))
-    .withPartitionBy("input", "name")
+    .withTableArgument(
+        TableArgument.forArgument("input")
+            .type(DataTypes.of("ROW<name STRING, value INT>"))
+            .partitionBy("name")
+            .build())
     .build()) {
 
     harness.processElement(Row.of("Alice", 10));
@@ -2365,8 +2387,11 @@ void testWithInitialState() throws Exception {
 
   try (ProcessTableFunctionTestHarness<Row> harness =
     ProcessTableFunctionTestHarness.ofClass(StatefulPTF.class)
-    .withTableArgument("input", DataTypes.of("ROW<name STRING, value INT>"))
-    .withPartitionBy("input", "name")
+    .withTableArgument(
+        TableArgument.forArgument("input")
+            .type(DataTypes.of("ROW<name STRING, value INT>"))
+            .partitionBy("name")
+            .build())
     // Initial state is set per partition key
     .withInitialStateForKey("valueState", Row.of("Alice"), initialValue)
     .withInitialStateForKey("rowState", Row.of("Alice"), initialRow)
@@ -2394,8 +2419,11 @@ inspect state during tests:
 void testStateIntrospection() throws Exception {
   try (ProcessTableFunctionTestHarness<Row> harness =
     ProcessTableFunctionTestHarness.ofClass(StatefulPTF.class)
-    .withTableArgument("input", DataTypes.of("ROW<name STRING, value INT>"))
-    .withPartitionBy("input", "name")
+    .withTableArgument(
+        TableArgument.forArgument("input")
+            .type(DataTypes.of("ROW<name STRING, value INT>"))
+            .partitionBy("name")
+            .build())
     .build()) {
 
     harness.processElement(Row.of("Alice", 10));
@@ -2442,8 +2470,11 @@ modify state during tests:
 void testStateMutation() throws Exception {
   try (ProcessTableFunctionTestHarness<Row> harness =
     ProcessTableFunctionTestHarness.ofClass(StatefulPTF.class)
-    .withTableArgument("input", DataTypes.of("ROW<name STRING, value INT>"))
-    .withPartitionBy("input", "name")
+    .withTableArgument(
+        TableArgument.forArgument("input")
+            .type(DataTypes.of("ROW<name STRING, value INT>"))
+            .partitionBy("name")
+            .build())
     .build()) {
 
     harness.processElement(Row.of("Alice", 10));
@@ -2499,9 +2530,11 @@ public class TimerPTF extends ProcessTableFunction<Row> {
 void testTimerRegistrationAndFiring() throws Exception {
   try (ProcessTableFunctionTestHarness<Row> harness =
       ProcessTableFunctionTestHarness.ofClass(TimerPTF.class)
-          .withTableArgument("input",
-              DataTypes.of("ROW<partition STRING, name STRING, ts TIMESTAMP(3)>"))
-          .withPartitionBy("input", "partition")
+          .withTableArgument(
+              TableArgument.forArgument("input")
+                  .type(DataTypes.of("ROW<partition STRING, name STRING, ts TIMESTAMP(3)>"))
+                  .partitionBy("partition")
+                  .build())
           .withOnTimeColumn("ts")
           .build()) {
 
@@ -2565,9 +2598,11 @@ public class TimerWithStatePTF extends ProcessTableFunction<Row> {
 void testTimerWithState() throws Exception {
   try (ProcessTableFunctionTestHarness<Row> harness =
       ProcessTableFunctionTestHarness.ofClass(TimerWithStatePTF.class)
-          .withTableArgument("input",
-              DataTypes.of("ROW<partition STRING, ts TIMESTAMP(3)>"))
-          .withPartitionBy("input", "partition")
+          .withTableArgument(
+              TableArgument.forArgument("input")
+                  .type(DataTypes.of("ROW<partition STRING, ts TIMESTAMP(3)>"))
+                  .partitionBy("partition")
+                  .build())
           .withOnTimeColumn("ts")
           .build()) {
 
@@ -2590,7 +2625,8 @@ void testTimerWithState() throws Exception {
 
 #### Optional Partitioning
 
-For PTFs with `OPTIONAL_PARTITION_BY`, you can omit `withPartitionBy()` during harness setup. The
+For PTFs with `OPTIONAL_PARTITION_BY`, you can omit `.partitionBy(...)` on `TableArgument.Builder`
+during harness setup. The
 harness executes the function as if it had a parallelism of 1, with the default `Row.of()` key,
 so all data is routed through the same function instance. Use `Row.of()` to access state:
 
@@ -2617,7 +2653,10 @@ public class GlobalCountPTF extends ProcessTableFunction<Row> {
 void testOptionalPartitionWithoutPartitionBy() throws Exception {
   try (ProcessTableFunctionTestHarness<Row> harness =
       ProcessTableFunctionTestHarness.ofClass(GlobalCountPTF.class)
-      .withTableArgument("input", DataTypes.of("ROW<key STRING, value INT>"))
+      .withTableArgument(
+          TableArgument.forArgument("input")
+              .type(DataTypes.of("ROW<key STRING, value INT>"))
+              .build())
       .build()) {
 
       harness.processElement(Row.of("A", 10));
@@ -2655,7 +2694,10 @@ public static class DoublePTF extends ProcessTableFunction<Row> {
 void testBuilderType() throws Exception {
   try (ProcessTableFunctionTestHarness<Row> harness =
     ProcessTableFunctionTestHarness.ofClass(DoublePTF.class)
-    .withTableArgument("input", DataTypes.of("ROW<value INT>"))
+    .withTableArgument(
+        TableArgument.forArgument("input")
+            .type(DataTypes.of("ROW<value INT>"))
+            .build())
     .build()) {
 
     harness.processElement(Row.of(5));
@@ -2687,7 +2729,10 @@ public class CustomerPTF extends ProcessTableFunction<Customer> {
 void testPOJO() throws Exception {
   try (ProcessTableFunctionTestHarness<Customer> harness =
       ProcessTableFunctionTestHarness.ofClass(CustomerPTF.class)
-          .withTableArgument("c", DataTypes.of(Customer.class))
+          .withTableArgument(
+              TableArgument.forArgument("c")
+                  .type(DataTypes.of(Customer.class))
+                  .build())
           .build()) {
 
     harness.processElement(Row.of(30, "Alice"));
@@ -2728,9 +2773,11 @@ public class PartitionedAtomicPTF extends ProcessTableFunction<Integer> {
 void testAtomicOutputFunctionOutput() throws Exception {
   try (ProcessTableFunctionTestHarness<Integer> harness =
       ProcessTableFunctionTestHarness.ofClass(PartitionedAtomicPTF.class)
-          .withTableArgument("input",
-              DataTypes.of("ROW<key STRING, value INT, ts TIMESTAMP(3)>"))
-          .withPartitionBy("input", "key")
+          .withTableArgument(
+              TableArgument.forArgument("input")
+                  .type(DataTypes.of("ROW<key STRING, value INT, ts TIMESTAMP(3)>"))
+                  .partitionBy("key")
+                  .build())
           .withOnTimeColumn("ts")
           .build()) {
 
diff --git a/flink-table/flink-table-test-utils/src/main/java/org/apache/flink/table/runtime/functions/ProcessTableFunctionTestHarness.java b/flink-table/flink-table-test-utils/src/main/java/org/apache/flink/table/runtime/functions/ProcessTableFunctionTestHarness.java
index 00c4d8213c155..d08b6a9701550 100644
--- a/flink-table/flink-table-test-utils/src/main/java/org/apache/flink/table/runtime/functions/ProcessTableFunctionTestHarness.java
+++ b/flink-table/flink-table-test-utils/src/main/java/org/apache/flink/table/runtime/functions/ProcessTableFunctionTestHarness.java
@@ -93,7 +93,10 @@
  * <pre>{@code
  * ProcessTableFunctionTestHarness<Row> harness =
  *     ProcessTableFunctionTestHarness.ofClass(MyPTF.class)
- *         .withTableArgument("input", DataTypes.of("ROW<id INT, name STRING>"))
+ *         .withTableArgument(
+ *             TableArgument.forArgument("input")
+ *                 .type(DataTypes.of("ROW<id INT, name STRING>"))
+ *                 .build())
  *         .withScalarArgument("threshold", 100)
  *         .build();
  *
@@ -1035,6 +1038,72 @@ private static DataType toRowDataType(StructuredType structuredType) {
                 new RowType(structuredType.isNullable(), rowFields));
     }
 
+    /**
+     * Configuration for a single named table argument.
+     *
+     * @see ProcessTableFunctionTestHarness.Builder#withTableArgument(TableArgument)
+     */
+    @PublicEvolving
+    public static final class TableArgument {
+        private final String name;
+        @Nullable private final AbstractDataType<?> type;
+        @Nullable private final String[] partitionColumns;
+
+        private TableArgument(Builder builder) {
+            this.name = builder.name;
+            this.type = builder.type;
+            this.partitionColumns = builder.partitionColumns;
+        }
+
+        public static Builder forArgument(String argumentName) {
+            return new Builder(argumentName);
+        }
+
+        /** Builder for {@link TableArgument}. */
+        public static final class Builder {
+            private final String name;
+            @Nullable private AbstractDataType<?> type;
+            @Nullable private String[] partitionColumns;
+
+            private Builder(String argumentName) {
+                this.name = checkNotNull(argumentName, "argumentName must not be null");
+            }
+
+            /**
+             * Configures the schema of this table argument.
+             *
+             * <p>Use this for dynamic tables that receive elements during the test. Elements are
+             * provided via {@link ProcessTableFunctionTestHarness#processElement(Row)} or {@link
+             * ProcessTableFunctionTestHarness#processElementForTable(String, Row)}.
+             *
+             * <p>Omit for structured type arguments whose type can be inferred from the PTF's
+             * eval() signature.
+             *
+             * @param dataType The schema/structure of the table
+             */
+            public Builder type(AbstractDataType<?> dataType) {
+                this.type = checkNotNull(dataType, "dataType must not be null");
+                return this;
+            }
+
+            /**
+             * Specifies partition columns for a set semantic table.
+             *
+             * @param columnNames The partition column names
+             */
+            public Builder partitionBy(String... columnNames) {
+                checkNotNull(columnNames, "columnNames must not be null");
+                checkArgument(columnNames.length > 0, "Must specify at least one column");
+                this.partitionColumns = columnNames;
+                return this;
+            }
+
+            public TableArgument build() {
+                return new TableArgument(this);
+            }
+        }
+    }
+
     /**
      * Builder for {@link ProcessTableFunctionTestHarness}.
      *
@@ -1065,40 +1134,22 @@ private void validateArgumentNotYetConfigured(String argumentName) {
         // ---------------------------------------------------------------------
 
         /**
-         * Configures a table argument with its schema (named argument).
-         *
-         * <p>Use this for dynamic tables that receive elements during the test. Elements are
-         * provided via {@link #processElement(Row)} or {@link #processElementForTable(String,
-         * Row)}.
+         * Configures a table argument.
          *
-         * @param argumentName The table argument name
-         * @param dataType The schema/structure of the table
+         * @param tableArgument The table argument configuration, built via {@link
+         *     TableArgument#forArgument(String)}
          */
-        public Builder<OUT> withTableArgument(String argumentName, AbstractDataType<?> dataType) {
-            checkNotNull(argumentName, "argumentName must not be null");
-            checkNotNull(dataType, "dataType must not be null");
+        public Builder<OUT> withTableArgument(TableArgument tableArgument) {
+            checkNotNull(tableArgument, "tableArgument must not be null");
 
-            validateArgumentNotYetConfigured(argumentName);
+            validateArgumentNotYetConfigured(tableArgument.name);
 
-            tableArgs.put(argumentName, new TableArgumentConfiguration(dataType));
-            return this;
-        }
-
-        /**
-         * Configures a table argument without an explicit schema.
-         *
-         * <p>Use this for structured type arguments where the type can be inferred from the PTF's
-         * eval() signature. For Row arguments, use {@link #withTableArgument(String,
-         * AbstractDataType)} with an explicit schema.
-         *
-         * @param argumentName The table argument name
-         */
-        public Builder<OUT> withTableArgument(String argumentName) {
-            checkNotNull(argumentName, "argumentName must not be null");
-
-            validateArgumentNotYetConfigured(argumentName);
-
-            tableArgs.put(argumentName, new TableArgumentConfiguration(null));
+            tableArgs.put(tableArgument.name, new TableArgumentConfiguration(tableArgument.type));
+            if (tableArgument.partitionColumns != null) {
+                partitionConfigs.put(
+                        tableArgument.name,
+                        new PartitionConfiguration(tableArgument.partitionColumns));
+            }
             return this;
         }
 
@@ -1135,32 +1186,6 @@ public Builder<OUT> withInitialStateForKey(
             return this;
         }
 
-        // ---------------------------------------------------------------------
-        // Partitioning
-        // ---------------------------------------------------------------------
-
-        /**
-         * Specifies partition columns for a set semantic table.
-         *
-         * @param argumentName The table argument name
-         * @param columnNames The partition column names
-         * @return This builder
-         */
-        public Builder<OUT> withPartitionBy(String argumentName, String... columnNames) {
-            checkNotNull(argumentName, "argumentName must not be null");
-            checkNotNull(columnNames, "columnNames must not be null");
-            checkArgument(columnNames.length > 0, "Must specify at least one column");
-
-            if (partitionConfigs.containsKey(argumentName)) {
-                throw new IllegalArgumentException(
-                        "Partition config already exists for: " + argumentName);
-            }
-
-            PartitionConfiguration config = new PartitionConfiguration(columnNames);
-            partitionConfigs.put(argumentName, config);
-            return this;
-        }
-
         // ---------------------------------------------------------------------
         // Timer & Watermark Configuration
         // ---------------------------------------------------------------------
@@ -1855,7 +1880,8 @@ private ArgumentInfo buildArgumentInfo(StaticArgument staticArg) {
                         throw new IllegalStateException(
                                 String.format(
                                         "Table argument '%s' requires explicit type configuration. "
-                                                + "Use .withTableArgument(\"%s\", DataTypes.of(\"ROW<...>\")) "
+                                                + "Use .withTableArgument(TableArgument.forArgument(\"%s\")"
+                                                + ".type(DataTypes.of(\"ROW<...>\")).build()) "
                                                 + "to explicitly declare it.",
                                         name, name));
                     }
@@ -1906,7 +1932,8 @@ private String[] extractAndValidatePartitionColumns(
                 throw new IllegalStateException(
                         String.format(
                                 "No partition configuration found for SET_SEMANTIC_TABLE argument '%s'. "
-                                        + "Use withPartitionBy(\"%s\", ...) to configure partitioning.",
+                                        + "Use TableArgument.forArgument(\"%s\").partitionBy(...) to "
+                                        + "configure partitioning.",
                                 name, name));
             }
 
diff --git a/flink-table/flink-table-test-utils/src/test/java/org/apache/flink/table/runtime/functions/ProcessTableFunctionTestHarnessTest.java b/flink-table/flink-table-test-utils/src/test/java/org/apache/flink/table/runtime/functions/ProcessTableFunctionTestHarnessTest.java
index 5041911d051cb..ef6e036db8289 100644
--- a/flink-table/flink-table-test-utils/src/test/java/org/apache/flink/table/runtime/functions/ProcessTableFunctionTestHarnessTest.java
+++ b/flink-table/flink-table-test-utils/src/test/java/org/apache/flink/table/runtime/functions/ProcessTableFunctionTestHarnessTest.java
@@ -30,6 +30,7 @@
 import org.apache.flink.table.connector.ChangelogMode;
 import org.apache.flink.table.functions.ProcessTableFunction;
 import org.apache.flink.table.functions.TableSemantics;
+import org.apache.flink.table.runtime.functions.ProcessTableFunctionTestHarness.TableArgument;
 import org.apache.flink.types.Row;
 import org.apache.flink.types.RowKind;
 
@@ -453,7 +454,10 @@ void testBuilderRejectsDuplicateScalarArguments() {
                         IllegalArgumentException.class,
                         () -> {
                             ProcessTableFunctionTestHarness.ofClass(FilterPTF.class)
-                                    .withTableArgument("input", DataTypes.of("ROW<value INT>"))
+                                    .withTableArgument(
+                                            TableArgument.forArgument("input")
+                                                    .type(DataTypes.of("ROW<value INT>"))
+                                                    .build())
                                     .withScalarArgument("threshold", 50)
                                     .withScalarArgument("threshold", 100);
                         });
@@ -469,9 +473,13 @@ void testBuilderRejectsDuplicateTableArguments() {
                         () -> {
                             ProcessTableFunctionTestHarness.ofClass(MultiTableUnionPTF.class)
                                     .withTableArgument(
-                                            "leftTable", DataTypes.of("ROW<id INT, name STRING>"))
+                                            TableArgument.forArgument("leftTable")
+                                                    .type(DataTypes.of("ROW<id INT, name STRING>"))
+                                                    .build())
                                     .withTableArgument(
-                                            "leftTable", DataTypes.of("ROW<id INT, value INT>"));
+                                            TableArgument.forArgument("leftTable")
+                                                    .type(DataTypes.of("ROW<id INT, value INT>"))
+                                                    .build());
                         });
 
         assertThat(exception.getMessage()).contains("leftTable");
@@ -484,7 +492,10 @@ void testBuilderRejectsMixedDuplicateArguments() {
                         IllegalArgumentException.class,
                         () -> {
                             ProcessTableFunctionTestHarness.ofClass(FilterPTF.class)
-                                    .withTableArgument("input", DataTypes.of("ROW<value INT>"))
+                                    .withTableArgument(
+                                            TableArgument.forArgument("input")
+                                                    .type(DataTypes.of("ROW<value INT>"))
+                                                    .build())
                                     .withScalarArgument("input", 42);
                         });
 
@@ -496,7 +507,10 @@ void testBuilderRejectsReservedArgumentOnTime() {
         // We should reject PTFs that use reserved argument name "on_time"
         ProcessTableFunctionTestHarness.Builder harnessBuilder =
                 ProcessTableFunctionTestHarness.ofClass(InvalidReservedArgOnTimePTF.class)
-                        .withTableArgument("on_time", DataTypes.of("ROW<id INT>"));
+                        .withTableArgument(
+                                TableArgument.forArgument("on_time")
+                                        .type(DataTypes.of("ROW<id INT>"))
+                                        .build());
 
         ValidationException exception =
                 assertThrows(
@@ -515,7 +529,10 @@ void testBuilderRejectsReservedArgumentUid() {
         // We should reject PTFs that use reserved argument name "uid"
         ProcessTableFunctionTestHarness.Builder harnessBuilder =
                 ProcessTableFunctionTestHarness.ofClass(InvalidReservedArgUidPTF.class)
-                        .withTableArgument("input", DataTypes.of("ROW<id INT>"))
+                        .withTableArgument(
+                                TableArgument.forArgument("input")
+                                        .type(DataTypes.of("ROW<id INT>"))
+                                        .build())
                         .withScalarArgument("uid", "my-id");
 
         ValidationException exception =
@@ -541,7 +558,10 @@ void testExplicitNameTakesPrecedence() throws Exception {
 
         try (ProcessTableFunctionTestHarness<Row> harness =
                 ProcessTableFunctionTestHarness.ofClass(ExplicitNamePTF.class)
-                        .withTableArgument("customName", DataTypes.of("ROW<value INT>"))
+                        .withTableArgument(
+                                TableArgument.forArgument("customName")
+                                        .type(DataTypes.of("ROW<value INT>"))
+                                        .build())
                         .build()) {
 
             harness.processElement(Row.of(42));
@@ -628,7 +648,10 @@ void testInvokeRejectsTableArguments() throws Exception {
         // Verify that invoke() rejects PTFs with table arguments
         try (ProcessTableFunctionTestHarness<Row> harness =
                 ProcessTableFunctionTestHarness.ofClass(FilterPTF.class)
-                        .withTableArgument("input", DataTypes.of("ROW<value INT>"))
+                        .withTableArgument(
+                                TableArgument.forArgument("input")
+                                        .type(DataTypes.of("ROW<value INT>"))
+                                        .build())
                         .withScalarArgument("threshold", 50)
                         .build()) {
 
@@ -648,7 +671,10 @@ void testTableProcessingWithScalarArgument() throws Exception {
         // Test a PTF that uses a scalar parameter
         try (ProcessTableFunctionTestHarness<Row> harness =
                 ProcessTableFunctionTestHarness.ofClass(FilterPTF.class)
-                        .withTableArgument("input", DataTypes.of("ROW<value INT>"))
+                        .withTableArgument(
+                                TableArgument.forArgument("input")
+                                        .type(DataTypes.of("ROW<value INT>"))
+                                        .build())
                         .withScalarArgument("threshold", 50) // Scalar argument: threshold = 50
                         .build()) {
 
@@ -674,7 +700,10 @@ void testTableProcessingWithScalarArgumentWrongType() {
                         IllegalStateException.class,
                         () -> {
                             ProcessTableFunctionTestHarness.ofClass(FilterPTF.class)
-                                    .withTableArgument("input", DataTypes.of("ROW<value INT>"))
+                                    .withTableArgument(
+                                            TableArgument.forArgument("input")
+                                                    .type(DataTypes.of("ROW<value INT>"))
+                                                    .build())
                                     .withScalarArgument("threshold", "not_an_integer")
                                     .build();
                         });
@@ -693,7 +722,10 @@ void testProcessElementWithRowKind() throws Exception {
         // Verify RowKind is preserved through processing (ROW_SEMANTIC_TABLE)
         try (ProcessTableFunctionTestHarness<Row> harness =
                 ProcessTableFunctionTestHarness.ofClass(PassthroughPTF.class)
-                        .withTableArgument("input", DataTypes.of("ROW<value INT>"))
+                        .withTableArgument(
+                                TableArgument.forArgument("input")
+                                        .type(DataTypes.of("ROW<value INT>"))
+                                        .build())
                         .build()) {
 
             harness.processElement(RowKind.INSERT, 10);
@@ -719,8 +751,11 @@ void testPassColumnsThroughTrait() throws Exception {
         // Verify PASS_COLUMNS_THROUGH prepends ALL input columns (not just partition keys)
         try (ProcessTableFunctionTestHarness<Row> harness =
                 ProcessTableFunctionTestHarness.ofClass(PassColumnsThroughPTF.class)
-                        .withTableArgument("input", DataTypes.of("ROW<key STRING, value INT>"))
-                        .withPartitionBy("input", "key")
+                        .withTableArgument(
+                                TableArgument.forArgument("input")
+                                        .type(DataTypes.of("ROW<key STRING, value INT>"))
+                                        .partitionBy("key")
+                                        .build())
                         .build()) {
 
             harness.processElement(Row.of("A", 10));
@@ -739,7 +774,10 @@ void testOptionalPartitionByWithoutPartition() throws Exception {
         // Verify OPTIONAL_PARTITION_BY allows omitting partition configuration
         try (ProcessTableFunctionTestHarness<Row> harness =
                 ProcessTableFunctionTestHarness.ofClass(OptionalPartitionPTF.class)
-                        .withTableArgument("input", DataTypes.of("ROW<key STRING, value INT>"))
+                        .withTableArgument(
+                                TableArgument.forArgument("input")
+                                        .type(DataTypes.of("ROW<key STRING, value INT>"))
+                                        .build())
                         .build()) {
 
             harness.processElement(Row.of("A", 10));
@@ -760,8 +798,11 @@ void testOptionalPartitionByWithPartition() throws Exception {
         // Verify OPTIONAL_PARTITION_BY still works when partition is configured
         try (ProcessTableFunctionTestHarness<Row> harness =
                 ProcessTableFunctionTestHarness.ofClass(OptionalPartitionPTF.class)
-                        .withTableArgument("input", DataTypes.of("ROW<key STRING, value INT>"))
-                        .withPartitionBy("input", "key")
+                        .withTableArgument(
+                                TableArgument.forArgument("input")
+                                        .type(DataTypes.of("ROW<key STRING, value INT>"))
+                                        .partitionBy("key")
+                                        .build())
                         .build()) {
 
             harness.processElement(Row.of("A", 10));
@@ -781,7 +822,10 @@ void testOptionalPartitionByWithPartition() throws Exception {
     void testOptionalPartitionByWithStateNoPartition() throws Exception {
         try (ProcessTableFunctionTestHarness<Row> harness =
                 ProcessTableFunctionTestHarness.ofClass(StatefulOptionalPartitionPTF.class)
-                        .withTableArgument("input", DataTypes.of("ROW<key STRING, value INT>"))
+                        .withTableArgument(
+                                TableArgument.forArgument("input")
+                                        .type(DataTypes.of("ROW<key STRING, value INT>"))
+                                        .build())
                         .build()) {
 
             harness.processElement(Row.of("A", 10));
@@ -804,8 +848,11 @@ void testOptionalPartitionByWithStateNoPartition() throws Exception {
     void testOptionalPartitionByWithStateAndPartition() throws Exception {
         try (ProcessTableFunctionTestHarness<Row> harness =
                 ProcessTableFunctionTestHarness.ofClass(StatefulOptionalPartitionPTF.class)
-                        .withTableArgument("input", DataTypes.of("ROW<key STRING, value INT>"))
-                        .withPartitionBy("input", "key")
+                        .withTableArgument(
+                                TableArgument.forArgument("input")
+                                        .type(DataTypes.of("ROW<key STRING, value INT>"))
+                                        .partitionBy("key")
+                                        .build())
                         .build()) {
 
             harness.processElement(Row.of("A", 10));
@@ -835,7 +882,10 @@ void testOptionalPartitionByWithInitialStateNoPartition() throws Exception {
 
         try (ProcessTableFunctionTestHarness<Row> harness =
                 ProcessTableFunctionTestHarness.ofClass(StatefulOptionalPartitionPTF.class)
-                        .withTableArgument("input", DataTypes.of("ROW<key STRING, value INT>"))
+                        .withTableArgument(
+                                TableArgument.forArgument("input")
+                                        .type(DataTypes.of("ROW<key STRING, value INT>"))
+                                        .build())
                         .withInitialStateForKey("state", Row.of(), initialState)
                         .build()) {
 
@@ -872,7 +922,10 @@ void testNamedRowFieldOrdering() throws Exception {
         // Test what happens when Row field order differs from DataType schema order
         try (ProcessTableFunctionTestHarness<Row> harness =
                 ProcessTableFunctionTestHarness.ofClass(UserValuePassthroughPTF.class)
-                        .withTableArgument("input", DataTypes.of("ROW<user STRING, value INT>"))
+                        .withTableArgument(
+                                TableArgument.forArgument("input")
+                                        .type(DataTypes.of("ROW<user STRING, value INT>"))
+                                        .build())
                         .build()) {
 
             Row rowA = Row.withNames();
@@ -897,7 +950,10 @@ void testPositionalRowWithWrongTypeOrder() throws Exception {
         // Verify that type mismatches are caught when Row values don't match schema types
         try (ProcessTableFunctionTestHarness<Row> harness =
                 ProcessTableFunctionTestHarness.ofClass(UserValuePassthroughPTF.class)
-                        .withTableArgument("input", DataTypes.of("ROW<user STRING, value INT>"))
+                        .withTableArgument(
+                                TableArgument.forArgument("input")
+                                        .type(DataTypes.of("ROW<user STRING, value INT>"))
+                                        .build())
                         .build()) {
 
             Row wrongOrderRow = Row.of(10, "Alice");
@@ -919,7 +975,7 @@ void testStructuredTypeInput() throws Exception {
         // that type.
         try (ProcessTableFunctionTestHarness<Row> harness =
                 ProcessTableFunctionTestHarness.ofClass(UserPTF.class)
-                        .withTableArgument("user")
+                        .withTableArgument(TableArgument.forArgument("user").build())
                         .build()) {
 
             harness.processElement(Row.of("Alice", 25));
@@ -939,7 +995,7 @@ void testStructuredTypeInputAndOutput() throws Exception {
         // Test PTF with structured type inputs and outputs
         try (ProcessTableFunctionTestHarness<User> harness =
                 ProcessTableFunctionTestHarness.ofClass(UserTransformPTF.class)
-                        .withTableArgument("user")
+                        .withTableArgument(TableArgument.forArgument("user").build())
                         .build()) {
 
             harness.processElement(Row.of("Alice", 25));
@@ -980,7 +1036,10 @@ void testInlineTypeMatchesBuilderConfig() throws Exception {
         // types, the harness builds successfully
         try (ProcessTableFunctionTestHarness<Row> harness =
                 ProcessTableFunctionTestHarness.ofClass(InlineTypePTF.class)
-                        .withTableArgument("input", DataTypes.of("ROW<value INT>"))
+                        .withTableArgument(
+                                TableArgument.forArgument("input")
+                                        .type(DataTypes.of("ROW<value INT>"))
+                                        .build())
                         .build()) {
 
             harness.processElement(Row.of(7));
@@ -1000,8 +1059,11 @@ void testSetSemanticWithPartitionByName() throws Exception {
         // Verify set-semantic table with partition configuration by column name
         try (ProcessTableFunctionTestHarness<Row> harness =
                 ProcessTableFunctionTestHarness.ofClass(PartitionedPTF.class)
-                        .withTableArgument("input", DataTypes.of("ROW<key STRING, value INT>"))
-                        .withPartitionBy("input", "key") // Partition by "key" column name
+                        .withTableArgument(
+                                TableArgument.forArgument("input")
+                                        .type(DataTypes.of("ROW<key STRING, value INT>"))
+                                        .partitionBy("key")
+                                        .build())
                         .build()) {
 
             harness.processElement(Row.of("X", 10));
@@ -1022,9 +1084,12 @@ void testSetSemanticWithMultiplePartitionColumns() throws Exception {
         try (ProcessTableFunctionTestHarness<Row> harness =
                 ProcessTableFunctionTestHarness.ofClass(PartitionedPTF.class)
                         .withTableArgument(
-                                "input",
-                                DataTypes.of("ROW<region STRING, country STRING, value INT>"))
-                        .withPartitionBy("input", "region", "country")
+                                TableArgument.forArgument("input")
+                                        .type(
+                                                DataTypes.of(
+                                                        "ROW<region STRING, country STRING, value INT>"))
+                                        .partitionBy("region", "country")
+                                        .build())
                         .build()) {
 
             harness.processElement(Row.of("EU", "DE", 100));
@@ -1047,10 +1112,12 @@ void testSetSemanticWithSelectivePartitioning() throws Exception {
         try (ProcessTableFunctionTestHarness<Row> harness =
                 ProcessTableFunctionTestHarness.ofClass(PartitionedPTF.class)
                         .withTableArgument(
-                                "input",
-                                DataTypes.of(
-                                        "ROW<id INT, region STRING, country STRING, city STRING, value INT>"))
-                        .withPartitionBy("input", "region")
+                                TableArgument.forArgument("input")
+                                        .type(
+                                                DataTypes.of(
+                                                        "ROW<id INT, region STRING, country STRING, city STRING, value INT>"))
+                                        .partitionBy("region")
+                                        .build())
                         .build()) {
 
             harness.processElement(Row.of(1, "EU", "DE", "Berlin", 100));
@@ -1067,11 +1134,16 @@ void testSetSemanticWithSelectivePartitioning() throws Exception {
     void testMultipleSetSemanticTablesWithMatchingPartitionKeys() throws Exception {
         try (ProcessTableFunctionTestHarness<Row> harness =
                 ProcessTableFunctionTestHarness.ofClass(MultiTableUnionPTF.class)
-                        .withTableArgument("leftTable", DataTypes.of("ROW<name STRING, score INT>"))
-                        .withPartitionBy("leftTable", "name")
                         .withTableArgument(
-                                "rightTable", DataTypes.of("ROW<name STRING, city STRING>"))
-                        .withPartitionBy("rightTable", "name")
+                                TableArgument.forArgument("leftTable")
+                                        .type(DataTypes.of("ROW<name STRING, score INT>"))
+                                        .partitionBy("name")
+                                        .build())
+                        .withTableArgument(
+                                TableArgument.forArgument("rightTable")
+                                        .type(DataTypes.of("ROW<name STRING, city STRING>"))
+                                        .partitionBy("name")
+                                        .build())
                         .build()) {
 
             harness.processElementForTable("leftTable", Row.of("Alice", 100));
@@ -1107,11 +1179,13 @@ void testMultipleSetSemanticTablesWithStructuredTypePartitioning() throws Except
         // and another is a Row, both partitioned by the same field type
         try (ProcessTableFunctionTestHarness<Row> harness =
                 ProcessTableFunctionTestHarness.ofClass(MixedTypeMultiTablePTF.class)
-                        .withTableArgument("userTable")
-                        .withPartitionBy("userTable", "age")
                         .withTableArgument(
-                                "rowTable", DataTypes.of("ROW<name STRING, age INT NOT NULL>"))
-                        .withPartitionBy("rowTable", "age")
+                                TableArgument.forArgument("userTable").partitionBy("age").build())
+                        .withTableArgument(
+                                TableArgument.forArgument("rowTable")
+                                        .type(DataTypes.of("ROW<name STRING, age INT NOT NULL>"))
+                                        .partitionBy("age")
+                                        .build())
                         .build()) {
 
             harness.processElementForTable("userTable", Row.of("Alice", 25));
@@ -1133,12 +1207,17 @@ void testMultipleSetSemanticTablesWithMismatchedPartitionTypes() {
                         () -> {
                             ProcessTableFunctionTestHarness.ofClass(MultiTableUnionPTF.class)
                                     .withTableArgument(
-                                            "leftTable", DataTypes.of("ROW<id INT, name STRING>"))
-                                    .withPartitionBy("leftTable", "id")
+                                            TableArgument.forArgument("leftTable")
+                                                    .type(DataTypes.of("ROW<id INT, name STRING>"))
+                                                    .partitionBy("id")
+                                                    .build())
                                     .withTableArgument(
-                                            "rightTable",
-                                            DataTypes.of("ROW<key STRING, city STRING>"))
-                                    .withPartitionBy("rightTable", "key")
+                                            TableArgument.forArgument("rightTable")
+                                                    .type(
+                                                            DataTypes.of(
+                                                                    "ROW<key STRING, city STRING>"))
+                                                    .partitionBy("key")
+                                                    .build())
                                     .build();
                         });
 
@@ -1154,12 +1233,17 @@ void testMultipleSetSemanticTablesWithMismatchedPartitionColumnCount() {
                         () -> {
                             ProcessTableFunctionTestHarness.ofClass(MultiTableUnionPTF.class)
                                     .withTableArgument(
-                                            "leftTable",
-                                            DataTypes.of("ROW<id INT, region STRING, name STRING>"))
-                                    .withPartitionBy("leftTable", "id", "region")
+                                            TableArgument.forArgument("leftTable")
+                                                    .type(
+                                                            DataTypes.of(
+                                                                    "ROW<id INT, region STRING, name STRING>"))
+                                                    .partitionBy("id", "region")
+                                                    .build())
                                     .withTableArgument(
-                                            "rightTable", DataTypes.of("ROW<id INT, city STRING>"))
-                                    .withPartitionBy("rightTable", "id")
+                                            TableArgument.forArgument("rightTable")
+                                                    .type(DataTypes.of("ROW<id INT, city STRING>"))
+                                                    .partitionBy("id")
+                                                    .build())
                                     .build();
                         });
 
@@ -1175,8 +1259,14 @@ void testPassColumnsThroughWithMultipleTablesRejected() {
                         () -> {
                             ProcessTableFunctionTestHarness.ofClass(
                                             InvalidPassColumnsThroughMultiTablePTF.class)
-                                    .withTableArgument("leftTable", DataTypes.of("ROW<a INT>"))
-                                    .withTableArgument("rightTable", DataTypes.of("ROW<b INT>"))
+                                    .withTableArgument(
+                                            TableArgument.forArgument("leftTable")
+                                                    .type(DataTypes.of("ROW<a INT>"))
+                                                    .build())
+                                    .withTableArgument(
+                                            TableArgument.forArgument("rightTable")
+                                                    .type(DataTypes.of("ROW<b INT>"))
+                                                    .build())
                                     .build();
                         });
 
@@ -1193,10 +1283,16 @@ void testPassColumnsThroughWithMultipleTablesRejected() {
     void testProcessElementOnMultiTableThrows() throws Exception {
         try (ProcessTableFunctionTestHarness<Row> harness =
                 ProcessTableFunctionTestHarness.ofClass(MultiTableUnionPTF.class)
-                        .withTableArgument("leftTable", DataTypes.of("ROW<id INT, name STRING>"))
-                        .withTableArgument("rightTable", DataTypes.of("ROW<id INT, value STRING>"))
-                        .withPartitionBy("leftTable", "id")
-                        .withPartitionBy("rightTable", "id")
+                        .withTableArgument(
+                                TableArgument.forArgument("leftTable")
+                                        .type(DataTypes.of("ROW<id INT, name STRING>"))
+                                        .partitionBy("id")
+                                        .build())
+                        .withTableArgument(
+                                TableArgument.forArgument("rightTable")
+                                        .type(DataTypes.of("ROW<id INT, value STRING>"))
+                                        .partitionBy("id")
+                                        .build())
                         .build()) {
 
             Exception exception =
@@ -1217,7 +1313,10 @@ void testProcessElementOnMultiTableThrows() throws Exception {
     void testClearOutput() throws Exception {
         try (ProcessTableFunctionTestHarness<Row> harness =
                 ProcessTableFunctionTestHarness.ofClass(PassthroughPTF.class)
-                        .withTableArgument("input", DataTypes.of("ROW<value INT>"))
+                        .withTableArgument(
+                                TableArgument.forArgument("input")
+                                        .type(DataTypes.of("ROW<value INT>"))
+                                        .build())
                         .build()) {
 
             harness.processElement(Row.of(10));
@@ -1243,7 +1342,10 @@ void testClearOutput() throws Exception {
     void testFunctionOutputReturnsUnwrappedAtomicValue() throws Exception {
         try (ProcessTableFunctionTestHarness<Integer> harness =
                 ProcessTableFunctionTestHarness.ofClass(AtomicOutputPTF.class)
-                        .withTableArgument("input", DataTypes.of("ROW<value INT>"))
+                        .withTableArgument(
+                                TableArgument.forArgument("input")
+                                        .type(DataTypes.of("ROW<value INT>"))
+                                        .build())
                         .build()) {
 
             harness.processElement(Row.of(21));
@@ -1257,8 +1359,11 @@ void testFunctionOutputReturnsUnwrappedAtomicValue() throws Exception {
     void testFunctionOutputExcludesPrependedPartitionKey() throws Exception {
         try (ProcessTableFunctionTestHarness<Row> harness =
                 ProcessTableFunctionTestHarness.ofClass(PartitionedPTF.class)
-                        .withTableArgument("input", DataTypes.of("ROW<key STRING, value INT>"))
-                        .withPartitionBy("input", "key")
+                        .withTableArgument(
+                                TableArgument.forArgument("input")
+                                        .type(DataTypes.of("ROW<key STRING, value INT>"))
+                                        .partitionBy("key")
+                                        .build())
                         .build()) {
 
             harness.processElement(Row.of("X", 10));
@@ -1277,7 +1382,10 @@ void testFunctionOutputExcludesPrependedPartitionKey() throws Exception {
     void testProcessElementForTableWithInvalidName() throws Exception {
         try (ProcessTableFunctionTestHarness<Row> harness =
                 ProcessTableFunctionTestHarness.ofClass(PassthroughPTF.class)
-                        .withTableArgument("input", DataTypes.of("ROW<value INT>"))
+                        .withTableArgument(
+                                TableArgument.forArgument("input")
+                                        .type(DataTypes.of("ROW<value INT>"))
+                                        .build())
                         .build()) {
 
             Exception exception =
@@ -1296,12 +1404,16 @@ void testSetSemanticMissingPartitionConfigThrows() {
                         () -> {
                             ProcessTableFunctionTestHarness.ofClass(PartitionedPTF.class)
                                     .withTableArgument(
-                                            "input", DataTypes.of("ROW<key STRING, value INT>"))
+                                            TableArgument.forArgument("input")
+                                                    .type(
+                                                            DataTypes.of(
+                                                                    "ROW<key STRING, value INT>"))
+                                                    .build())
                                     .build();
                         });
 
         assertThat(exception.getMessage()).contains("No partition configuration found");
-        assertThat(exception.getMessage()).contains("withPartitionBy");
+        assertThat(exception.getMessage()).contains("TableArgument.forArgument");
     }
 
     @Test
@@ -1312,8 +1424,12 @@ void testPartitionByInvalidColumnName() {
                         () -> {
                             ProcessTableFunctionTestHarness.ofClass(PartitionedPTF.class)
                                     .withTableArgument(
-                                            "input", DataTypes.of("ROW<key STRING, value INT>"))
-                                    .withPartitionBy("input", "nonexistent")
+                                            TableArgument.forArgument("input")
+                                                    .type(
+                                                            DataTypes.of(
+                                                                    "ROW<key STRING, value INT>"))
+                                                    .partitionBy("nonexistent")
+                                                    .build())
                                     .build();
                         });
 
@@ -1322,19 +1438,31 @@ void testPartitionByInvalidColumnName() {
     }
 
     @Test
-    void testPartitionByDuplicateConfigThrows() {
+    void testBuilderRejectsDuplicateTableArgumentWithPartitioning() {
+        // Configuring the same table argument twice is rejected, even when the second
+        // configuration only differs in its partitioning.
         Exception exception =
                 assertThrows(
                         IllegalArgumentException.class,
                         () -> {
                             ProcessTableFunctionTestHarness.ofClass(PartitionedPTF.class)
-                                    .withTableArgument(
-                                            "input", DataTypes.of("ROW<key STRING, value INT>"))
-                                    .withPartitionBy("input", "key") // First config
-                                    .withPartitionBy("input", "key"); // Duplicate - should fail
+                                    .withTableArgument( // First config
+                                            TableArgument.forArgument("input")
+                                                    .type(
+                                                            DataTypes.of(
+                                                                    "ROW<key STRING, value INT>"))
+                                                    .partitionBy("key")
+                                                    .build())
+                                    .withTableArgument( // Duplicate - should fail
+                                            TableArgument.forArgument("input")
+                                                    .type(
+                                                            DataTypes.of(
+                                                                    "ROW<key STRING, value INT>"))
+                                                    .partitionBy("key")
+                                                    .build());
                         });
 
-        assertThat(exception.getMessage()).contains("Partition config already exists");
+        assertThat(exception.getMessage()).contains("Argument already configured");
     }
 
     // -------------------------------------------------------------------------
@@ -1345,8 +1473,11 @@ void testPartitionByDuplicateConfigThrows() {
     void testValueState() throws Exception {
         ProcessTableFunctionTestHarness<Row> harness =
                 ProcessTableFunctionTestHarness.ofClass(PTFWithValueState.class)
-                        .withTableArgument("input", DataTypes.of("ROW<name STRING, value INT>"))
-                        .withPartitionBy("input", "name")
+                        .withTableArgument(
+                                TableArgument.forArgument("input")
+                                        .type(DataTypes.of("ROW<name STRING, value INT>"))
+                                        .partitionBy("name")
+                                        .build())
                         .build();
 
         harness.processElementForTable("input", Row.of("Alice", 10));
@@ -1368,8 +1499,11 @@ void testValueState() throws Exception {
     void testValueStatePartitionIsolation() throws Exception {
         ProcessTableFunctionTestHarness<Row> harness =
                 ProcessTableFunctionTestHarness.ofClass(PTFWithValueState.class)
-                        .withTableArgument("input", DataTypes.of("ROW<name STRING, value INT>"))
-                        .withPartitionBy("input", "name")
+                        .withTableArgument(
+                                TableArgument.forArgument("input")
+                                        .type(DataTypes.of("ROW<name STRING, value INT>"))
+                                        .partitionBy("name")
+                                        .build())
                         .build();
 
         harness.processElementForTable("input", Row.of("Alice", 10));
@@ -1393,8 +1527,11 @@ void testValueStateWithInitialState() throws Exception {
 
         ProcessTableFunctionTestHarness<Row> harness =
                 ProcessTableFunctionTestHarness.ofClass(PTFWithValueState.class)
-                        .withTableArgument("input", DataTypes.of("ROW<id INT>"))
-                        .withPartitionBy("input", "id")
+                        .withTableArgument(
+                                TableArgument.forArgument("input")
+                                        .type(DataTypes.of("ROW<id INT>"))
+                                        .partitionBy("id")
+                                        .build())
                         .withInitialStateForKey("state", Row.of(1), initialState)
                         .build();
 
@@ -1414,8 +1551,11 @@ void testValueStateWithInitialState() throws Exception {
     void testGetStateKeys() throws Exception {
         ProcessTableFunctionTestHarness<Row> harness =
                 ProcessTableFunctionTestHarness.ofClass(PTFWithValueState.class)
-                        .withTableArgument("input", DataTypes.of("ROW<name STRING, value INT>"))
-                        .withPartitionBy("input", "name")
+                        .withTableArgument(
+                                TableArgument.forArgument("input")
+                                        .type(DataTypes.of("ROW<name STRING, value INT>"))
+                                        .partitionBy("name")
+                                        .build())
                         .build();
 
         harness.processElementForTable("input", Row.of("Alice", 10));
@@ -1433,8 +1573,11 @@ void testGetStateKeys() throws Exception {
     void testGetAllState() throws Exception {
         ProcessTableFunctionTestHarness<Row> harness =
                 ProcessTableFunctionTestHarness.ofClass(PTFWithValueState.class)
-                        .withTableArgument("input", DataTypes.of("ROW<name STRING, value INT>"))
-                        .withPartitionBy("input", "name")
+                        .withTableArgument(
+                                TableArgument.forArgument("input")
+                                        .type(DataTypes.of("ROW<name STRING, value INT>"))
+                                        .partitionBy("name")
+                                        .build())
                         .build();
 
         harness.processElementForTable("input", Row.of("Alice", 10));
@@ -1454,8 +1597,11 @@ void testGetAllState() throws Exception {
     void testListViewState() throws Exception {
         ProcessTableFunctionTestHarness<Row> harness =
                 ProcessTableFunctionTestHarness.ofClass(PTFWithListViewState.class)
-                        .withTableArgument("input", DataTypes.of("ROW<key STRING, value INT>"))
-                        .withPartitionBy("input", "key")
+                        .withTableArgument(
+                                TableArgument.forArgument("input")
+                                        .type(DataTypes.of("ROW<key STRING, value INT>"))
+                                        .partitionBy("key")
+                                        .build())
                         .build();
 
         harness.processElementForTable("input", Row.of("A", 1));
@@ -1475,8 +1621,10 @@ void testMapViewState() throws Exception {
         ProcessTableFunctionTestHarness<Row> harness =
                 ProcessTableFunctionTestHarness.ofClass(PTFWithMapViewState.class)
                         .withTableArgument(
-                                "input", DataTypes.of("ROW<partition STRING, key STRING>"))
-                        .withPartitionBy("input", "partition")
+                                TableArgument.forArgument("input")
+                                        .type(DataTypes.of("ROW<partition STRING, key STRING>"))
+                                        .partitionBy("partition")
+                                        .build())
                         .build();
 
         harness.processElementForTable("input", Row.of("P1", "foo"));
@@ -1499,8 +1647,11 @@ void testMapViewState() throws Exception {
     void testRowState() throws Exception {
         ProcessTableFunctionTestHarness<Row> harness =
                 ProcessTableFunctionTestHarness.ofClass(PTFWithRowState.class)
-                        .withTableArgument("input", DataTypes.of("ROW<name STRING, value INT>"))
-                        .withPartitionBy("input", "name")
+                        .withTableArgument(
+                                TableArgument.forArgument("input")
+                                        .type(DataTypes.of("ROW<name STRING, value INT>"))
+                                        .partitionBy("name")
+                                        .build())
                         .build();
 
         harness.processElementForTable("input", Row.of("Alice", 10));
@@ -1519,8 +1670,11 @@ void testRowState() throws Exception {
     void testEmptyState() throws Exception {
         ProcessTableFunctionTestHarness<Row> harness =
                 ProcessTableFunctionTestHarness.ofClass(PTFWithValueState.class)
-                        .withTableArgument("input", DataTypes.of("ROW<name STRING, value INT>"))
-                        .withPartitionBy("input", "name")
+                        .withTableArgument(
+                                TableArgument.forArgument("input")
+                                        .type(DataTypes.of("ROW<name STRING, value INT>"))
+                                        .partitionBy("name")
+                                        .build())
                         .build();
 
         PTFWithValueState.CounterState state = harness.getStateForKey("state", Row.of("Alice"));
@@ -1534,8 +1688,11 @@ void testEmptyState() throws Exception {
     void testClearAllStatesForKey() throws Exception {
         ProcessTableFunctionTestHarness<Row> harness =
                 ProcessTableFunctionTestHarness.ofClass(PTFWithValueState.class)
-                        .withTableArgument("input", DataTypes.of("ROW<name STRING, value INT>"))
-                        .withPartitionBy("input", "name")
+                        .withTableArgument(
+                                TableArgument.forArgument("input")
+                                        .type(DataTypes.of("ROW<name STRING, value INT>"))
+                                        .partitionBy("name")
+                                        .build())
                         .build();
 
         harness.processElementForTable("input", Row.of("Alice", 10));
@@ -1560,8 +1717,11 @@ void testClearAllStatesForKey() throws Exception {
     void testClearStateForKey() throws Exception {
         ProcessTableFunctionTestHarness<Row> harness =
                 ProcessTableFunctionTestHarness.ofClass(PTFWithValueState.class)
-                        .withTableArgument("input", DataTypes.of("ROW<name STRING, value INT>"))
-                        .withPartitionBy("input", "name")
+                        .withTableArgument(
+                                TableArgument.forArgument("input")
+                                        .type(DataTypes.of("ROW<name STRING, value INT>"))
+                                        .partitionBy("name")
+                                        .build())
                         .build();
 
         harness.processElementForTable("input", Row.of("Alice", 10));
@@ -1586,8 +1746,11 @@ void testClearStateForKey() throws Exception {
     void testMultipleStateParameters() throws Exception {
         ProcessTableFunctionTestHarness<Row> harness =
                 ProcessTableFunctionTestHarness.ofClass(PTFWithMultipleStates.class)
-                        .withTableArgument("input", DataTypes.of("ROW<key STRING, value INT>"))
-                        .withPartitionBy("input", "key")
+                        .withTableArgument(
+                                TableArgument.forArgument("input")
+                                        .type(DataTypes.of("ROW<key STRING, value INT>"))
+                                        .partitionBy("key")
+                                        .build())
                         .build();
 
         harness.processElementForTable("input", Row.of("A", 10));
@@ -1615,8 +1778,11 @@ void testInitialStateWithListView() throws Exception {
 
         ProcessTableFunctionTestHarness<Row> harness =
                 ProcessTableFunctionTestHarness.ofClass(PTFWithListViewState.class)
-                        .withTableArgument("input", DataTypes.of("ROW<key STRING, value INT>"))
-                        .withPartitionBy("input", "key")
+                        .withTableArgument(
+                                TableArgument.forArgument("input")
+                                        .type(DataTypes.of("ROW<key STRING, value INT>"))
+                                        .partitionBy("key")
+                                        .build())
                         .withInitialStateForKey("listState", Row.of("A"), initialList)
                         .build();
 
@@ -1637,8 +1803,10 @@ void testInitialStateWithMapView() throws Exception {
         ProcessTableFunctionTestHarness<Row> harness =
                 ProcessTableFunctionTestHarness.ofClass(PTFWithMapViewState.class)
                         .withTableArgument(
-                                "input", DataTypes.of("ROW<partition STRING, key STRING>"))
-                        .withPartitionBy("input", "partition")
+                                TableArgument.forArgument("input")
+                                        .type(DataTypes.of("ROW<partition STRING, key STRING>"))
+                                        .partitionBy("partition")
+                                        .build())
                         .withInitialStateForKey("mapState", Row.of("P1"), initialMap)
                         .build();
 
@@ -1659,9 +1827,12 @@ void testInitialStateKeyArityMismatch() {
                         () ->
                                 ProcessTableFunctionTestHarness.ofClass(PTFWithValueState.class)
                                         .withTableArgument(
-                                                "input",
-                                                DataTypes.of("ROW<name STRING, value INT>"))
-                                        .withPartitionBy("input", "name")
+                                                TableArgument.forArgument("input")
+                                                        .type(
+                                                                DataTypes.of(
+                                                                        "ROW<name STRING, value INT>"))
+                                                        .partitionBy("name")
+                                                        .build())
                                         .withInitialStateForKey(
                                                 "state",
                                                 Row.of("Alice", 42),
@@ -1681,9 +1852,12 @@ void testInitialStateKeyTypeMismatch() {
                         () ->
                                 ProcessTableFunctionTestHarness.ofClass(PTFWithValueState.class)
                                         .withTableArgument(
-                                                "input",
-                                                DataTypes.of("ROW<name STRING, value INT>"))
-                                        .withPartitionBy("input", "name")
+                                                TableArgument.forArgument("input")
+                                                        .type(
+                                                                DataTypes.of(
+                                                                        "ROW<name STRING, value INT>"))
+                                                        .partitionBy("name")
+                                                        .build())
                                         .withInitialStateForKey(
                                                 "state",
                                                 Row.of(42),
@@ -1700,8 +1874,11 @@ void testInitialStateKeyTypeMismatch() {
     void testSetStateForKey() throws Exception {
         ProcessTableFunctionTestHarness<Row> harness =
                 ProcessTableFunctionTestHarness.ofClass(PTFWithValueState.class)
-                        .withTableArgument("input", DataTypes.of("ROW<name STRING, value INT>"))
-                        .withPartitionBy("input", "name")
+                        .withTableArgument(
+                                TableArgument.forArgument("input")
+                                        .type(DataTypes.of("ROW<name STRING, value INT>"))
+                                        .partitionBy("name")
+                                        .build())
                         .build();
 
         harness.processElementForTable("input", Row.of("Alice", 10));
@@ -1730,8 +1907,11 @@ void testInvalidStateNameInWithInitialState() {
                         IllegalArgumentException.class,
                         () ->
                                 ProcessTableFunctionTestHarness.ofClass(PTFWithValueState.class)
-                                        .withTableArgument("input", DataTypes.of("ROW<id INT>"))
-                                        .withPartitionBy("input", "id")
+                                        .withTableArgument(
+                                                TableArgument.forArgument("input")
+                                                        .type(DataTypes.of("ROW<id INT>"))
+                                                        .partitionBy("id")
+                                                        .build())
                                         .withInitialStateForKey(
                                                 "nonExistentState", Row.of(1), "value")
                                         .build());
@@ -1750,8 +1930,11 @@ void testInvalidStateNameInWithInitialState() {
     void testPartitionKeyValidationWrongArity() throws Exception {
         ProcessTableFunctionTestHarness<Row> harness =
                 ProcessTableFunctionTestHarness.ofClass(PTFWithValueState.class)
-                        .withTableArgument("input", DataTypes.of("ROW<name STRING, value INT>"))
-                        .withPartitionBy("input", "name")
+                        .withTableArgument(
+                                TableArgument.forArgument("input")
+                                        .type(DataTypes.of("ROW<name STRING, value INT>"))
+                                        .partitionBy("name")
+                                        .build())
                         .build();
 
         Exception exception =
@@ -1768,8 +1951,11 @@ void testPartitionKeyValidationWrongArity() throws Exception {
     void testPartitionKeyValidationWrongType() throws Exception {
         ProcessTableFunctionTestHarness<Row> harness =
                 ProcessTableFunctionTestHarness.ofClass(PTFWithValueState.class)
-                        .withTableArgument("input", DataTypes.of("ROW<name STRING, value INT>"))
-                        .withPartitionBy("input", "name")
+                        .withTableArgument(
+                                TableArgument.forArgument("input")
+                                        .type(DataTypes.of("ROW<name STRING, value INT>"))
+                                        .partitionBy("name")
+                                        .build())
                         .build();
 
         Exception exception =
@@ -1787,8 +1973,11 @@ void testPartitionKeyValidationWrongType() throws Exception {
     void testPartitionKeyValidationOnSetState() throws Exception {
         ProcessTableFunctionTestHarness<Row> harness =
                 ProcessTableFunctionTestHarness.ofClass(PTFWithValueState.class)
-                        .withTableArgument("input", DataTypes.of("ROW<name STRING, value INT>"))
-                        .withPartitionBy("input", "name")
+                        .withTableArgument(
+                                TableArgument.forArgument("input")
+                                        .type(DataTypes.of("ROW<name STRING, value INT>"))
+                                        .partitionBy("name")
+                                        .build())
                         .build();
 
         PTFWithValueState.CounterState state = new PTFWithValueState.CounterState();
@@ -1805,8 +1994,11 @@ void testPartitionKeyValidationOnSetState() throws Exception {
     void testPartitionKeyValidationOnClearAllStates() throws Exception {
         ProcessTableFunctionTestHarness<Row> harness =
                 ProcessTableFunctionTestHarness.ofClass(PTFWithValueState.class)
-                        .withTableArgument("input", DataTypes.of("ROW<name STRING, value INT>"))
-                        .withPartitionBy("input", "name")
+                        .withTableArgument(
+                                TableArgument.forArgument("input")
+                                        .type(DataTypes.of("ROW<name STRING, value INT>"))
+                                        .partitionBy("name")
+                                        .build())
                         .build();
 
         assertThrows(
@@ -1820,8 +2012,11 @@ void testPartitionKeyValidationOnClearAllStates() throws Exception {
     void testPartitionKeyValidationOnClearState() throws Exception {
         ProcessTableFunctionTestHarness<Row> harness =
                 ProcessTableFunctionTestHarness.ofClass(PTFWithValueState.class)
-                        .withTableArgument("input", DataTypes.of("ROW<name STRING, value INT>"))
-                        .withPartitionBy("input", "name")
+                        .withTableArgument(
+                                TableArgument.forArgument("input")
+                                        .type(DataTypes.of("ROW<name STRING, value INT>"))
+                                        .partitionBy("name")
+                                        .build())
                         .build();
 
         assertThrows(
@@ -2130,12 +2325,19 @@ void testPerTableWatermarkTimerFiring() throws Exception {
         try (ProcessTableFunctionTestHarness<Row> harness =
                 ProcessTableFunctionTestHarness.ofClass(MultiTableTimerPTF.class)
                         .withTableArgument(
-                                "leftTable", DataTypes.of("ROW<partition STRING, ts TIMESTAMP(3)>"))
+                                TableArgument.forArgument("leftTable")
+                                        .type(
+                                                DataTypes.of(
+                                                        "ROW<partition STRING, ts TIMESTAMP(3)>"))
+                                        .partitionBy("partition")
+                                        .build())
                         .withTableArgument(
-                                "rightTable",
-                                DataTypes.of("ROW<partition STRING, ts TIMESTAMP(3)>"))
-                        .withPartitionBy("leftTable", "partition")
-                        .withPartitionBy("rightTable", "partition")
+                                TableArgument.forArgument("rightTable")
+                                        .type(
+                                                DataTypes.of(
+                                                        "ROW<partition STRING, ts TIMESTAMP(3)>"))
+                                        .partitionBy("partition")
+                                        .build())
                         .withOnTimeColumn("ts")
                         .build()) {
 
@@ -2172,10 +2374,12 @@ void testContextClearState() throws Exception {
         try (ProcessTableFunctionTestHarness<Row> harness =
                 ProcessTableFunctionTestHarness.ofClass(ContextClearStatePTF.class)
                         .withTableArgument(
-                                "input",
-                                DataTypes.of(
-                                        "ROW<partition STRING, action STRING, ts TIMESTAMP(3)>"))
-                        .withPartitionBy("input", "partition")
+                                TableArgument.forArgument("input")
+                                        .type(
+                                                DataTypes.of(
+                                                        "ROW<partition STRING, action STRING, ts TIMESTAMP(3)>"))
+                                        .partitionBy("partition")
+                                        .build())
                         .withOnTimeColumn("ts")
                         .build()) {
 
@@ -2205,10 +2409,12 @@ void testContextClearAllState() throws Exception {
         try (ProcessTableFunctionTestHarness<Row> harness =
                 ProcessTableFunctionTestHarness.ofClass(ContextClearStatePTF.class)
                         .withTableArgument(
-                                "input",
-                                DataTypes.of(
-                                        "ROW<partition STRING, action STRING, ts TIMESTAMP(3)>"))
-                        .withPartitionBy("input", "partition")
+                                TableArgument.forArgument("input")
+                                        .type(
+                                                DataTypes.of(
+                                                        "ROW<partition STRING, action STRING, ts TIMESTAMP(3)>"))
+                                        .partitionBy("partition")
+                                        .build())
                         .withOnTimeColumn("ts")
                         .build()) {
 
@@ -2232,10 +2438,12 @@ void testContextClearAll() throws Exception {
         try (ProcessTableFunctionTestHarness<Row> harness =
                 ProcessTableFunctionTestHarness.ofClass(ContextClearStatePTF.class)
                         .withTableArgument(
-                                "input",
-                                DataTypes.of(
-                                        "ROW<partition STRING, action STRING, ts TIMESTAMP(3)>"))
-                        .withPartitionBy("input", "partition")
+                                TableArgument.forArgument("input")
+                                        .type(
+                                                DataTypes.of(
+                                                        "ROW<partition STRING, action STRING, ts TIMESTAMP(3)>"))
+                                        .partitionBy("partition")
+                                        .build())
                         .withOnTimeColumn("ts")
                         .build()) {
 
@@ -2268,8 +2476,12 @@ void testContextTableSemanticsAndChangelogMode() throws Exception {
         try (ProcessTableFunctionTestHarness<Row> harness =
                 ProcessTableFunctionTestHarness.ofClass(ContextSemanticsIntrospectionPTF.class)
                         .withTableArgument(
-                                "input", DataTypes.of("ROW<partition STRING, ts TIMESTAMP(3)>"))
-                        .withPartitionBy("input", "partition")
+                                TableArgument.forArgument("input")
+                                        .type(
+                                                DataTypes.of(
+                                                        "ROW<partition STRING, ts TIMESTAMP(3)>"))
+                                        .partitionBy("partition")
+                                        .build())
                         .withOnTimeColumn("ts")
                         .build()) {
 
@@ -2300,8 +2512,8 @@ void testContextTableSemanticsAndChangelogMode() throws Exception {
     void testPojoInputWithOnTimeColumn() throws Exception {
         try (ProcessTableFunctionTestHarness<Row> harness =
                 ProcessTableFunctionTestHarness.ofClass(PojoTimerPTF.class)
-                        .withTableArgument("input")
-                        .withPartitionBy("input", "key")
+                        .withTableArgument(
+                                TableArgument.forArgument("input").partitionBy("key").build())
                         .withOnTimeColumn("ts")
                         .build()) {
 
@@ -2329,9 +2541,12 @@ void testOnTimeColumnAppendedToEvalOutput() throws Exception {
         try (ProcessTableFunctionTestHarness<Row> harness =
                 ProcessTableFunctionTestHarness.ofClass(OnTimePTF.class)
                         .withTableArgument(
-                                "input",
-                                DataTypes.of("ROW<partition STRING, name STRING, ts TIMESTAMP(3)>"))
-                        .withPartitionBy("input", "partition")
+                                TableArgument.forArgument("input")
+                                        .type(
+                                                DataTypes.of(
+                                                        "ROW<partition STRING, name STRING, ts TIMESTAMP(3)>"))
+                                        .partitionBy("partition")
+                                        .build())
                         .withOnTimeColumn("ts")
                         .build()) {
 
@@ -2347,9 +2562,12 @@ void testWatermarkAdvancesWithoutTimers() throws Exception {
         try (ProcessTableFunctionTestHarness<Row> harness =
                 ProcessTableFunctionTestHarness.ofClass(OnTimePTF.class)
                         .withTableArgument(
-                                "input",
-                                DataTypes.of("ROW<partition STRING, name STRING, ts TIMESTAMP(3)>"))
-                        .withPartitionBy("input", "partition")
+                                TableArgument.forArgument("input")
+                                        .type(
+                                                DataTypes.of(
+                                                        "ROW<partition STRING, name STRING, ts TIMESTAMP(3)>"))
+                                        .partitionBy("partition")
+                                        .build())
                         .withOnTimeColumn("ts")
                         .build()) {
 
@@ -2365,7 +2583,10 @@ void testWatermarkAdvancesWithoutTimers() throws Exception {
     void testWatermarkAdvancesWithoutOnTimeColumn() throws Exception {
         try (ProcessTableFunctionTestHarness<Row> harness =
                 ProcessTableFunctionTestHarness.ofClass(PassthroughPTF.class)
-                        .withTableArgument("input", DataTypes.of("ROW<value INT>"))
+                        .withTableArgument(
+                                TableArgument.forArgument("input")
+                                        .type(DataTypes.of("ROW<value INT>"))
+                                        .build())
                         .build()) {
 
             harness.processElement(Row.of(42));
@@ -2383,9 +2604,12 @@ void testNamedTimerRegistrationAndFiring() throws Exception {
         try (ProcessTableFunctionTestHarness<Row> harness =
                 ProcessTableFunctionTestHarness.ofClass(TimerPTF.class)
                         .withTableArgument(
-                                "input",
-                                DataTypes.of("ROW<partition STRING, name STRING, ts TIMESTAMP(3)>"))
-                        .withPartitionBy("input", "partition")
+                                TableArgument.forArgument("input")
+                                        .type(
+                                                DataTypes.of(
+                                                        "ROW<partition STRING, name STRING, ts TIMESTAMP(3)>"))
+                                        .partitionBy("partition")
+                                        .build())
                         .withOnTimeColumn("ts")
                         .build()) {
 
@@ -2425,8 +2649,12 @@ void testUnnamedTimerFiring() throws Exception {
         try (ProcessTableFunctionTestHarness<Row> harness =
                 ProcessTableFunctionTestHarness.ofClass(UnnamedTimerPTF.class)
                         .withTableArgument(
-                                "input", DataTypes.of("ROW<partition STRING, ts TIMESTAMP(3)>"))
-                        .withPartitionBy("input", "partition")
+                                TableArgument.forArgument("input")
+                                        .type(
+                                                DataTypes.of(
+                                                        "ROW<partition STRING, ts TIMESTAMP(3)>"))
+                                        .partitionBy("partition")
+                                        .build())
                         .withOnTimeColumn("ts")
                         .build()) {
 
@@ -2448,9 +2676,12 @@ void testTimerReplacementSemantics() throws Exception {
         try (ProcessTableFunctionTestHarness<Row> harness =
                 ProcessTableFunctionTestHarness.ofClass(TimerPTF.class)
                         .withTableArgument(
-                                "input",
-                                DataTypes.of("ROW<partition STRING, name STRING, ts TIMESTAMP(3)>"))
-                        .withPartitionBy("input", "partition")
+                                TableArgument.forArgument("input")
+                                        .type(
+                                                DataTypes.of(
+                                                        "ROW<partition STRING, name STRING, ts TIMESTAMP(3)>"))
+                                        .partitionBy("partition")
+                                        .build())
                         .withOnTimeColumn("ts")
                         .build()) {
 
@@ -2485,9 +2716,12 @@ void testMultipleTimersFiringOrder() throws Exception {
         try (ProcessTableFunctionTestHarness<Row> harness =
                 ProcessTableFunctionTestHarness.ofClass(TimerPTF.class)
                         .withTableArgument(
-                                "input",
-                                DataTypes.of("ROW<partition STRING, name STRING, ts TIMESTAMP(3)>"))
-                        .withPartitionBy("input", "partition")
+                                TableArgument.forArgument("input")
+                                        .type(
+                                                DataTypes.of(
+                                                        "ROW<partition STRING, name STRING, ts TIMESTAMP(3)>"))
+                                        .partitionBy("partition")
+                                        .build())
                         .withOnTimeColumn("ts")
                         .build()) {
 
@@ -2519,8 +2753,12 @@ void testStateAccessInOnTimer() throws Exception {
         try (ProcessTableFunctionTestHarness<Row> harness =
                 ProcessTableFunctionTestHarness.ofClass(TimerWithStatePTF.class)
                         .withTableArgument(
-                                "input", DataTypes.of("ROW<partition STRING, ts TIMESTAMP(3)>"))
-                        .withPartitionBy("input", "partition")
+                                TableArgument.forArgument("input")
+                                        .type(
+                                                DataTypes.of(
+                                                        "ROW<partition STRING, ts TIMESTAMP(3)>"))
+                                        .partitionBy("partition")
+                                        .build())
                         .withOnTimeColumn("ts")
                         .build()) {
 
@@ -2549,9 +2787,12 @@ void testWatermarkCannotMoveBackward() throws Exception {
         try (ProcessTableFunctionTestHarness<Row> harness =
                 ProcessTableFunctionTestHarness.ofClass(TimerPTF.class)
                         .withTableArgument(
-                                "input",
-                                DataTypes.of("ROW<partition STRING, name STRING, ts TIMESTAMP(3)>"))
-                        .withPartitionBy("input", "partition")
+                                TableArgument.forArgument("input")
+                                        .type(
+                                                DataTypes.of(
+                                                        "ROW<partition STRING, name STRING, ts TIMESTAMP(3)>"))
+                                        .partitionBy("partition")
+                                        .build())
                         .withOnTimeColumn("ts")
                         .build()) {
 
@@ -2577,8 +2818,12 @@ void testPassThroughWithTimersIsRejected() throws Exception {
         try (ProcessTableFunctionTestHarness<Row> harness =
                 ProcessTableFunctionTestHarness.ofClass(PassThroughTimerPTF.class)
                         .withTableArgument(
-                                "input", DataTypes.of("ROW<partition STRING, ts TIMESTAMP(3)>"))
-                        .withPartitionBy("input", "partition")
+                                TableArgument.forArgument("input")
+                                        .type(
+                                                DataTypes.of(
+                                                        "ROW<partition STRING, ts TIMESTAMP(3)>"))
+                                        .partitionBy("partition")
+                                        .build())
                         .withOnTimeColumn("ts")
                         .build()) {
 
@@ -2597,8 +2842,12 @@ void testNoOnTimerMethodThrows() throws Exception {
         try (ProcessTableFunctionTestHarness<Row> harness =
                 ProcessTableFunctionTestHarness.ofClass(NoOnTimerPTF.class)
                         .withTableArgument(
-                                "input", DataTypes.of("ROW<partition STRING, ts TIMESTAMP(3)>"))
-                        .withPartitionBy("input", "partition")
+                                TableArgument.forArgument("input")
+                                        .type(
+                                                DataTypes.of(
+                                                        "ROW<partition STRING, ts TIMESTAMP(3)>"))
+                                        .partitionBy("partition")
+                                        .build())
                         .withOnTimeColumn("ts")
                         .build()) {
 
@@ -2620,10 +2869,12 @@ void testRequireOnTimeWithoutOnTimeColumnIsRejected() {
                         () ->
                                 ProcessTableFunctionTestHarness.ofClass(TimerPTF.class)
                                         .withTableArgument(
-                                                "input",
-                                                DataTypes.of(
-                                                        "ROW<partition STRING, ts TIMESTAMP(3)>"))
-                                        .withPartitionBy("input", "partition")
+                                                TableArgument.forArgument("input")
+                                                        .type(
+                                                                DataTypes.of(
+                                                                        "ROW<partition STRING, ts TIMESTAMP(3)>"))
+                                                        .partitionBy("partition")
+                                                        .build())
                                         .build());
         assertThat(e.getMessage()).contains("requires a time attribute", "on_time");
     }
@@ -2633,8 +2884,12 @@ void testMultipleOnTimerMethods() throws Exception {
         try (ProcessTableFunctionTestHarness<Row> harness =
                 ProcessTableFunctionTestHarness.ofClass(MultipleOnTimerPTF.class)
                         .withTableArgument(
-                                "input", DataTypes.of("ROW<partition STRING, ts TIMESTAMP(3)>"))
-                        .withPartitionBy("input", "partition")
+                                TableArgument.forArgument("input")
+                                        .type(
+                                                DataTypes.of(
+                                                        "ROW<partition STRING, ts TIMESTAMP(3)>"))
+                                        .partitionBy("partition")
+                                        .build())
                         .withOnTimeColumn("ts")
                         .build()) {
 
@@ -2653,8 +2908,12 @@ void testCascadingTimerFromOnTimer() throws Exception {
         try (ProcessTableFunctionTestHarness<Row> harness =
                 ProcessTableFunctionTestHarness.ofClass(CascadingTimerPTF.class)
                         .withTableArgument(
-                                "input", DataTypes.of("ROW<partition STRING, ts TIMESTAMP(3)>"))
-                        .withPartitionBy("input", "partition")
+                                TableArgument.forArgument("input")
+                                        .type(
+                                                DataTypes.of(
+                                                        "ROW<partition STRING, ts TIMESTAMP(3)>"))
+                                        .partitionBy("partition")
+                                        .build())
                         .withOnTimeColumn("ts")
                         .build()) {
 
@@ -2683,10 +2942,12 @@ void testClearTimerByName() throws Exception {
         try (ProcessTableFunctionTestHarness<Row> harness =
                 ProcessTableFunctionTestHarness.ofClass(ClearTimerPTF.class)
                         .withTableArgument(
-                                "input",
-                                DataTypes.of(
-                                        "ROW<partition STRING, action STRING, ts TIMESTAMP(3)>"))
-                        .withPartitionBy("input", "partition")
+                                TableArgument.forArgument("input")
+                                        .type(
+                                                DataTypes.of(
+                                                        "ROW<partition STRING, action STRING, ts TIMESTAMP(3)>"))
+                                        .partitionBy("partition")
+                                        .build())
                         .withOnTimeColumn("ts")
                         .build()) {
 
@@ -2709,10 +2970,12 @@ void testClearTimerByTimestamp() throws Exception {
         try (ProcessTableFunctionTestHarness<Row> harness =
                 ProcessTableFunctionTestHarness.ofClass(ClearTimerPTF.class)
                         .withTableArgument(
-                                "input",
-                                DataTypes.of(
-                                        "ROW<partition STRING, action STRING, ts TIMESTAMP(3)>"))
-                        .withPartitionBy("input", "partition")
+                                TableArgument.forArgument("input")
+                                        .type(
+                                                DataTypes.of(
+                                                        "ROW<partition STRING, action STRING, ts TIMESTAMP(3)>"))
+                                        .partitionBy("partition")
+                                        .build())
                         .withOnTimeColumn("ts")
                         .build()) {
 
@@ -2737,10 +3000,12 @@ void testClearTimerByTimestampDoesNotClearNamedTimers() throws Exception {
         try (ProcessTableFunctionTestHarness<Row> harness =
                 ProcessTableFunctionTestHarness.ofClass(ClearTimerPTF.class)
                         .withTableArgument(
-                                "input",
-                                DataTypes.of(
-                                        "ROW<partition STRING, action STRING, ts TIMESTAMP(3)>"))
-                        .withPartitionBy("input", "partition")
+                                TableArgument.forArgument("input")
+                                        .type(
+                                                DataTypes.of(
+                                                        "ROW<partition STRING, action STRING, ts TIMESTAMP(3)>"))
+                                        .partitionBy("partition")
+                                        .build())
                         .withOnTimeColumn("ts")
                         .build()) {
 
@@ -2766,10 +3031,12 @@ void testClearAllTimers() throws Exception {
         try (ProcessTableFunctionTestHarness<Row> harness =
                 ProcessTableFunctionTestHarness.ofClass(ClearTimerPTF.class)
                         .withTableArgument(
-                                "input",
-                                DataTypes.of(
-                                        "ROW<partition STRING, action STRING, ts TIMESTAMP(3)>"))
-                        .withPartitionBy("input", "partition")
+                                TableArgument.forArgument("input")
+                                        .type(
+                                                DataTypes.of(
+                                                        "ROW<partition STRING, action STRING, ts TIMESTAMP(3)>"))
+                                        .partitionBy("partition")
+                                        .build())
                         .withOnTimeColumn("ts")
                         .build()) {
 
@@ -2793,8 +3060,10 @@ void testTimersWithoutOnTimeColumn() throws Exception {
         try (ProcessTableFunctionTestHarness<Row> harness =
                 ProcessTableFunctionTestHarness.ofClass(WatermarkOnlyTimerPTF.class)
                         .withTableArgument(
-                                "input", DataTypes.of("ROW<partition STRING, value INT>"))
-                        .withPartitionBy("input", "partition")
+                                TableArgument.forArgument("input")
+                                        .type(DataTypes.of("ROW<partition STRING, value INT>"))
+                                        .partitionBy("partition")
+                                        .build())
                         .build()) {
 
             harness.setWatermark(Instant.ofEpochMilli(1000));