private void assertStatelessRegions(
final OperatorDef stateless1,
final OperatorDef stateless2,
final Map<String, OperatorDef> optimizedOperators,
final Collection<Entry<Port, Port>> optimizedConnections,
final List<RegionDef> statelessRegions) {
for (RegionDef region : statelessRegions) {
if (region.getOperatorCount() != 2) {
continue;
}
final List<OperatorDef> operators = region.getOperators();
final OperatorDef optimizedStateless1 = operators.get(0);
final OperatorDef optimizedStateless2 = operators.get(1);
assertEquals(StatelessOperator2.class, optimizedStateless1.operatorClazz());
assertEquals(StatelessOperator2.class, optimizedStateless2.operatorClazz());
assertTrue(optimizedOperators.containsKey(optimizedStateless1.id()));
assertTrue(optimizedOperators.containsKey(optimizedStateless2.id()));
assertTrue(
optimizedConnections.contains(
new SimpleEntry<>(
new Port(optimizedStateless1.id(), 0), new Port(optimizedStateless2.id(), 0))));
}
assertFalse(optimizedOperators.containsKey(stateless1.id()));
assertFalse(optimizedOperators.containsKey(stateless2.id()));
}
@Test
public void shouldMergePartitionedStatefulAndStatelessRegionsAfterOptimization() {
final OperatorDef partitionedStateful1 =
OperatorDefBuilder.newInstance("partitionedStateful1", PartitionedStatefulOperator.class)
.setExtendingSchema(schema)
.setPartitionFieldNames(partitionFieldNames)
.build();
final OperatorDef partitionedStateful2 =
OperatorDefBuilder.newInstance("partitionedStateful2", PartitionedStatefulOperator.class)
.setExtendingSchema(schema)
.setPartitionFieldNames(partitionFieldNames)
.build();
final OperatorDef stateless =
OperatorDefBuilder.newInstance("stateless", StatelessOperator.class)
.setExtendingSchema(schema)
.build();
final FlowDef flow =
new FlowDefBuilder()
.add(partitionedStateful1)
.add(partitionedStateful2)
.add(stateless)
.connect("partitionedStateful1", "stateless")
.connect("partitionedStateful2", "stateless")
.build();
final List<RegionDef> regions = regionDefFormer.createRegions(flow);
final Map<String, OperatorDef> operators = flow.getOperatorsMap();
final Collection<Entry<Port, Port>> connections = flow.getConnections();
flowOptimizer.duplicateStatelessRegions(operators, connections, regions);
flowOptimizer.mergeRegions(operators, connections, regions);
assertEquals(2, regions.size());
final RegionDef region1 = regions.get(0);
assertEquals(PARTITIONED_STATEFUL, region1.getRegionType());
assertEquals(
asList(PartitionedStatefulOperator.class, StatelessOperator.class),
region1.getOperators().stream().map(OperatorDef::operatorClazz).collect(toList()));
final RegionDef region2 = regions.get(1);
assertEquals(PARTITIONED_STATEFUL, region2.getRegionType());
assertEquals(
asList(PartitionedStatefulOperator.class, StatelessOperator.class),
region2.getOperators().stream().map(OperatorDef::operatorClazz).collect(toList()));
}
@Test
public void shouldMergeStatefulAndStatelessRegions() {
final OperatorDef stateful =
OperatorDefBuilder.newInstance("stateful", StatefulOperator.class).build();
final OperatorDef stateless =
OperatorDefBuilder.newInstance("stateless", StatelessOperator.class).build();
final FlowDef flow =
new FlowDefBuilder().add(stateful).add(stateless).connect("stateful", "stateless").build();
final List<RegionDef> regions = regionDefFormer.createRegions(flow);
flowOptimizer.mergeRegions(flow.getOperatorsMap(), flow.getConnections(), regions);
assertEquals(1, regions.size());
final RegionDef region = regions.get(0);
assertEquals(STATEFUL, region.getRegionType());
assertEquals(asList(stateful, stateless), region.getOperators());
}
private void assertStatelessRegions(
final OperatorDef stateless1,
final OperatorDef stateless2,
final OperatorDef stateless3,
final Map<String, OperatorDef> optimizedOperators,
final Collection<Entry<Port, Port>> optimizedConnections,
final List<RegionDef> statelessRegions,
final boolean expectedNonMatchingStatelessRegion) {
boolean stateful1ConnectionExists = false, stateful2ConnectionExists = false;
for (RegionDef region : statelessRegions) {
if (region.getOperatorCount() != 3) {
assertTrue(expectedNonMatchingStatelessRegion);
continue;
}
final List<OperatorDef> operators = region.getOperators();
final OperatorDef optimizedStateless1 = operators.get(0);
final OperatorDef optimizedStateless2 = operators.get(1);
final OperatorDef optimizedStateless3 = operators.get(2);
assertEquals(StatelessOperator.class, optimizedStateless1.operatorClazz());
assertEquals(StatelessOperator.class, optimizedStateless2.operatorClazz());
assertEquals(StatelessOperator.class, optimizedStateless3.operatorClazz());
assertTrue(optimizedOperators.containsKey(optimizedStateless1.id()));
assertTrue(optimizedOperators.containsKey(optimizedStateless2.id()));
assertTrue(optimizedOperators.containsKey(optimizedStateless3.id()));
final Collection<Port> stateful1DownstreamConnections =
getDownstream(optimizedConnections, new Port("stateful1", 0));
final Collection<Port> stateful2DownstreamConnections =
getDownstream(optimizedConnections, new Port("stateful2", 0));
assertEquals(1, stateful1DownstreamConnections.size());
assertEquals(1, stateful2DownstreamConnections.size());
final Port stateful1DownstreamPort = stateful1DownstreamConnections.iterator().next();
if (stateful1DownstreamPort.operatorId.equals(optimizedStateless1.id())) {
stateful1ConnectionExists = true;
} else {
final Port stateful2DownstreamPort = stateful2DownstreamConnections.iterator().next();
if (stateful2DownstreamPort.operatorId.equals(optimizedStateless1.id())) {
stateful2ConnectionExists = true;
}
}
assertTrue(
optimizedConnections.contains(
new SimpleEntry<>(
new Port(optimizedStateless1.id(), 0), new Port(optimizedStateless2.id(), 0))));
assertTrue(
optimizedConnections.contains(
new SimpleEntry<>(
new Port(optimizedStateless2.id(), 0), new Port(optimizedStateless3.id(), 0))));
assertTrue(
optimizedConnections.contains(
new SimpleEntry<>(new Port(optimizedStateless3.id(), 0), new Port("stateful3", 0))));
assertTrue(
optimizedConnections.contains(
new SimpleEntry<>(new Port(optimizedStateless3.id(), 0), new Port("stateful4", 0))));
}
assertTrue(stateful1ConnectionExists);
assertTrue(stateful2ConnectionExists);
assertFalse(optimizedOperators.containsKey(stateless1.id()));
assertFalse(optimizedOperators.containsKey(stateless2.id()));
assertFalse(optimizedOperators.containsKey(stateless3.id()));
}