Introduce TableFunctionProvider

to simplify adding table functions, we are not tightly coupled with concrete table connector function implementations in GlobalSystemConnector and GlobalFunctionCatalog now.
trinodb · Aug 21, 2023 · ef99a66 · ef99a66
1 parent cac5b0f
commit ef99a66
Show file tree

Hide file tree

Showing 13 changed files with 191 additions and 53 deletions.
diff --git a/core/trino-main/src/main/java/io/trino/connector/system/GlobalSystemConnector.java b/core/trino-main/src/main/java/io/trino/connector/system/GlobalSystemConnector.java
@@ -15,7 +15,7 @@
 
 import com.google.common.collect.ImmutableSet;
 import com.google.inject.Inject;
-import io.trino.operator.table.Sequence.SequenceFunctionHandle;
+import io.trino.metadata.TableFunctionProvider;
 import io.trino.spi.connector.CatalogHandle;
 import io.trino.spi.connector.CatalogHandle.CatalogVersion;
 import io.trino.spi.connector.ConnectorMetadata;
@@ -31,9 +31,12 @@
 import io.trino.transaction.InternalConnector;
 import io.trino.transaction.TransactionId;
 
+import java.util.Optional;
 import java.util.Set;
 
-import static io.trino.operator.table.Sequence.getSequenceFunctionSplitSource;
+import static com.google.common.base.Preconditions.checkArgument;
+import static com.google.common.collect.ImmutableSet.toImmutableSet;
+import static io.trino.metadata.TableFunctionProvider.TableFunction;
 import static io.trino.spi.connector.CatalogHandle.createRootCatalogHandle;
 import static java.util.Objects.requireNonNull;
 
@@ -46,13 +49,15 @@ public class GlobalSystemConnector
     private final Set<SystemTable> systemTables;
     private final Set<Procedure> procedures;
     private final Set<ConnectorTableFunction> tableFunctions;
+    private final Set<TableFunctionProvider> tableFunctionProviders;
 
     @Inject
-    public GlobalSystemConnector(Set<SystemTable> systemTables, Set<Procedure> procedures, Set<ConnectorTableFunction> tableFunctions)
+    public GlobalSystemConnector(Set<SystemTable> systemTables, Set<Procedure> procedures, Set<TableFunctionProvider> tableFunctionProviders, Set<ConnectorTableFunction> tableFunctions)
     {
         this.systemTables = ImmutableSet.copyOf(requireNonNull(systemTables, "systemTables is null"));
         this.procedures = ImmutableSet.copyOf(requireNonNull(procedures, "procedures is null"));
         this.tableFunctions = ImmutableSet.copyOf(requireNonNull(tableFunctions, "tableFunctions is null"));
+        this.tableFunctionProviders = ImmutableSet.copyOf(requireNonNull(tableFunctionProviders, "tableFunctionFeaturesFactories is null"));
     }
 
     @Override
@@ -93,11 +98,18 @@ public ConnectorSplitManager getSplitManager()
             @Override
             public ConnectorSplitSource getSplits(ConnectorTransactionHandle transaction, ConnectorSession session, ConnectorTableFunctionHandle functionHandle)
             {
-                if (functionHandle instanceof SequenceFunctionHandle sequenceFunctionHandle) {
-                    return getSequenceFunctionSplitSource(sequenceFunctionHandle);
-                }
+                Set<TableFunction> tableFunctionAdditionalFeatures = tableFunctionProviders.stream()
+                        .map(factory -> factory.get(functionHandle))
+                        .filter(Optional::isPresent)
+                        .map(Optional::get)
+                        .collect(toImmutableSet());
 
-                throw new UnsupportedOperationException();
+                checkArgument(tableFunctionAdditionalFeatures.size() <= 1, "we should have only one implementation for concrete function handle or nothing");
+
+                return tableFunctionAdditionalFeatures.stream()
+                        .map(TableFunction::getSplitSource)
+                        .findFirst()
+                        .orElseThrow(UnsupportedOperationException::new);
             }
         };
     }

diff --git a/core/trino-main/src/main/java/io/trino/connector/system/SystemConnectorModule.java b/core/trino-main/src/main/java/io/trino/connector/system/SystemConnectorModule.java
@@ -30,12 +30,15 @@
 import io.trino.connector.system.jdbc.TableTypeJdbcTable;
 import io.trino.connector.system.jdbc.TypesJdbcTable;
 import io.trino.connector.system.jdbc.UdtJdbcTable;
+import io.trino.metadata.TableFunctionProvider;
 import io.trino.operator.table.ExcludeColumns;
 import io.trino.operator.table.Sequence;
 import io.trino.spi.connector.SystemTable;
 import io.trino.spi.function.table.ConnectorTableFunction;
 import io.trino.spi.procedure.Procedure;
 
+import static io.trino.operator.table.Sequence.SequenceTableFunctionProvider;
+
 public class SystemConnectorModule
         implements Module
 {
@@ -77,6 +80,8 @@ public void configure(Binder binder)
         binder.bind(KillQueryProcedure.class).in(Scopes.SINGLETON);
 
         binder.bind(GlobalSystemConnector.class).in(Scopes.SINGLETON);
+        Multibinder<TableFunctionProvider> tableFunctionProviders = Multibinder.newSetBinder(binder, TableFunctionProvider.class);
+        tableFunctionProviders.addBinding().to(SequenceTableFunctionProvider.class).in(Scopes.SINGLETON);
 
         Multibinder<ConnectorTableFunction> tableFunctions = Multibinder.newSetBinder(binder, ConnectorTableFunction.class);
         tableFunctions.addBinding().toProvider(ExcludeColumns.class).in(Scopes.SINGLETON);

diff --git a/core/trino-main/src/main/java/io/trino/metadata/FunctionManager.java b/core/trino-main/src/main/java/io/trino/metadata/FunctionManager.java
@@ -14,6 +14,7 @@
 package io.trino.metadata;
 
 import com.google.common.cache.CacheBuilder;
+import com.google.common.collect.ImmutableSet;
 import com.google.common.util.concurrent.UncheckedExecutionException;
 import com.google.inject.Inject;
 import io.trino.FeaturesConfig;
@@ -340,7 +341,7 @@ public String toString()
     public static FunctionManager createTestingFunctionManager()
     {
         TypeOperators typeOperators = new TypeOperators();
-        GlobalFunctionCatalog functionCatalog = new GlobalFunctionCatalog();
+        GlobalFunctionCatalog functionCatalog = new GlobalFunctionCatalog(ImmutableSet.of());
         functionCatalog.addFunctions(SystemFunctionBundle.create(new FeaturesConfig(), typeOperators, new BlockTypeOperators(typeOperators), UNKNOWN));
         functionCatalog.addFunctions(new InternalFunctionBundle(new LiteralFunction(new InternalBlockEncodingSerde(new BlockEncodingManager(), TESTING_TYPE_MANAGER))));
         return new FunctionManager(CatalogServiceProvider.fail(), functionCatalog);

diff --git a/core/trino-main/src/main/java/io/trino/metadata/GlobalFunctionCatalog.java b/core/trino-main/src/main/java/io/trino/metadata/GlobalFunctionCatalog.java
@@ -16,10 +16,10 @@
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableListMultimap;
 import com.google.common.collect.ImmutableMap;
+import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.Multimap;
 import com.google.errorprone.annotations.ThreadSafe;
-import io.trino.operator.table.ExcludeColumns.ExcludeColumnsFunctionHandle;
-import io.trino.operator.table.Sequence.SequenceFunctionHandle;
+import com.google.inject.Inject;
 import io.trino.spi.function.AggregationFunctionMetadata;
 import io.trino.spi.function.AggregationImplementation;
 import io.trino.spi.function.BoundSignature;
@@ -42,27 +42,37 @@
 import java.util.Collections;
 import java.util.List;
 import java.util.Map;
+import java.util.Optional;
+import java.util.Set;
 import java.util.function.Function;
+import java.util.stream.Collectors;
 
 import static com.google.common.base.Preconditions.checkArgument;
 import static com.google.common.base.Preconditions.checkState;
 import static com.google.common.collect.ImmutableMap.toImmutableMap;
 import static io.trino.metadata.OperatorNameUtil.isOperatorName;
 import static io.trino.metadata.OperatorNameUtil.unmangleOperator;
-import static io.trino.operator.table.ExcludeColumns.getExcludeColumnsFunctionProcessorProvider;
-import static io.trino.operator.table.Sequence.getSequenceFunctionProcessorProvider;
+import static io.trino.metadata.TableFunctionProvider.TableFunction;
 import static io.trino.spi.function.FunctionKind.AGGREGATE;
 import static io.trino.spi.type.BigintType.BIGINT;
 import static io.trino.spi.type.BooleanType.BOOLEAN;
 import static io.trino.spi.type.IntegerType.INTEGER;
 import static java.util.Locale.ENGLISH;
+import static java.util.Objects.requireNonNull;
 
 @ThreadSafe
 public class GlobalFunctionCatalog
         implements FunctionProvider
 {
     public static final String BUILTIN_SCHEMA = "builtin";
     private volatile FunctionMap functions = new FunctionMap();
+    private final Set<TableFunctionProvider> tableFunctionProviders;
+
+    @Inject
+    public GlobalFunctionCatalog(Set<TableFunctionProvider> tableFunctionProviders)
+    {
+        this.tableFunctionProviders = ImmutableSet.copyOf(requireNonNull(tableFunctionProviders, "tableFunctionFeaturesFactories is null"));
+    }
 
     public final synchronized void addFunctions(FunctionBundle functionBundle)
     {
@@ -180,14 +190,18 @@ public ScalarFunctionImplementation getScalarFunctionImplementation(
     @Override
     public TableFunctionProcessorProvider getTableFunctionProcessorProvider(ConnectorTableFunctionHandle functionHandle)
     {
-        if (functionHandle instanceof ExcludeColumnsFunctionHandle) {
-            return getExcludeColumnsFunctionProcessorProvider();
-        }
-        if (functionHandle instanceof SequenceFunctionHandle) {
-            return getSequenceFunctionProcessorProvider();
-        }
+        Set<TableFunction> tableFunctions = tableFunctionProviders.stream()
+                .map(factory -> factory.get(functionHandle))
+                .filter(Optional::isPresent)
+                .map(Optional::get)
+                .collect(Collectors.toSet());
+
+        checkArgument(tableFunctions.size() <= 1, "we should have only one implementation for concrete function handle or nothing");
 
-        return null;
+        return tableFunctions.stream()
+                .map(TableFunction::getTableFunctionProcessorProvider)
+                .findFirst()
+                .orElse(null);
     }
 
     private static class FunctionMap

diff --git a/core/trino-main/src/main/java/io/trino/metadata/MetadataManager.java b/core/trino-main/src/main/java/io/trino/metadata/MetadataManager.java
@@ -2817,7 +2817,7 @@ public MetadataManager build()
 
             GlobalFunctionCatalog globalFunctionCatalog = this.globalFunctionCatalog;
             if (globalFunctionCatalog == null) {
-                globalFunctionCatalog = new GlobalFunctionCatalog();
+                globalFunctionCatalog = new GlobalFunctionCatalog(ImmutableSet.of());
                 TypeOperators typeOperators = new TypeOperators();
                 globalFunctionCatalog.addFunctions(SystemFunctionBundle.create(new FeaturesConfig(), typeOperators, new BlockTypeOperators(typeOperators), UNKNOWN));
                 globalFunctionCatalog.addFunctions(new InternalFunctionBundle(new LiteralFunction(new InternalBlockEncodingSerde(new BlockEncodingManager(), typeManager))));

diff --git a/core/trino-main/src/main/java/io/trino/metadata/TableFunctionProvider.java b/core/trino-main/src/main/java/io/trino/metadata/TableFunctionProvider.java
@@ -0,0 +1,43 @@
+/*
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package io.trino.metadata;
+
+import io.trino.spi.connector.ConnectorSplitSource;
+import io.trino.spi.function.table.ConnectorTableFunctionHandle;
+import io.trino.spi.function.table.TableFunctionProcessorProvider;
+
+import java.util.Optional;
+
+public interface TableFunctionProvider
+{
+    /**
+     * Provides specific TableFunction
+     * with TableFunctionProcessorProvider and ConnectorSplitSource implementation
+     * for concrete function handle
+     */
+    Optional<TableFunction> get(ConnectorTableFunctionHandle functionHandle);
+
+    interface TableFunction
+    {
+        default TableFunctionProcessorProvider getTableFunctionProcessorProvider()
+        {
+            throw new UnsupportedOperationException("%s does not provide table function processor provider".formatted(getClass().getName()));
+        }
+
+        default ConnectorSplitSource getSplitSource()
+        {
+            throw new UnsupportedOperationException("%s does not provide connector split source".formatted(getClass().getName()));
+        }
+    }
+}
diff --git a/core/trino-main/src/main/java/io/trino/operator/table/ExcludeColumns.java b/core/trino-main/src/main/java/io/trino/operator/table/ExcludeColumns.java
@@ -16,6 +16,7 @@
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.Sets;
 import com.google.inject.Provider;
+import io.trino.metadata.TableFunctionProvider;
 import io.trino.plugin.base.classloader.ClassLoaderSafeConnectorTableFunction;
 import io.trino.spi.TrinoException;
 import io.trino.spi.connector.ConnectorAccessControl;
@@ -43,6 +44,7 @@
 import static com.google.common.collect.ImmutableSet.toImmutableSet;
 import static com.google.common.collect.Iterables.getOnlyElement;
 import static io.trino.metadata.GlobalFunctionCatalog.BUILTIN_SCHEMA;
+import static io.trino.metadata.TableFunctionProvider.TableFunction;
 import static io.trino.spi.StandardErrorCode.INVALID_FUNCTION_ARGUMENT;
 import static io.trino.spi.function.table.DescriptorArgument.NULL_DESCRIPTOR;
 import static io.trino.spi.function.table.ReturnTypeSpecification.GenericTable.GENERIC_TABLE;
@@ -149,26 +151,44 @@ public TableFunctionAnalysis analyze(
         }
     }
 
-    public static TableFunctionProcessorProvider getExcludeColumnsFunctionProcessorProvider()
+    public record ExcludeColumnsFunctionHandle()
+            implements ConnectorTableFunctionHandle
     {
-        return new TableFunctionProcessorProvider()
+        // there's no information to remember. All logic is effectively delegated to the engine via `requiredColumns`.
+    }
+
+    public static class ExcludeColumnsTableFunctionProvider
+            implements TableFunctionProvider
+    {
+        @Override
+        public Optional<TableFunction> get(ConnectorTableFunctionHandle functionHandle)
         {
-            @Override
-            public TableFunctionDataProcessor getDataProcessor(ConnectorTableFunctionHandle handle)
-            {
-                return input -> {
-                    if (input == null) {
-                        return FINISHED;
-                    }
-                    return usedInputAndProduced(getOnlyElement(input).orElseThrow());
-                };
+            if (!(functionHandle instanceof ExcludeColumnsFunctionHandle)) {
+                return Optional.empty();
             }
-        };
+            return Optional.of(new ExcludeColumnsTableFunctionAdditionalFeature());
+        }
     }
 
-    public record ExcludeColumnsFunctionHandle()
-            implements ConnectorTableFunctionHandle
+    private static class ExcludeColumnsTableFunctionAdditionalFeature
+            implements TableFunction
     {
-        // there's no information to remember. All logic is effectively delegated to the engine via `requiredColumns`.
+        @Override
+        public TableFunctionProcessorProvider getTableFunctionProcessorProvider()
+        {
+            return new TableFunctionProcessorProvider()
+            {
+                @Override
+                public TableFunctionDataProcessor getDataProcessor(ConnectorTableFunctionHandle handle)
+                {
+                    return input -> {
+                        if (input == null) {
+                            return FINISHED;
+                        }
+                        return usedInputAndProduced(getOnlyElement(input).orElseThrow());
+                    };
+                }
+            };
+        }
     }
 }