Modified: 
incubator/pig/branches/types/test/org/apache/pig/test/utils/LogicalPlanTester.java
URL: 
http://svn.apache.org/viewvc/incubator/pig/branches/types/test/org/apache/pig/test/utils/LogicalPlanTester.java?rev=672801&r1=672800&r2=672801&view=diff
==============================================================================
--- 
incubator/pig/branches/types/test/org/apache/pig/test/utils/LogicalPlanTester.java
 (original)
+++ 
incubator/pig/branches/types/test/org/apache/pig/test/utils/LogicalPlanTester.java
 Mon Jun 30 08:56:41 2008
@@ -19,12 +19,14 @@
 package org.apache.pig.test.utils;
 
 import org.apache.pig.impl.logicalLayer.*;
+import org.apache.pig.impl.logicalLayer.optimizer.LogicalOptimizer;
 import org.apache.pig.impl.logicalLayer.validators.TypeCheckingValidator;
 import org.apache.pig.impl.PigContext;
 import org.apache.pig.impl.plan.PlanValidationException;
 import org.apache.pig.impl.plan.CompilationMessageCollector;
 import org.apache.pig.impl.plan.OperatorKey;
 import org.apache.pig.impl.plan.NodeIdGenerator;
+import org.apache.pig.impl.plan.optimizer.OptimizerException;
 import org.apache.pig.ExecType;
 import static org.apache.pig.test.utils.TypeCheckingTestUtil.* ;
 import org.apache.pig.test.utils.dotGraph.LogicalPlanLoader;
@@ -87,7 +89,14 @@
         TypeCheckingValidator typeValidator = new TypeCheckingValidator() ;
         typeValidator.validate(plan, collector) ;
         printMessageCollector(collector) ;
-        System.out.println("Actual plan:") ;
+        System.out.println("Actual plan after type check:") ;
+        printTypeGraph(plan) ;
+    }
+
+    public void optimizePlan(LogicalPlan plan) throws OptimizerException {
+        LogicalOptimizer optimizer = new LogicalOptimizer(plan);
+        optimizer.optimize();
+        System.out.println("Actual plan after after optimization:") ;
         printTypeGraph(plan) ;
     }
 
@@ -98,11 +107,30 @@
      * @param file
      * @throws PlanValidationException
      */
-    public void typeCheckAgainstDotFile(LogicalPlan plan, String file)
-                                            throws PlanValidationException {
+    public void typeCheckAgainstDotFile(
+            LogicalPlan plan,
+            String file) throws PlanValidationException, OptimizerException {
+        typeCheckAgainstDotFile(plan, file, false);
+    }
+
+    /***
+     * Run type checking and compare the result with  plan structure
+     * stored in Dot file
+     * @param plan
+     * @param file
+     * @param optimize if true, the plan will be run through the optimizer
+     * @throws PlanValidationException
+     */
+    public void typeCheckAgainstDotFile(
+            LogicalPlan plan,
+            String file,
+            boolean optimize) throws PlanValidationException,
+                                     OptimizerException {
         // validate the given plan
         typeCheckPlan(plan);
 
+        if (optimize) optimizePlan(plan);
+
         // load the expected plan from file
         LogicalPlanLoader planLoader = new LogicalPlanLoader() ;
         LogicalPlan expectedPlan = planLoader.loadFromFile(file, 
LogicalPlan.class) ;
@@ -120,8 +148,15 @@
         System.out.println("Checking DONE!") ;
     }
 
-    public void typeCheckUsingDotFile(String file)
-                                            throws PlanValidationException {
+    public void typeCheckUsingDotFile(
+            String file) throws PlanValidationException, OptimizerException {
+        typeCheckUsingDotFile(file, false);
+    }
+
+    public void typeCheckUsingDotFile(
+            String file,
+            boolean optimize) throws PlanValidationException,
+                                     OptimizerException {
         DotGraphReader reader = new DotGraphReader() ;
         DotGraph graph = reader.loadFromFile(file) ;
         if (!graph.attributes.containsKey("pigScript")) {
@@ -138,7 +173,7 @@
                 plan = buildPlan(query + ";") ;
             }
         }
-        typeCheckAgainstDotFile(plan, file) ;
+        typeCheckAgainstDotFile(plan, file, optimize) ;
 
     }
 


Reply via email to