@Override public ListenableFuture<SplitBatch> getNextBatch(ConnectorPartitionHandle partitionHandle, Lifespan lifespan, int maxSize) { ListenableFuture<ConnectorSplitBatch> nextBatch = toListenableFuture(source.getNextBatch(partitionHandle, maxSize)); return Futures.transform(nextBatch, splitBatch -> { ImmutableList.Builder<Split> result = ImmutableList.builder(); for (ConnectorSplit connectorSplit : splitBatch.getSplits()) { result.add(new Split(connectorId, transactionHandle, connectorSplit, lifespan)); } return new SplitBatch(result.build(), splitBatch.isNoMoreSplits()); }, directExecutor()); }
@Test public void testScheduleLocal() { Split split = new Split(CONNECTOR_ID, TestingTransactionHandle.create(), new TestSplitLocal()); Set<Split> splits = ImmutableSet.of(split); Map.Entry<Node, Split> assignment = Iterables.getOnlyElement(nodeSelector.computeAssignments(splits, ImmutableList.copyOf(taskMap.values())).getAssignments().entries()); assertEquals(assignment.getKey().getHostAndPort(), split.getAddresses().get(0)); assertEquals(assignment.getValue(), split); }
@Override public ConnectorPageSource createPageSource(Session session, Split split, List<ColumnHandle> columns) { requireNonNull(split, "split is null"); requireNonNull(columns, "columns is null"); ConnectorSession connectorSession = session.toConnectorSession(split.getConnectorId()); return getPageSourceProvider(split).createPageSource(split.getTransactionHandle(), connectorSession, split.getConnectorSplit(), columns); }
@Override public Supplier<Optional<UpdatablePageSource>> addSplit(Split split) { requireNonNull(split, "split is null"); checkArgument(split.getConnectorId().equals(REMOTE_CONNECTOR_ID), "split is not a remote split"); URI location = ((RemoteSplit) split.getConnectorSplit()).getLocation(); exchangeClient.addLocation(location); return Optional::empty; }
@Override public Supplier<Optional<UpdatablePageSource>> addSplit(Split split) { requireNonNull(split, "split is null"); checkState(this.split == null, "Table scan split already set"); if (finishing) { return Optional::empty; } this.split = split; Object splitInfo = split.getInfo(); if (splitInfo != null) { operatorContext.setInfoSupplier(() -> new SplitOperatorInfo(splitInfo)); } blocked.set(null); if (split.getConnectorSplit() instanceof EmptySplit) { pageSource = new EmptySplitPageSource(); } return () -> { if (pageSource instanceof UpdatablePageSource) { return Optional.of((UpdatablePageSource) pageSource); } return Optional.empty(); }; }
boolean splitWaitingForAnyNode = false; for (Split split : splits) { if (!split.isRemotelyAccessible()) { List<Node> candidateNodes = selectExactNodes(nodeMap, split.getAddresses(), includeCoordinator); if (candidateNodes.isEmpty()) { log.debug("No nodes available to schedule %s. Available nodes %s", split, nodeMap.getNodesByHost().keys()); int chosenDepth = 0; Set<NetworkLocation> locations = new HashSet<>(); for (HostAddress host : split.getAddresses()) { locations.add(networkLocationCache.get(host));
@Override public Supplier<Optional<UpdatablePageSource>> addSplit(Split split) { requireNonNull(split, "split is null"); checkState(this.split == null, "Table scan split already set"); if (finished) { return Optional::empty; } this.split = (TestingSplit) split.getConnectorSplit(); blocked.set(null); return Optional::empty; }
private ToIntFunction<Split> getSplitToBucket(Session session, PartitioningHandle partitioningHandle) { ConnectorNodePartitioningProvider partitioningProvider = partitioningProviders.get(partitioningHandle.getConnectorId().get()); checkArgument(partitioningProvider != null, "No partitioning provider for connector %s", partitioningHandle.getConnectorId().get()); ToIntFunction<ConnectorSplit> splitBucketFunction = partitioningProvider.getSplitBucketFunction( partitioningHandle.getTransactionHandle().orElse(null), session.toConnectorSession(), partitioningHandle.getConnectorHandle()); checkArgument(splitBucketFunction != null, "No partitioning %s", partitioningHandle); return split -> { int bucket; if (split.getConnectorSplit() instanceof EmptySplit) { bucket = split.getLifespan().isTaskWide() ? 0 : split.getLifespan().getId(); } else { bucket = splitBucketFunction.applyAsInt(split.getConnectorSplit()); } if (!split.getLifespan().isTaskWide()) { checkArgument(split.getLifespan().getId() == bucket); } return bucket; }; }
@Override public String getInfo() { return (partitionedSplit == null) ? "" : partitionedSplit.getSplit().getInfo().toString(); }
private ConnectorPageSourceProvider getPageSourceProvider(Split split) { ConnectorPageSourceProvider provider = pageSourceProviders.get(split.getConnectorId()); checkArgument(provider != null, "No page stream provider for '%s", split.getConnectorId()); return provider; } }
if (!split.isRemotelyAccessible()) { candidateNodes = selectExactNodes(nodeMap, split.getAddresses(), includeCoordinator); if (split.isRemotelyAccessible()) { splitWaitingForAnyNode = true;
@Override public Supplier<Optional<UpdatablePageSource>> addSplit(Split split) { requireNonNull(split, "split is null"); checkState(this.split == null, "Table scan split already set"); if (finished) { return Optional::empty; } this.split = split; Object splitInfo = split.getInfo(); if (splitInfo != null) { operatorContext.setInfoSupplier(() -> new SplitOperatorInfo(splitInfo)); } blocked.set(null); if (split.getConnectorSplit() instanceof EmptySplit) { source = new EmptySplitPageSource(); } return () -> { if (source instanceof UpdatablePageSource) { return Optional.of((UpdatablePageSource) source); } return Optional.empty(); }; }
@Override public Supplier<Optional<UpdatablePageSource>> addSplit(Split split) { requireNonNull(split, "split is null"); checkArgument(split.getConnectorSplit() instanceof RemoteSplit, "split is not a remote split"); checkState(!blockedOnSplits.isDone(), "noMoreSplits has been called already"); URI location = ((RemoteSplit) split.getConnectorSplit()).getLocation(); ExchangeClient exchangeClient = closer.register(exchangeClientSupplier.get(operatorContext.localSystemMemoryContext())); exchangeClient.addLocation(location); exchangeClient.noMoreLocations(); pageProducers.add(exchangeClient.pages() .map(serializedPage -> { operatorContext.recordRawInput(serializedPage.getSizeInBytes()); return pagesSerde.deserialize(serializedPage); })); return Optional::empty; }
@Override public Supplier<Optional<UpdatablePageSource>> addSplit(Split split) { requireNonNull(split, "split is null"); checkArgument(split.getConnectorId().equals("remote"), "split is not a remote split"); URI location = ((RemoteSplit) split.getConnectorSplit()).getLocation(); exchangeClient.addLocation(location); return Optional::empty; }
@Override public String getInfo() { return (partitionedSplit == null) ? "" : partitionedSplit.getSplit().getInfo().toString(); }
private ConnectorPageSourceProvider getPageSourceProvider(Split split) { ConnectorPageSourceProvider provider = pageSourceProviders.get(split.getConnectorId()); checkArgument(provider != null, "No page stream provider for '%s", split.getConnectorId()); return provider; } }
@Test public void testScheduleRemote() { Set<Split> splits = new HashSet<>(); splits.add(new Split(CONNECTOR_ID, TestingTransactionHandle.create(), new TestSplitRemote())); Multimap<Node, Split> assignments = nodeSelector.computeAssignments(splits, ImmutableList.copyOf(taskMap.values())).getAssignments(); assertEquals(assignments.size(), 1); }
nonRackLocalBuilder.add(new Split(CONNECTOR_ID, transactionHandle, new TestSplitRemote(HostAddress.fromParts("data.other_rack", 1)))); HostAddress dataHost2 = HostAddress.fromParts("data.rack2", 1); for (int i = 0; i < 6 * 2; i++) { rackLocalSplits.add(new Split(CONNECTOR_ID, transactionHandle, new TestSplitRemote(dataHost1))); rackLocalSplits.add(new Split(CONNECTOR_ID, transactionHandle, new TestSplitRemote(dataHost2))); int rack2 = 0; for (Split split : unassigned) { String rack = topology.locate(split.getAddresses().get(0)).getSegments().get(0); switch (rack) { case "rack1": localSplits.add(new Split(CONNECTOR_ID, transactionHandle, new TestSplitRemote(HostAddress.fromParts("host1.rack1", 1)))); localSplits.add(new Split(CONNECTOR_ID, transactionHandle, new TestSplitRemote(HostAddress.fromParts("host2.rack1", 1)))); localSplits.add(new Split(CONNECTOR_ID, transactionHandle, new TestSplitRemote(HostAddress.fromParts("host3.rack2", 1)))); assignments = nodeSelector.computeAssignments(localSplits.build(), ImmutableList.copyOf(taskMap.values())).getAssignments(); assertEquals(assignments.size(), 3);
Set<NetworkLocation> filledLocations = new HashSet<>(); for (Split split : splits) { if (!split.isRemotelyAccessible()) { List<Node> candidateNodes = selectExactNodes(nodeMap, split.getAddresses(), includeCoordinator); if (candidateNodes.isEmpty()) { log.debug("No nodes available to schedule %s. Available nodes %s", split, nodeMap.getNodesByHost().keys()); int chosenDepth = 0; Set<NetworkLocation> locations = new HashSet<>(); for (HostAddress host : split.getAddresses()) { locations.add(networkLocationCache.get(host));
@Override public ConnectorPageSource createPageSource(Session session, Split split, List<ColumnHandle> columns) { requireNonNull(split, "split is null"); requireNonNull(columns, "columns is null"); // assumes connectorId and catalog are the same ConnectorSession connectorSession = session.toConnectorSession(split.getConnectorId()); return getPageSourceProvider(split).createPageSource(split.getTransactionHandle(), connectorSession, split.getConnectorSplit(), columns); }