Search in sources :

Example 11 with NodePartitioningManager

use of io.trino.sql.planner.NodePartitioningManager in project trino by trinodb.

the class TestLocalExchange method testPartitionCustomPartitioning.

@Test(dataProvider = "executionStrategy")
public void testPartitionCustomPartitioning(PipelineExecutionStrategy executionStrategy) {
    ConnectorPartitioningHandle connectorPartitioningHandle = new ConnectorPartitioningHandle() {
    };
    ConnectorNodePartitioningProvider connectorNodePartitioningProvider = new ConnectorNodePartitioningProvider() {

        @Override
        public ConnectorBucketNodeMap getBucketNodeMap(ConnectorTransactionHandle transactionHandle, ConnectorSession session, ConnectorPartitioningHandle partitioningHandle) {
            return createBucketNodeMap(2);
        }

        @Override
        public ToIntFunction<ConnectorSplit> getSplitBucketFunction(ConnectorTransactionHandle transactionHandle, ConnectorSession session, ConnectorPartitioningHandle partitioningHandle) {
            throw new UnsupportedOperationException();
        }

        @Override
        public BucketFunction getBucketFunction(ConnectorTransactionHandle transactionHandle, ConnectorSession session, ConnectorPartitioningHandle partitioningHandle, List<Type> partitionChannelTypes, int bucketCount) {
            return (page, position) -> {
                long rowValue = BIGINT.getLong(page.getBlock(0), position);
                if (rowValue == 42) {
                    return 0;
                }
                return 1;
            };
        }
    };
    List<Type> types = ImmutableList.of(VARCHAR, BIGINT);
    nodePartitioningManager.addPartitioningProvider(new CatalogName("foo"), connectorNodePartitioningProvider);
    PartitioningHandle partitioningHandle = new PartitioningHandle(Optional.of(new CatalogName("foo")), Optional.of(TestingTransactionHandle.create()), connectorPartitioningHandle);
    LocalExchangeFactory localExchangeFactory = new LocalExchangeFactory(nodePartitioningManager, SESSION, partitioningHandle, 2, types, ImmutableList.of(1), Optional.empty(), executionStrategy, LOCAL_EXCHANGE_MAX_BUFFERED_BYTES, TYPE_OPERATOR_FACTORY);
    LocalExchangeSinkFactoryId localExchangeSinkFactoryId = localExchangeFactory.newSinkFactoryId();
    localExchangeFactory.noMoreSinkFactories();
    run(localExchangeFactory, executionStrategy, exchange -> {
        assertEquals(exchange.getBufferCount(), 2);
        assertExchangeTotalBufferedBytes(exchange, 0);
        LocalExchangeSinkFactory sinkFactory = exchange.getSinkFactory(localExchangeSinkFactoryId);
        LocalExchangeSink sink = sinkFactory.createSink();
        assertSinkCanWrite(sink);
        sinkFactory.close();
        sinkFactory.noMoreSinkFactories();
        LocalExchangeSource sourceA = exchange.getSource(1);
        assertSource(sourceA, 0);
        LocalExchangeSource sourceB = exchange.getSource(0);
        assertSource(sourceB, 0);
        Page pageA = SequencePageBuilder.createSequencePage(types, 1, 100, 42);
        sink.addPage(pageA);
        assertSource(sourceA, 1);
        assertSource(sourceB, 0);
        assertRemovePage(types, sourceA, pageA);
        assertSource(sourceA, 0);
        Page pageB = SequencePageBuilder.createSequencePage(types, 100, 100, 43);
        sink.addPage(pageB);
        assertSource(sourceA, 0);
        assertSource(sourceB, 1);
        assertRemovePage(types, sourceB, pageB);
        assertSource(sourceB, 0);
    });
}
Also used : BucketFunction(io.trino.spi.connector.BucketFunction) NodeTaskMap(io.trino.execution.NodeTaskMap) FIXED_HASH_DISTRIBUTION(io.trino.sql.planner.SystemPartitioningHandle.FIXED_HASH_DISTRIBUTION) Test(org.testng.annotations.Test) TypeOperators(io.trino.spi.type.TypeOperators) SequencePageBuilder(io.trino.SequencePageBuilder) CatalogName(io.trino.connector.CatalogName) NodeScheduler(io.trino.execution.scheduler.NodeScheduler) FinalizerService(io.trino.util.FinalizerService) SINGLE_DISTRIBUTION(io.trino.sql.planner.SystemPartitioningHandle.SINGLE_DISTRIBUTION) Assert.assertFalse(org.testng.Assert.assertFalse) ConnectorNodePartitioningProvider(io.trino.spi.connector.ConnectorNodePartitioningProvider) FIXED_BROADCAST_DISTRIBUTION(io.trino.sql.planner.SystemPartitioningHandle.FIXED_BROADCAST_DISTRIBUTION) NodeSchedulerConfig(io.trino.execution.scheduler.NodeSchedulerConfig) BeforeMethod(org.testng.annotations.BeforeMethod) PartitioningHandle(io.trino.sql.planner.PartitioningHandle) LocalExchangeSinkFactory(io.trino.operator.exchange.LocalExchange.LocalExchangeSinkFactory) Assert.assertNotNull(org.testng.Assert.assertNotNull) UniformNodeSelectorFactory(io.trino.execution.scheduler.UniformNodeSelectorFactory) ConnectorBucketNodeMap(io.trino.spi.connector.ConnectorBucketNodeMap) FIXED_ARBITRARY_DISTRIBUTION(io.trino.sql.planner.SystemPartitioningHandle.FIXED_ARBITRARY_DISTRIBUTION) DataSize(io.airlift.units.DataSize) List(java.util.List) TestingSession.testSessionBuilder(io.trino.testing.TestingSession.testSessionBuilder) BIGINT(io.trino.spi.type.BigintType.BIGINT) ConnectorPartitioningHandle(io.trino.spi.connector.ConnectorPartitioningHandle) InMemoryNodeManager(io.trino.metadata.InMemoryNodeManager) FIXED_PASSTHROUGH_DISTRIBUTION(io.trino.sql.planner.SystemPartitioningHandle.FIXED_PASSTHROUGH_DISTRIBUTION) Optional(java.util.Optional) Session(io.trino.Session) ListenableFuture(com.google.common.util.concurrent.ListenableFuture) DataProvider(org.testng.annotations.DataProvider) Assert.assertNull(org.testng.Assert.assertNull) Type(io.trino.spi.type.Type) ConnectorBucketNodeMap.createBucketNodeMap(io.trino.spi.connector.ConnectorBucketNodeMap.createBucketNodeMap) Page(io.trino.spi.Page) Assert.assertEquals(org.testng.Assert.assertEquals) NodePartitioningManager(io.trino.sql.planner.NodePartitioningManager) LocalExchangeFactory(io.trino.operator.exchange.LocalExchange.LocalExchangeFactory) VARCHAR(io.trino.spi.type.VarcharType.VARCHAR) ImmutableList(com.google.common.collect.ImmutableList) GROUPED_EXECUTION(io.trino.operator.PipelineExecutionStrategy.GROUPED_EXECUTION) Assertions.assertThatThrownBy(org.assertj.core.api.Assertions.assertThatThrownBy) PipelineExecutionStrategy(io.trino.operator.PipelineExecutionStrategy) Lifespan(io.trino.execution.Lifespan) InterpretedHashGenerator(io.trino.operator.InterpretedHashGenerator) BlockTypeOperators(io.trino.type.BlockTypeOperators) UNGROUPED_EXECUTION(io.trino.operator.PipelineExecutionStrategy.UNGROUPED_EXECUTION) ConnectorSplit(io.trino.spi.connector.ConnectorSplit) PageAssertions(io.trino.operator.PageAssertions) ToIntFunction(java.util.function.ToIntFunction) LocalExchangeSinkFactoryId(io.trino.operator.exchange.LocalExchange.LocalExchangeSinkFactoryId) TestingTransactionHandle(io.trino.testing.TestingTransactionHandle) ConnectorSession(io.trino.spi.connector.ConnectorSession) Consumer(java.util.function.Consumer) Assert.assertTrue(org.testng.Assert.assertTrue) ConnectorTransactionHandle(io.trino.spi.connector.ConnectorTransactionHandle) ConnectorTransactionHandle(io.trino.spi.connector.ConnectorTransactionHandle) Page(io.trino.spi.Page) LocalExchangeSinkFactory(io.trino.operator.exchange.LocalExchange.LocalExchangeSinkFactory) Type(io.trino.spi.type.Type) ConnectorPartitioningHandle(io.trino.spi.connector.ConnectorPartitioningHandle) ConnectorNodePartitioningProvider(io.trino.spi.connector.ConnectorNodePartitioningProvider) ConnectorSession(io.trino.spi.connector.ConnectorSession) List(java.util.List) ImmutableList(com.google.common.collect.ImmutableList) CatalogName(io.trino.connector.CatalogName) LocalExchangeFactory(io.trino.operator.exchange.LocalExchange.LocalExchangeFactory) PartitioningHandle(io.trino.sql.planner.PartitioningHandle) ConnectorPartitioningHandle(io.trino.spi.connector.ConnectorPartitioningHandle) ConnectorSplit(io.trino.spi.connector.ConnectorSplit) LocalExchangeSinkFactoryId(io.trino.operator.exchange.LocalExchange.LocalExchangeSinkFactoryId) Test(org.testng.annotations.Test)

Aggregations

NodeScheduler (io.trino.execution.scheduler.NodeScheduler)11 NodeSchedulerConfig (io.trino.execution.scheduler.NodeSchedulerConfig)11 NodePartitioningManager (io.trino.sql.planner.NodePartitioningManager)11 BlockTypeOperators (io.trino.type.BlockTypeOperators)11 FinalizerService (io.trino.util.FinalizerService)11 NodeTaskMap (io.trino.execution.NodeTaskMap)10 UniformNodeSelectorFactory (io.trino.execution.scheduler.UniformNodeSelectorFactory)10 InMemoryNodeManager (io.trino.metadata.InMemoryNodeManager)10 TypeOperators (io.trino.spi.type.TypeOperators)10 ImmutableList (com.google.common.collect.ImmutableList)9 List (java.util.List)9 BeforeMethod (org.testng.annotations.BeforeMethod)9 DataSize (io.airlift.units.DataSize)8 BIGINT (io.trino.spi.type.BigintType.BIGINT)8 Type (io.trino.spi.type.Type)8 VARCHAR (io.trino.spi.type.VarcharType.VARCHAR)8 Optional (java.util.Optional)8 ListenableFuture (com.google.common.util.concurrent.ListenableFuture)7 Threads.daemonThreadsNamed (io.airlift.concurrent.Threads.daemonThreadsNamed)7 Preconditions.checkState (com.google.common.base.Preconditions.checkState)6