PARQUET-116: Pass a filter object to user defined predicate in filter2 api

parquet-column/src/main/java/parquet/filter2/predicate/FilterApi.java

@@ -1,5 +1,7 @@
 package parquet.filter2.predicate;
 
+import java.io.Serializable;
+
 import parquet.common.schema.ColumnPath;
 import parquet.filter2.predicate.Operators.And;
 import parquet.filter2.predicate.Operators.BinaryColumn;
@@ -20,6 +22,8 @@
 import parquet.filter2.predicate.Operators.SupportsEqNotEq;
 import parquet.filter2.predicate.Operators.SupportsLtGt;
 import parquet.filter2.predicate.Operators.UserDefined;
+import parquet.filter2.predicate.Operators.UserDefinedByClass;
+import parquet.filter2.predicate.Operators.UserDefinedByInstance;
 
 /**
  * The Filter API is expressed through these static methods.
@@ -144,10 +148,23 @@ public static <T extends Comparable<T>, C extends Column<T> & SupportsLtGt> GtEq
 
   /**
    * Keeps records that pass the provided {@link UserDefinedPredicate}
+   *
+   * The provided class must have a default constructor. To use an instance
+   * of a UserDefinedPredicate instead, see {@link #userDefined(column, udp)} below.
    */
   public static <T extends Comparable<T>, U extends UserDefinedPredicate<T>>
     UserDefined<T, U> userDefined(Column<T> column, Class<U> clazz) {
-    return new UserDefined<T, U>(column, clazz);
+    return new UserDefinedByClass<T, U>(column, clazz);
+  }
+
+  /**
+   * Keeps records that pass the provided {@link UserDefinedPredicate}
+   *
+   * The provided instance of UserDefinedPredicate must be serializable.
+   */
+  public static <T extends Comparable<T>, U extends UserDefinedPredicate<T> & Serializable>
+    UserDefined<T, U> userDefined(Column<T> column, U udp) {
+    return new UserDefinedByInstance<T, U>(column, udp);
   }
 
   /**

parquet-column/src/main/java/parquet/filter2/predicate/Operators.java

@@ -340,15 +340,33 @@ public int hashCode() {
     }
   }
 
-  public static final class UserDefined<T extends Comparable<T>, U extends UserDefinedPredicate<T>> implements FilterPredicate, Serializable {
-    private final Column<T> column;
+  public static abstract class UserDefined<T extends Comparable<T>, U extends UserDefinedPredicate<T>> implements FilterPredicate, Serializable {
+    protected final Column<T> column;
+
+    UserDefined(Column<T> column) {
+      this.column = checkNotNull(column, "column");
+    }
+
+    public Column<T> getColumn() {
+      return column;
+    }
+
+    public abstract U getUserDefinedPredicate();
+
+    @Override
+    public <R> R accept(Visitor<R> visitor) {
+      return visitor.visit(this);
+    }
+  }
+
+  public static final class UserDefinedByClass<T extends Comparable<T>, U extends UserDefinedPredicate<T>> extends UserDefined<T, U> {
     private final Class<U> udpClass;
     private final String toString;
     private static final String INSTANTIATION_ERROR_MESSAGE =
         "Could not instantiate custom filter: %s. User defined predicates must be static classes with a default constructor.";
 
-    UserDefined(Column<T> column, Class<U> udpClass) {
-      this.column = checkNotNull(column, "column");
+    UserDefinedByClass(Column<T> column, Class<U> udpClass) {
+      super(column);
       this.udpClass = checkNotNull(udpClass, "udpClass");
       String name = getClass().getSimpleName().toLowerCase();
       this.toString = name + "(" + column.getColumnPath().toDotString() + ", " + udpClass.getName() + ")";
@@ -357,14 +375,11 @@ public static final class UserDefined<T extends Comparable<T>, U extends UserDef
       getUserDefinedPredicate();
     }
 
-    public Column<T> getColumn() {
-      return column;
-    }
-
     public Class<U> getUserDefinedPredicateClass() {
       return udpClass;
     }
 
+    @Override
     public U getUserDefinedPredicate() {
       try {
         return udpClass.newInstance();
@@ -376,8 +391,46 @@ public U getUserDefinedPredicate() {
     }
 
     @Override
-    public <R> R accept(Visitor<R> visitor) {
-      return visitor.visit(this);
+    public String toString() {
+      return toString;
+    }
+
+    @Override
+    public boolean equals(Object o) {
+      if (this == o) return true;
+      if (o == null || getClass() != o.getClass()) return false;
+
+      UserDefinedByClass that = (UserDefinedByClass) o;
+
+      if (!column.equals(that.column)) return false;
+      if (!udpClass.equals(that.udpClass)) return false;
+
+      return true;
+    }
+
+    @Override
+    public int hashCode() {
+      int result = column.hashCode();
+      result = 31 * result + udpClass.hashCode();
+      result = result * 31 + getClass().hashCode();
+      return result;
+    }
+  }
+
+  public static final class UserDefinedByInstance<T extends Comparable<T>, U extends UserDefinedPredicate<T> & Serializable> extends UserDefined<T, U> {
+    private final String toString;
+    private final U udpInstance;
+
+    UserDefinedByInstance(Column<T> column, U udpInstance) {
+      super(column);
+      this.udpInstance = checkNotNull(udpInstance, "udpInstance");
+      String name = getClass().getSimpleName().toLowerCase();
+      this.toString = name + "(" + column.getColumnPath().toDotString() + ", " + udpInstance + ")";
+    }
+
+    @Override
+    public U getUserDefinedPredicate() {
+      return udpInstance;
     }
 
     @Override
@@ -390,18 +443,18 @@ public boolean equals(Object o) {
       if (this == o) return true;
       if (o == null || getClass() != o.getClass()) return false;
 
-      UserDefined that = (UserDefined) o;
+      UserDefinedByInstance that = (UserDefinedByInstance) o;
 
       if (!column.equals(that.column)) return false;
-      if (!udpClass.equals(that.udpClass)) return false;
+      if (!udpInstance.equals(that.udpInstance)) return false;
 
       return true;
     }
 
     @Override
     public int hashCode() {
       int result = column.hashCode();
-      result = 31 * result + udpClass.hashCode();
+      result = 31 * result + udpInstance.hashCode();
       result = result * 31 + getClass().hashCode();
       return result;
     }

parquet-column/src/main/java/parquet/filter2/predicate/UserDefinedPredicate.java

@@ -87,4 +87,4 @@ public UserDefinedPredicate() { }
    * }
    */
   public abstract boolean inverseCanDrop(Statistics<T> statistics);
-}
+}

parquet-column/src/test/java/parquet/filter2/predicate/TestFilterApiMethods.java

@@ -4,6 +4,7 @@
 import java.io.ByteArrayOutputStream;
 import java.io.ObjectInputStream;
 import java.io.ObjectOutputStream;
+import java.io.Serializable;
 
 import org.junit.Test;
 
@@ -14,9 +15,11 @@
 import parquet.filter2.predicate.Operators.Eq;
 import parquet.filter2.predicate.Operators.Gt;
 import parquet.filter2.predicate.Operators.IntColumn;
+import parquet.filter2.predicate.Operators.LongColumn;
 import parquet.filter2.predicate.Operators.Not;
 import parquet.filter2.predicate.Operators.Or;
 import parquet.filter2.predicate.Operators.UserDefined;
+import parquet.filter2.predicate.Operators.UserDefinedByClass;
 import parquet.io.api.Binary;
 
 import static org.junit.Assert.assertEquals;
@@ -27,6 +30,7 @@
 import static parquet.filter2.predicate.FilterApi.eq;
 import static parquet.filter2.predicate.FilterApi.gt;
 import static parquet.filter2.predicate.FilterApi.intColumn;
+import static parquet.filter2.predicate.FilterApi.longColumn;
 import static parquet.filter2.predicate.FilterApi.not;
 import static parquet.filter2.predicate.FilterApi.notEq;
 import static parquet.filter2.predicate.FilterApi.or;
@@ -36,6 +40,7 @@
 public class TestFilterApiMethods {
 
   private static final IntColumn intColumn = intColumn("a.b.c");
+  private static final LongColumn longColumn = longColumn("a.b.l");
   private static final DoubleColumn doubleColumn = doubleColumn("x.y.z");
   private static final BinaryColumn binColumn = binaryColumn("a.string.column");
 
@@ -82,15 +87,15 @@ public void testUdp() {
     FilterPredicate predicate = or(eq(doubleColumn, 12.0), userDefined(intColumn, DummyUdp.class));
     assertTrue(predicate instanceof Or);
     FilterPredicate ud = ((Or) predicate).getRight();
-    assertTrue(ud instanceof UserDefined);
-    assertEquals(DummyUdp.class, ((UserDefined) ud).getUserDefinedPredicateClass());
+    assertTrue(ud instanceof UserDefinedByClass);
+    assertEquals(DummyUdp.class, ((UserDefinedByClass) ud).getUserDefinedPredicateClass());
     assertTrue(((UserDefined) ud).getUserDefinedPredicate() instanceof DummyUdp);
   }
 
   @Test
-  public void testSerializable() throws Exception {
+  public void testSerializable() throws Exception {    
     BinaryColumn binary = binaryColumn("foo");
-    FilterPredicate p = or(and(userDefined(intColumn, DummyUdp.class), predicate), eq(binary, Binary.fromString("hi")));
+    FilterPredicate p = and(or(and(userDefined(intColumn, DummyUdp.class), predicate), eq(binary, Binary.fromString("hi"))), userDefined(longColumn, new IsMultipleOf(7)));
     ByteArrayOutputStream baos = new ByteArrayOutputStream();
     ObjectOutputStream oos = new ObjectOutputStream(baos);
     oos.writeObject(p);
@@ -100,4 +105,50 @@ public void testSerializable() throws Exception {
     FilterPredicate read = (FilterPredicate) is.readObject();
     assertEquals(p, read);
   }
+
+  public static class IsMultipleOf extends UserDefinedPredicate<Long> implements Serializable {
+
+    private long of; 
+
+    public IsMultipleOf(long of) {
+      this.of = of;
+    }
+
+    @Override
+    public boolean keep(Long value) {
+      if (value == null) {
+        return false;
+      }
+      return value % of == 0;
+    }
+
+    @Override
+    public boolean canDrop(Statistics<Long> statistics) {
+      return false;
+    }
+
+    @Override
+    public boolean inverseCanDrop(Statistics<Long> statistics) {
+      return false;
+    }
+
+    @Override
+    public boolean equals(Object o) {
+      if (this == o) return true;
+      if (o == null || getClass() != o.getClass()) return false;
+
+      IsMultipleOf that = (IsMultipleOf) o;
+      return this.of == that.of;
+    }
+
+    @Override
+    public int hashCode() {
+      return new Long(of).hashCode();
+    }
+
+    @Override
+    public String toString() {
+      return "IsMultipleOf(" + of + ")";
+    }
+  }
 }

parquet-hadoop/src/test/java/parquet/filter2/recordlevel/TestRecordLevelFilters.java

@@ -2,10 +2,13 @@
 
 import java.io.File;
 import java.io.IOException;
+import java.io.Serializable;
 import java.util.ArrayList;
 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 +18,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 +30,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 +165,34 @@ public boolean inverseCanDrop(Statistics<Binary> statistics) {
       return false;
     }
   }
+
+  public static class SetInFilter extends UserDefinedPredicate<Long> implements Serializable {
+
+    private HashSet<Long> hSet;
+
+    public SetInFilter(HashSet<Long> phSet) {
+      hSet = phSet;
+    }
+
+    @Override
+    public boolean keep(Long value) {
+      if (value == null) {
+        return false;
+      }
+
+      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 +209,27 @@ public boolean keep(User u) {
       }
     });
   }
+
+  @Test
+  public void testUserDefinedByInstance() throws Exception {
+    LongColumn name = longColumn("id");
+
+    final HashSet<Long> h = new HashSet<Long>();
+    h.add(20L); 
+    h.add(27L);
+    h.add(28L);
+
+    FilterPredicate pred = userDefined(name, new SetInFilter(h));
+
+    List<Group> found = PhoneBookWriter.readFile(phonebookFile, FilterCompat.get(pred));
+
+    assertFilter(found, new UserFilter() {
+      @Override
+      public boolean keep(User u) {
+        return u != null && h.contains(u.getId());
+      }
+    });
+  }
 
   @Test
   public void testComplex() throws Exception {

parquet-scala/src/main/scala/parquet/filter2/dsl/Dsl.scala

@@ -1,6 +1,7 @@
 package parquet.filter2.dsl
 
 import java.lang.{Boolean => JBoolean, Double => JDouble, Float => JFloat, Integer => JInt, Long => JLong}
+import java.io.Serializable
 
 import parquet.filter2.predicate.{FilterApi, FilterPredicate, Operators, UserDefinedPredicate}
 import parquet.io.api.Binary
@@ -30,6 +31,8 @@ object Dsl {
     val javaColumn: C
 
     def filterBy[U <: UserDefinedPredicate[T]](clazz: Class[U]) = FilterApi.userDefined(javaColumn, clazz)
+
+    def filterBy[U <: UserDefinedPredicate[T] with Serializable](udp: U) = FilterApi.userDefined(javaColumn, udp)
 
     // this is not supported because it allows for easy mistakes. For example:
     // val pred = IntColumn("foo") == "hello"