PartitionerFactory { public static Partitioner create(Config config, JavaPairRDD<Row, Row> rdd) { String partitionerType = config.getString(TYPE_CONFIG_NAME); if (!config.hasPath(TYPE_CONFIG_NAME)) { throw new RuntimeException("Partitioner type not specified"); } Partitioner partitioner; switch (partitionerType) { case "hash": partitioner = new HashPartitioner(rdd.getNumPartitions()); break; case "range": Ordering<Row> rowOrdering = Ordering$.MODULE$.<Row>comparatorToOrdering(new RowComparator()); ClassTag<Row> rowClassTag = ClassTag$.MODULE$.<Row>apply(Row.class); partitioner = new RangePartitioner<Row, Row>(rdd.getNumPartitions(), rdd.rdd(), true, rowOrdering, rowClassTag); break; default: partitioner = ComponentFactory.create(ConfigurablePartitioner.class, config, true); ((ConfigurablePartitioner) partitioner).configureRDD(rdd); } return partitioner; }  static Partitioner create(Config config, JavaPairRDD<Row, Row> rdd); static final String TYPE_CONFIG_NAME; }
@Test public void testHash() { Map<String, Object> configMap = Maps.newHashMap(); configMap.put("type", "hash"); JavaPairRDD<Row, Row> base = getDummyRDD(10); Config config = ConfigFactory.parseMap(configMap); Partitioner p = PartitionerFactory.create(config, base); assertTrue(p instanceof HashPartitioner); assertEquals(p.numPartitions(), 10); }
@Test public void testRange() { Map<String, Object> configMap = Maps.newHashMap(); configMap.put("type", "range"); JavaPairRDD<Row, Row> base = getDummyRDD(10); Config config = ConfigFactory.parseMap(configMap); Partitioner p = PartitionerFactory.create(config, base); assertTrue(p instanceof RangePartitioner); assertEquals(p.numPartitions(), 10); }
@Test public void testCustom() { Map<String, Object> configMap = Maps.newHashMap(); configMap.put("type", "com.cloudera.labs.envelope.partition.DummyPartitioner"); JavaPairRDD<Row, Row> base = getDummyRDD(10); Config config = ConfigFactory.parseMap(configMap); DummyPartitioner p = (DummyPartitioner)PartitionerFactory.create(config, base); assertTrue(p instanceof DummyPartitioner); assertTrue(p.hasBeenConfigured()); }