PIG-4974: A simple map reference fail to cast (knoguchi)


git-svn-id: https://svn.apache.org/repos/asf/pig/trunk@1759893 13f79535-47bb-0310-9956-ffa450edef68
diff --git a/CHANGES.txt b/CHANGES.txt
index 5574f37..a503ab2 100644
--- a/CHANGES.txt
+++ b/CHANGES.txt
@@ -44,6 +44,8 @@
  
 BUG FIXES
 
+PIG-4974: A simple map reference fail to cast (knoguchi)
+
 PIG-4975 Map schema shows "Type: null Uid: null" in explain (knoguchi)
 
 PIG-4973: Bigdecimal divison fails (szita via daijy)
diff --git a/src/org/apache/pig/newplan/logical/visitor/LineageFindRelVisitor.java b/src/org/apache/pig/newplan/logical/visitor/LineageFindRelVisitor.java
index b5c44ee..ec86129 100644
--- a/src/org/apache/pig/newplan/logical/visitor/LineageFindRelVisitor.java
+++ b/src/org/apache/pig/newplan/logical/visitor/LineageFindRelVisitor.java
@@ -729,6 +729,15 @@
             }
         }
 
+        @Override
+        public void visit(UserFuncExpression op) throws FrontendException {
+            if(op.getFuncSpec().getClassName().equals(IdentityColumn.class.getName())) {
+                // IdentityColumn only expects one arg
+                FuncSpec funcSpec = uid2LoadFuncMap.get(op.getArguments().get(0).getFieldSchema().uid);
+                addUidLoadFuncToMap(op.getFieldSchema().uid, funcSpec);
+            }
+        }
+
         /**
          * if there is a null constant under casts, return it
          * @param rel
diff --git a/test/org/apache/pig/test/TestLineageFindRelVisitor.java b/test/org/apache/pig/test/TestLineageFindRelVisitor.java
index 0e5db9b..2b1f388 100644
--- a/test/org/apache/pig/test/TestLineageFindRelVisitor.java
+++ b/test/org/apache/pig/test/TestLineageFindRelVisitor.java
@@ -24,11 +24,16 @@
 
 import java.io.IOException;
 import java.lang.reflect.Method;
+import java.util.List;
 
 import org.apache.pig.FuncSpec;
 import org.apache.pig.LoadCaster;
+import org.apache.pig.PigServer;
 import org.apache.pig.builtin.PigStorage;
 import org.apache.pig.builtin.Utf8StorageConverter;
+import org.apache.pig.builtin.mock.Storage;
+import org.apache.pig.data.DataByteArray;
+import org.apache.pig.data.Tuple;
 import org.apache.pig.impl.io.FileSpec;
 import org.apache.pig.newplan.logical.relational.LOLoad;
 import org.apache.pig.newplan.logical.relational.LogicalPlan;
@@ -42,6 +47,12 @@
 
 public class TestLineageFindRelVisitor {
 
+    private PigServer pig ;
+
+    public TestLineageFindRelVisitor() throws Throwable {
+        pig = new PigServer(Util.getLocalTestMode()) ;
+    }
+
     public static class SillyLoadCasterWithExtraConstructor extends Utf8StorageConverter {
         public SillyLoadCasterWithExtraConstructor(String ignored) {
             super();
@@ -125,4 +136,32 @@
 
         Assert.assertEquals("Loader should be instantiated at most once.", SillyLoaderWithLoadCasterWithExtraConstructor.counter, 1);
     }
+
+    @Test
+    public void testIdenticalColumnUDFForwardingLoadCaster() throws Exception {
+        Storage.Data data = Storage.resetData(pig);
+        data.set("input",
+                Storage.tuple(Storage.map(
+                                 "key1",new DataByteArray("aaa"),
+                                 "key2",new DataByteArray("bbb"),
+                                 "key3",new DataByteArray("ccc"))),
+                Storage.tuple(Storage.map(
+                                 "key1",new DataByteArray("zzz"),
+                                 "key2",new DataByteArray("yyy"),
+                                 "key3",new DataByteArray("xxx"))));
+        pig.setBatchOn();
+        pig.registerQuery("A = load 'input' using mock.Storage() as (m:[bytearray]);");
+        pig.registerQuery("B = foreach A GENERATE m#'key1' as key1, m#'key2' as key2; "
+                // this equal comparison creates implicit typecast to chararray
+                // which requires loadcaster
+                + "C = FILTER B by key1 == 'aaa' and key2 == 'bbb';");
+        pig.registerQuery("store C into 'output' using mock.Storage();");
+
+        pig.executeBatch();
+
+        List<Tuple> actualResults = data.get("output");
+        List<Tuple> expectedResults = Util.getTuplesFromConstantTupleStrings(
+                new String[] {"('aaa', 'bbb')"});
+        Util.checkQueryOutputs(actualResults.iterator(), expectedResults);
+    }
 }