Return fully qualified functionName and cleanup table function resolution logic

Author: pdabre12

presto-main-base/src/main/java/com/facebook/presto/metadata/BuiltInTypeAndFunctionNamespaceManager.java

@@ -236,8 +236,6 @@
 import com.facebook.presto.spi.function.SqlFunctionVisibility;
 import com.facebook.presto.spi.function.SqlInvokedFunction;
 import com.facebook.presto.spi.function.SqlInvokedScalarFunctionImplementation;
-import com.facebook.presto.spi.function.table.TableFunctionMetadata;
-import com.facebook.presto.spi.tvf.TVFProvider;
 import com.facebook.presto.sql.analyzer.FunctionsConfig;
 import com.facebook.presto.type.BigintOperators;
 import com.facebook.presto.type.BooleanOperators;
@@ -530,7 +528,7 @@
 
 @ThreadSafe
 public class BuiltInTypeAndFunctionNamespaceManager
-        implements FunctionNamespaceManager<SqlFunction>, TVFProvider, TypeManager
+        implements FunctionNamespaceManager<SqlFunction>, TypeManager
 {
     public static final CatalogSchemaName JAVA_BUILTIN_NAMESPACE = new CatalogSchemaName("presto", "default");
     public static final String ID = "builtin";
@@ -1400,12 +1398,6 @@ public SpecializedFunctionKey doGetSpecializedFunctionKeyForMagicLiteralFunction
         throw new PrestoException(FUNCTION_IMPLEMENTATION_MISSING, format("%s not found", signature));
     }
 
-    @Override
-    public TableFunctionMetadata resolveTableFunction(String functionName)
-    {
-        return null;
-    }
-
     private static class EmptyTransactionHandle
             implements FunctionNamespaceTransactionHandle
     {

presto-main-base/src/main/java/com/facebook/presto/metadata/FunctionAndTypeManager.java

@@ -39,7 +39,6 @@
 import com.facebook.presto.spi.StandardErrorCode;
 import com.facebook.presto.spi.function.AggregationFunctionImplementation;
 import com.facebook.presto.spi.function.AlterRoutineCharacteristics;
-import com.facebook.presto.spi.function.CatalogSchemaFunctionName;
 import com.facebook.presto.spi.function.FunctionHandle;
 import com.facebook.presto.spi.function.FunctionMetadata;
 import com.facebook.presto.spi.function.FunctionMetadataManager;
@@ -108,7 +107,6 @@
 import static com.facebook.presto.metadata.FunctionSignatureMatcher.constructFunctionNotFoundErrorMessage;
 import static com.facebook.presto.metadata.SessionFunctionHandle.SESSION_NAMESPACE;
 import static com.facebook.presto.metadata.SignatureBinder.applyBoundVariables;
-import static com.facebook.presto.metadata.TableFunctionRegistry.toPath;
 import static com.facebook.presto.spi.StandardErrorCode.ALREADY_EXISTS;
 import static com.facebook.presto.spi.StandardErrorCode.AMBIGUOUS_FUNCTION_CALL;
 import static com.facebook.presto.spi.StandardErrorCode.FUNCTION_IMPLEMENTATION_MISSING;
@@ -167,7 +165,6 @@ public class FunctionAndTypeManager
     private final ConcurrentHashMap<ConnectorId, Function<ConnectorTableFunctionHandle, TableFunctionProcessorProvider>> tableFunctionProcessorProviderMap = new ConcurrentHashMap<>();
     private final Map<String, TVFProviderFactory> tvfProviderFactories = new ConcurrentHashMap<>();
     private final Map<ConnectorId, TVFProvider> tvfProviders = new ConcurrentHashMap<>();
-    private final AtomicReference<TVFProvider> servingTableFunctionsProvider;
 
     @Inject
     public FunctionAndTypeManager(
@@ -204,7 +201,6 @@ public FunctionAndTypeManager(
         this.servingTypeManager = new AtomicReference<>(builtInTypeAndFunctionNamespaceManager);
         this.servingTypeManagerParametricTypesSupplier = new AtomicReference<>(this::getServingTypeManagerParametricTypes);
         this.builtInPluginFunctionNamespaceManager = new BuiltInPluginFunctionNamespaceManager(this);
-        this.servingTableFunctionsProvider = new AtomicReference<>(builtInTypeAndFunctionNamespaceManager);
     }
 
     public static FunctionAndTypeManager createTestFunctionAndTypeManager()
@@ -373,12 +369,11 @@ public void loadTVFProvider(String tvfProviderName, NodeManager nodeManager)
         requireNonNull(tvfProviderName, "tvfProviderName is null");
         TVFProviderFactory factory = tvfProviderFactories.get(tvfProviderName);
         checkState(factory != null, "No factory for tvf provider %s", tvfProviderName);
-        TVFProvider tvfProvider = factory.createTVFProvider(tvfProviderName, ImmutableMap.of(), new TVFProviderContext(nodeManager, this));
+        TVFProvider tvfProvider = factory.createTVFProvider(ImmutableMap.of(), new TVFProviderContext(nodeManager, this));
 
         if (tvfProviders.putIfAbsent(new ConnectorId(tvfProviderName), tvfProvider) != null) {
             throw new IllegalArgumentException(format("TVF provider [%s] is already registered", tvfProvider));
         }
-        servingTableFunctionsProvider.compareAndSet(servingTableFunctionsProvider.get(), tvfProvider);
     }
 
     public void loadTVFProviders(NodeManager nodeManager)
@@ -498,19 +493,19 @@ public void addTypeManagerFactory(TypeManagerFactory factory)
 
     public TableFunctionMetadata resolveTableFunction(Session session, QualifiedName qualifiedName)
     {
-        // Fetch if it's a builtin function first
-        TableFunctionMetadata tableFunctionMetadata =
-                 servingTableFunctionsProvider.get()
-                         .resolveTableFunction(
-                                 getFunctionAndTypeResolver()
-                                         .qualifyObjectName(qualifiedName).getObjectName());
-        if (tableFunctionMetadata == null) {
-            // populate the registry before trying to resolve the table functions from table functions provider
-            List<CatalogSchemaFunctionName> name = toPath(session, qualifiedName);
-            ConnectorId connectorId = new ConnectorId(name.get(0).getCatalogName());
-            return tableFunctionRegistry.resolve(connectorId, name.get(0));
+        // Before resolving the table function, add all the TVF provider's table functions to the function registry.
+        if (!tableFunctionRegistry.areTvfProviderFunctionsLoaded()) {
+            for (ConnectorId connectorId : tvfProviders.keySet()) {
+                // In terms of the NativeTVFProvider, you want it to act similarly to the system connector table functions, hence we replace the Java loaded system connector table functions.
+                // This is only enforced when native execution is enabled and  presto-native-tvf module is loaded.
+                if (connectorId.getCatalogName().equals("system")) {
+                    tableFunctionRegistry.removeTableFunctions(connectorId);
+                }
+                tableFunctionRegistry.addTableFunctions(connectorId, tvfProviders.get(connectorId).getTableFunctions());
+            }
+            tableFunctionRegistry.updateTvfProviderFunctionsLoaded();
         }
-        return tableFunctionMetadata;
+        return tableFunctionRegistry.resolve(session, qualifiedName);
     }
 
     public TransactionManager getTransactionManager()

presto-main-base/src/main/java/com/facebook/presto/metadata/TableFunctionRegistry.java

@@ -36,14 +36,13 @@
 import java.util.Map;
 import java.util.Set;
 import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.atomic.AtomicBoolean;
 
-import static com.facebook.presto.spi.StandardErrorCode.GENERIC_USER_ERROR;
 import static com.facebook.presto.spi.StandardErrorCode.MISSING_CATALOG_NAME;
 import static com.facebook.presto.spi.function.table.Preconditions.checkArgument;
 import static com.facebook.presto.sql.analyzer.SemanticErrorCode.CATALOG_NOT_SPECIFIED;
 import static com.facebook.presto.sql.analyzer.SemanticErrorCode.SCHEMA_NOT_SPECIFIED;
 import static com.google.common.base.Preconditions.checkState;
-import static java.lang.String.format;
 import static java.util.Locale.ENGLISH;
 import static java.util.Objects.requireNonNull;
 
@@ -52,6 +51,7 @@ public class TableFunctionRegistry
 {
     // catalog name in the original case; schema and function name in lowercase
     private final Map<ConnectorId, Map<SchemaFunctionName, TableFunctionMetadata>> tableFunctions = new ConcurrentHashMap<>();
+    private final AtomicBoolean tvfProviderFunctionsLoaded = new AtomicBoolean(false);
 
     public void addTableFunctions(ConnectorId catalogName, Collection<ConnectorTableFunction> functions)
     {
@@ -77,9 +77,14 @@ public void removeTableFunctions(ConnectorId catalogName)
         tableFunctions.remove(catalogName);
     }
 
-    public boolean areTableFunctionsLoaded(ConnectorId catalogName)
+    public boolean areTvfProviderFunctionsLoaded()
     {
-        return tableFunctions.containsKey(catalogName);
+        return tvfProviderFunctionsLoaded.get();
+    }
+
+    public void updateTvfProviderFunctionsLoaded()
+    {
+        tvfProviderFunctionsLoaded.compareAndSet(false, true);
     }
 
     public static List<CatalogSchemaFunctionName> toPath(Session session, QualifiedName name)
@@ -117,19 +122,22 @@ public static List<CatalogSchemaFunctionName> toPath(Session session, QualifiedN
      * Resolve table function with given qualified name.
      * Table functions are resolved case-insensitive for consistency with existing scalar function resolution.
      */
-    public TableFunctionMetadata resolve(ConnectorId connectorId, CatalogSchemaFunctionName name)
+    public TableFunctionMetadata resolve(Session session, QualifiedName qualifiedName)
     {
-        Map<SchemaFunctionName, TableFunctionMetadata> catalogFunctions = tableFunctions.get(connectorId);
-        if (catalogFunctions != null) {
-            String lowercasedSchemaName = name.getSchemaFunctionName().getSchemaName().toLowerCase(ENGLISH);
-            String lowercasedFunctionName = name.getSchemaFunctionName().getFunctionName().toLowerCase(ENGLISH);
-            TableFunctionMetadata function = catalogFunctions.get(new SchemaFunctionName(lowercasedSchemaName, lowercasedFunctionName));
-            if (function != null) {
-                return function;
+        for (CatalogSchemaFunctionName name : toPath(session, qualifiedName)) {
+            ConnectorId connectorId = new ConnectorId(name.getCatalogName());
+            Map<SchemaFunctionName, TableFunctionMetadata> catalogFunctions = tableFunctions.get(connectorId);
+            if (catalogFunctions != null) {
+                String lowercasedSchemaName = name.getSchemaFunctionName().getSchemaName().toLowerCase(ENGLISH);
+                String lowercasedFunctionName = name.getSchemaFunctionName().getFunctionName().toLowerCase(ENGLISH);
+                TableFunctionMetadata function = catalogFunctions.get(new SchemaFunctionName(lowercasedSchemaName, lowercasedFunctionName));
+                if (function != null) {
+                    return function;
+                }
             }
         }
 
-        throw new PrestoException(GENERIC_USER_ERROR, format("Table functions for catalog %s could not be resolved.", connectorId.getCatalogName()));
+        return null;
     }
 
     private static void validateTableFunction(ConnectorTableFunction tableFunction)

presto-native-execution/presto_cpp/main/types/FunctionMetadata.cpp

@@ -423,19 +423,18 @@ json getTableValuedFunctionsMetadata() {
   // Get metadata for all registered table valued functions in velox.
   const auto signatures = tableFunctions();
   for (const auto& entry : signatures) {
+    // Just a check to see if it's a fully qualified name
     const auto parts = getFunctionNameParts(entry.first);
-    const auto schema = parts[1];
     const auto functionName = parts[2];
 
-    protocol::AbstractConnectorTableFunction function;
+    protocol::JsonBasedTableFunctionMetadata jsonBasedTableFunctionMetadata;
     json tj;
-    function.name = functionName;
-    function.schema = schema;
-    function.returnTypeSpecification =
+    jsonBasedTableFunctionMetadata.functionName = entry.first;
+    jsonBasedTableFunctionMetadata.returnTypeSpecification =
         buildReturnTypeSpecification(getTableFunctionReturnType(entry.first));
-    function.arguments =
+    jsonBasedTableFunctionMetadata.arguments =
         buildArgumentSpecsList(getTableFunctionArgumentSpecs(entry.first));
-    protocol::to_json(tj, function);
+    protocol::to_json(tj, jsonBasedTableFunctionMetadata);
     j[functionName] = tj;
   }
   return j;
@@ -456,9 +455,10 @@ getRequiredColumns(const tvf::TableFunctionAnalysis* tableFunctionAnalysis) {
 }
 
 protocol::NativeTableFunctionHandle buildNativeTableFunctionHandle(
-    const TableFunctionHandlePtr tableFunctionHandle) {
+    const TableFunctionHandlePtr tableFunctionHandle,
+    std::string functionName) {
   protocol::NativeTableFunctionHandle handle;
-  handle.functionName = tableFunctionHandle->name();
+  handle.functionName = functionName;
   handle.serializedTableFunctionHandle =
       folly::toJson(tableFunctionHandle->serialize());
   return handle;
@@ -475,7 +475,7 @@ protocol::NativeTableFunctionAnalysis getNativeTableFunctionAnalysis(
       std::make_shared<protocol::NativeDescriptor>(
           buildNativeDescriptor(*tableFunctionAnalysis->returnType()));
   nativeTableFunctionAnalysis.handle = buildNativeTableFunctionHandle(
-      tableFunctionAnalysis->tableFunctionHandle());
+      tableFunctionAnalysis->tableFunctionHandle(), functionName);
   return nativeTableFunctionAnalysis;
 }
 
@@ -484,7 +484,7 @@ json getAnalyzedTableValueFunction(
     velox::memory::MemoryPool* pool) {
   TypeParser parser;
   VeloxExprConverter exprConverter_{pool, &parser};
-  protocol::ConnectorTableMetadata1 connectorTableMetadata =
+  protocol::ConnectorTableMetadata connectorTableMetadata =
       json::parse(connectorTableMetadataJson);
   std::unordered_map<std::string, std::shared_ptr<tvf::Argument>> args;
   for (const auto& entry : connectorTableMetadata.arguments) {
@@ -527,11 +527,11 @@ json getAnalyzedTableValueFunction(
       functionArg = std::make_shared<tvf::Descriptor>(
           std::move(fieldNames), std::move(fieldTypes));
     } else {
-      VELOX_FAIL("Failed to convert to a valid Argument");
+      VELOX_UNSUPPORTED("Failed to convert to a valid Argument");
     }
     args[entry.first] = functionArg;
   }
-  return json(
-      getNativeTableFunctionAnalysis(connectorTableMetadata.name, args));
+  return json(getNativeTableFunctionAnalysis(
+      connectorTableMetadata.functionName, args));
 }
 } // namespace facebook::presto