PARQUET-116: Add a test case for passing a filter object to user defi…

…ned predicate
rdblue · Oct 18, 2014 · dfd0478 · dfd0478
1 parent 4ab46ec
commit dfd0478
Showing 1 changed file with 49 additions and 0 deletions.
diff --git a/parquet-hadoop/src/test/java/parquet/filter2/recordlevel/TestRecordLevelFilters.java b/parquet-hadoop/src/test/java/parquet/filter2/recordlevel/TestRecordLevelFilters.java
@@ -6,6 +6,8 @@
 import java.util.Arrays;
 import java.util.Iterator;
 import java.util.List;
+import java.util.Set;
+import java.util.HashSet;
 
 import org.junit.BeforeClass;
 import org.junit.Test;
@@ -15,6 +17,7 @@
 import parquet.filter2.predicate.FilterPredicate;
 import parquet.filter2.predicate.Operators.BinaryColumn;
 import parquet.filter2.predicate.Operators.DoubleColumn;
+import parquet.filter2.predicate.Operators.LongColumn;
 import parquet.filter2.predicate.Statistics;
 import parquet.filter2.predicate.UserDefinedPredicate;
 import parquet.filter2.recordlevel.PhoneBookWriter.Location;
@@ -26,6 +29,7 @@
 import static parquet.filter2.predicate.FilterApi.and;
 import static parquet.filter2.predicate.FilterApi.binaryColumn;
 import static parquet.filter2.predicate.FilterApi.doubleColumn;
+import static parquet.filter2.predicate.FilterApi.longColumn;
 import static parquet.filter2.predicate.FilterApi.eq;
 import static parquet.filter2.predicate.FilterApi.gt;
 import static parquet.filter2.predicate.FilterApi.not;
@@ -160,6 +164,29 @@ public boolean inverseCanDrop(Statistics<Binary> statistics) {
       return false;
     }
   }
+
+  public static class SetInFilter extends UserDefinedPredicate<Long> {
+
+    @Override
+    public boolean keep(Long value, Object o) {
+      if (value == null) {
+        return false;
+      }
+
+      Set<Long> hSet = (HashSet<Long>) o;
+      return hSet.contains(value);
+    }
+
+    @Override
+    public boolean canDrop(Statistics<Long> statistics) {
+      return false;
+    }
+
+    @Override
+    public boolean inverseCanDrop(Statistics<Long> statistics) {
+      return false;
+    }
+  }
 
   @Test
   public void testNameNotStartWithP() throws Exception {
@@ -176,6 +203,28 @@ public boolean keep(User u) {
       }
     });
   }
+
+  @Test
+  public void testIdIn() throws Exception {
+    LongColumn name = longColumn("id");
+
+    Set<Long> h = new HashSet<Long>() {{
+          add(20L); add(27L); add(28L);
+    }}; 
+    FilterPredicate pred = userDefined(name, SetInFilter.class, h);
+
+    List<Group> found = PhoneBookWriter.readFile(phonebookFile, FilterCompat.get(pred));
+
+    assertFilter(found, new UserFilter() {
+      @Override
+      public boolean keep(User u) {
+        Set<Long> h = new HashSet<Long>() {{
+          add(20L); add(27L); add(28L);
+        }}; 
+        return h.contains(u.getId());
+      }
+    });
+  }
 
   @Test
   public void testComplex() throws Exception {