From 2bce2ce4d946731e4d2f883bb65f80b9ff5a23b2 Mon Sep 17 00:00:00 2001
From: manuzhang <owenzhang1990@gmail.com>
Date: Wed, 23 Nov 2016 09:24:05 +0800
Subject: [PATCH] [BEAM-800] add getFn to DoFnInvoker

---
 .../sdk/transforms/reflect/ByteBuddyDoFnInvokerFactory.java  | 5 +++++
 .../org/apache/beam/sdk/transforms/reflect/DoFnInvoker.java  | 3 +++
 .../org/apache/beam/sdk/transforms/reflect/DoFnInvokers.java | 5 +++++
 3 files changed, 13 insertions(+)

diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/reflect/ByteBuddyDoFnInvokerFactory.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/reflect/ByteBuddyDoFnInvokerFactory.java
index 9998c9d78801..408ea98e2d51 100644
--- a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/reflect/ByteBuddyDoFnInvokerFactory.java
+++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/reflect/ByteBuddyDoFnInvokerFactory.java
@@ -175,6 +175,11 @@ public void invokeOnTimer(
                 timerId, delegate.getClass().getName()));
       }
     }
+
+    @Override
+    public DoFn<InputT, OutputT> getFn() {
+      return delegate;
+    }
   }
 
   /** @return the {@link DoFnInvoker} for the given {@link DoFn}. */
diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/reflect/DoFnInvoker.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/reflect/DoFnInvoker.java
index d899207a6f34..5e61bddbe4d9 100644
--- a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/reflect/DoFnInvoker.java
+++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/reflect/DoFnInvoker.java
@@ -84,6 +84,9 @@ <RestrictionT> void invokeSplitRestriction(
   <RestrictionT, TrackerT extends RestrictionTracker<RestrictionT>> TrackerT invokeNewTracker(
       RestrictionT restriction);
 
+  /** Get the bound {@link DoFn}. */
+  DoFn<InputT, OutputT> getFn();
+
   /**
    * Interface for runner implementors to provide implementations of extra context information.
    *
diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/reflect/DoFnInvokers.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/reflect/DoFnInvokers.java
index 15ba1987eaab..4ad7dad30e95 100644
--- a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/reflect/DoFnInvokers.java
+++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/reflect/DoFnInvokers.java
@@ -180,5 +180,10 @@ public <RestrictionT> void invokeSplitRestriction(
         TrackerT invokeNewTracker(RestrictionT restriction) {
       throw new UnsupportedOperationException("OldDoFn is not splittable");
     }
+
+    @Override
+    public DoFn<InputT, OutputT> getFn() {
+      throw new UnsupportedOperationException("getFn is not supported for OldDoFn");
+    }
   }
 }