Add lookup function API to StandardFunctionResolution

FunctionAndTypeResolver is in presto-main, and for code which has no access to presto-main, it relies on the interface StandardFunctionResolution in SPI to get function handle. In this PR, I add a API lookUpFunction so that it can work with internal functions which does not exist in OSS
unix280 · Nov 21, 2024 · 0d60c68 · 0d60c68
1 parent c3c7713
commit 0d60c68
Show file tree

Hide file tree

Showing 3 changed files with 15 additions and 0 deletions.
diff --git a/presto-main/src/main/java/com/facebook/presto/sql/relational/FunctionResolution.java b/presto-main/src/main/java/com/facebook/presto/sql/relational/FunctionResolution.java
@@ -371,4 +371,10 @@ public boolean isWindowValueFunction(FunctionHandle functionHandle)
     {
         return windowValueFunctions.contains(functionAndTypeResolver.getFunctionMetadata(functionHandle).getName());
     }
+
+    @Override
+    public FunctionHandle lookupBuiltInFunction(String functionName, List<Type> inputTypes)
+    {
+        return functionAndTypeResolver.lookupFunction(functionName, fromTypes(inputTypes));
+    }
 }
diff --git a/presto-main/src/test/java/com/facebook/presto/sql/relational/TestFunctionResolution.java b/presto-main/src/test/java/com/facebook/presto/sql/relational/TestFunctionResolution.java
@@ -16,14 +16,17 @@
 import com.facebook.presto.common.type.ArrayType;
 import com.facebook.presto.metadata.FunctionAndTypeManager;
 import com.facebook.presto.spi.function.StandardFunctionResolution;
+import com.google.common.collect.ImmutableList;
 import org.testng.annotations.BeforeClass;
 import org.testng.annotations.Test;
 
 import static com.facebook.presto.common.function.OperatorType.ADD;
 import static com.facebook.presto.common.function.OperatorType.GREATER_THAN;
 import static com.facebook.presto.common.type.BigintType.BIGINT;
+import static com.facebook.presto.common.type.BooleanType.BOOLEAN;
 import static com.facebook.presto.common.type.DoubleType.DOUBLE;
 import static com.facebook.presto.metadata.FunctionAndTypeManager.createTestFunctionAndTypeManager;
+import static org.testng.Assert.assertEquals;
 import static org.testng.Assert.assertFalse;
 import static org.testng.Assert.assertTrue;
 
@@ -65,5 +68,9 @@ public void testStandardFunctionResolution()
         // subscript
         assertTrue(standardFunctionResolution.isSubscriptFunction(standardFunctionResolution.subscriptFunction(new ArrayType(DOUBLE), BIGINT)));
         assertFalse(standardFunctionResolution.isBetweenFunction(standardFunctionResolution.subscriptFunction(new ArrayType(DOUBLE), BIGINT)));
+
+        // BuiltInFunction
+        assertEquals(standardFunctionResolution.notFunction(), standardFunctionResolution.lookupBuiltInFunction("not", ImmutableList.of(BOOLEAN)));
+        assertEquals(standardFunctionResolution.countFunction(), standardFunctionResolution.lookupBuiltInFunction("count", ImmutableList.of()));
     }
 }
diff --git a/presto-spi/src/main/java/com/facebook/presto/spi/function/StandardFunctionResolution.java b/presto-spi/src/main/java/com/facebook/presto/spi/function/StandardFunctionResolution.java
@@ -85,4 +85,6 @@ public interface StandardFunctionResolution
     boolean isApproximateSetFunction(FunctionHandle functionHandle);
 
     FunctionHandle approximateSetFunction(Type valueType);
+
+    FunctionHandle lookupBuiltInFunction(String functionName, List<Type> inputTypes);
 }