@Override public ListenableFuture<SplitBatch> getNextBatch(ConnectorPartitionHandle partitionHandle, Lifespan lifespan, int maxSize) { ListenableFuture<SplitBatch> batch = splitSource.getNextBatch(partitionHandle, lifespan, maxSize); return Futures.transform(batch, splitBatch -> new SplitBatch( splitBatch.getSplits().stream() .filter(input -> ThreadLocalRandom.current().nextDouble() < sampleRatio) .collect(toImmutableList()), splitBatch.isLastBatch()), directExecutor()); }
@Override public boolean isFinished() { return source.isFinished(); }
@Override public void close() { splitSource.close(); }
scheduleGroup.nextSplitBatchFuture = splitSource.getNextBatch(scheduleGroup.partitionHandle, lifespan, splitBatchSize - pendingSplits.size()); splitSource.getConnectorId(), splitSource.getTransactionHandle(), new EmptySplit(splitSource.getConnectorId()), lifespan)); if ((state == State.NO_MORE_SPLITS || state == State.FINISHED) || (noMoreScheduleGroups && scheduleGroups.isEmpty() && splitSource.isFinished())) { switch (state) { case INITIALIZED: case SPLITS_ADDED: state = State.NO_MORE_SPLITS; splitSource.close();
private Split getLocalQuerySplit(Session session, TableLayoutHandle handle) { SplitSource splitSource = splitManager.getSplits(session, handle); List<Split> splits = new ArrayList<>(); splits.addAll(getFutureValue(splitSource.getNextBatch(1000))); while (!splitSource.isFinished()) { splits.addAll(getFutureValue(splitSource.getNextBatch(1000))); } checkArgument(splits.size() == 1, "Expected only one split for a local query, but got %s splits", splits.size()); return splits.get(0); }
if (splitSource.isFinished()) { splitSource.close(); return new ScheduleResult(true, ImmutableSet.of(), CompletableFuture.completedFuture(null)); batchFuture = splitSource.getNextBatch(splitBatchSize); batchFuture.thenRun(() -> stage.recordGetSplitTime(start)); boolean finished = splitSource.isFinished(); if (finished) { splitSource.close();
@Override public ConnectorId getConnectorId() { return source.getConnectorId(); }
@Nullable @Override public String getDataSourceName() { return splitSource.getDataSourceName(); }
@Override public ConnectorTransactionHandle getTransactionHandle() { return source.getTransactionHandle(); }
while (!splitSource.isFinished()) { for (Split split : getFutureValue(splitSource.getNextBatch(1000))) { scheduledSplits.add(new ScheduledSplit(sequenceId++, split));
@Nullable @Override public ConnectorId getConnectorId() { return splitSource.getConnectorId(); }
NodeSelector nodeSelector = nodeScheduler.createNodeSelector(splitSource.getDataSourceName()); SplitPlacementPolicy placementPolicy = new DynamicSplitPlacementPolicy(nodeSelector, stage::getAllTasks); stageSchedulers.put(stageId, new SourcePartitionedScheduler(stage, splitSource, placementPolicy, splitBatchSize));
@Override public ConnectorTransactionHandle getTransactionHandle() { return splitSource.getTransactionHandle(); }
@Override public boolean isFinished() { return splitSource.isFinished(); } }
private ListenableFuture<?> fetchSplits() { if (splits.size() >= min) { return immediateFuture(null); } ListenableFuture<SplitBatch> future = splitSource.getNextBatch(partitionHandle, lifespan, max - splits.size()); return Futures.transformAsync(future, splitBatch -> { splits.addAll(splitBatch.getSplits()); if (splitBatch.isLastBatch()) { noMoreSplits = true; return immediateFuture(null); } return fetchSplits(); }, directExecutor()); } }
@Override public void close() { splitSource.close(); }
PlanNodeId planNodeId = entry.getKey(); SplitSource splitSource = entry.getValue(); ConnectorId connectorId = splitSource.getConnectorId(); if (isInternalSystemConnector(connectorId)) { connectorId = null;
private static SourcePartitionedScheduler getSourcePartitionedScheduler( StageExecutionPlan plan, SqlStageExecution stage, NodeManager nodeManager, NodeTaskMap nodeTaskMap, int splitBatchSize) { NodeSchedulerConfig nodeSchedulerConfig = new NodeSchedulerConfig() .setIncludeCoordinator(false) .setMaxSplitsPerNode(20) .setMaxPendingSplitsPerNodePerTask(0); NodeScheduler nodeScheduler = new NodeScheduler(new LegacyNetworkTopology(), nodeManager, nodeSchedulerConfig, nodeTaskMap); SplitSource splitSource = plan.getDataSource().get(); SplitPlacementPolicy placementPolicy = new DynamicSplitPlacementPolicy(nodeScheduler.createNodeSelector(splitSource.getDataSourceName()), stage::getAllTasks); return new SourcePartitionedScheduler(stage, splitSource, placementPolicy, splitBatchSize); }
private Split getLocalQuerySplit(Session session, TableLayoutHandle handle) { SplitSource splitSource = localQueryRunner.getSplitManager().getSplits(session, handle, UNGROUPED_SCHEDULING); List<Split> splits = new ArrayList<>(); while (!splitSource.isFinished()) { splits.addAll(getNextBatch(splitSource)); } checkArgument(splits.size() == 1, "Expected only one split for a local query, but got %s splits", splits.size()); return splits.get(0); }
private static ListenableFuture<NextBatchResult> getNextBatch(SplitSource splitSource, int maxSize) { ListenableFuture<SplitBatch> future = splitSource.getNextBatch(NOT_PARTITIONED, Lifespan.taskWide(), maxSize); return Futures.transform(future, NextBatchResult::new, directExecutor()); }