Examples with IdentityPartitionerMapper org.apache.flink.optimizer.testfunctions.IdentityPartitionerMapper used on opensource projects

Example 1 with IdentityPartitionerMapper

use of org.apache.flink.optimizer.testfunctions.IdentityPartitionerMapper in project flink by apache.

the class CustomPartitioningTest method testPartitionPojo.

@Test
public void testPartitionPojo() {
    try {
        final Partitioner<Integer> part = new TestPartitionerInt();
        final int parallelism = 4;
        ExecutionEnvironment env = ExecutionEnvironment.getExecutionEnvironment();
        env.setParallelism(parallelism);
        DataSet<Pojo> data = env.fromElements(new Pojo()).rebalance();
        data.partitionCustom(part, "a").mapPartition(new IdentityPartitionerMapper<Pojo>()).output(new DiscardingOutputFormat<Pojo>());
        Plan p = env.createProgramPlan();
        OptimizedPlan op = compileNoStats(p);
        SinkPlanNode sink = op.getDataSinks().iterator().next();
        SingleInputPlanNode mapper = (SingleInputPlanNode) sink.getInput().getSource();
        SingleInputPlanNode partitioner = (SingleInputPlanNode) mapper.getInput().getSource();
        SingleInputPlanNode balancer = (SingleInputPlanNode) partitioner.getInput().getSource();
        assertEquals(ShipStrategyType.FORWARD, sink.getInput().getShipStrategy());
        assertEquals(parallelism, sink.getParallelism());
        assertEquals(ShipStrategyType.FORWARD, mapper.getInput().getShipStrategy());
        assertEquals(parallelism, mapper.getParallelism());
        assertEquals(ShipStrategyType.PARTITION_CUSTOM, partitioner.getInput().getShipStrategy());
        assertEquals(part, partitioner.getInput().getPartitioner());
        assertEquals(parallelism, partitioner.getParallelism());
        assertEquals(ShipStrategyType.PARTITION_FORCED_REBALANCE, balancer.getInput().getShipStrategy());
        assertEquals(parallelism, balancer.getParallelism());
    } catch (Exception e) {
        e.printStackTrace();
        fail(e.getMessage());
    }
}

Example 2 with IdentityPartitionerMapper

use of org.apache.flink.optimizer.testfunctions.IdentityPartitionerMapper in project flink by apache.

the class CustomPartitioningTest method testPartitionTuples.

@Test
public void testPartitionTuples() {
    try {
        final Partitioner<Integer> part = new TestPartitionerInt();
        final int parallelism = 4;
        ExecutionEnvironment env = ExecutionEnvironment.getExecutionEnvironment();
        env.setParallelism(parallelism);
        DataSet<Tuple2<Integer, Integer>> data = env.fromElements(new Tuple2<Integer, Integer>(0, 0)).rebalance();
        data.partitionCustom(part, 0).mapPartition(new IdentityPartitionerMapper<Tuple2<Integer, Integer>>()).output(new DiscardingOutputFormat<Tuple2<Integer, Integer>>());
        Plan p = env.createProgramPlan();
        OptimizedPlan op = compileNoStats(p);
        SinkPlanNode sink = op.getDataSinks().iterator().next();
        SingleInputPlanNode mapper = (SingleInputPlanNode) sink.getInput().getSource();
        SingleInputPlanNode partitioner = (SingleInputPlanNode) mapper.getInput().getSource();
        SingleInputPlanNode balancer = (SingleInputPlanNode) partitioner.getInput().getSource();
        assertEquals(ShipStrategyType.FORWARD, sink.getInput().getShipStrategy());
        assertEquals(parallelism, sink.getParallelism());
        assertEquals(ShipStrategyType.FORWARD, mapper.getInput().getShipStrategy());
        assertEquals(parallelism, mapper.getParallelism());
        assertEquals(ShipStrategyType.PARTITION_CUSTOM, partitioner.getInput().getShipStrategy());
        assertEquals(part, partitioner.getInput().getPartitioner());
        assertEquals(parallelism, partitioner.getParallelism());
        assertEquals(ShipStrategyType.PARTITION_FORCED_REBALANCE, balancer.getInput().getShipStrategy());
        assertEquals(parallelism, balancer.getParallelism());
    } catch (Exception e) {
        e.printStackTrace();
        fail(e.getMessage());
    }
}

Example 3 with IdentityPartitionerMapper

use of org.apache.flink.optimizer.testfunctions.IdentityPartitionerMapper in project flink by apache.

the class CustomPartitioningTest method testPartitionKeySelector.

@Test
public void testPartitionKeySelector() {
    try {
        final Partitioner<Integer> part = new TestPartitionerInt();
        final int parallelism = 4;
        ExecutionEnvironment env = ExecutionEnvironment.getExecutionEnvironment();
        env.setParallelism(parallelism);
        DataSet<Pojo> data = env.fromElements(new Pojo()).rebalance();
        data.partitionCustom(part, new TestKeySelectorInt<Pojo>()).mapPartition(new IdentityPartitionerMapper<Pojo>()).output(new DiscardingOutputFormat<Pojo>());
        Plan p = env.createProgramPlan();
        OptimizedPlan op = compileNoStats(p);
        SinkPlanNode sink = op.getDataSinks().iterator().next();
        SingleInputPlanNode mapper = (SingleInputPlanNode) sink.getInput().getSource();
        SingleInputPlanNode keyRemover = (SingleInputPlanNode) mapper.getInput().getSource();
        SingleInputPlanNode partitioner = (SingleInputPlanNode) keyRemover.getInput().getSource();
        SingleInputPlanNode keyExtractor = (SingleInputPlanNode) partitioner.getInput().getSource();
        SingleInputPlanNode balancer = (SingleInputPlanNode) keyExtractor.getInput().getSource();
        assertEquals(ShipStrategyType.FORWARD, sink.getInput().getShipStrategy());
        assertEquals(parallelism, sink.getParallelism());
        assertEquals(ShipStrategyType.FORWARD, mapper.getInput().getShipStrategy());
        assertEquals(parallelism, mapper.getParallelism());
        assertEquals(ShipStrategyType.FORWARD, keyRemover.getInput().getShipStrategy());
        assertEquals(parallelism, keyRemover.getParallelism());
        assertEquals(ShipStrategyType.PARTITION_CUSTOM, partitioner.getInput().getShipStrategy());
        assertEquals(part, partitioner.getInput().getPartitioner());
        assertEquals(parallelism, partitioner.getParallelism());
        assertEquals(ShipStrategyType.FORWARD, keyExtractor.getInput().getShipStrategy());
        assertEquals(parallelism, keyExtractor.getParallelism());
        assertEquals(ShipStrategyType.PARTITION_FORCED_REBALANCE, balancer.getInput().getShipStrategy());
        assertEquals(parallelism, balancer.getParallelism());
    } catch (Exception e) {
        e.printStackTrace();
        fail(e.getMessage());
    }
}