Search in sources :

Example 1 with PageConsumerOutputFactory

use of io.trino.testing.PageConsumerOperator.PageConsumerOutputFactory in project trino by trinodb.

the class LocalQueryRunner method executeInternal.

private MaterializedResultWithPlan executeInternal(Session session, @Language("SQL") String sql) {
    lock.readLock().lock();
    try (Closer closer = Closer.create()) {
        accessControl.checkCanExecuteQuery(session.getIdentity());
        AtomicReference<MaterializedResult.Builder> builder = new AtomicReference<>();
        PageConsumerOutputFactory outputFactory = new PageConsumerOutputFactory(types -> {
            builder.compareAndSet(null, MaterializedResult.resultBuilder(session, types));
            return builder.get()::page;
        });
        TaskContext taskContext = TestingTaskContext.builder(notificationExecutor, yieldExecutor, session).setMaxSpillSize(nodeSpillConfig.getMaxSpillPerNode()).setQueryMaxSpillSize(nodeSpillConfig.getQueryMaxSpillPerNode()).build();
        Plan plan = createPlan(session, sql, WarningCollector.NOOP);
        List<Driver> drivers = createDrivers(session, plan, outputFactory, taskContext);
        drivers.forEach(closer::register);
        boolean done = false;
        while (!done) {
            boolean processed = false;
            for (Driver driver : drivers) {
                if (alwaysRevokeMemory) {
                    driver.getDriverContext().getOperatorContexts().stream().filter(operatorContext -> operatorContext.getNestedOperatorStats().stream().mapToLong(stats -> stats.getRevocableMemoryReservation().toBytes()).sum() > 0).forEach(OperatorContext::requestMemoryRevoking);
                }
                if (!driver.isFinished()) {
                    driver.process();
                    processed = true;
                }
            }
            done = !processed;
        }
        verify(builder.get() != null, "Output operator was not created");
        return new MaterializedResultWithPlan(builder.get().build(), plan);
    } catch (IOException e) {
        throw new UncheckedIOException(e);
    } finally {
        lock.readLock().unlock();
    }
}
Also used : Closer(com.google.common.io.Closer) ConnectorManager(io.trino.connector.ConnectorManager) PagesIndexPageSorter(io.trino.operator.PagesIndexPageSorter) TransactionManager(io.trino.transaction.TransactionManager) TransactionManagerConfig(io.trino.transaction.TransactionManagerConfig) FilterStatsCalculator(io.trino.cost.FilterStatsCalculator) FinalizerService(io.trino.util.FinalizerService) PlanNodeId(io.trino.sql.planner.plan.PlanNodeId) MaterializedViewPropertiesSystemTable(io.trino.connector.system.MaterializedViewPropertiesSystemTable) HandleResolver(io.trino.metadata.HandleResolver) Map(java.util.Map) MetadataManager(io.trino.metadata.MetadataManager) CostCalculatorUsingExchanges(io.trino.cost.CostCalculatorUsingExchanges) TableScanNode(io.trino.sql.planner.plan.TableScanNode) InMemoryTransactionManager(io.trino.transaction.InMemoryTransactionManager) InjectedFailureType(io.trino.execution.FailureInjector.InjectedFailureType) SplitSource(io.trino.split.SplitSource) SystemSessionPropertiesProvider(io.trino.SystemSessionPropertiesProvider) PreparedQuery(io.trino.execution.QueryPreparer.PreparedQuery) PlanSanityChecker(io.trino.sql.planner.sanity.PlanSanityChecker) Plugin(io.trino.spi.Plugin) NOT_PARTITIONED(io.trino.spi.connector.NotPartitionedPartitionHandle.NOT_PARTITIONED) InternalBlockEncodingSerde(io.trino.metadata.InternalBlockEncodingSerde) SubPlan(io.trino.sql.planner.SubPlan) ShowStatsRewrite(io.trino.sql.rewrite.ShowStatsRewrite) Split(io.trino.metadata.Split) InMemoryNodeManager(io.trino.metadata.InMemoryNodeManager) Session(io.trino.Session) Analyzer(io.trino.sql.analyzer.Analyzer) MaterializedViewSystemTable(io.trino.connector.system.MaterializedViewSystemTable) NodePartitioningManager(io.trino.sql.planner.NodePartitioningManager) PluginManager(io.trino.server.PluginManager) RuleStatsRecorder(io.trino.sql.planner.RuleStatsRecorder) ColumnPropertiesSystemTable(io.trino.connector.system.ColumnPropertiesSystemTable) ComposableStatsCalculator(io.trino.cost.ComposableStatsCalculator) PageIndexerFactory(io.trino.spi.PageIndexerFactory) ScheduledExecutorService(java.util.concurrent.ScheduledExecutorService) Lifespan(io.trino.execution.Lifespan) StageExecutionDescriptor(io.trino.operator.StageExecutionDescriptor) LocalExecutionPlan(io.trino.sql.planner.LocalExecutionPlanner.LocalExecutionPlan) SpillerStats(io.trino.spiller.SpillerStats) IOException(java.io.IOException) LocalExecutionPlanner(io.trino.sql.planner.LocalExecutionPlanner) FunctionBundle(io.trino.metadata.FunctionBundle) SessionPropertyDefaults(io.trino.server.SessionPropertyDefaults) TableHandle(io.trino.metadata.TableHandle) QualifiedObjectName(io.trino.metadata.QualifiedObjectName) EventListenerManager(io.trino.eventlistener.EventListenerManager) InternalFunctionBundle(io.trino.metadata.InternalFunctionBundle) StatementAnalyzerFactory(io.trino.sql.analyzer.StatementAnalyzerFactory) TablePropertyManager(io.trino.metadata.TablePropertyManager) TransactionId(io.trino.transaction.TransactionId) Driver(io.trino.operator.Driver) AllowAllSystemAccessControl(io.trino.plugin.base.security.AllowAllSystemAccessControl) OptimizerConfig(io.trino.sql.planner.OptimizerConfig) AnalyzePropertiesSystemTable(io.trino.connector.system.AnalyzePropertiesSystemTable) ParsingUtil.createParsingOptions(io.trino.sql.ParsingUtil.createParsingOptions) StatementRewrite(io.trino.sql.rewrite.StatementRewrite) TransactionBuilder.transaction(io.trino.transaction.TransactionBuilder.transaction) NodeTaskMap(io.trino.execution.NodeTaskMap) JoinFilterFunctionCompiler(io.trino.sql.gen.JoinFilterFunctionCompiler) InternalTypeManager(io.trino.type.InternalTypeManager) NodeSpillConfig(io.trino.spiller.NodeSpillConfig) Duration(io.airlift.units.Duration) StatsCalculator(io.trino.cost.StatsCalculator) Preconditions.checkArgument(com.google.common.base.Preconditions.checkArgument) DisabledSystemSecurityMetadata(io.trino.metadata.DisabledSystemSecurityMetadata) TreeAssertions.assertFormattedSql(io.trino.sql.testing.TreeAssertions.assertFormattedSql) PageSorter(io.trino.spi.PageSorter) LiteralFunction(io.trino.metadata.LiteralFunction) DescribeInputRewrite(io.trino.sql.rewrite.DescribeInputRewrite) ImmutableSet(com.google.common.collect.ImmutableSet) GroupProviderManager(io.trino.security.GroupProviderManager) NodeSystemTable(io.trino.connector.system.NodeSystemTable) CostComparator(io.trino.cost.CostComparator) Constraint.alwaysTrue(io.trino.spi.connector.Constraint.alwaysTrue) PlanOptimizers(io.trino.sql.planner.PlanOptimizers) ColumnPropertyManager(io.trino.metadata.ColumnPropertyManager) GlobalSystemConnectorFactory(io.trino.connector.system.GlobalSystemConnectorFactory) PlannerContext(io.trino.sql.PlannerContext) Analysis(io.trino.sql.analyzer.Analysis) GlobalFunctionCatalog(io.trino.metadata.GlobalFunctionCatalog) TableCommentSystemTable(io.trino.connector.system.TableCommentSystemTable) FileSingleStreamSpillerFactory(io.trino.spiller.FileSingleStreamSpillerFactory) OrderingCompiler(io.trino.sql.gen.OrderingCompiler) ExchangeHandleResolver(io.trino.metadata.ExchangeHandleResolver) Function(java.util.function.Function) ExpressionCompiler(io.trino.sql.gen.ExpressionCompiler) GROUPED_SCHEDULING(io.trino.spi.connector.ConnectorSplitManager.SplitSchedulingStrategy.GROUPED_SCHEDULING) ImmutableList(com.google.common.collect.ImmutableList) SystemSessionProperties(io.trino.SystemSessionProperties) TaskContext(io.trino.operator.TaskContext) DynamicFilterConfig(io.trino.execution.DynamicFilterConfig) IndexManager(io.trino.index.IndexManager) ExecutorService(java.util.concurrent.ExecutorService) PagesIndex(io.trino.operator.PagesIndex) TaskCountEstimator(io.trino.cost.TaskCountEstimator) SchemaPropertyManager(io.trino.metadata.SchemaPropertyManager) MetadataUtil(io.trino.metadata.MetadataUtil) TableExecuteContextManager(io.trino.execution.TableExecuteContextManager) GenericSpillerFactory(io.trino.spiller.GenericSpillerFactory) PageSourceManager(io.trino.split.PageSourceManager) FunctionManager(io.trino.metadata.FunctionManager) FeaturesConfig(io.trino.FeaturesConfig) CatalogManager(io.trino.metadata.CatalogManager) WarningCollector(io.trino.execution.warnings.WarningCollector) Plan(io.trino.sql.planner.Plan) TypeManager(io.trino.spi.type.TypeManager) QueryExplainerFactory(io.trino.sql.analyzer.QueryExplainerFactory) SplitAssignment(io.trino.execution.SplitAssignment) ExchangeManagerRegistry(io.trino.exchange.ExchangeManagerRegistry) EMPTY(io.trino.spi.connector.DynamicFilter.EMPTY) TableProceduresPropertyManager(io.trino.metadata.TableProceduresPropertyManager) SchemaPropertiesSystemTable(io.trino.connector.system.SchemaPropertiesSystemTable) NoOpResourceGroupManager(io.trino.execution.resourcegroups.NoOpResourceGroupManager) NodeInfo(io.airlift.node.NodeInfo) TypeOperators(io.trino.spi.type.TypeOperators) AnalyzerFactory(io.trino.sql.analyzer.AnalyzerFactory) OperatorContext(io.trino.operator.OperatorContext) PlanNode(io.trino.sql.planner.plan.PlanNode) CatalogName(io.trino.connector.CatalogName) NodeScheduler(io.trino.execution.scheduler.NodeScheduler) Executors.newScheduledThreadPool(java.util.concurrent.Executors.newScheduledThreadPool) ReadWriteLock(java.util.concurrent.locks.ReadWriteLock) StatsRulesProvider(io.trino.cost.StatsCalculatorModule.StatsRulesProvider) NodeSchedulerConfig(io.trino.execution.scheduler.NodeSchedulerConfig) PageSinkManager(io.trino.split.PageSinkManager) PasswordAuthenticatorConfig(io.trino.server.security.PasswordAuthenticatorConfig) LogicalPlanner(io.trino.sql.planner.LogicalPlanner) Set(java.util.Set) TrinoOperatorFactories(io.trino.operator.TrinoOperatorFactories) QueryPreparer(io.trino.execution.QueryPreparer) IndexJoinLookupStats(io.trino.operator.index.IndexJoinLookupStats) UncheckedIOException(java.io.UncheckedIOException) TaskManagerConfig(io.trino.execution.TaskManagerConfig) EmbedVersion.testingVersionEmbedder(io.trino.version.EmbedVersion.testingVersionEmbedder) SpillerFactory(io.trino.spiller.SpillerFactory) GlobalSystemConnector(io.trino.connector.system.GlobalSystemConnector) PlanPrinter(io.trino.sql.planner.planprinter.PlanPrinter) PlanNodeSearcher.searchFrom(io.trino.sql.planner.optimizations.PlanNodeSearcher.searchFrom) ExplainRewrite(io.trino.sql.rewrite.ExplainRewrite) ShowQueriesRewrite(io.trino.sql.rewrite.ShowQueriesRewrite) JoinCompiler(io.trino.sql.gen.JoinCompiler) ReentrantReadWriteLock(java.util.concurrent.locks.ReentrantReadWriteLock) GroupByHashPageIndexerFactory(io.trino.operator.GroupByHashPageIndexerFactory) ArrayList(java.util.ArrayList) TypeRegistry(io.trino.metadata.TypeRegistry) Closer(com.google.common.io.Closer) UNGROUPED_SCHEDULING(io.trino.spi.connector.ConnectorSplitManager.SplitSchedulingStrategy.UNGROUPED_SCHEDULING) SystemFunctionBundle(io.trino.metadata.SystemFunctionBundle) PlanOptimizer(io.trino.sql.planner.optimizations.PlanOptimizer) EventListenerConfig(io.trino.eventlistener.EventListenerConfig) ConnectorFactory(io.trino.spi.connector.ConnectorFactory) BlockTypeOperators(io.trino.type.BlockTypeOperators) Language(org.intellij.lang.annotations.Language) HeaderAuthenticatorConfig(io.trino.server.security.HeaderAuthenticatorConfig) MoreFutures.getFutureValue(io.airlift.concurrent.MoreFutures.getFutureValue) PasswordAuthenticatorManager(io.trino.server.security.PasswordAuthenticatorManager) Lock(java.util.concurrent.locks.Lock) TypeAnalyzer(io.trino.sql.planner.TypeAnalyzer) MaterializedViewPropertyManager(io.trino.metadata.MaterializedViewPropertyManager) PlanNodeIdAllocator(io.trino.sql.planner.PlanNodeIdAllocator) PlanOptimizersFactory(io.trino.sql.planner.PlanOptimizersFactory) NodeMemoryConfig(io.trino.memory.NodeMemoryConfig) DriverFactory(io.trino.operator.DriverFactory) OPTIMIZED_AND_VALIDATED(io.trino.sql.planner.LogicalPlanner.Stage.OPTIMIZED_AND_VALIDATED) SplitManager(io.trino.split.SplitManager) ParameterUtils.parameterExtractor(io.trino.sql.ParameterUtils.parameterExtractor) TransactionsSystemTable(io.trino.connector.system.TransactionsSystemTable) ErrorType(io.trino.spi.ErrorType) SqlParser(io.trino.sql.parser.SqlParser) StatsNormalizer(io.trino.cost.StatsNormalizer) CertificateAuthenticatorManager(io.trino.server.security.CertificateAuthenticatorManager) ImmutableMap(com.google.common.collect.ImmutableMap) TablePropertiesSystemTable(io.trino.connector.system.TablePropertiesSystemTable) PageFunctionCompiler(io.trino.sql.gen.PageFunctionCompiler) PropertyMetadata(io.trino.spi.session.PropertyMetadata) OutputFactory(io.trino.operator.OutputFactory) ScalarStatsCalculator(io.trino.cost.ScalarStatsCalculator) ScheduledSplit(io.trino.execution.ScheduledSplit) UniformNodeSelectorFactory(io.trino.execution.scheduler.UniformNodeSelectorFactory) HeaderAuthenticatorManager(io.trino.server.security.HeaderAuthenticatorManager) Preconditions.checkState(com.google.common.base.Preconditions.checkState) SessionPropertyManager(io.trino.metadata.SessionPropertyManager) QueryManagerConfig(io.trino.execution.QueryManagerConfig) List(java.util.List) DriverContext(io.trino.operator.DriverContext) Optional(java.util.Optional) OperatorFactories(io.trino.operator.OperatorFactories) PlanFragmenter(io.trino.sql.planner.PlanFragmenter) HashMap(java.util.HashMap) AtomicReference(java.util.concurrent.atomic.AtomicReference) MemoryManagerConfig(io.trino.memory.MemoryManagerConfig) BlockEncodingManager(io.trino.metadata.BlockEncodingManager) Verify.verify(com.google.common.base.Verify.verify) QueryExplainer(io.trino.sql.analyzer.QueryExplainer) Threads.daemonThreadsNamed(io.airlift.concurrent.Threads.daemonThreadsNamed) DescribeOutputRewrite(io.trino.sql.rewrite.DescribeOutputRewrite) Objects.requireNonNull(java.util.Objects.requireNonNull) PageConsumerOutputFactory(io.trino.testing.PageConsumerOperator.PageConsumerOutputFactory) GenericPartitioningSpillerFactory(io.trino.spiller.GenericPartitioningSpillerFactory) TableProceduresRegistry(io.trino.metadata.TableProceduresRegistry) QualifiedTablePrefix(io.trino.metadata.QualifiedTablePrefix) TimeUnit(java.util.concurrent.TimeUnit) PartitioningSpillerFactory(io.trino.spiller.PartitioningSpillerFactory) CatalogSystemTable(io.trino.connector.system.CatalogSystemTable) Executors.newCachedThreadPool(java.util.concurrent.Executors.newCachedThreadPool) CostCalculator(io.trino.cost.CostCalculator) ProcedureRegistry(io.trino.metadata.ProcedureRegistry) CostCalculatorWithEstimatedExchanges(io.trino.cost.CostCalculatorWithEstimatedExchanges) Metadata(io.trino.metadata.Metadata) AnalyzePropertyManager(io.trino.metadata.AnalyzePropertyManager) TaskContext(io.trino.operator.TaskContext) PageConsumerOutputFactory(io.trino.testing.PageConsumerOperator.PageConsumerOutputFactory) Driver(io.trino.operator.Driver) AtomicReference(java.util.concurrent.atomic.AtomicReference) UncheckedIOException(java.io.UncheckedIOException) IOException(java.io.IOException) UncheckedIOException(java.io.UncheckedIOException) SubPlan(io.trino.sql.planner.SubPlan) LocalExecutionPlan(io.trino.sql.planner.LocalExecutionPlanner.LocalExecutionPlan) Plan(io.trino.sql.planner.Plan) OperatorContext(io.trino.operator.OperatorContext)

Example 2 with PageConsumerOutputFactory

use of io.trino.testing.PageConsumerOperator.PageConsumerOutputFactory in project trino by trinodb.

the class TestMemoryPools method setupConsumeRevocableMemory.

private RevocableMemoryOperator setupConsumeRevocableMemory(DataSize reservedPerPage, long numberOfPages) {
    AtomicReference<RevocableMemoryOperator> createOperator = new AtomicReference<>();
    setUp(() -> {
        DriverContext driverContext = taskContext.addPipelineContext(0, false, false, false).addDriverContext();
        OperatorContext revokableOperatorContext = driverContext.addOperatorContext(Integer.MAX_VALUE, new PlanNodeId("revokable_operator"), TableScanOperator.class.getSimpleName());
        OutputFactory outputFactory = new PageConsumerOutputFactory(types -> (page -> {
        }));
        Operator outputOperator = outputFactory.createOutputOperator(2, new PlanNodeId("output"), ImmutableList.of(), Function.identity(), new TestingPagesSerdeFactory()).createOperator(driverContext);
        RevocableMemoryOperator revocableMemoryOperator = new RevocableMemoryOperator(revokableOperatorContext, reservedPerPage, numberOfPages);
        createOperator.set(revocableMemoryOperator);
        Driver driver = Driver.createDriver(driverContext, revocableMemoryOperator, outputOperator);
        return ImmutableList.of(driver);
    });
    return createOperator.get();
}
Also used : QueryId(io.trino.spi.QueryId) SpillSpaceTracker(io.trino.spiller.SpillSpaceTracker) ListenableFuture(com.google.common.util.concurrent.ListenableFuture) MEGABYTE(io.airlift.units.DataSize.Unit.MEGABYTE) Page(io.trino.spi.Page) Assert.assertEquals(org.testng.Assert.assertEquals) Test(org.testng.annotations.Test) AfterMethod(org.testng.annotations.AfterMethod) AtomicReference(java.util.concurrent.atomic.AtomicReference) Function(java.util.function.Function) Supplier(java.util.function.Supplier) OperatorContext(io.trino.operator.OperatorContext) GIGABYTE(io.airlift.units.DataSize.Unit.GIGABYTE) ImmutableList(com.google.common.collect.ImmutableList) Assertions.assertThatThrownBy(org.assertj.core.api.Assertions.assertThatThrownBy) PlanNodeId(io.trino.sql.planner.plan.PlanNodeId) TpchConnectorFactory(io.trino.plugin.tpch.TpchConnectorFactory) LocalQueryRunner(io.trino.testing.LocalQueryRunner) Map(java.util.Map) TaskContext(io.trino.operator.TaskContext) PageConsumerOutputFactory(io.trino.testing.PageConsumerOperator.PageConsumerOutputFactory) Assert.assertFalse(org.testng.Assert.assertFalse) TestingPagesSerdeFactory(io.trino.execution.buffer.TestingPagesSerdeFactory) TestingTaskContext.createTaskContext(io.trino.testing.TestingTaskContext.createTaskContext) Futures.immediateVoidFuture(com.google.common.util.concurrent.Futures.immediateVoidFuture) Operator(io.trino.operator.Operator) ImmutableMap(com.google.common.collect.ImmutableMap) TableScanOperator(io.trino.operator.TableScanOperator) Predicate(java.util.function.Predicate) OutputFactory(io.trino.operator.OutputFactory) String.format(java.lang.String.format) Preconditions.checkState(com.google.common.base.Preconditions.checkState) AtomicLong(java.util.concurrent.atomic.AtomicLong) DataSize(io.airlift.units.DataSize) List(java.util.List) TestingSession.testSessionBuilder(io.trino.testing.TestingSession.testSessionBuilder) DriverContext(io.trino.operator.DriverContext) Assert.assertTrue(org.testng.Assert.assertTrue) LocalMemoryContext(io.trino.memory.context.LocalMemoryContext) Driver(io.trino.operator.Driver) TestingGcMonitor(io.airlift.stats.TestingGcMonitor) Session(io.trino.Session) Operator(io.trino.operator.Operator) TableScanOperator(io.trino.operator.TableScanOperator) DriverContext(io.trino.operator.DriverContext) TableScanOperator(io.trino.operator.TableScanOperator) PageConsumerOutputFactory(io.trino.testing.PageConsumerOperator.PageConsumerOutputFactory) TestingPagesSerdeFactory(io.trino.execution.buffer.TestingPagesSerdeFactory) Driver(io.trino.operator.Driver) AtomicReference(java.util.concurrent.atomic.AtomicReference) PlanNodeId(io.trino.sql.planner.plan.PlanNodeId) OperatorContext(io.trino.operator.OperatorContext) PageConsumerOutputFactory(io.trino.testing.PageConsumerOperator.PageConsumerOutputFactory) OutputFactory(io.trino.operator.OutputFactory)

Aggregations

Preconditions.checkState (com.google.common.base.Preconditions.checkState)2 ImmutableList (com.google.common.collect.ImmutableList)2 ImmutableMap (com.google.common.collect.ImmutableMap)2 Session (io.trino.Session)2 Preconditions.checkArgument (com.google.common.base.Preconditions.checkArgument)1 Verify.verify (com.google.common.base.Verify.verify)1 ImmutableSet (com.google.common.collect.ImmutableSet)1 Closer (com.google.common.io.Closer)1 Futures.immediateVoidFuture (com.google.common.util.concurrent.Futures.immediateVoidFuture)1 ListenableFuture (com.google.common.util.concurrent.ListenableFuture)1 MoreFutures.getFutureValue (io.airlift.concurrent.MoreFutures.getFutureValue)1 Threads.daemonThreadsNamed (io.airlift.concurrent.Threads.daemonThreadsNamed)1 NodeInfo (io.airlift.node.NodeInfo)1 TestingGcMonitor (io.airlift.stats.TestingGcMonitor)1 DataSize (io.airlift.units.DataSize)1 GIGABYTE (io.airlift.units.DataSize.Unit.GIGABYTE)1 MEGABYTE (io.airlift.units.DataSize.Unit.MEGABYTE)1 Duration (io.airlift.units.Duration)1 FeaturesConfig (io.trino.FeaturesConfig)1 SystemSessionProperties (io.trino.SystemSessionProperties)1