Add getSplits

Author: xin-zhang2

presto-hive/pom.xml

@@ -110,11 +110,6 @@
             <artifactId>presto-rcfile</artifactId>
         </dependency>
 
-        <dependency>
-            <groupId>com.facebook.presto</groupId>
-            <artifactId>presto-native-tvf</artifactId>
-        </dependency>
-
         <dependency>
             <groupId>com.facebook.presto.hadoop</groupId>
             <artifactId>hadoop-apache2</artifactId>

presto-main-base/src/main/java/com/facebook/presto/connector/system/GlobalSystemConnector.java

@@ -15,6 +15,7 @@
 
 import com.facebook.presto.common.RuntimeStats;
 import com.facebook.presto.common.transaction.TransactionId;
+import com.facebook.presto.metadata.FunctionAndTypeManager;
 import com.facebook.presto.operator.table.ExcludeColumns;
 import com.facebook.presto.operator.table.Sequence;
 import com.facebook.presto.operator.table.Sequence.SequenceFunctionHandle;
@@ -30,6 +31,7 @@
 import com.facebook.presto.spi.ConnectorTableLayoutResult;
 import com.facebook.presto.spi.ConnectorTableMetadata;
 import com.facebook.presto.spi.Constraint;
+import com.facebook.presto.spi.NodeManager;
 import com.facebook.presto.spi.SchemaTableName;
 import com.facebook.presto.spi.SchemaTablePrefix;
 import com.facebook.presto.spi.SplitContext;
@@ -54,7 +56,6 @@
 import java.util.Set;
 import java.util.function.Function;
 
-import static com.facebook.presto.operator.table.Sequence.getSequenceFunctionSplitSource;
 import static java.util.Objects.requireNonNull;
 
 public class GlobalSystemConnector
@@ -66,13 +67,17 @@ public class GlobalSystemConnector
     private final Set<SystemTable> systemTables;
     private final Set<Procedure> procedures;
     private final Set<ConnectorTableFunction> tableFunctions;
+    private final NodeManager nodeManager;
+    private final FunctionAndTypeManager functionAndTypeManager;
 
-    public GlobalSystemConnector(String connectorId, Set<SystemTable> systemTables, Set<Procedure> procedures, Set<ConnectorTableFunction> tableFunctions)
+    public GlobalSystemConnector(String connectorId, Set<SystemTable> systemTables, Set<Procedure> procedures, Set<ConnectorTableFunction> tableFunctions, NodeManager nodeManager, FunctionAndTypeManager functionAndTypeManager)
     {
         this.connectorId = requireNonNull(connectorId, "connectorId is null");
         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.nodeManager = requireNonNull(nodeManager, "nodeManager is null");
+        this.functionAndTypeManager = requireNonNull(functionAndTypeManager, "functionAndTypeManager is null");
     }
 
     @Override
@@ -160,11 +165,7 @@ public ConnectorSplitSource getSplits(ConnectorTransactionHandle transactionHand
             @Override
             public ConnectorSplitSource getSplits(ConnectorTransactionHandle transaction, ConnectorSession session, ConnectorTableFunctionHandle function)
             {
-                if (function instanceof SequenceFunctionHandle) {
-                    SequenceFunctionHandle sequenceFunctionHandle = (SequenceFunctionHandle) function;
-                    return getSequenceFunctionSplitSource(sequenceFunctionHandle);
-                }
-                throw new UnsupportedOperationException();
+                return function.getSplits(transaction, session, nodeManager, functionAndTypeManager);
             }
         };
     }

presto-main-base/src/main/java/com/facebook/presto/connector/system/GlobalSystemConnectorFactory.java

@@ -13,7 +13,9 @@
  */
 package com.facebook.presto.connector.system;
 
+import com.facebook.presto.metadata.FunctionAndTypeManager;
 import com.facebook.presto.spi.ConnectorHandleResolver;
+import com.facebook.presto.spi.NodeManager;
 import com.facebook.presto.spi.SystemTable;
 import com.facebook.presto.spi.connector.Connector;
 import com.facebook.presto.spi.connector.ConnectorContext;
@@ -34,13 +36,17 @@ public class GlobalSystemConnectorFactory
     private final Set<SystemTable> tables;
     private final Set<Procedure> procedures;
     private final Set<ConnectorTableFunction> tableFunctions;
+    private final NodeManager nodeManager;
+    private final FunctionAndTypeManager functionAndTypeManager;
 
     @Inject
-    public GlobalSystemConnectorFactory(Set<SystemTable> tables, Set<Procedure> procedures, Set<ConnectorTableFunction> tableFunctions)
+    public GlobalSystemConnectorFactory(Set<SystemTable> tables, Set<Procedure> procedures, Set<ConnectorTableFunction> tableFunctions, NodeManager nodeManager, FunctionAndTypeManager functionAndTypeManager)
     {
         this.tables = ImmutableSet.copyOf(requireNonNull(tables, "tables is null"));
         this.procedures = ImmutableSet.copyOf(requireNonNull(procedures, "procedures is null"));
         this.tableFunctions = ImmutableSet.copyOf(requireNonNull(tableFunctions, "tableFunctions is null"));
+        this.nodeManager = requireNonNull(nodeManager, "nodeManager is null");
+        this.functionAndTypeManager = requireNonNull(functionAndTypeManager, "functionAndTypeManager is null");
     }
 
     @Override
@@ -58,6 +64,6 @@ public ConnectorHandleResolver getHandleResolver()
     @Override
     public Connector create(String catalogName, Map<String, String> config, ConnectorContext context)
     {
-        return new GlobalSystemConnector(catalogName, tables, procedures, tableFunctions);
+        return new GlobalSystemConnector(catalogName, tables, procedures, tableFunctions, nodeManager, functionAndTypeManager);
     }
 }

presto-main-base/src/main/java/com/facebook/presto/connector/system/SystemConnectorModule.java

@@ -27,8 +27,11 @@
 import com.facebook.presto.connector.system.jdbc.TableTypeJdbcTable;
 import com.facebook.presto.connector.system.jdbc.TypesJdbcTable;
 import com.facebook.presto.connector.system.jdbc.UdtJdbcTable;
+import com.facebook.presto.metadata.FunctionAndTypeManager;
+import com.facebook.presto.nodeManager.PluginNodeManager;
 import com.facebook.presto.operator.table.ExcludeColumns;
 import com.facebook.presto.operator.table.Sequence;
+import com.facebook.presto.spi.NodeManager;
 import com.facebook.presto.spi.SystemTable;
 import com.facebook.presto.spi.function.table.ConnectorTableFunction;
 import com.facebook.presto.spi.procedure.Procedure;
@@ -80,6 +83,9 @@ public void configure(Binder binder)
 
         binder.bind(GlobalSystemConnectorFactory.class).in(Scopes.SINGLETON);
         binder.bind(SystemConnectorRegistrar.class).asEagerSingleton();
+        binder.bind(PluginNodeManager.class).in(Scopes.SINGLETON);
+        binder.bind(NodeManager.class).to(PluginNodeManager.class).in(Scopes.SINGLETON);
+        binder.bind(FunctionAndTypeManager.class).in(Scopes.SINGLETON);
 
         Multibinder<ConnectorTableFunction> tableFunctions = Multibinder.newSetBinder(binder, ConnectorTableFunction.class);
         tableFunctions.addBinding().toProvider(ExcludeColumns.class).in(Scopes.SINGLETON);

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

@@ -385,6 +385,12 @@ public void addTVFProviderFactory(TVFProviderFactory factory)
             throw new IllegalArgumentException(format("TVF provider '%s' is already registered", factory.getName()));
         }
         handleResolver.addTableFunctionNamespace(factory.getName(), factory.getTableFunctionHandleResolver());
+        handleResolver.addTableFunctionSplitNamespace(factory.getName(), factory.getTableFunctionSplitResolver());
+    }
+
+    public HandleResolver getHandleResolver()
+    {
+        return handleResolver;
     }
 
     @Override

presto-main-base/src/main/java/com/facebook/presto/operator/PageBuffer.java

@@ -14,11 +14,11 @@
 package com.facebook.presto.operator;
 
 import com.facebook.presto.common.Page;
+import com.facebook.presto.operator.WorkProcessor.ProcessState;
 import jakarta.annotation.Nullable;
 
 import static com.facebook.presto.operator.WorkProcessor.ProcessState.finished;
 import static com.facebook.presto.operator.WorkProcessor.ProcessState.ofResult;
-import static com.facebook.presto.operator.WorkProcessor.ProcessState.yield;
 import static com.google.common.base.Preconditions.checkState;
 import static java.util.Objects.requireNonNull;
 
@@ -41,7 +41,7 @@ public WorkProcessor<Page> pages()
                 return ofResult(result);
             }
 
-            return yield();
+            return ProcessState.yield();
         });
     }
 

presto-main-base/src/main/java/com/facebook/presto/operator/project/PageProcessor.java

@@ -43,7 +43,6 @@
 import static com.facebook.presto.common.block.DictionaryId.randomDictionaryId;
 import static com.facebook.presto.operator.WorkProcessor.ProcessState.finished;
 import static com.facebook.presto.operator.WorkProcessor.ProcessState.ofResult;
-import static com.facebook.presto.operator.WorkProcessor.ProcessState.yield;
 import static com.facebook.presto.operator.project.SelectedPositions.positionsRange;
 import static com.google.common.base.Preconditions.checkArgument;
 import static com.google.common.base.Verify.verify;
@@ -204,7 +203,7 @@ public ProcessState<Page> process()
                     lastComputeYielded = true;
                     lastComputeBatchSize = batchSize;
                     updateRetainedSize();
-                    return yield();
+                    return ProcessState.yield();
                 }
 
                 if (result.isPageTooLarge()) {

presto-main-base/src/main/java/com/facebook/presto/operator/table/Sequence.java

@@ -20,6 +20,7 @@
 import com.facebook.presto.spi.ConnectorSplitSource;
 import com.facebook.presto.spi.FixedSplitSource;
 import com.facebook.presto.spi.HostAddress;
+import com.facebook.presto.spi.NodeManager;
 import com.facebook.presto.spi.NodeProvider;
 import com.facebook.presto.spi.PrestoException;
 import com.facebook.presto.spi.connector.ConnectorTransactionHandle;
@@ -167,6 +168,12 @@ public long step()
         {
             return step;
         }
+
+        @Override
+        public ConnectorSplitSource getSplits(ConnectorTransactionHandle transaction, ConnectorSession session, NodeManager nodeManager, Object functionAndTypeManager)
+        {
+            return getSequenceFunctionSplitSource(this);
+        }
     }
 
     public static ConnectorSplitSource getSequenceFunctionSplitSource(SequenceFunctionHandle handle)

presto-main-base/src/main/java/com/facebook/presto/testing/LocalQueryRunner.java

@@ -522,7 +522,9 @@ private LocalQueryRunner(Session defaultSession, FeaturesConfig featuresConfig,
                 new AnalyzePropertiesSystemTable(transactionManager, metadata),
                 new TransactionsSystemTable(metadata.getFunctionAndTypeManager(), transactionManager)),
                 ImmutableSet.of(),
-                ImmutableSet.of(new ExcludeColumnsFunction()));
+                ImmutableSet.of(new ExcludeColumnsFunction()),
+                null,
+                getFunctionAndTypeManager());
 
         BuiltInQueryAnalyzer queryAnalyzer = new BuiltInQueryAnalyzer(metadata, sqlParser, accessControl, Optional.empty(), metadataExtractorExecutor);
         BuiltInAnalyzerProvider analyzerProvider = new BuiltInAnalyzerProvider(queryAnalyzer);

presto-native-execution/presto_cpp/main/PrestoServer.cpp

@@ -406,6 +406,18 @@ void PrestoServer::run() {
                 util::extractMessageBody(body),
                 server->nativeWorkerPool_.get()));
       });
+  httpServer_->registerPost(
+    "/v1/tvf/splits",
+    [server = this](
+        proxygen::HTTPMessage* message,
+        const std::vector<std::unique_ptr<folly::IOBuf>>& body,
+        proxygen::ResponseHandler* downstream) {
+      http::sendOkResponse(
+          downstream,
+          getSplits(
+              util::extractMessageBody(body),
+              server->nativeWorkerPool_.get()));
+    });
 
   if (systemConfig->enableRuntimeMetricsCollection()) {
     enableWorkerStatsReporting();