Examples with Split io.trino.metadata.Split used on opensource projects

Example 21 with Split

use of io.trino.metadata.Split in project trino by trinodb.

the class TestNodeScheduler method testEmptyAssignmentWithFullNodes.

@Test
public void testEmptyAssignmentWithFullNodes() {
    InternalNode node1 = new InternalNode("node1", URI.create("http://10.0.0.1:11"), NodeVersion.UNKNOWN, false);
    nodeManager.addNode(CONNECTOR_ID, node1);
    InternalNode node2 = new InternalNode("node2", URI.create("http://10.0.0.1:12"), NodeVersion.UNKNOWN, false);
    nodeManager.addNode(CONNECTOR_ID, node2);
    Set<Split> splits = new LinkedHashSet<>();
    // 20 splits with node1 as local node to be assigned in the first iteration of computeAssignments
    for (int i = 0; i < (20 + 10 + 5) * 2; i++) {
        splits.add(new Split(CONNECTOR_ID, new TestSplitLocal(), Lifespan.taskWide()));
    }
    // computeAssignments just returns a mapping of nodes with splits to be assigned, it does not assign splits
    Multimap<InternalNode, Split> assignments1 = nodeSelector.computeAssignments(splits, ImmutableList.copyOf(taskMap.values())).getAssignments();
    assertEquals(assignments1.size(), 40);
    assertEquals(assignments1.keySet().size(), 2);
    assertEquals(assignments1.get(node1).size(), 20);
    assertEquals(assignments1.get(node2).size(), 20);
    MockRemoteTaskFactory remoteTaskFactory = new MockRemoteTaskFactory(remoteTaskExecutor, remoteTaskScheduledExecutor);
    int task = 0;
    for (InternalNode node : assignments1.keySet()) {
        TaskId taskId = new TaskId(new StageId("test", 1), task, 0);
        task++;
        MockRemoteTaskFactory.MockRemoteTask remoteTask = remoteTaskFactory.createTableScanTask(taskId, node, ImmutableList.copyOf(assignments1.get(node)), nodeTaskMap.createPartitionedSplitCountTracker(node, taskId));
        remoteTask.startSplits(20);
        nodeTaskMap.addTask(node, remoteTask);
        taskMap.put(node, remoteTask);
    }
    Set<Split> unassignedSplits = Sets.difference(splits, new HashSet<>(assignments1.values()));
    assertEquals(unassignedSplits.size(), 30);
    Multimap<InternalNode, Split> assignments2 = nodeSelector.computeAssignments(unassignedSplits, ImmutableList.copyOf(taskMap.values())).getAssignments();
    for (InternalNode node : assignments2.keySet()) {
        RemoteTask remoteTask = taskMap.get(node);
        remoteTask.addSplits(ImmutableMultimap.<PlanNodeId, Split>builder().putAll(new PlanNodeId("sourceId"), assignments2.get(node)).build());
    }
    unassignedSplits = Sets.difference(unassignedSplits, new HashSet<>(assignments2.values()));
    assertEquals(unassignedSplits.size(), 10);
    Multimap<InternalNode, Split> assignments3 = nodeSelector.computeAssignments(unassignedSplits, ImmutableList.copyOf(taskMap.values())).getAssignments();
    assertTrue(assignments3.isEmpty());
}

Example 22 with Split

use of io.trino.metadata.Split in project trino by trinodb.

the class TestStageTaskSourceFactory method testHashDistributionTaskSource.

@Test
public void testHashDistributionTaskSource() {
    TaskSource taskSource = createHashDistributionTaskSource(ImmutableMap.of(), ImmutableListMultimap.of(), ImmutableListMultimap.of(), 1, new int[] { 0, 1, 2, 3 }, Optional.empty(), 0, DataSize.of(3, BYTE));
    assertFalse(taskSource.isFinished());
    assertEquals(taskSource.getMoreTasks(), ImmutableList.of());
    assertTrue(taskSource.isFinished());
    taskSource = createHashDistributionTaskSource(ImmutableMap.of(), ImmutableListMultimap.of(PLAN_NODE_1, new TestingExchangeSourceHandle(0, 1), PLAN_NODE_1, new TestingExchangeSourceHandle(1, 1), PLAN_NODE_2, new TestingExchangeSourceHandle(0, 1), PLAN_NODE_2, new TestingExchangeSourceHandle(3, 1)), ImmutableListMultimap.of(PLAN_NODE_3, new TestingExchangeSourceHandle(0, 1)), 1, new int[] { 0, 1, 2, 3 }, Optional.empty(), 0, DataSize.of(0, BYTE));
    assertFalse(taskSource.isFinished());
    assertEquals(taskSource.getMoreTasks(), ImmutableList.of(new TaskDescriptor(0, ImmutableListMultimap.of(), ImmutableListMultimap.of(PLAN_NODE_1, new TestingExchangeSourceHandle(0, 1), PLAN_NODE_2, new TestingExchangeSourceHandle(0, 1), PLAN_NODE_3, new TestingExchangeSourceHandle(0, 1)), new NodeRequirements(Optional.of(CATALOG), ImmutableSet.of(), DataSize.of(4, GIGABYTE))), new TaskDescriptor(1, ImmutableListMultimap.of(), ImmutableListMultimap.of(PLAN_NODE_1, new TestingExchangeSourceHandle(1, 1), PLAN_NODE_3, new TestingExchangeSourceHandle(0, 1)), new NodeRequirements(Optional.of(CATALOG), ImmutableSet.of(), DataSize.of(4, GIGABYTE))), new TaskDescriptor(2, ImmutableListMultimap.of(), ImmutableListMultimap.of(PLAN_NODE_2, new TestingExchangeSourceHandle(3, 1), PLAN_NODE_3, new TestingExchangeSourceHandle(0, 1)), new NodeRequirements(Optional.of(CATALOG), ImmutableSet.of(), DataSize.of(4, GIGABYTE)))));
    assertTrue(taskSource.isFinished());
    Split bucketedSplit1 = createBucketedSplit(0, 0);
    Split bucketedSplit2 = createBucketedSplit(0, 2);
    Split bucketedSplit3 = createBucketedSplit(0, 3);
    Split bucketedSplit4 = createBucketedSplit(0, 1);
    taskSource = createHashDistributionTaskSource(ImmutableMap.of(PLAN_NODE_4, new TestingSplitSource(CATALOG, ImmutableList.of(bucketedSplit1, bucketedSplit2, bucketedSplit3)), PLAN_NODE_5, new TestingSplitSource(CATALOG, ImmutableList.of(bucketedSplit4))), ImmutableListMultimap.of(), ImmutableListMultimap.of(PLAN_NODE_3, new TestingExchangeSourceHandle(0, 1)), 1, new int[] { 0, 1, 2, 3 }, Optional.of(getTestingBucketNodeMap(4)), 0, DataSize.of(0, BYTE));
    assertFalse(taskSource.isFinished());
    assertEquals(taskSource.getMoreTasks(), ImmutableList.of(new TaskDescriptor(0, ImmutableListMultimap.of(PLAN_NODE_4, bucketedSplit1), ImmutableListMultimap.of(PLAN_NODE_3, new TestingExchangeSourceHandle(0, 1)), new NodeRequirements(Optional.of(CATALOG), ImmutableSet.of(), DataSize.of(4, GIGABYTE))), new TaskDescriptor(1, ImmutableListMultimap.of(PLAN_NODE_5, bucketedSplit4), ImmutableListMultimap.of(PLAN_NODE_3, new TestingExchangeSourceHandle(0, 1)), new NodeRequirements(Optional.of(CATALOG), ImmutableSet.of(), DataSize.of(4, GIGABYTE))), new TaskDescriptor(2, ImmutableListMultimap.of(PLAN_NODE_4, bucketedSplit2), ImmutableListMultimap.of(PLAN_NODE_3, new TestingExchangeSourceHandle(0, 1)), new NodeRequirements(Optional.of(CATALOG), ImmutableSet.of(), DataSize.of(4, GIGABYTE))), new TaskDescriptor(3, ImmutableListMultimap.of(PLAN_NODE_4, bucketedSplit3), ImmutableListMultimap.of(PLAN_NODE_3, new TestingExchangeSourceHandle(0, 1)), new NodeRequirements(Optional.of(CATALOG), ImmutableSet.of(), DataSize.of(4, GIGABYTE)))));
    assertTrue(taskSource.isFinished());
    taskSource = createHashDistributionTaskSource(ImmutableMap.of(PLAN_NODE_4, new TestingSplitSource(CATALOG, ImmutableList.of(bucketedSplit1, bucketedSplit2, bucketedSplit3)), PLAN_NODE_5, new TestingSplitSource(CATALOG, ImmutableList.of(bucketedSplit4))), ImmutableListMultimap.of(PLAN_NODE_1, new TestingExchangeSourceHandle(0, 1), PLAN_NODE_1, new TestingExchangeSourceHandle(1, 1), PLAN_NODE_2, new TestingExchangeSourceHandle(0, 1), PLAN_NODE_2, new TestingExchangeSourceHandle(3, 1)), ImmutableListMultimap.of(PLAN_NODE_3, new TestingExchangeSourceHandle(0, 1)), 1, new int[] { 0, 1, 2, 3 }, Optional.of(getTestingBucketNodeMap(4)), 0, DataSize.of(0, BYTE));
    assertFalse(taskSource.isFinished());
    assertEquals(taskSource.getMoreTasks(), ImmutableList.of(new TaskDescriptor(0, ImmutableListMultimap.of(PLAN_NODE_4, bucketedSplit1), ImmutableListMultimap.of(PLAN_NODE_1, new TestingExchangeSourceHandle(0, 1), PLAN_NODE_2, new TestingExchangeSourceHandle(0, 1), PLAN_NODE_3, new TestingExchangeSourceHandle(0, 1)), new NodeRequirements(Optional.of(CATALOG), ImmutableSet.of(), DataSize.of(4, GIGABYTE))), new TaskDescriptor(1, ImmutableListMultimap.of(PLAN_NODE_5, bucketedSplit4), ImmutableListMultimap.of(PLAN_NODE_1, new TestingExchangeSourceHandle(1, 1), PLAN_NODE_3, new TestingExchangeSourceHandle(0, 1)), new NodeRequirements(Optional.of(CATALOG), ImmutableSet.of(), DataSize.of(4, GIGABYTE))), new TaskDescriptor(2, ImmutableListMultimap.of(PLAN_NODE_4, bucketedSplit2), ImmutableListMultimap.of(PLAN_NODE_3, new TestingExchangeSourceHandle(0, 1)), new NodeRequirements(Optional.of(CATALOG), ImmutableSet.of(), DataSize.of(4, GIGABYTE))), new TaskDescriptor(3, ImmutableListMultimap.of(PLAN_NODE_4, bucketedSplit3), ImmutableListMultimap.of(PLAN_NODE_2, new TestingExchangeSourceHandle(3, 1), PLAN_NODE_3, new TestingExchangeSourceHandle(0, 1)), new NodeRequirements(Optional.of(CATALOG), ImmutableSet.of(), DataSize.of(4, GIGABYTE)))));
    assertTrue(taskSource.isFinished());
    taskSource = createHashDistributionTaskSource(ImmutableMap.of(PLAN_NODE_4, new TestingSplitSource(CATALOG, ImmutableList.of(bucketedSplit1, bucketedSplit2, bucketedSplit3)), PLAN_NODE_5, new TestingSplitSource(CATALOG, ImmutableList.of(bucketedSplit4))), ImmutableListMultimap.of(PLAN_NODE_1, new TestingExchangeSourceHandle(0, 1), PLAN_NODE_1, new TestingExchangeSourceHandle(1, 1), PLAN_NODE_2, new TestingExchangeSourceHandle(0, 1)), ImmutableListMultimap.of(PLAN_NODE_3, new TestingExchangeSourceHandle(0, 1)), 2, new int[] { 0, 1, 0, 1 }, Optional.of(getTestingBucketNodeMap(4)), 0, DataSize.of(0, BYTE));
    assertFalse(taskSource.isFinished());
    assertEquals(taskSource.getMoreTasks(), ImmutableList.of(new TaskDescriptor(0, ImmutableListMultimap.of(PLAN_NODE_4, bucketedSplit1, PLAN_NODE_4, bucketedSplit2), ImmutableListMultimap.of(PLAN_NODE_1, new TestingExchangeSourceHandle(0, 1), PLAN_NODE_2, new TestingExchangeSourceHandle(0, 1), PLAN_NODE_3, new TestingExchangeSourceHandle(0, 1)), new NodeRequirements(Optional.of(CATALOG), ImmutableSet.of(), DataSize.of(4, GIGABYTE))), new TaskDescriptor(1, ImmutableListMultimap.of(PLAN_NODE_4, bucketedSplit3, PLAN_NODE_5, bucketedSplit4), ImmutableListMultimap.of(PLAN_NODE_1, new TestingExchangeSourceHandle(1, 1), PLAN_NODE_3, new TestingExchangeSourceHandle(0, 1)), new NodeRequirements(Optional.of(CATALOG), ImmutableSet.of(), DataSize.of(4, GIGABYTE)))));
    assertTrue(taskSource.isFinished());
    // join based on split target split weight
    taskSource = createHashDistributionTaskSource(ImmutableMap.of(PLAN_NODE_4, new TestingSplitSource(CATALOG, ImmutableList.of(bucketedSplit1, bucketedSplit2, bucketedSplit3)), PLAN_NODE_5, new TestingSplitSource(CATALOG, ImmutableList.of(bucketedSplit4))), ImmutableListMultimap.of(PLAN_NODE_1, new TestingExchangeSourceHandle(0, 1), PLAN_NODE_1, new TestingExchangeSourceHandle(1, 1), PLAN_NODE_2, new TestingExchangeSourceHandle(1, 1), PLAN_NODE_2, new TestingExchangeSourceHandle(2, 1), PLAN_NODE_2, new TestingExchangeSourceHandle(3, 1)), ImmutableListMultimap.of(PLAN_NODE_3, new TestingExchangeSourceHandle(17, 1)), 2, new int[] { 0, 1, 2, 3 }, Optional.of(getTestingBucketNodeMap(4)), 2 * STANDARD_WEIGHT, DataSize.of(100, GIGABYTE));
    assertFalse(taskSource.isFinished());
    assertEquals(taskSource.getMoreTasks(), ImmutableList.of(new TaskDescriptor(0, ImmutableListMultimap.of(PLAN_NODE_4, bucketedSplit1, PLAN_NODE_5, bucketedSplit4), ImmutableListMultimap.of(PLAN_NODE_1, new TestingExchangeSourceHandle(0, 1), PLAN_NODE_1, new TestingExchangeSourceHandle(1, 1), PLAN_NODE_2, new TestingExchangeSourceHandle(1, 1), PLAN_NODE_3, new TestingExchangeSourceHandle(17, 1)), new NodeRequirements(Optional.of(CATALOG), ImmutableSet.of(), DataSize.of(4, GIGABYTE))), new TaskDescriptor(1, ImmutableListMultimap.of(PLAN_NODE_4, bucketedSplit2, PLAN_NODE_4, bucketedSplit3), ImmutableListMultimap.of(PLAN_NODE_2, new TestingExchangeSourceHandle(2, 1), PLAN_NODE_2, new TestingExchangeSourceHandle(3, 1), PLAN_NODE_3, new TestingExchangeSourceHandle(17, 1)), new NodeRequirements(Optional.of(CATALOG), ImmutableSet.of(), DataSize.of(4, GIGABYTE)))));
    assertTrue(taskSource.isFinished());
    // join based on target exchange size
    taskSource = createHashDistributionTaskSource(ImmutableMap.of(PLAN_NODE_4, new TestingSplitSource(CATALOG, ImmutableList.of(bucketedSplit1, bucketedSplit2, bucketedSplit3)), PLAN_NODE_5, new TestingSplitSource(CATALOG, ImmutableList.of(bucketedSplit4))), ImmutableListMultimap.of(PLAN_NODE_1, new TestingExchangeSourceHandle(0, 20), PLAN_NODE_1, new TestingExchangeSourceHandle(1, 30), PLAN_NODE_2, new TestingExchangeSourceHandle(1, 20), PLAN_NODE_2, new TestingExchangeSourceHandle(2, 99), PLAN_NODE_2, new TestingExchangeSourceHandle(3, 30)), ImmutableListMultimap.of(PLAN_NODE_3, new TestingExchangeSourceHandle(17, 1)), 2, new int[] { 0, 1, 2, 3 }, Optional.of(getTestingBucketNodeMap(4)), 100 * STANDARD_WEIGHT, DataSize.of(100, BYTE));
    assertFalse(taskSource.isFinished());
    assertEquals(taskSource.getMoreTasks(), ImmutableList.of(new TaskDescriptor(0, ImmutableListMultimap.of(PLAN_NODE_4, bucketedSplit1, PLAN_NODE_5, bucketedSplit4), ImmutableListMultimap.of(PLAN_NODE_1, new TestingExchangeSourceHandle(0, 20), PLAN_NODE_1, new TestingExchangeSourceHandle(1, 30), PLAN_NODE_2, new TestingExchangeSourceHandle(1, 20), PLAN_NODE_3, new TestingExchangeSourceHandle(17, 1)), new NodeRequirements(Optional.of(CATALOG), ImmutableSet.of(), DataSize.of(4, GIGABYTE))), new TaskDescriptor(1, ImmutableListMultimap.of(PLAN_NODE_4, bucketedSplit2), ImmutableListMultimap.of(PLAN_NODE_2, new TestingExchangeSourceHandle(2, 99), PLAN_NODE_3, new TestingExchangeSourceHandle(17, 1)), new NodeRequirements(Optional.of(CATALOG), ImmutableSet.of(), DataSize.of(4, GIGABYTE))), new TaskDescriptor(2, ImmutableListMultimap.of(PLAN_NODE_4, bucketedSplit3), ImmutableListMultimap.of(PLAN_NODE_2, new TestingExchangeSourceHandle(3, 30), PLAN_NODE_3, new TestingExchangeSourceHandle(17, 1)), new NodeRequirements(Optional.of(CATALOG), ImmutableSet.of(), DataSize.of(4, GIGABYTE)))));
    assertTrue(taskSource.isFinished());
}

Example 23 with Split

use of io.trino.metadata.Split in project trino by trinodb.

the class TestStageTaskSourceFactory method testSourceDistributionTaskSourceWithWeights.

@Test
public void testSourceDistributionTaskSourceWithWeights() {
    Split split1 = createWeightedSplit(1, STANDARD_WEIGHT);
    long heavyWeight = 2 * STANDARD_WEIGHT;
    Split heavySplit1 = createWeightedSplit(11, heavyWeight);
    Split heavySplit2 = createWeightedSplit(12, heavyWeight);
    Split heavySplit3 = createWeightedSplit(13, heavyWeight);
    long lightWeight = (long) (0.5 * STANDARD_WEIGHT);
    Split lightSplit1 = createWeightedSplit(21, lightWeight);
    Split lightSplit2 = createWeightedSplit(22, lightWeight);
    Split lightSplit3 = createWeightedSplit(23, lightWeight);
    Split lightSplit4 = createWeightedSplit(24, lightWeight);
    // no limits
    TaskSource taskSource = createSourceDistributionTaskSource(ImmutableList.of(lightSplit1, lightSplit2, split1, heavySplit1, heavySplit2, lightSplit4), ImmutableListMultimap.of(), // single split per batch for predictable results
    1, 0, (long) (1.9 * STANDARD_WEIGHT), 1000);
    List<TaskDescriptor> tasks = readAllTasks(taskSource);
    assertThat(tasks).hasSize(4);
    assertThat(tasks).allMatch(task -> getOnlyElement(task.getSplits().keySet()).equals(PLAN_NODE_1));
    assertThat(tasks.get(0).getSplits().values()).containsExactlyInAnyOrder(lightSplit1, lightSplit2, split1);
    assertThat(tasks.get(1).getSplits().values()).containsExactlyInAnyOrder(heavySplit1);
    assertThat(tasks.get(2).getSplits().values()).containsExactlyInAnyOrder(heavySplit2);
    // remainder
    assertThat(tasks.get(3).getSplits().values()).containsExactlyInAnyOrder(lightSplit4);
    assertTrue(taskSource.isFinished());
    // min splits == 2
    taskSource = createSourceDistributionTaskSource(ImmutableList.of(heavySplit1, heavySplit2, heavySplit3, lightSplit1, lightSplit2, lightSplit3, lightSplit4), ImmutableListMultimap.of(), // single split per batch for predictable results
    1, 2, 2 * STANDARD_WEIGHT, 1000);
    tasks = readAllTasks(taskSource);
    assertThat(tasks).hasSize(3);
    assertThat(tasks).allMatch(task -> getOnlyElement(task.getSplits().keySet()).equals(PLAN_NODE_1));
    assertThat(tasks.get(0).getSplits().values()).containsExactlyInAnyOrder(heavySplit1, heavySplit2);
    assertThat(tasks.get(1).getSplits().values()).containsExactlyInAnyOrder(heavySplit3, lightSplit1);
    assertThat(tasks.get(2).getSplits().values()).containsExactlyInAnyOrder(lightSplit2, lightSplit3, lightSplit4);
    assertTrue(taskSource.isFinished());
    // max splits == 3
    taskSource = createSourceDistributionTaskSource(ImmutableList.of(lightSplit1, lightSplit2, lightSplit3, heavySplit1, lightSplit4), ImmutableListMultimap.of(), // single split per batch for predictable results
    1, 0, 2 * STANDARD_WEIGHT, 3);
    tasks = readAllTasks(taskSource);
    assertThat(tasks).hasSize(3);
    assertThat(tasks).allMatch(task -> getOnlyElement(task.getSplits().keySet()).equals(PLAN_NODE_1));
    assertThat(tasks.get(0).getSplits().values()).containsExactlyInAnyOrder(lightSplit1, lightSplit2, lightSplit3);
    assertThat(tasks.get(1).getSplits().values()).containsExactlyInAnyOrder(heavySplit1);
    assertThat(tasks.get(2).getSplits().values()).containsExactlyInAnyOrder(lightSplit4);
    assertTrue(taskSource.isFinished());
    // with addresses
    Split split1a1 = createWeightedSplit(1, STANDARD_WEIGHT, "host1:8080");
    Split split2a2 = createWeightedSplit(2, STANDARD_WEIGHT, "host2:8080");
    Split split3a1 = createWeightedSplit(3, STANDARD_WEIGHT, "host1:8080");
    Split split3a12 = createWeightedSplit(3, STANDARD_WEIGHT, "host1:8080", "host2:8080");
    Split heavySplit2a2 = createWeightedSplit(12, heavyWeight, "host2:8080");
    Split lightSplit1a1 = createWeightedSplit(21, lightWeight, "host1:8080");
    taskSource = createSourceDistributionTaskSource(ImmutableList.of(split1a1, heavySplit2a2, split3a1, lightSplit1a1), ImmutableListMultimap.of(), // single split per batch for predictable results
    1, 0, 2 * STANDARD_WEIGHT, 3);
    tasks = readAllTasks(taskSource);
    assertThat(tasks).hasSize(3);
    assertThat(tasks).allMatch(task -> getOnlyElement(task.getSplits().keySet()).equals(PLAN_NODE_1));
    assertThat(tasks.get(0).getSplits().values()).containsExactlyInAnyOrder(heavySplit2a2);
    assertThat(tasks.get(1).getSplits().values()).containsExactlyInAnyOrder(split1a1, split3a1);
    assertThat(tasks.get(2).getSplits().values()).containsExactlyInAnyOrder(lightSplit1a1);
    assertTrue(taskSource.isFinished());
    // with addresses with multiple matching
    taskSource = createSourceDistributionTaskSource(ImmutableList.of(split1a1, split3a12, split2a2), ImmutableListMultimap.of(), // single split per batch for predictable results
    1, 0, 2 * STANDARD_WEIGHT, 3);
    tasks = readAllTasks(taskSource);
    assertThat(tasks).hasSize(2);
    assertThat(tasks).allMatch(task -> getOnlyElement(task.getSplits().keySet()).equals(PLAN_NODE_1));
    assertThat(tasks.get(0).getSplits().values()).containsExactlyInAnyOrder(split1a1, split3a12);
    assertThat(tasks.get(1).getSplits().values()).containsExactlyInAnyOrder(split2a2);
    assertTrue(taskSource.isFinished());
}

Example 24 with Split

use of io.trino.metadata.Split in project trino by trinodb.

the class TestMemoryBlocking method testTableScanMemoryBlocking.

@Test
public void testTableScanMemoryBlocking() {
    PlanNodeId sourceId = new PlanNodeId("source");
    List<Type> types = ImmutableList.of(VARCHAR);
    TableScanOperator source = new TableScanOperator(driverContext.addOperatorContext(1, new PlanNodeId("test"), "values"), sourceId, (session, split, table, columns, dynamicFilter) -> new FixedPageSource(rowPagesBuilder(types).addSequencePage(10, 1).addSequencePage(10, 1).addSequencePage(10, 1).addSequencePage(10, 1).addSequencePage(10, 1).build()), TEST_TABLE_HANDLE, ImmutableList.of(), DynamicFilter.EMPTY);
    PageConsumerOperator sink = createSinkOperator(types);
    Driver driver = Driver.createDriver(driverContext, source, sink);
    assertSame(driver.getDriverContext(), driverContext);
    assertFalse(driver.isFinished());
    Split testSplit = new Split(new CatalogName("test"), new TestSplit(), Lifespan.taskWide());
    driver.updateSplitAssignment(new SplitAssignment(sourceId, ImmutableSet.of(new ScheduledSplit(0, sourceId, testSplit)), true));
    ListenableFuture<Void> blocked = driver.processFor(new Duration(1, NANOSECONDS));
    // the driver shouldn't block in the first call as it will be able to move a page between source and the sink operator
    // but the operator should be blocked
    assertTrue(blocked.isDone());
    assertFalse(source.getOperatorContext().isWaitingForMemory().isDone());
    // and they should stay blocked until more memory becomes available
    for (int i = 0; i < 10; i++) {
        blocked = driver.processFor(new Duration(1, NANOSECONDS));
        assertFalse(blocked.isDone());
        assertFalse(source.getOperatorContext().isWaitingForMemory().isDone());
    }
    // free up some memory
    memoryPool.free(QUERY_ID, "test", memoryPool.getReservedBytes());
    // the operator should be unblocked
    assertTrue(source.getOperatorContext().isWaitingForMemory().isDone());
    // the driver shouldn't be blocked
    blocked = driver.processFor(new Duration(1, NANOSECONDS));
    assertTrue(blocked.isDone());
}

Example 25 with Split

use of io.trino.metadata.Split in project trino by trinodb.

the class TestKuduIntegrationDynamicFilter method testIncompleteDynamicFilterTimeout.

@Test(timeOut = 30_000)
public void testIncompleteDynamicFilterTimeout() throws Exception {
    QueryRunner runner = getQueryRunner();
    TransactionManager transactionManager = runner.getTransactionManager();
    TransactionId transactionId = transactionManager.beginTransaction(false);
    Session session = Session.builder(getSession()).setCatalogSessionProperty("kudu", "dynamic_filtering_wait_timeout", "1s").build().beginTransactionId(transactionId, transactionManager, new AllowAllAccessControl());
    QualifiedObjectName tableName = new QualifiedObjectName("kudu", "tpch", "orders");
    Optional<TableHandle> tableHandle = runner.getMetadata().getTableHandle(session, tableName);
    assertTrue(tableHandle.isPresent());
    SplitSource splitSource = runner.getSplitManager().getSplits(session, tableHandle.get(), UNGROUPED_SCHEDULING, new IncompleteDynamicFilter(), alwaysTrue());
    List<Split> splits = new ArrayList<>();
    while (!splitSource.isFinished()) {
        splits.addAll(splitSource.getNextBatch(NOT_PARTITIONED, Lifespan.taskWide(), 1000).get().getSplits());
    }
    splitSource.close();
    assertFalse(splits.isEmpty());
}