Beam: Moved instantiation to Setup annotated method

nielsbasjes · Jul 19, 2017 · 1458afb · 1458afb
1 parent dfacd6b
commit 1458afb
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 6 deletions.
diff --git a/udfs/beam/src/main/java/nl/basjes/parse/useragent/beam/UserAgentAnalysisDoFn.java b/udfs/beam/src/main/java/nl/basjes/parse/useragent/beam/UserAgentAnalysisDoFn.java
@@ -28,13 +28,14 @@ public abstract class UserAgentAnalysisDoFn<T extends Serializable> extends DoFn
     implements UseragentAnnotationMapper<T>, Serializable {
     private transient UserAgentAnnotationAnalyzer<T> userAgentAnalyzer = null;
 
+    @Setup
+    public void initialize() {
+        userAgentAnalyzer = new UserAgentAnnotationAnalyzer<>();
+        userAgentAnalyzer.initialize(this);
+    }
+
     @ProcessElement
     public void processElement(ProcessContext c) {
-        if (userAgentAnalyzer == null) {
-            userAgentAnalyzer = new UserAgentAnnotationAnalyzer<>();
-            userAgentAnalyzer.initialize(this);
-        }
-
         // Currently Beam does not allow changing the input instance.
         // So unfortunately we must clone the entire thing :(
         // See also: https://issues.apache.org/jira/browse/BEAM-1164

diff --git a/udfs/beam/src/test/java/nl/basjes/parse/useragent/beam/TestUserAgentAnalysisDoFnRaw.java b/udfs/beam/src/test/java/nl/basjes/parse/useragent/beam/TestUserAgentAnalysisDoFnRaw.java
@@ -21,6 +21,7 @@
 import nl.basjes.parse.useragent.annonate.YauaaField;
 import org.apache.beam.sdk.transforms.DoFn;
 import org.apache.beam.sdk.transforms.DoFnTester;
+import org.apache.beam.sdk.util.UserCodeException;
 import org.junit.Test;
 
 import java.io.Serializable;
@@ -82,7 +83,7 @@ public void testUserAgentAnalysisDoFn() throws Exception {
         assertNull(record.shouldRemainNull);
     }
 
-    @Test(expected = InvalidParserConfigurationException.class)
+    @Test(expected = UserCodeException.class)
     public void testImpossibleField() throws Exception {
         DoFn<TestRecord, TestRecord> fn = new UserAgentAnalysisDoFn<TestRecord>() {
             @Override