Apply DomainTranslator to STARTS_WITH function

trinodb · Aug 4, 2020 · c56f90a · c56f90a
1 parent aa7b555
commit c56f90a
Show file tree

Hide file tree

Showing 2 changed files with 100 additions and 1 deletion.
diff --git a/presto-main/src/main/java/io/prestosql/sql/planner/DomainTranslator.java b/presto-main/src/main/java/io/prestosql/sql/planner/DomainTranslator.java
@@ -46,6 +46,7 @@
 import io.prestosql.sql.tree.Cast;
 import io.prestosql.sql.tree.ComparisonExpression;
 import io.prestosql.sql.tree.Expression;
+import io.prestosql.sql.tree.FunctionCall;
 import io.prestosql.sql.tree.InListExpression;
 import io.prestosql.sql.tree.InPredicate;
 import io.prestosql.sql.tree.IsNotNullPredicate;
@@ -928,7 +929,58 @@ private Optional<ExtractionResult> tryVisitLikePredicate(LikePredicate node, Boo
             }
 
             Slice constantPrefix = LikeFunctions.unescapeLiteralLikePattern(pattern.slice(0, patternConstantPrefixBytes), escape);
+            return createRangeDomain(type, constantPrefix).map(domain -> new ExtractionResult(TupleDomain.withColumnDomains(ImmutableMap.of(symbol, domain)), node));
+        }
+
+        @Override
+        protected ExtractionResult visitFunctionCall(FunctionCall node, Boolean complement)
+        {
+            String name = ResolvedFunction.extractFunctionName(node.getName());
+            if (name.equals("starts_with")) {
+                Optional<ExtractionResult> result = tryVisitStartsWithFunction(node, complement);
+                if (result.isPresent()) {
+                    return result.get();
+                }
+            }
+            return visitExpression(node, complement);
+        }
 
+        private Optional<ExtractionResult> tryVisitStartsWithFunction(FunctionCall node, Boolean complement)
+        {
+            List<Expression> args = node.getArguments();
+            if (args.size() != 2) {
+                return Optional.empty();
+            }
+
+            Expression target = args.get(0);
+            if (!(target instanceof SymbolReference)) {
+                // Target is not a symbol
+                return Optional.empty();
+            }
+
+            Expression prefix = args.get(1);
+            if (!(prefix instanceof StringLiteral)) {
+                // dynamic pattern
+                return Optional.empty();
+            }
+
+            Type type = typeAnalyzer.getType(session, types, target);
+            if (!(type instanceof VarcharType)) {
+                // TODO support CharType
+                return Optional.empty();
+            }
+            if (complement) {
+                return Optional.empty();
+            }
+
+            Symbol symbol = Symbol.from(target);
+            Slice constantPrefix = ((StringLiteral) prefix).getSlice();
+
+            return createRangeDomain(type, constantPrefix).map(domain -> new ExtractionResult(TupleDomain.withColumnDomains(ImmutableMap.of(symbol, domain)), node));
+        }
+
+        private Optional<Domain> createRangeDomain(Type type, Slice constantPrefix)
+        {
             int lastIncrementable = -1;
             for (int position = 0; position < constantPrefix.length(); position += lengthOfCodePoint(constantPrefix, position)) {
                 // Get last ASCII character to increment, so that character length in bytes does not change.
@@ -948,7 +1000,7 @@ private Optional<ExtractionResult> tryVisitLikePredicate(LikePredicate node, Boo
             setCodePointAt(getCodePointAt(constantPrefix, lastIncrementable) + 1, upperBound, lastIncrementable);
 
             Domain domain = Domain.create(ValueSet.ofRanges(Range.range(type, lowerBound, true, upperBound, false)), false);
-            return Optional.of(new ExtractionResult(TupleDomain.withColumnDomains(ImmutableMap.of(symbol, domain)), node));
+            return Optional.of(domain);
         }
 
         @Override

diff --git a/presto-main/src/test/java/io/prestosql/sql/planner/TestDomainTranslator.java b/presto-main/src/test/java/io/prestosql/sql/planner/TestDomainTranslator.java
@@ -1582,6 +1582,48 @@ public void testLikePredicate()
         assertUnsupportedPredicate(not(like(C_VARCHAR, stringLiteral("abc\\_def"))));
     }
 
+    @Test
+    public void testStartsWithFunction()
+    {
+        Type varcharType = createUnboundedVarcharType();
+
+        // constant
+        testSimpleComparison(
+                startsWith(C_VARCHAR, stringLiteral("abc")),
+                C_VARCHAR,
+                startsWith(C_VARCHAR, stringLiteral("abc")),
+                Domain.create(ValueSet.ofRanges(Range.range(varcharType, utf8Slice("abc"), true, utf8Slice("abd"), false)), false));
+
+        testSimpleComparison(
+                startsWith(C_VARCHAR, stringLiteral("_abc")),
+                C_VARCHAR,
+                startsWith(C_VARCHAR, stringLiteral("_abc")),
+                Domain.create(ValueSet.ofRanges(Range.range(varcharType, utf8Slice("_abc"), true, utf8Slice("_abd"), false)), false));
+
+        // empty
+        assertUnsupportedPredicate(startsWith(C_VARCHAR, stringLiteral("")));
+        // complement
+        assertUnsupportedPredicate(not(startsWith(C_VARCHAR, stringLiteral("abc"))));
+
+        // non-ASCII
+        testSimpleComparison(
+                startsWith(C_VARCHAR, stringLiteral("abc\u0123\ud83d\ude80def\u007e\u007f\u00ff\u0123\uccf0")),
+                C_VARCHAR,
+                startsWith(C_VARCHAR, stringLiteral("abc\u0123\ud83d\ude80def\u007e\u007f\u00ff\u0123\uccf0")),
+                Domain.create(
+                        ValueSet.ofRanges(Range.range(varcharType,
+                                utf8Slice("abc\u0123\ud83d\ude80def\u007e\u007f\u00ff\u0123\uccf0"), true,
+                                utf8Slice("abc\u0123\ud83d\ude80def\u007f"), false)),
+                        false));
+    }
+
+    @Test
+    public void testUnsupportedFunctions()
+    {
+        assertUnsupportedPredicate(new FunctionCall(QualifiedName.of("LENGTH"), ImmutableList.of(C_VARCHAR.toSymbolReference())));
+        assertUnsupportedPredicate(new FunctionCall(QualifiedName.of("REPLACE"), ImmutableList.of(C_VARCHAR.toSymbolReference(), stringLiteral("abc"))));
+    }
+
     @Test
     public void testCharComparedToVarcharExpression()
     {
@@ -1703,6 +1745,11 @@ private static LikePredicate like(Symbol symbol, Expression expression, Expressi
         return new LikePredicate(symbol.toSymbolReference(), expression, Optional.of(escape));
     }
 
+    private static FunctionCall startsWith(Symbol symbol, Expression expression)
+    {
+        return new FunctionCall(QualifiedName.of("STARTS_WITH"), ImmutableList.of(symbol.toSymbolReference(), expression));
+    }
+
     private static Expression isNotNull(Symbol symbol)
     {
         return isNotNull(symbol.toSymbolReference());