org.apache.flink.api.common.functions.Partitioner Java Examples

@Test
public void testCustomPartitioningPreserved2() {

	SingleInputSemanticProperties sprops = new SingleInputSemanticProperties();
	SemanticPropUtil.getSemanticPropsSingleFromString(sprops, new String[]{"0->1; 1->2; 4->3"}, null, null, tupleInfo, tupleInfo);

	GlobalProperties gprops = new GlobalProperties();
	Partitioner<Tuple2<Long, Integer>> myP = new MockPartitioner();
	gprops.setCustomPartitioned(new FieldList(0, 4), myP);

	GlobalProperties result = gprops.filterBySemanticProperties(sprops, 0);

	assertEquals(PartitioningProperty.CUSTOM_PARTITIONING, result.getPartitioning());
	FieldList pFields = result.getPartitioningFields();
	assertEquals(2, pFields.size());
	assertTrue(pFields.contains(1));
	assertTrue(pFields.contains(3));
	assertEquals(myP, result.getCustomPartitioner());
}
 

@Override
public <E> void validateCustomPartitioner(Partitioner<E> partitioner, TypeInformation<E> typeInfo) {

	if (keyFields.size() != 1) {
		throw new InvalidProgramException("Custom partitioners can only be used with keys that have one key field.");
	}
	
	if (typeInfo == null) {
		// try to extract key type from partitioner
		try {
			typeInfo = TypeExtractor.getPartitionerTypes(partitioner);
		}
		catch (Throwable t) {
			// best effort check, so we ignore exceptions
		}
	}

	// only check if type is known and not a generic type
	if (typeInfo != null && !(typeInfo instanceof GenericTypeInfo)) {
		// check equality of key and partitioner type
		if (!keyType.equals(typeInfo)) {
			throw new InvalidProgramException("The partitioner is incompatible with the key type. "
				+ "Partitioner type: " + typeInfo + " , key type: " + keyType);
		}
	}
}
 

public static <X, K> KeySelector<X, K> getSelectorForOneKey(
		Keys<X> keys, Partitioner<K> partitioner, TypeInformation<X> typeInfo, ExecutionConfig executionConfig) {
	if (!(typeInfo instanceof CompositeType)) {
		throw new InvalidTypesException(
				"This key operation requires a composite type such as Tuples, POJOs, case classes, etc");
	}
	if (partitioner != null) {
		keys.validateCustomPartitioner(partitioner, null);
	}

	CompositeType<X> compositeType = (CompositeType<X>) typeInfo;
	int[] logicalKeyPositions = keys.computeLogicalKeyPositions();
	if (logicalKeyPositions.length != 1) {
		throw new IllegalArgumentException("There must be exactly 1 key specified");
	}

	TypeComparator<X> comparator = compositeType.createComparator(
			logicalKeyPositions, new boolean[] { true }, 0, executionConfig);
	return new OneKeySelector<>(comparator);
}
 

public GroupReduceWithCombineProperties(FieldSet groupKeys, Ordering additionalOrderKeys, Partitioner<?> customPartitioner) {
	super(groupKeys);
	
	// if we have an additional ordering, construct the ordering to have primarily the grouping fields
	if (additionalOrderKeys != null) {
		this.ordering = new Ordering();
		for (Integer key : this.keyList) {
			this.ordering.appendOrdering(key, null, Order.ANY);
		}
	
		// and next the additional order fields
		for (int i = 0; i < additionalOrderKeys.getNumberOfFields(); i++) {
			Integer field = additionalOrderKeys.getFieldNumber(i);
			Order order = additionalOrderKeys.getOrder(i);
			this.ordering.appendOrdering(field, additionalOrderKeys.getType(i), order);
		}
	} else {
		this.ordering = null;
	}
	
	this.customPartitioner = customPartitioner;
}
 

public GroupReduceProperties(FieldSet groupKeys, Ordering additionalOrderKeys, Partitioner<?> customPartitioner) {
	super(groupKeys);
	
	// if we have an additional ordering, construct the ordering to have primarily the grouping fields
	if (additionalOrderKeys != null) {
		this.ordering = new Ordering();
		for (Integer key : this.keyList) {
			this.ordering.appendOrdering(key, null, Order.ANY);
		}
	
		// and next the additional order fields
		for (int i = 0; i < additionalOrderKeys.getNumberOfFields(); i++) {
			Integer field = additionalOrderKeys.getFieldNumber(i);
			Order order = additionalOrderKeys.getOrder(i);
			this.ordering.appendOrdering(field, additionalOrderKeys.getType(i), order);
		}
	}
	else {
		this.ordering = null;
	}
	
	this.customPartitioner = customPartitioner;
}
 

private List<OperatorDescriptorDual> initializeDataProperties(Partitioner<?> customPartitioner) {
	Ordering groupOrder1 = null;
	Ordering groupOrder2 = null;
	
	CoGroupOperatorBase<?, ?, ?, ?> cgc = getOperator();
	groupOrder1 = cgc.getGroupOrderForInputOne();
	groupOrder2 = cgc.getGroupOrderForInputTwo();
		
	if (groupOrder1 != null && groupOrder1.getNumberOfFields() == 0) {
		groupOrder1 = null;
	}
	if (groupOrder2 != null && groupOrder2.getNumberOfFields() == 0) {
		groupOrder2 = null;
	}
	
	CoGroupDescriptor descr = new CoGroupDescriptor(this.keys1, this.keys2, groupOrder1, groupOrder2);
	if (customPartitioner != null) {
		descr.setCustomPartitioner(customPartitioner);
	}
	
	return Collections.<OperatorDescriptorDual>singletonList(descr);
}
 

private <P> PartitionOperator(DataSet<T> input, PartitionMethod pMethod, Keys<T> pKeys, Partitioner<P> customPartitioner,
		TypeInformation<P> partitionerTypeInfo, DataDistribution distribution, String partitionLocationName) {
	super(input, input.getType());

	Preconditions.checkNotNull(pMethod);
	Preconditions.checkArgument(pKeys != null || pMethod == PartitionMethod.REBALANCE, "Partitioning requires keys");
	Preconditions.checkArgument(pMethod != PartitionMethod.CUSTOM || customPartitioner != null, "Custom partioning requires a partitioner.");
	Preconditions.checkArgument(distribution == null || pMethod == PartitionMethod.RANGE, "Customized data distribution is only neccessary for range partition.");

	if (distribution != null) {
		Preconditions.checkArgument(pKeys.getNumberOfKeyFields() <= distribution.getNumberOfFields(), "The distribution must provide at least as many fields as flat key fields are specified.");
		Preconditions.checkArgument(Arrays.equals(pKeys.getKeyFieldTypes(), Arrays.copyOfRange(distribution.getKeyTypes(), 0, pKeys.getNumberOfKeyFields())),
				"The types of the flat key fields must be equal to the types of the fields of the distribution.");
	}

	if (customPartitioner != null) {
		pKeys.validateCustomPartitioner(customPartitioner, partitionerTypeInfo);
	}

	this.pMethod = pMethod;
	this.pKeys = pKeys;
	this.partitionLocationName = partitionLocationName;
	this.customPartitioner = customPartitioner;
	this.distribution = distribution;
}
 

@Test
public void testCustomPartitioningErased() {

	SingleInputSemanticProperties sprops = new SingleInputSemanticProperties();
	SemanticPropUtil.getSemanticPropsSingleFromString(sprops, new String[]{"0;1"}, null, null, tupleInfo, tupleInfo);

	GlobalProperties gprops = new GlobalProperties();
	Partitioner<Tuple2<Long, Integer>> myP = new MockPartitioner();
	gprops.setCustomPartitioned(new FieldList(0, 4), myP);

	GlobalProperties result = gprops.filterBySemanticProperties(sprops, 0);

	assertEquals(PartitioningProperty.RANDOM_PARTITIONED, result.getPartitioning());
	assertNull(result.getPartitioningFields());
	assertNull(result.getCustomPartitioner());
}
 

private List<OperatorDescriptorDual> initializeDataProperties(Partitioner<?> customPartitioner) {
	Ordering groupOrder1 = null;
	Ordering groupOrder2 = null;
	
	CoGroupOperatorBase<?, ?, ?, ?> cgc = getOperator();
	groupOrder1 = cgc.getGroupOrderForInputOne();
	groupOrder2 = cgc.getGroupOrderForInputTwo();
		
	if (groupOrder1 != null && groupOrder1.getNumberOfFields() == 0) {
		groupOrder1 = null;
	}
	if (groupOrder2 != null && groupOrder2.getNumberOfFields() == 0) {
		groupOrder2 = null;
	}
	
	CoGroupDescriptor descr = new CoGroupDescriptor(this.keys1, this.keys2, groupOrder1, groupOrder2);
	if (customPartitioner != null) {
		descr.setCustomPartitioner(customPartitioner);
	}
	
	return Collections.<OperatorDescriptorDual>singletonList(descr);
}
 

public GroupReduceProperties(FieldSet groupKeys, Ordering additionalOrderKeys, Partitioner<?> customPartitioner) {
	super(groupKeys);
	
	// if we have an additional ordering, construct the ordering to have primarily the grouping fields
	if (additionalOrderKeys != null) {
		this.ordering = new Ordering();
		for (Integer key : this.keyList) {
			this.ordering.appendOrdering(key, null, Order.ANY);
		}
	
		// and next the additional order fields
		for (int i = 0; i < additionalOrderKeys.getNumberOfFields(); i++) {
			Integer field = additionalOrderKeys.getFieldNumber(i);
			Order order = additionalOrderKeys.getOrder(i);
			this.ordering.appendOrdering(field, additionalOrderKeys.getType(i), order);
		}
	}
	else {
		this.ordering = null;
	}
	
	this.customPartitioner = customPartitioner;
}
 

@PublicEvolving
public static <T> TypeInformation<T> getPartitionerTypes(
	Partitioner<T> partitioner,
	String functionName,
	boolean allowMissing) {

	return getUnaryOperatorReturnType(
		partitioner,
		Partitioner.class,
		-1,
		0,
		new int[]{0},
		null,
		functionName,
		allowMissing);
}
 

@SuppressWarnings("unchecked")
private static <T, K> org.apache.flink.api.common.operators.SingleInputOperator<?, T, ?> translateSelectorFunctionPartitioner(
	SelectorFunctionKeys<T, ?> rawKeys,
	PartitionMethod pMethod,
	String name,
	Operator<T> input,
	int partitionDop,
	Partitioner<?> customPartitioner,
	Order[] orders) {
	final SelectorFunctionKeys<T, K> keys = (SelectorFunctionKeys<T, K>) rawKeys;
	TypeInformation<Tuple2<K, T>> typeInfoWithKey = KeyFunctions.createTypeWithKey(keys);

	Operator<Tuple2<K, T>> keyedInput = KeyFunctions.appendKeyExtractor(input, keys);

	PartitionOperatorBase<Tuple2<K, T>> keyedPartitionedInput =
		new PartitionOperatorBase<>(new UnaryOperatorInformation<>(typeInfoWithKey, typeInfoWithKey), pMethod, new int[]{0}, name);
	keyedPartitionedInput.setInput(keyedInput);
	keyedPartitionedInput.setCustomPartitioner(customPartitioner);
	keyedPartitionedInput.setParallelism(partitionDop);
	keyedPartitionedInput.setOrdering(new Ordering(0, null, orders != null ? orders[0] : Order.ASCENDING));

	return KeyFunctions.appendKeyRemover(keyedPartitionedInput, keys);
}
 

private List<OperatorDescriptorSingle> initPossibleProperties(Partitioner<?> customPartitioner) {
	// see if an internal hint dictates the strategy to use
	final Configuration conf = getOperator().getParameters();
	final String localStrategy = conf.getString(Optimizer.HINT_LOCAL_STRATEGY, null);

	final boolean useCombiner;
	if (localStrategy != null) {
		if (Optimizer.HINT_LOCAL_STRATEGY_SORT.equals(localStrategy)) {
			useCombiner = false;
		}
		else if (Optimizer.HINT_LOCAL_STRATEGY_COMBINING_SORT.equals(localStrategy)) {
			if (!isCombineable()) {
				Optimizer.LOG.warn("Strategy hint for GroupReduce '" + getOperator().getName() +
					"' requires combinable reduce, but user function is not marked combinable.");
			}
			useCombiner = true;
		} else {
			throw new CompilerException("Invalid local strategy hint for match contract: " + localStrategy);
		}
	} else {
		useCombiner = isCombineable();
	}
	
	// check if we can work with a grouping (simple reducer), or if we need ordering because of a group order
	Ordering groupOrder = null;
	if (getOperator() != null) {
		groupOrder = getOperator().getGroupOrder();
		if (groupOrder != null && groupOrder.getNumberOfFields() == 0) {
			groupOrder = null;
		}
	}
	
	OperatorDescriptorSingle props = useCombiner ?
		(this.keys == null ? new AllGroupWithPartialPreGroupProperties() : new GroupReduceWithCombineProperties(this.keys, groupOrder, customPartitioner)) :
		(this.keys == null ? new AllGroupReduceProperties() : new GroupReduceProperties(this.keys, groupOrder, customPartitioner));

	return Collections.singletonList(props);
}
 

public void setCustomPartitioned(FieldList partitionedFields, Partitioner<?> partitioner) {
	if (partitionedFields == null || partitioner == null) {
		throw new NullPointerException();
	}
	
	this.partitioning = PartitioningProperty.CUSTOM_PARTITIONING;
	this.partitioningFields = partitionedFields;
	this.ordering = null;
	this.customPartitioner = partitioner;
}
 

public void setOutputPartitioner(Partitioner<?> partitioner, int outputNum) {
	try {
		InstantiationUtil.writeObjectToConfig(partitioner, config, OUTPUT_PARTITIONER + outputNum);
	}
	catch (Throwable t) {
		throw new RuntimeException("Could not serialize custom partitioner.", t);
	}
}
 

@Test
public void testCoGroupWithKeySelectors() {
	try {
		final Partitioner<Integer> partitioner = new TestPartitionerInt();
		
		ExecutionEnvironment env = ExecutionEnvironment.getExecutionEnvironment();
		
		DataSet<Pojo2> input1 = env.fromElements(new Pojo2());
		DataSet<Pojo3> input2 = env.fromElements(new Pojo3());
		
		input1
			.coGroup(input2)
			.where(new Pojo2KeySelector()).equalTo(new Pojo3KeySelector())
			.withPartitioner(partitioner)
			.with(new DummyCoGroupFunction<Pojo2, Pojo3>())
			.output(new DiscardingOutputFormat<Tuple2<Pojo2, Pojo3>>());
		
		Plan p = env.createProgramPlan();
		OptimizedPlan op = compileNoStats(p);
		
		SinkPlanNode sink = op.getDataSinks().iterator().next();
		DualInputPlanNode join = (DualInputPlanNode) sink.getInput().getSource();
		
		assertEquals(ShipStrategyType.PARTITION_CUSTOM, join.getInput1().getShipStrategy());
		assertEquals(ShipStrategyType.PARTITION_CUSTOM, join.getInput2().getShipStrategy());
		assertEquals(partitioner, join.getInput1().getPartitioner());
		assertEquals(partitioner, join.getInput2().getPartitioner());
	}
	catch (Exception e) {
		e.printStackTrace();
		fail(e.getMessage());
	}
}
 

public Partitioner<?> getOutputPartitioner(int outputNum, final ClassLoader cl) throws ClassNotFoundException {
	try {
		return (Partitioner<?>) InstantiationUtil.readObjectFromConfig(config, OUTPUT_PARTITIONER + outputNum, cl);
	}
	catch (ClassNotFoundException e) {
		throw e;
	}
	catch (Throwable t) {
		throw new RuntimeException("Could not deserialize custom partitioner.", t);
	}
}
 

public PartitionDescriptor(PartitionMethod pMethod, FieldSet pKeys, Ordering ordering, Partitioner<?>
		customPartitioner, DataDistribution distribution) {
	super(pKeys);

	Preconditions.checkArgument(pMethod != PartitionMethod.RANGE
			|| pKeys.equals(new FieldSet(ordering.getFieldPositions())),
			"Partition keys must match the given ordering.");

	this.pMethod = pMethod;
	this.customPartitioner = customPartitioner;
	this.distribution = distribution;
	this.ordering = ordering;
}
 

private List<OperatorDescriptorDual> getDataProperties() {
	OuterJoinOperatorBase<?, ?, ?, ?> operator = getOperator();

	OuterJoinType type = operator.getOuterJoinType();

	JoinHint joinHint = operator.getJoinHint();
	joinHint = joinHint == null ? JoinHint.OPTIMIZER_CHOOSES : joinHint;

	List<OperatorDescriptorDual> list;
	switch (type) {
		case LEFT:
			list = createLeftOuterJoinDescriptors(joinHint);
			break;
		case RIGHT:
			list = createRightOuterJoinDescriptors(joinHint);
			break;
		case FULL:
			list = createFullOuterJoinDescriptors(joinHint);
			break;
		default:
			throw new CompilerException("Unknown outer join type: " + type);
	}

	Partitioner<?> customPartitioner = operator.getCustomPartitioner();
	if (customPartitioner != null) {
		for (OperatorDescriptorDual desc : list) {
			((AbstractJoinDescriptor) desc).setCustomPartitioner(customPartitioner);
		}
	}
	return list;
}
 

public void setShipStrategy(ShipStrategyType strategy, FieldList keys,
							boolean[] sortDirection, Partitioner<?> partitioner,
							DataExchangeMode dataExchangeMode) {
	this.shipStrategy = strategy;
	this.shipKeys = keys;
	this.shipSortOrder = sortDirection;
	this.partitioner = partitioner;
	this.dataExchangeMode = dataExchangeMode;
	this.globalProps = null;		// reset the global properties
}
 

@Test
public void testPartitionCustomOperatorPreservesFields() {
	try {
		ExecutionEnvironment env = ExecutionEnvironment.getExecutionEnvironment();
		
		DataSet<Tuple2<Long, Long>> data = env.fromCollection(Collections.singleton(new Tuple2<>(0L, 0L)));
		
		data.partitionCustom(new Partitioner<Long>() {
				public int partition(Long key, int numPartitions) { return key.intValue(); }
			}, 1)
			.groupBy(1)
			.reduceGroup(new IdentityGroupReducerCombinable<Tuple2<Long, Long>>())
			.output(new DiscardingOutputFormat<Tuple2<Long, Long>>());
		
		Plan p = env.createProgramPlan();
		OptimizedPlan op = compileNoStats(p);
		
		SinkPlanNode sink = op.getDataSinks().iterator().next();
		SingleInputPlanNode reducer = (SingleInputPlanNode) sink.getInput().getSource();
		SingleInputPlanNode partitioner = (SingleInputPlanNode) reducer.getInput().getSource();

		assertEquals(ShipStrategyType.FORWARD, reducer.getInput().getShipStrategy());
		assertEquals(ShipStrategyType.PARTITION_CUSTOM, partitioner.getInput().getShipStrategy());
	}
	catch (Exception e) {
		e.printStackTrace();
		fail(e.getMessage());
	}
}
 

@Test
public void testCoGroupWithPojos() {
	try {
		final Partitioner<Integer> partitioner = new TestPartitionerInt();
		
		ExecutionEnvironment env = ExecutionEnvironment.getExecutionEnvironment();
		
		DataSet<Pojo2> input1 = env.fromElements(new Pojo2());
		DataSet<Pojo3> input2 = env.fromElements(new Pojo3());
		
		input1
			.coGroup(input2)
			.where("b").equalTo("a")
			.withPartitioner(partitioner)
			.with(new DummyCoGroupFunction<Pojo2, Pojo3>())
			.output(new DiscardingOutputFormat<Tuple2<Pojo2, Pojo3>>());
		
		Plan p = env.createProgramPlan();
		OptimizedPlan op = compileNoStats(p);
		
		SinkPlanNode sink = op.getDataSinks().iterator().next();
		DualInputPlanNode join = (DualInputPlanNode) sink.getInput().getSource();
		
		assertEquals(ShipStrategyType.PARTITION_CUSTOM, join.getInput1().getShipStrategy());
		assertEquals(ShipStrategyType.PARTITION_CUSTOM, join.getInput2().getShipStrategy());
		assertEquals(partitioner, join.getInput1().getPartitioner());
		assertEquals(partitioner, join.getInput2().getPartitioner());
	}
	catch (Exception e) {
		e.printStackTrace();
		fail(e.getMessage());
	}
}
 

@Test
public void testCoGroupWithKeySelectors() {
	try {
		final Partitioner<Integer> partitioner = new TestPartitionerInt();
		
		ExecutionEnvironment env = ExecutionEnvironment.getExecutionEnvironment();
		
		DataSet<Pojo2> input1 = env.fromElements(new Pojo2());
		DataSet<Pojo3> input2 = env.fromElements(new Pojo3());
		
		input1
			.coGroup(input2)
			.where(new Pojo2KeySelector()).equalTo(new Pojo3KeySelector())
			.withPartitioner(partitioner)
			.with(new DummyCoGroupFunction<Pojo2, Pojo3>())
			.output(new DiscardingOutputFormat<Tuple2<Pojo2, Pojo3>>());
		
		Plan p = env.createProgramPlan();
		OptimizedPlan op = compileNoStats(p);
		
		SinkPlanNode sink = op.getDataSinks().iterator().next();
		DualInputPlanNode join = (DualInputPlanNode) sink.getInput().getSource();
		
		assertEquals(ShipStrategyType.PARTITION_CUSTOM, join.getInput1().getShipStrategy());
		assertEquals(ShipStrategyType.PARTITION_CUSTOM, join.getInput2().getShipStrategy());
		assertEquals(partitioner, join.getInput1().getPartitioner());
		assertEquals(partitioner, join.getInput2().getPartitioner());
	}
	catch (Exception e) {
		e.printStackTrace();
		fail(e.getMessage());
	}
}
 

private List<OperatorDescriptorDual> getDataProperties() {
	OuterJoinOperatorBase<?, ?, ?, ?> operator = getOperator();

	OuterJoinType type = operator.getOuterJoinType();

	JoinHint joinHint = operator.getJoinHint();
	joinHint = joinHint == null ? JoinHint.OPTIMIZER_CHOOSES : joinHint;

	List<OperatorDescriptorDual> list;
	switch (type) {
		case LEFT:
			list = createLeftOuterJoinDescriptors(joinHint);
			break;
		case RIGHT:
			list = createRightOuterJoinDescriptors(joinHint);
			break;
		case FULL:
			list = createFullOuterJoinDescriptors(joinHint);
			break;
		default:
			throw new CompilerException("Unknown outer join type: " + type);
	}

	Partitioner<?> customPartitioner = operator.getCustomPartitioner();
	if (customPartitioner != null) {
		for (OperatorDescriptorDual desc : list) {
			((AbstractJoinDescriptor) desc).setCustomPartitioner(customPartitioner);
		}
	}
	return list;
}
 

@Test
public void testPartitionerLambda() {
	Partitioner<Tuple2<Integer, String>> partitioner = (key, numPartitions) -> key.f1.length() % numPartitions;
	final TypeInformation<?> ti = TypeExtractor.getPartitionerTypes(partitioner, null, true);

	if (!(ti instanceof MissingTypeInfo)) {
		assertTrue(ti.isTupleType());
		assertEquals(2, ti.getArity());
		assertEquals(((TupleTypeInfo<?>) ti).getTypeAt(0), BasicTypeInfo.INT_TYPE_INFO);
		assertEquals(((TupleTypeInfo<?>) ti).getTypeAt(1), BasicTypeInfo.STRING_TYPE_INFO);
	}
}
 

@Override
public <E> void validateCustomPartitioner(Partitioner<E> partitioner, TypeInformation<E> typeInfo) {

	if (keyFields.size() != 1) {
		throw new InvalidProgramException("Custom partitioners can only be used with keys that have one key field.");
	}

	if (typeInfo == null) {
		// try to extract key type from partitioner
		try {
			typeInfo = TypeExtractor.getPartitionerTypes(partitioner);
		}
		catch (Throwable t) {
			// best effort check, so we ignore exceptions
		}
	}

	if (typeInfo != null && !(typeInfo instanceof GenericTypeInfo)) {
		// only check type compatibility if type is known and not a generic type

		TypeInformation<?> keyType = keyFields.get(0).getType();
		if (!keyType.equals(typeInfo)) {
			throw new InvalidProgramException("The partitioner is incompatible with the key type. "
								+ "Partitioner type: " + typeInfo + " , key type: " + keyType);
		}
	}
}
 

private <K> DataStream<T> partitionCustom(Partitioner<K> partitioner, Keys<T> keys) {
	KeySelector<T, K> keySelector = KeySelectorUtil.getSelectorForOneKey(keys, partitioner, getType(), getExecutionConfig());

	return setConnectionType(
			new CustomPartitionerWrapper<>(
					clean(partitioner),
					clean(keySelector)));
}
 

private <K> DataStream<T> partitionCustom(Partitioner<K> partitioner, Keys<T> keys) {
	KeySelector<T, K> keySelector = KeySelectorUtil.getSelectorForOneKey(keys, partitioner, getType(), getExecutionConfig());

	return setConnectionType(
			new CustomPartitionerWrapper<>(
					clean(partitioner),
					clean(keySelector)));
}
 

@Test
public void testPartitionCustomOperatorPreservesFields() {
	try {
		ExecutionEnvironment env = ExecutionEnvironment.getExecutionEnvironment();
		
		DataSet<Tuple2<Long, Long>> data = env.fromCollection(Collections.singleton(new Tuple2<>(0L, 0L)));
		
		data.partitionCustom(new Partitioner<Long>() {
				public int partition(Long key, int numPartitions) { return key.intValue(); }
			}, 1)
			.groupBy(1)
			.reduceGroup(new IdentityGroupReducerCombinable<Tuple2<Long, Long>>())
			.output(new DiscardingOutputFormat<Tuple2<Long, Long>>());
		
		Plan p = env.createProgramPlan();
		OptimizedPlan op = compileNoStats(p);
		
		SinkPlanNode sink = op.getDataSinks().iterator().next();
		SingleInputPlanNode reducer = (SingleInputPlanNode) sink.getInput().getSource();
		SingleInputPlanNode partitioner = (SingleInputPlanNode) reducer.getInput().getSource();

		assertEquals(ShipStrategyType.FORWARD, reducer.getInput().getShipStrategy());
		assertEquals(ShipStrategyType.PARTITION_CUSTOM, partitioner.getInput().getShipStrategy());
	}
	catch (Exception e) {
		e.printStackTrace();
		fail(e.getMessage());
	}
}
 

/**
 * Uses a custom partitioner for the grouping.
 *
 * @param partitioner The custom partitioner.
 * @return The grouping object itself, to allow for method chaining.
 */
public UnsortedGrouping<T> withPartitioner(Partitioner<?> partitioner) {
	Preconditions.checkNotNull(partitioner);
	getKeys().validateCustomPartitioner(partitioner, null);

	this.customPartitioner = partitioner;
	return this;
}