Merge pull request twitter#95 from rangadi/allow_enum_keys

Allow enum keys in thrift maps
jcoveney · Nov 4, 2011 · 9126c07 · 9126c07
2 parents 4dd770c + 014d9ca
commit 9126c07
Show file tree

Hide file tree

Showing 4 changed files with 52 additions and 7 deletions.
diff --git a/src/java/com/twitter/elephantbird/pig/util/PigToThrift.java b/src/java/com/twitter/elephantbird/pig/util/PigToThrift.java
@@ -143,14 +143,17 @@ private static Object toStringType(Object value) {
     return null;
   }
 
-  private static Map<String, Object> toThriftMap(Field field, Map<String, Object> map) {
-    if (field.getMapKeyField().getType() != TType.STRING) {
-      throw new IllegalArgumentException("TStructs's map key is not a String");
-    }
-    HashMap<String, Object> out = new HashMap<String, Object>(map.size());
+  private static Map<Object, Object> toThriftMap(Field field, Map<String, Object> map) {
+    Field keyField   = field.getMapKeyField();
     Field valueField = field.getMapValueField();
+    if (keyField.getType() != TType.STRING
+        && keyField.getType() != TType.ENUM) {
+      throw new IllegalArgumentException("TStructs's map key should be a STRING or an ENUM");
+    }
+    HashMap<Object, Object> out = new HashMap<Object, Object>(map.size());
     for(Entry<String, Object> e : map.entrySet()) {
-      out.put(e.getKey(), toThriftValue(valueField, e.getValue()));
+      out.put(toThriftValue(keyField,   e.getKey()),
+              toThriftValue(valueField, e.getValue()));
     }
     return out;
   }

diff --git a/src/java/com/twitter/elephantbird/pig/util/ThriftToPig.java b/src/java/com/twitter/elephantbird/pig/util/ThriftToPig.java
@@ -245,7 +245,8 @@ private static FieldSchema singleFieldToFieldSchema(String fieldName, Field fiel
         return new FieldSchema(fieldName, singleFieldToTupleSchema(fieldName + "_tuple", field.getSetElemField()), DataType.BAG);
       case TType.MAP:
         // can not specify types for maps in Pig.
-        if (field.getMapKeyField().getType() != TType.STRING) {
+        if (field.getMapKeyField().getType() != TType.STRING
+            && field.getMapKeyField().getType() != TType.ENUM) {
           LOG.warn("Using a map with non-string key for field " + field.getName()
               + ". while converting to PIG Tuple, toString() is used for the key."
               + " It could result in incorrect maps.");

diff --git a/src/test/com/twitter/elephantbird/pig/piggybank/TestThriftToPig.java b/src/test/com/twitter/elephantbird/pig/piggybank/TestThriftToPig.java
@@ -1,6 +1,7 @@
 package com.twitter.elephantbird.pig.piggybank;
 
 import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertTrue;
 
 import java.io.IOException;
 import java.nio.ByteBuffer;
@@ -18,6 +19,7 @@
 import thrift.test.Nesting;
 import thrift.test.OneOfEach;
 
+import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.Lists;
 import com.twitter.data.proto.tutorial.thrift.Name;
 import com.twitter.data.proto.tutorial.thrift.Person;
@@ -26,6 +28,9 @@
 import com.twitter.elephantbird.mapreduce.io.ThriftConverter;
 import com.twitter.elephantbird.pig.util.ThriftToPig;
 import com.twitter.elephantbird.pig.util.PigToThrift;
+import com.twitter.elephantbird.thrift.test.TestName;
+import com.twitter.elephantbird.thrift.test.TestPerson;
+import com.twitter.elephantbird.thrift.test.TestPhoneType;
 import com.twitter.elephantbird.util.TypeRef;
 
 public class TestThriftToPig {
@@ -132,5 +137,19 @@ private void tupleTest(TestType type) throws Exception {
 
     Person person = new Person(new Name("bob", "jenkins"), 42, "foo@bar.com", Lists.newArrayList(ph));
     assertEquals("(bob,jenkins),42,foo@bar.com,{(415-555-5555,HOME)}", toTuple(type, person).toDelimitedString(","));
+
+    // test Enum map
+    TestPerson testPerson = new TestPerson(
+                                  new TestName("bob", "jenkins"),
+                                  ImmutableMap.of(
+                                      TestPhoneType.HOME,   "408-555-5555",
+                                      TestPhoneType.MOBILE, "650-555-5555",
+                                      TestPhoneType.WORK,   "415-555-5555"));
+
+    String tupleString = toTuple(type, testPerson).toDelimitedString("-");
+    assertTrue( // the order of elements in map could vary because of HashMap
+        tupleString.equals("(bob,jenkins)-{MOBILE=650-555-5555, WORK=415-555-5555, HOME=408-555-5555}") ||
+        tupleString.equals("(bob,jenkins)-{MOBILE=650-555-5555, HOME=408-555-5555, WORK=415-555-5555}")
+    );
   }
 }
diff --git a/src/thrift/test.thrift b/src/thrift/test.thrift
@@ -0,0 +1,22 @@
+namespace java com.twitter.elephantbird.thrift.test
+
+/**
+ * various thrift classes used in unit tests
+ */
+
+enum TestPhoneType {
+  MOBILE = 0,
+  HOME = 1,
+  WORK = 2
+}
+
+struct TestName {
+  1: string first_name,
+  2: string last_name
+}
+
+struct TestPerson {
+  1: TestName                     name,
+  2: map<TestPhoneType, string>   phones, // for testing enum keys in maps.
+}
+