@Override
public ConnectorPageSource createPageSource(
Session session, Split split, List<ColumnHandle> columns) {
assertInstanceOf(split.getConnectorSplit(), FunctionAssertions.TestSplit.class);
FunctionAssertions.TestSplit testSplit =
(FunctionAssertions.TestSplit) split.getConnectorSplit();
if (testSplit.isRecordSet()) {
RecordSet records =
InMemoryRecordSet.builder(
ImmutableList.<Type>of(
BIGINT,
VARCHAR,
DOUBLE,
BOOLEAN,
BIGINT,
VARCHAR,
VARCHAR,
TIMESTAMP_WITH_TIME_ZONE))
.addRow(
1234L,
"hello",
12.34,
true,
new DateTime(2001, 8, 22, 3, 4, 5, 321, DateTimeZone.UTC).getMillis(),
"%el%",
null,
packDateTimeWithZone(
new DateTime(1970, 1, 1, 0, 1, 0, 999, DateTimeZone.UTC).getMillis(),
TimeZoneKey.getTimeZoneKey("Z")))
.build();
return new RecordPageSource(records);
} else {
return new FixedPageSource(ImmutableList.of(SOURCE_PAGE));
}
}
@Override
public Supplier<Optional<UpdatablePageSource>> addSplit(Split split) {
checkNotNull(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;
}
private List<Node> selectCandidateNodes(NodeMap nodeMap, Split split) {
Set<Node> chosen = new LinkedHashSet<>(minCandidates);
// first look for nodes that match the hint
for (HostAddress hint : split.getAddresses()) {
for (Node node : nodeMap.getNodesByHostAndPort().get(hint)) {
if (chosen.add(node)) {
scheduleLocal.incrementAndGet();
}
}
InetAddress address;
try {
address = hint.toInetAddress();
} catch (UnknownHostException e) {
// skip addresses that don't resolve
continue;
}
// consider a split with a host hint without a port as being accessible
// by all nodes in that host
if (!hint.hasPort() || split.isRemotelyAccessible()) {
for (Node node : nodeMap.getNodesByHost().get(address)) {
if (chosen.add(node)) {
scheduleLocal.incrementAndGet();
}
}
}
}
// add nodes in same rack, if below the minimum count
if (split.isRemotelyAccessible() && chosen.size() < minCandidates) {
for (HostAddress hint : split.getAddresses()) {
InetAddress address;
try {
address = hint.toInetAddress();
} catch (UnknownHostException e) {
// skip addresses that don't resolve
continue;
}
for (Node node : nodeMap.getNodesByRack().get(Rack.of(address))) {
if (chosen.add(node)) {
scheduleRack.incrementAndGet();
}
if (chosen.size() == minCandidates) {
break;
}
}
if (chosen.size() == minCandidates) {
break;
}
}
}
// add some random nodes if below the minimum count
if (split.isRemotelyAccessible()) {
if (chosen.size() < minCandidates) {
for (Node node : lazyShuffle(nodeMap.getNodesByHost().values())) {
if (chosen.add(node)) {
scheduleRandom.incrementAndGet();
}
if (chosen.size() == minCandidates) {
break;
}
}
}
}
return ImmutableList.copyOf(chosen);
}
