Add Projections

apache · Jul 29, 2019 · 863c8e9 · 863c8e9
1 parent 0d53c82
commit 863c8e9
Show file tree

Hide file tree

Showing 4 changed files with 60 additions and 4 deletions.
diff --git a/api/src/main/java/org/apache/iceberg/transforms/Bucket.java b/api/src/main/java/org/apache/iceberg/transforms/Bucket.java
@@ -31,10 +31,13 @@
 import java.util.UUID;
 import org.apache.iceberg.expressions.BoundPredicate;
 import org.apache.iceberg.expressions.Expressions;
+import org.apache.iceberg.expressions.Literal;
 import org.apache.iceberg.expressions.UnboundPredicate;
 import org.apache.iceberg.types.Type;
 import org.apache.iceberg.types.Types;
 
+import static org.apache.iceberg.expressions.Expression.Operation.STARTS_WITH;
+import static org.apache.iceberg.expressions.Expressions.startsWith;
 import static org.apache.iceberg.types.Type.TypeID;
 
 abstract class Bucket<T> implements Transform<T, Integer> {
@@ -113,6 +116,7 @@ public UnboundPredicate<Integer> project(String name, BoundPredicate<T> predicat
             predicate.op(), name, apply(predicate.literal().value()));
 //      case IN:
 //        return Expressions.predicate();
+      case STARTS_WITH:
       default:
         // comparison predicates can't be projected, notEq can't be projected
         // TODO: small ranges can be projected.

diff --git a/api/src/main/java/org/apache/iceberg/transforms/ProjectionUtil.java b/api/src/main/java/org/apache/iceberg/transforms/ProjectionUtil.java
@@ -197,6 +197,8 @@ static <S, T> UnboundPredicate<T> truncateArray(
         return predicate(Expression.Operation.GT_EQ, name, transform.apply(boundary));
       case EQ:
         return predicate(Expression.Operation.EQ, name, transform.apply(boundary));
+      case STARTS_WITH:
+        return predicate(Expression.Operation.STARTS_WITH, name, transform.apply(boundary));
 //        case IN: // TODO
 //          return Expressions.predicate(Operation.IN, name, transform.apply(boundary));
       default:

diff --git a/api/src/main/java/org/apache/iceberg/transforms/Truncate.java b/api/src/main/java/org/apache/iceberg/transforms/Truncate.java
@@ -243,11 +243,20 @@ public boolean canTransform(Type type) {
 
     @Override
     public UnboundPredicate<CharSequence> project(String name,
-                                                  BoundPredicate<CharSequence> pred) {
-      if (pred.op() == NOT_NULL || pred.op() == IS_NULL) {
-        return Expressions.predicate(pred.op(), name);
+                                                  BoundPredicate<CharSequence> predicate) {
+      switch(predicate.op()){
+        case NOT_NULL:
+        case IS_NULL:
+          return Expressions.predicate(predicate.op(), name);
+        case STARTS_WITH:
+          if(predicate.literal().value().length() < width()) {
+              Expressions.predicate(predicate.op(), name, predicate.literal().value());
+          } else {
+              return ProjectionUtil.truncateArray(name, predicate, this);
+          }
+        default:
+          return ProjectionUtil.truncateArray(name, predicate, this);
       }
-      return ProjectionUtil.truncateArray(name, pred, this);
     }
 
     @Override

diff --git a/api/src/test/java/org/apache/iceberg/transforms/TestStartsWithProjection.java b/api/src/test/java/org/apache/iceberg/transforms/TestStartsWithProjection.java
@@ -0,0 +1,41 @@
+package org.apache.iceberg.transforms;
+
+import org.apache.iceberg.PartitionSpec;
+import org.apache.iceberg.Schema;
+import org.apache.iceberg.expressions.Expression;
+import org.apache.iceberg.expressions.Literal;
+import org.apache.iceberg.expressions.Projections;
+import org.apache.iceberg.expressions.UnboundPredicate;
+import org.apache.iceberg.types.Types;
+import org.junit.Assert;
+import org.junit.Test;
+
+import static org.apache.iceberg.TestHelpers.assertAndUnwrapUnbound;
+import static org.apache.iceberg.expressions.Expressions.startsWith;
+import static org.apache.iceberg.types.Types.NestedField.optional;
+
+public class TestStartsWithProjection {
+    private static final Schema SCHEMA = new Schema(optional(1, "someStringCol", Types.StringType.get()));
+    @Test
+    public void assertTruncateProjections(){
+        PartitionSpec spec = PartitionSpec.builderFor(SCHEMA).truncate("someStringCol", 4).build();
+
+        assertTruncateProjectionInclusive(spec, startsWith("someStringCol", "ab"), "ab");
+        assertTruncateProjectionInclusive(spec, startsWith("someStringCol", "abab"), "abab");
+        assertTruncateProjectionInclusive(spec, startsWith("someStringCol", "ababab"), "abab");
+
+    }
+
+    private void assertTruncateProjectionInclusive(PartitionSpec spec, UnboundPredicate<?> filter,
+                                                   String expectedLiteral) {
+        Expression projection = Projections.inclusive(spec).project(filter);
+        UnboundPredicate<?> predicate = assertAndUnwrapUnbound(projection);
+
+        Assert.assertEquals(predicate.op(), Expression.Operation.STARTS_WITH);
+
+        Literal literal = predicate.literal();
+        Truncate<CharSequence> transform = (Truncate<CharSequence>) spec.getFieldsBySourceId(1).get(0).transform();
+        String output = transform.toHumanString((String) literal.value());
+        Assert.assertEquals(expectedLiteral, output);
+    }
+}