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

Example 26 with Split

use of io.prestosql.metadata.Split in project hetu-core by openlookeng.

the class NodeScheduler method selectDistributionNodes.

public static SplitPlacementResult selectDistributionNodes(NodeMap nodeMap, NodeTaskMap nodeTaskMap, int maxSplitsPerNode, int maxPendingSplitsPerTask, Set<Split> splits, List<RemoteTask> existingTasks, BucketNodeMap bucketNodeMap) {
    Multimap<InternalNode, Split> assignments = HashMultimap.create();
    NodeAssignmentStats assignmentStats = new NodeAssignmentStats(nodeTaskMap, nodeMap, existingTasks);
    Set<InternalNode> blockedNodes = new HashSet<>();
    for (Split split : splits) {
        // node placement is forced by the bucket to node map
        InternalNode node = bucketNodeMap.getAssignedNode(split).get();
        // if node is full, don't schedule now, which will push back on the scheduling of splits
        if (assignmentStats.getTotalSplitCount(node) < maxSplitsPerNode || assignmentStats.getQueuedSplitCountForStage(node) < maxPendingSplitsPerTask) {
            assignments.put(node, split);
            assignmentStats.addAssignedSplit(node);
        } else {
            blockedNodes.add(node);
        }
    }
    ListenableFuture<?> blocked = toWhenHasSplitQueueSpaceFuture(blockedNodes, existingTasks, calculateLowWatermark(maxPendingSplitsPerTask));
    return new SplitPlacementResult(blocked, ImmutableMultimap.copyOf(assignments));
}

Example 27 with Split

use of io.prestosql.metadata.Split in project hetu-core by openlookeng.

the class SimpleFixedNodeSelector method computeAssignments.

@Override
public SplitPlacementResult computeAssignments(Set<Split> splits, List<RemoteTask> existingTasks, Optional<SqlStageExecution> stage) {
    Multimap<InternalNode, Split> assignment = HashMultimap.create();
    NodeMap nodeMap = this.nodeMap.get().get();
    NodeAssignmentStats assignmentStats = new NodeAssignmentStats(nodeTaskMap, nodeMap, existingTasks);
    List<InternalNode> candidateNodes = new ArrayList<>();
    if (!stage.isPresent()) {
        log.error("Cant schedule as stage missing");
        throw new PrestoException(GENERIC_INTERNAL_ERROR, "stage is empty");
    }
    PlanNodeId planNodeId = stage.get().getFragment().getFeederCTEParentId().get();
    // if still feeder has not been scheduled then no point in scheduling this also
    if (!feederScheduledNodes.containsKey(planNodeId)) {
        return new SplitPlacementResult(immediateFuture(null), assignment);
    }
    // Find max number of splits consumer can schedule in current cycle.
    int maxSplitsToSchedule = feederScheduledNodes.get(planNodeId).getSplitCount() - consumedNodes.getSplitCount();
    // find list of nodes where still consumer has not been scheduled.
    if (feederScheduledNodes.get(planNodeId).getAssignedNodes().equals(consumedNodes.getAssignedNodes())) {
        candidateNodes = new ArrayList<>(consumedNodes.getAssignedNodes());
    } else {
        for (InternalNode node : feederScheduledNodes.get(planNodeId).getAssignedNodes()) {
            if (!consumedNodes.getAssignedNodes().contains(node)) {
                candidateNodes.add(node);
                consumedNodes.getAssignedNodes().add(node);
            }
        }
    }
    // schedule derived number of splits on derived list of nodes.
    // It is expected that splits count should be at-least equal to number of nodes so that each node gets at-least
    // one split.
    int index = 0;
    int totalNodes = candidateNodes.size();
    for (Split split : Iterables.limit(splits, maxSplitsToSchedule)) {
        InternalNode chosenNode = candidateNodes.get(index % totalNodes);
        assignment.put(chosenNode, split);
        assignmentStats.addAssignedSplit(chosenNode);
        index++;
    }
    consumedNodes.updateSplitCount(maxSplitsToSchedule);
    return new SplitPlacementResult(immediateFuture(null), assignment);
}

Example 28 with Split

use of io.prestosql.metadata.Split in project hetu-core by openlookeng.

the class IndexLoader method streamIndexDataForSingleKey.

public IndexedData streamIndexDataForSingleKey(UpdateRequest updateRequest) {
    Page indexKeyTuple = updateRequest.getPage().getRegion(0, 1);
    PageBuffer pageBuffer = new PageBuffer(100);
    DriverFactory driverFactory = indexBuildDriverFactoryProvider.createStreaming(pageBuffer, indexKeyTuple);
    Driver driver = driverFactory.createDriver(pipelineContext.addDriverContext());
    PageRecordSet pageRecordSet = new PageRecordSet(keyTypes, indexKeyTuple);
    PlanNodeId planNodeId = driverFactory.getSourceId().get();
    ScheduledSplit split = new ScheduledSplit(0, planNodeId, new Split(INDEX_CONNECTOR_ID, new IndexSplit(pageRecordSet), Lifespan.taskWide()));
    driver.updateSource(new TaskSource(planNodeId, ImmutableSet.of(split), true));
    return new StreamingIndexedData(outputTypes, keyTypes, indexKeyTuple, pageBuffer, driver);
}

Example 29 with Split

use of io.prestosql.metadata.Split in project hetu-core by openlookeng.

the class ExtractSpatialJoins method loadKdbTree.

private static KdbTree loadKdbTree(String tableName, Session session, Metadata metadata, SplitManager splitManager, PageSourceManager pageSourceManager, PlanNodeId nodeId) {
    QualifiedObjectName name = toQualifiedObjectName(tableName, session.getCatalog().get(), session.getSchema().get());
    TableHandle tableHandle = metadata.getTableHandle(session, name).orElseThrow(() -> new PrestoException(INVALID_SPATIAL_PARTITIONING, format("Table not found: %s", name)));
    Map<String, ColumnHandle> columnHandles = metadata.getColumnHandles(session, tableHandle);
    List<ColumnHandle> visibleColumnHandles = columnHandles.values().stream().filter(handle -> !metadata.getColumnMetadata(session, tableHandle, handle).isHidden()).collect(toImmutableList());
    checkSpatialPartitioningTable(visibleColumnHandles.size() == 1, "Expected single column for table %s, but found %s columns", name, columnHandles.size());
    ColumnHandle kdbTreeColumn = Iterables.getOnlyElement(visibleColumnHandles);
    Optional<KdbTree> kdbTree = Optional.empty();
    try (SplitSource splitSource = splitManager.getSplits(session, tableHandle, UNGROUPED_SCHEDULING, null, Optional.empty(), Collections.emptyMap(), ImmutableSet.of(), false, nodeId)) {
        while (!Thread.currentThread().isInterrupted()) {
            SplitBatch splitBatch = getFutureValue(splitSource.getNextBatch(NOT_PARTITIONED, Lifespan.taskWide(), 1000));
            List<Split> splits = splitBatch.getSplits();
            for (Split split : splits) {
                try (ConnectorPageSource pageSource = pageSourceManager.createPageSource(session, split, tableHandle, ImmutableList.of(kdbTreeColumn), Optional.empty())) {
                    do {
                        getFutureValue(pageSource.isBlocked());
                        Page page = pageSource.getNextPage();
                        if (page != null && page.getPositionCount() > 0) {
                            checkSpatialPartitioningTable(!kdbTree.isPresent(), "Expected exactly one row for table %s, but found more", name);
                            checkSpatialPartitioningTable(page.getPositionCount() == 1, "Expected exactly one row for table %s, but found %s rows", name, page.getPositionCount());
                            String kdbTreeJson = VARCHAR.getSlice(page.getBlock(0), 0).toStringUtf8();
                            try {
                                kdbTree = Optional.of(KdbTreeUtils.fromJson(kdbTreeJson));
                            } catch (IllegalArgumentException e) {
                                checkSpatialPartitioningTable(false, "Invalid JSON string for KDB tree: %s", e.getMessage());
                            }
                        }
                    } while (!pageSource.isFinished());
                } catch (IOException e) {
                    throw new UncheckedIOException(e);
                }
            }
            if (splitBatch.isLastBatch()) {
                break;
            }
        }
    }
    checkSpatialPartitioningTable(kdbTree.isPresent(), "Expected exactly one row for table %s, but got none", name);
    return kdbTree.get();
}

Example 30 with Split

use of io.prestosql.metadata.Split in project hetu-core by openlookeng.

the class ConnectorAwareSplitSource method groupSmallSplits.

public List<Split> groupSmallSplits(List<Split> pendingSplits, Lifespan lifespan, int maxGroupSize) {
    List<ConnectorSplit> connectorSplits = new ArrayList<>();
    for (Split split : pendingSplits) {
        connectorSplits.add(split.getConnectorSplit());
    }
    List<ConnectorSplit> connectorSplits1 = source.groupSmallSplits(connectorSplits, maxGroupSize);
    ImmutableList.Builder<Split> result = ImmutableList.builder();
    for (ConnectorSplit connectorSplit : connectorSplits1) {
        result.add(new Split(catalogName, connectorSplit, lifespan));
    }
    return result.build();
}