Search in sources :

Example 1 with SqlStage

use of io.trino.execution.SqlStage in project trino by trinodb.

the class TestSourcePartitionedScheduler method createStageExecution.

private StageExecution createStageExecution(PlanFragment fragment, NodeTaskMap nodeTaskMap) {
    StageId stageId = new StageId(QUERY_ID, 0);
    SqlStage stage = SqlStage.createSqlStage(stageId, fragment, ImmutableMap.of(TABLE_SCAN_NODE_ID, new TableInfo(new QualifiedObjectName("test", "test", "test"), TupleDomain.all())), new MockRemoteTaskFactory(queryExecutor, scheduledExecutor), TEST_SESSION, true, nodeTaskMap, queryExecutor, new SplitSchedulerStats());
    ImmutableMap.Builder<PlanFragmentId, OutputBufferManager> outputBuffers = ImmutableMap.builder();
    outputBuffers.put(fragment.getId(), new PartitionedOutputBufferManager(FIXED_HASH_DISTRIBUTION, 1));
    fragment.getRemoteSourceNodes().stream().flatMap(node -> node.getSourceFragmentIds().stream()).forEach(fragmentId -> outputBuffers.put(fragmentId, new PartitionedOutputBufferManager(FIXED_HASH_DISTRIBUTION, 10)));
    return createPipelinedStageExecution(stage, outputBuffers.buildOrThrow(), TaskLifecycleListener.NO_OP, new NoOpFailureDetector(), queryExecutor, Optional.of(new int[] { 0 }), 0);
}
Also used : PlanFragment(io.trino.sql.planner.PlanFragment) PartitionedSplitsInfo(io.trino.execution.PartitionedSplitsInfo) FIXED_HASH_DISTRIBUTION(io.trino.sql.planner.SystemPartitioningHandle.FIXED_HASH_DISTRIBUTION) Test(org.testng.annotations.Test) TypeOperators(io.trino.spi.type.TypeOperators) TrinoExceptionAssert.assertTrinoExceptionThrownBy(io.trino.testing.assertions.TrinoExceptionAssert.assertTrinoExceptionThrownBy) FixedSplitSource(io.trino.spi.connector.FixedSplitSource) RemoteSourceNode(io.trino.sql.planner.plan.RemoteSourceNode) CatalogName(io.trino.connector.CatalogName) FinalizerService(io.trino.util.FinalizerService) ConnectorPartitionHandle(io.trino.spi.connector.ConnectorPartitionHandle) DynamicFilters(io.trino.sql.DynamicFilters) PlanNodeId(io.trino.sql.planner.plan.PlanNodeId) Executors.newScheduledThreadPool(java.util.concurrent.Executors.newScheduledThreadPool) TEST_SESSION(io.trino.SessionTestUtils.TEST_SESSION) Assert.assertFalse(org.testng.Assert.assertFalse) TableScanNode(io.trino.sql.planner.plan.TableScanNode) SplitsBalancingPolicy(io.trino.execution.scheduler.NodeSchedulerConfig.SplitsBalancingPolicy) NOT_PARTITIONED(io.trino.spi.connector.NotPartitionedPartitionHandle.NOT_PARTITIONED) InMemoryNodeManager(io.trino.metadata.InMemoryNodeManager) NODE(io.trino.execution.scheduler.NodeSchedulerConfig.SplitsBalancingPolicy.NODE) TestingSession(io.trino.testing.TestingSession) PlanFragmentId(io.trino.sql.planner.plan.PlanFragmentId) Session(io.trino.Session) InternalNodeManager(io.trino.metadata.InternalNodeManager) StageId(io.trino.execution.StageId) ConnectorAwareSplitSource(io.trino.split.ConnectorAwareSplitSource) Supplier(java.util.function.Supplier) REPLICATE(io.trino.sql.planner.plan.ExchangeNode.Type.REPLICATE) StatsAndCosts(io.trino.cost.StatsAndCosts) DynamicFilterId(io.trino.sql.planner.plan.DynamicFilterId) ArrayList(java.util.ArrayList) STAGE(io.trino.execution.scheduler.NodeSchedulerConfig.SplitsBalancingPolicy.STAGE) VARCHAR(io.trino.spi.type.VarcharType.VARCHAR) ScheduledExecutorService(java.util.concurrent.ScheduledExecutorService) TestingSplit(io.trino.testing.TestingSplit) StageExecutionDescriptor.ungroupedExecution(io.trino.operator.StageExecutionDescriptor.ungroupedExecution) NodeVersion(io.trino.client.NodeVersion) MockRemoteTaskFactory(io.trino.execution.MockRemoteTaskFactory) NoOpFailureDetector(io.trino.failuredetector.NoOpFailureDetector) Integer.min(java.lang.Integer.min) AfterClass(org.testng.annotations.AfterClass) SqlStage(io.trino.execution.SqlStage) ConnectorSplit(io.trino.spi.connector.ConnectorSplit) ConnectorSplitSource(io.trino.spi.connector.ConnectorSplitSource) MockRemoteTask(io.trino.execution.MockRemoteTaskFactory.MockRemoteTask) QualifiedObjectName(io.trino.metadata.QualifiedObjectName) QueryId(io.trino.spi.QueryId) NodeTaskMap(io.trino.execution.NodeTaskMap) SymbolAllocator(io.trino.sql.planner.SymbolAllocator) SCHEDULING(io.trino.execution.scheduler.StageExecution.State.SCHEDULING) FilterNode(io.trino.sql.planner.plan.FilterNode) Duration(io.airlift.units.Duration) Preconditions.checkArgument(com.google.common.base.Preconditions.checkArgument) URI(java.net.URI) SINGLE_DISTRIBUTION(io.trino.sql.planner.SystemPartitioningHandle.SINGLE_DISTRIBUTION) JoinNode(io.trino.sql.planner.plan.JoinNode) DynamicFilterService(io.trino.server.DynamicFilterService) SourcePartitionedScheduler.newSourcePartitionedSchedulerAsStageScheduler(io.trino.execution.scheduler.SourcePartitionedScheduler.newSourcePartitionedSchedulerAsStageScheduler) ImmutableSet(com.google.common.collect.ImmutableSet) ImmutableMap(com.google.common.collect.ImmutableMap) SPLIT_QUEUES_FULL(io.trino.execution.scheduler.ScheduleResult.BlockedReason.SPLIT_QUEUES_FULL) BeforeClass(org.testng.annotations.BeforeClass) DynamicFilters.createDynamicFilterExpression(io.trino.sql.DynamicFilters.createDynamicFilterExpression) LinkedBlockingQueue(java.util.concurrent.LinkedBlockingQueue) List(java.util.List) BIGINT(io.trino.spi.type.BigintType.BIGINT) TableInfo(io.trino.execution.TableInfo) PLANNED(io.trino.execution.scheduler.StageExecution.State.PLANNED) DynamicFilter(io.trino.spi.connector.DynamicFilter) MetadataManager.createTestMetadataManager(io.trino.metadata.MetadataManager.createTestMetadataManager) Optional(java.util.Optional) NO_NODES_AVAILABLE(io.trino.spi.StandardErrorCode.NO_NODES_AVAILABLE) TestingColumnHandle(io.trino.testing.TestingMetadata.TestingColumnHandle) INNER(io.trino.sql.planner.plan.JoinNode.Type.INNER) Assert.assertEquals(org.testng.Assert.assertEquals) CompletableFuture(java.util.concurrent.CompletableFuture) RetryPolicy(io.trino.operator.RetryPolicy) Partitioning(io.trino.sql.planner.Partitioning) PartitioningScheme(io.trino.sql.planner.PartitioningScheme) FunctionManager.createTestingFunctionManager(io.trino.metadata.FunctionManager.createTestingFunctionManager) ImmutableList(com.google.common.collect.ImmutableList) Threads.daemonThreadsNamed(io.airlift.concurrent.Threads.daemonThreadsNamed) Objects.requireNonNull(java.util.Objects.requireNonNull) DynamicFilterConfig(io.trino.execution.DynamicFilterConfig) ExecutorService(java.util.concurrent.ExecutorService) Symbol(io.trino.sql.planner.Symbol) TableExecuteContextManager(io.trino.execution.TableExecuteContextManager) RemoteTask(io.trino.execution.RemoteTask) TupleDomain(io.trino.spi.predicate.TupleDomain) FunctionManager(io.trino.metadata.FunctionManager) InternalNode(io.trino.metadata.InternalNode) TEST_TABLE_HANDLE(io.trino.testing.TestingHandles.TEST_TABLE_HANDLE) Executors.newCachedThreadPool(java.util.concurrent.Executors.newCachedThreadPool) PipelinedStageExecution.createPipelinedStageExecution(io.trino.execution.scheduler.PipelinedStageExecution.createPipelinedStageExecution) SOURCE_DISTRIBUTION(io.trino.sql.planner.SystemPartitioningHandle.SOURCE_DISTRIBUTION) Metadata(io.trino.metadata.Metadata) Assert.assertTrue(org.testng.Assert.assertTrue) SECONDS(java.util.concurrent.TimeUnit.SECONDS) NoOpFailureDetector(io.trino.failuredetector.NoOpFailureDetector) SqlStage(io.trino.execution.SqlStage) StageId(io.trino.execution.StageId) QualifiedObjectName(io.trino.metadata.QualifiedObjectName) ImmutableMap(com.google.common.collect.ImmutableMap) TableInfo(io.trino.execution.TableInfo) PlanFragmentId(io.trino.sql.planner.plan.PlanFragmentId) MockRemoteTaskFactory(io.trino.execution.MockRemoteTaskFactory)

Aggregations

Preconditions.checkArgument (com.google.common.base.Preconditions.checkArgument)1 ImmutableList (com.google.common.collect.ImmutableList)1 ImmutableMap (com.google.common.collect.ImmutableMap)1 ImmutableSet (com.google.common.collect.ImmutableSet)1 Threads.daemonThreadsNamed (io.airlift.concurrent.Threads.daemonThreadsNamed)1 Duration (io.airlift.units.Duration)1 Session (io.trino.Session)1 TEST_SESSION (io.trino.SessionTestUtils.TEST_SESSION)1 NodeVersion (io.trino.client.NodeVersion)1 CatalogName (io.trino.connector.CatalogName)1 StatsAndCosts (io.trino.cost.StatsAndCosts)1 DynamicFilterConfig (io.trino.execution.DynamicFilterConfig)1 MockRemoteTaskFactory (io.trino.execution.MockRemoteTaskFactory)1 MockRemoteTask (io.trino.execution.MockRemoteTaskFactory.MockRemoteTask)1 NodeTaskMap (io.trino.execution.NodeTaskMap)1 PartitionedSplitsInfo (io.trino.execution.PartitionedSplitsInfo)1 RemoteTask (io.trino.execution.RemoteTask)1 SqlStage (io.trino.execution.SqlStage)1 StageId (io.trino.execution.StageId)1 TableExecuteContextManager (io.trino.execution.TableExecuteContextManager)1