Use a bimap for reverse lookups on injective maps (#5681)

* Use a bimap for reverse lookups on injective maps

- A BiMap provides constant-time lookups for mapping values to keys

* Address comments

* Fix Tests

processing/src/main/java/io/druid/query/extraction/MapLookupExtractor.java

@@ -23,9 +23,9 @@ import com.fasterxml.jackson.annotation.JsonCreator;
 import com.fasterxml.jackson.annotation.JsonProperty;
 import com.fasterxml.jackson.annotation.JsonTypeName;
 import com.google.common.base.Preconditions;
-import com.google.common.base.Predicate;
 import com.google.common.base.Strings;
 import com.google.common.base.Throwables;
+import com.google.common.collect.HashBiMap;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.Lists;
 import com.google.common.collect.Maps;
@@ -36,6 +36,7 @@ import javax.annotation.Nullable;
 import javax.validation.constraints.NotNull;
 import java.io.ByteArrayOutputStream;
 import java.io.IOException;
+import java.util.Collections;
 import java.util.List;
 import java.util.Map;
 
@@ -43,6 +44,7 @@ import java.util.Map;
 public class MapLookupExtractor extends LookupExtractor
 {
   private final Map<String, String> map;
+  private final Map<String, String> reverseMap;
 
   private final boolean isOneToOne;
 
@@ -52,8 +54,18 @@ public class MapLookupExtractor extends LookupExtractor
       @JsonProperty("isOneToOne") boolean isOneToOne
   )
   {
-    this.map = Preconditions.checkNotNull(map, "map");
+
+    Preconditions.checkNotNull(map, "map");
+
     this.isOneToOne = isOneToOne;
+
+    if (this.isOneToOne) {
+      this.map = HashBiMap.create(map);
+      this.reverseMap = ((HashBiMap<String, String>) this.map).inverse();
+    } else {
+      this.map = map;
+      this.reverseMap = null;
+    }
   }
 
   @JsonProperty
@@ -72,14 +84,14 @@ public class MapLookupExtractor extends LookupExtractor
   @Override
   public List<String> unapply(final String value)
   {
-    return Lists.newArrayList(Maps.filterKeys(map, new Predicate<String>()
-    {
-      @Override public boolean apply(@Nullable String key)
-      {
-        return map.get(key).equals(Strings.nullToEmpty(value));
-      }
-    }).keySet());
+    String valueToLookup = Strings.nullToEmpty(value);
 
+    if (this.reverseMap != null) {
+      String val = this.reverseMap.get(valueToLookup);
+      return (val != null) ? Collections.singletonList(val) : Collections.emptyList();
+    } else {
+      return Lists.newArrayList(Maps.filterKeys(map, key -> map.get(key).equals(valueToLookup)).keySet());
+    }
   }
 
   @Override

processing/src/main/java/io/druid/query/lookup/LookupExtractor.java

@@ -74,7 +74,7 @@ public abstract class LookupExtractor
    *              Null and empty are considered to be the same value = nullToEmpty(value)
    *
    * @return the list of keys that maps to value or empty list.
-   * Note that for the case of a none existing value in the lookup we have to cases either return an empty list OR list with null element.
+   * Note that for the case of a none existing value in the lookup we have two cases either return an empty list OR list with null element.
    * returning an empty list implies that user want to ignore such a lookup value.
    * In the other hand returning a list with the null element implies user want to map the none existing value to the key null.
    */

processing/src/test/java/io/druid/query/extraction/MapLookupExtractorTest.java

@@ -35,7 +35,7 @@ public class MapLookupExtractorTest
   private final MapLookupExtractor fn = new MapLookupExtractor(lookupMap, false);
 
   @Test
-  public void testUnApply()
+  public void testNonInjectiveUnApply()
   {
     Assert.assertEquals(Arrays.asList("foo"), fn.unapply("bar"));
     Assert.assertEquals(Sets.newHashSet("null", "empty String"), Sets.newHashSet(fn.unapply("")));
@@ -46,6 +46,23 @@ public class MapLookupExtractorTest
     Assert.assertEquals("not existing value returns empty list", Collections.EMPTY_LIST, fn.unapply("not There"));
   }
 
+  @Test
+  public void testInjectiveUnApply()
+  {
+    MapLookupExtractor injectiveFn = new MapLookupExtractor(
+        ImmutableMap.of("foo", "bar", "null", "", "", "empty_string"), true
+    );
+
+    Assert.assertEquals(Arrays.asList("foo"), injectiveFn.unapply("bar"));
+    Assert.assertEquals(
+        "Null value should be equal to empty string",
+        Sets.newHashSet("null"),
+        Sets.newHashSet(injectiveFn.unapply((String) null))
+    );
+    Assert.assertEquals(Sets.newHashSet(""), Sets.newHashSet(injectiveFn.unapply("empty_string")));
+    Assert.assertEquals("not existing value returns empty list", Collections.EMPTY_LIST, injectiveFn.unapply("not There"));
+  }
+
   @Test
   public void testGetMap()
   {