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

@SuppressWarnings("unchecked")
private static <IN, OUT, K1, K2> PlanUnwrappingSortedGroupCombineOperator<IN, OUT, K1, K2> translateSelectorFunctionSortedReducer(
		SelectorFunctionKeys<IN, ?> rawGroupingKey,
		SelectorFunctionKeys<IN, ?> rawSortingKeys,
		Ordering groupOrder,
		GroupCombineFunction<IN, OUT> function,
		TypeInformation<OUT> outputType,
		String name,
		Operator<IN> input) {
	final SelectorFunctionKeys<IN, K1> groupingKey = (SelectorFunctionKeys<IN, K1>) rawGroupingKey;
	final SelectorFunctionKeys<IN, K2> sortingKey = (SelectorFunctionKeys<IN, K2>) rawSortingKeys;
	TypeInformation<Tuple3<K1, K2, IN>> typeInfoWithKey = KeyFunctions.createTypeWithKey(groupingKey, sortingKey);

	Operator<Tuple3<K1, K2, IN>> inputWithKey = KeyFunctions.appendKeyExtractor(input, groupingKey, sortingKey);

	PlanUnwrappingSortedGroupCombineOperator<IN, OUT, K1, K2> reducer =
		new PlanUnwrappingSortedGroupCombineOperator<>(function, groupingKey, sortingKey, name, outputType, typeInfoWithKey);
	reducer.setInput(inputWithKey);
	reducer.setGroupOrder(groupOrder);

	return reducer;
}
 

@SuppressWarnings("unchecked")
private static <IN, OUT, K1, K2> PlanUnwrappingSortedGroupCombineOperator<IN, OUT, K1, K2> translateSelectorFunctionSortedReducer(
		SelectorFunctionKeys<IN, ?> rawGroupingKey,
		SelectorFunctionKeys<IN, ?> rawSortingKeys,
		Ordering groupOrder,
		GroupCombineFunction<IN, OUT> function,
		TypeInformation<OUT> outputType,
		String name,
		Operator<IN> input) {
	final SelectorFunctionKeys<IN, K1> groupingKey = (SelectorFunctionKeys<IN, K1>) rawGroupingKey;
	final SelectorFunctionKeys<IN, K2> sortingKey = (SelectorFunctionKeys<IN, K2>) rawSortingKeys;
	TypeInformation<Tuple3<K1, K2, IN>> typeInfoWithKey = KeyFunctions.createTypeWithKey(groupingKey, sortingKey);

	Operator<Tuple3<K1, K2, IN>> inputWithKey = KeyFunctions.appendKeyExtractor(input, groupingKey, sortingKey);

	PlanUnwrappingSortedGroupCombineOperator<IN, OUT, K1, K2> reducer =
		new PlanUnwrappingSortedGroupCombineOperator<>(function, groupingKey, sortingKey, name, outputType, typeInfoWithKey);
	reducer.setInput(inputWithKey);
	reducer.setGroupOrder(groupOrder);

	return reducer;
}
 

@SuppressWarnings("unchecked")
private static <IN, OUT, K> PlanUnwrappingGroupCombineOperator<IN, OUT, K> translateSelectorFunctionReducer(
		SelectorFunctionKeys<IN, ?> rawKeys,
		GroupCombineFunction<IN, OUT> function,
		TypeInformation<OUT> outputType,
		String name,
		Operator<IN> input) {
	final SelectorFunctionKeys<IN, K> keys = (SelectorFunctionKeys<IN, K>) rawKeys;

	TypeInformation<Tuple2<K, IN>> typeInfoWithKey = KeyFunctions.createTypeWithKey(keys);
	Operator<Tuple2<K, IN>> keyedInput = KeyFunctions.appendKeyExtractor(input, keys);

	PlanUnwrappingGroupCombineOperator<IN, OUT, K> reducer =
		new PlanUnwrappingGroupCombineOperator<>(function, keys, name, outputType, typeInfoWithKey);
	reducer.setInput(keyedInput);

	return reducer;
}
 

@SuppressWarnings("unchecked")
private static <IN, OUT, K> PlanUnwrappingGroupCombineOperator<IN, OUT, K> translateSelectorFunctionReducer(
		SelectorFunctionKeys<IN, ?> rawKeys,
		GroupCombineFunction<IN, OUT> function,
		TypeInformation<OUT> outputType,
		String name,
		Operator<IN> input) {
	final SelectorFunctionKeys<IN, K> keys = (SelectorFunctionKeys<IN, K>) rawKeys;

	TypeInformation<Tuple2<K, IN>> typeInfoWithKey = KeyFunctions.createTypeWithKey(keys);
	Operator<Tuple2<K, IN>> keyedInput = KeyFunctions.appendKeyExtractor(input, keys);

	PlanUnwrappingGroupCombineOperator<IN, OUT, K> reducer =
		new PlanUnwrappingGroupCombineOperator<>(function, keys, name, outputType, typeInfoWithKey);
	reducer.setInput(keyedInput);

	return reducer;
}
 

@SuppressWarnings("unchecked")
private static <IN, OUT, K1, K2> PlanUnwrappingSortedGroupCombineOperator<IN, OUT, K1, K2> translateSelectorFunctionSortedReducer(
		SelectorFunctionKeys<IN, ?> rawGroupingKey,
		SelectorFunctionKeys<IN, ?> rawSortingKeys,
		Ordering groupOrder,
		GroupCombineFunction<IN, OUT> function,
		TypeInformation<OUT> outputType,
		String name,
		Operator<IN> input) {
	final SelectorFunctionKeys<IN, K1> groupingKey = (SelectorFunctionKeys<IN, K1>) rawGroupingKey;
	final SelectorFunctionKeys<IN, K2> sortingKey = (SelectorFunctionKeys<IN, K2>) rawSortingKeys;
	TypeInformation<Tuple3<K1, K2, IN>> typeInfoWithKey = KeyFunctions.createTypeWithKey(groupingKey, sortingKey);

	Operator<Tuple3<K1, K2, IN>> inputWithKey = KeyFunctions.appendKeyExtractor(input, groupingKey, sortingKey);

	PlanUnwrappingSortedGroupCombineOperator<IN, OUT, K1, K2> reducer =
		new PlanUnwrappingSortedGroupCombineOperator<>(function, groupingKey, sortingKey, name, outputType, typeInfoWithKey);
	reducer.setInput(inputWithKey);
	reducer.setGroupOrder(groupOrder);

	return reducer;
}
 

@SuppressWarnings("unchecked")
@Override
public void combine(Iterable<Tuple2<K, IN>> values, Collector<Tuple2<K, IN>> out) throws Exception {

	iter.set(values.iterator());
	coll.set(out);
	((GroupCombineFunction<IN, IN>) this.wrappedFunction).combine(iter, coll);
}
 

@Override
public void setup(AbstractInvokable parent) {
	this.parent = parent;

	@SuppressWarnings("unchecked")
	final GroupCombineFunction<IN, OUT> combiner =
		BatchTask.instantiateUserCode(this.config, userCodeClassLoader, GroupCombineFunction.class);
	this.combiner = combiner;
	FunctionUtils.setFunctionRuntimeContext(combiner, getUdfRuntimeContext());
}
 

private TupleUnwrappingGroupCombinableGroupReducer(GroupReduceFunction<IN, OUT> wrapped) {
	super(wrapped);

	if (!GroupCombineFunction.class.isAssignableFrom(wrappedFunction.getClass())) {
		throw new IllegalArgumentException("Wrapped reduce function does not implement the GroupCombineFunction interface.");
	}

	this.iter = new TupleUnwrappingIterator<>();
	this.coll = new TupleWrappingCollector<>(this.iter);
}
 

/**
 * Constructor for a grouped reduce.
 *
 * @param input The grouped input to be processed group-wise by the groupReduce function.
 * @param function The user-defined GroupReduce function.
 */
public GroupCombineOperator(Grouping<IN> input, TypeInformation<OUT> resultType, GroupCombineFunction<IN, OUT> function, String defaultName) {
	super(input != null ? input.getInputDataSet() : null, resultType);

	this.function = function;
	this.grouper = input;
	this.defaultName = defaultName;
}
 

@Test
// check if parallelism of 1 results in the same data like a shuffle
public void testCheckPartitionShuffleDOP1() throws Exception {

	final ExecutionEnvironment env = ExecutionEnvironment.getExecutionEnvironment();

	env.setParallelism(1);

	// data
	DataSet<Tuple3<Integer, Long, String>> ds = CollectionDataSets.get3TupleDataSet(env);

	// partition and group data
	UnsortedGrouping<Tuple3<Integer, Long, String>> partitionedDS = ds.partitionByHash(0).groupBy(1);

	List<Tuple2<Long, Integer>> result = partitionedDS
			.combineGroup(
			new GroupCombineFunction<Tuple3<Integer, Long, String>, Tuple2<Long, Integer>>() {
				@Override
				public void combine(Iterable<Tuple3<Integer, Long, String>> values, Collector<Tuple2<Long, Integer>> out) throws Exception {
					int count = 0;
					long key = 0;
					for (Tuple3<Integer, Long, String> value : values) {
						key = value.f1;
						count++;
					}
					out.collect(new Tuple2<>(key, count));
				}
			}).collect();

	String expected = "6,6\n" +
			"5,5\n" +
			"4,4\n" +
			"3,3\n" +
			"2,2\n" +
			"1,1\n";

	compareResultAsTuples(result, expected);
}
 

@Test
// check if no shuffle is being executed
public void testCheckPartitionShuffleGroupBy() throws Exception {

	org.junit.Assume.assumeTrue(mode != TestExecutionMode.COLLECTION);

	final ExecutionEnvironment env = ExecutionEnvironment.getExecutionEnvironment();

	// data
	DataSet<Tuple3<Integer, Long, String>> ds = CollectionDataSets.get3TupleDataSet(env);

	// partition and group data
	UnsortedGrouping<Tuple3<Integer, Long, String>> partitionedDS = ds.partitionByHash(0).groupBy(1);

	List<Tuple2<Long, Integer>> result = partitionedDS
			.combineGroup(
					new GroupCombineFunction<Tuple3<Integer, Long, String>, Tuple2<Long, Integer>>() {
		@Override
		public void combine(Iterable<Tuple3<Integer, Long, String>> values, Collector<Tuple2<Long, Integer>> out) throws Exception {
			int count = 0;
			long key = 0;
			for (Tuple3<Integer, Long, String> value : values) {
				key = value.f1;
				count++;
			}
			out.collect(new Tuple2<>(key, count));
		}
	}).collect();

	String[] localExpected = new String[] { "(6,6)", "(5,5)" + "(4,4)", "(3,3)", "(2,2)", "(1,1)" };

	String[] resultAsStringArray = new String[result.size()];
	for (int i = 0; i < resultAsStringArray.length; ++i) {
		resultAsStringArray[i] = result.get(i).toString();
	}
	Arrays.sort(resultAsStringArray);

	Assert.assertEquals("The two arrays were identical.", false, Arrays.equals(localExpected, resultAsStringArray));
}
 

@Override
public void setup(AbstractInvokable parent) {
	this.parent = parent;

	@SuppressWarnings("unchecked")
	final GroupCombineFunction<IN, OUT> combiner =
		BatchTask.instantiateUserCode(this.config, userCodeClassLoader, GroupCombineFunction.class);
	this.combiner = combiner;
	FunctionUtils.setFunctionRuntimeContext(combiner, getUdfRuntimeContext());
}
 

private TupleUnwrappingGroupCombinableGroupReducer(GroupReduceFunction<IN, OUT> wrapped) {
	super(wrapped);

	if (!GroupCombineFunction.class.isAssignableFrom(wrappedFunction.getClass())) {
		throw new IllegalArgumentException("Wrapped reduce function does not implement the GroupCombineFunction interface.");
	}

	this.iter = new TupleUnwrappingIterator<>();
	this.coll = new TupleWrappingCollector<>(this.iter);
}
 

/**
 * Constructor for a grouped reduce.
 *
 * @param input The grouped input to be processed group-wise by the groupReduce function.
 * @param function The user-defined GroupReduce function.
 */
public GroupCombineOperator(Grouping<IN> input, TypeInformation<OUT> resultType, GroupCombineFunction<IN, OUT> function, String defaultName) {
	super(input != null ? input.getInputDataSet() : null, resultType);

	this.function = function;
	this.grouper = input;
	this.defaultName = defaultName;
}
 

public PlanUnwrappingSortedGroupCombineOperator(GroupCombineFunction<IN, OUT> udf, Keys.SelectorFunctionKeys<IN, K1> groupingKey, Keys.SelectorFunctionKeys<IN, K2> sortingKey, String name,
												TypeInformation<OUT> outType, TypeInformation<Tuple3<K1, K2, IN>> typeInfoWithKey) {
	super(new TupleUnwrappingGroupReducer<IN, OUT, K1, K2>(udf),
			new UnaryOperatorInformation<Tuple3<K1, K2, IN>, OUT>(typeInfoWithKey, outType),
			groupingKey.computeLogicalKeyPositions(),
			name);

}
 

/**
 * Marks the group reduce operation as combinable. Combinable operations may pre-reduce the
 * data before the actual group reduce operations. Combinable user-defined functions
 * must implement the interface {@link GroupCombineFunction}.
 * 
 * @param combinable Flag to mark the group reduce operation as combinable.
 */
public void setCombinable(boolean combinable) {
	// sanity check
	if (combinable && !GroupCombineFunction.class.isAssignableFrom(this.userFunction.getUserCodeClass())) {
		throw new IllegalArgumentException("Cannot set a UDF as combinable if it does not implement the interface " +
				GroupCombineFunction.class.getName());
	} else {
		this.combinable = combinable;
	}
}
 

@PublicEvolving
public static <IN, OUT> TypeInformation<OUT> getGroupCombineReturnTypes(GroupCombineFunction<IN, OUT> combineInterface, TypeInformation<IN> inType,
																		String functionName, boolean allowMissing)
{
	return getUnaryOperatorReturnType(
		(Function) combineInterface,
		GroupCombineFunction.class,
		0,
		1,
		new int[]{1, 0},
		inType,
		functionName,
		allowMissing);
}
 

private TupleUnwrappingGroupCombinableGroupReducer(GroupReduceFunction<IN, OUT> wrapped) {
	super(wrapped);

	if (!GroupCombineFunction.class.isAssignableFrom(wrappedFunction.getClass())) {
		throw new IllegalArgumentException("Wrapped reduce function does not implement the GroupCombineFunction interface.");
	}

	this.iter = new Tuple3UnwrappingIterator<>();
	this.coll = new Tuple3WrappingCollector<>(this.iter);
}
 

@SuppressWarnings("unchecked")
@Override
public void combine(Iterable<Tuple3<K1, K2, IN>> values, Collector<Tuple3<K1, K2, IN>> out) throws Exception {
	iter.set(values.iterator());
	coll.set(out);
	((GroupCombineFunction<IN, IN>) this.wrappedFunction).combine(iter, coll);
}
 

@Test
// check if parallelism of 1 results in the same data like a shuffle
public void testCheckPartitionShuffleDOP1() throws Exception {

	final ExecutionEnvironment env = ExecutionEnvironment.getExecutionEnvironment();

	env.setParallelism(1);

	// data
	DataSet<Tuple3<Integer, Long, String>> ds = CollectionDataSets.get3TupleDataSet(env);

	// partition and group data
	UnsortedGrouping<Tuple3<Integer, Long, String>> partitionedDS = ds.partitionByHash(0).groupBy(1);

	List<Tuple2<Long, Integer>> result = partitionedDS
			.combineGroup(
			new GroupCombineFunction<Tuple3<Integer, Long, String>, Tuple2<Long, Integer>>() {
				@Override
				public void combine(Iterable<Tuple3<Integer, Long, String>> values, Collector<Tuple2<Long, Integer>> out) throws Exception {
					int count = 0;
					long key = 0;
					for (Tuple3<Integer, Long, String> value : values) {
						key = value.f1;
						count++;
					}
					out.collect(new Tuple2<>(key, count));
				}
			}).collect();

	String expected = "6,6\n" +
			"5,5\n" +
			"4,4\n" +
			"3,3\n" +
			"2,2\n" +
			"1,1\n";

	compareResultAsTuples(result, expected);
}
 

/**
 * Marks the group reduce operation as combinable. Combinable operations may pre-reduce the
 * data before the actual group reduce operations. Combinable user-defined functions
 * must implement the interface {@link GroupCombineFunction}.
 * 
 * @param combinable Flag to mark the group reduce operation as combinable.
 */
public void setCombinable(boolean combinable) {
	// sanity check
	if (combinable && !GroupCombineFunction.class.isAssignableFrom(this.userFunction.getUserCodeClass())) {
		throw new IllegalArgumentException("Cannot set a UDF as combinable if it does not implement the interface " +
				GroupCombineFunction.class.getName());
	} else {
		this.combinable = combinable;
	}
}
 

@Override
public void prepare() throws Exception {
	final TaskConfig config = this.taskContext.getTaskConfig();
	this.strategy = config.getDriverStrategy();

	switch (this.strategy) {
		case ALL_GROUP_REDUCE_COMBINE:
			if (!(this.taskContext.getStub() instanceof GroupCombineFunction)) {
				throw new Exception("Using combiner on a UDF that does not implement the combiner interface " + GroupCombineFunction.class.getName());
			}
		case ALL_GROUP_REDUCE:
		case ALL_GROUP_COMBINE:
			break;
		default:
			throw new Exception("Unrecognized driver strategy for AllGroupReduce driver: " + this.strategy.name());
	}

	this.serializer = this.taskContext.<IT>getInputSerializer(0).getSerializer();
	this.input = this.taskContext.getInput(0);

	ExecutionConfig executionConfig = taskContext.getExecutionConfig();
	this.objectReuseEnabled = executionConfig.isObjectReuseEnabled();

	if (LOG.isDebugEnabled()) {
		LOG.debug("AllGroupReduceDriver object reuse: " + (this.objectReuseEnabled ? "ENABLED" : "DISABLED") + ".");
	}
}
 

@Override
public void setup(AbstractInvokable parent) {
	this.parent = parent;

	@SuppressWarnings("unchecked")
	final GroupCombineFunction<IN, OUT> combiner =
		BatchTask.instantiateUserCode(this.config, userCodeClassLoader, GroupCombineFunction.class);
	this.combiner = combiner;
	FunctionUtils.setFunctionRuntimeContext(combiner, getUdfRuntimeContext());
}
 

@Test
// check if no shuffle is being executed
public void testCheckPartitionShuffleGroupBy() throws Exception {

	org.junit.Assume.assumeTrue(mode != TestExecutionMode.COLLECTION);

	final ExecutionEnvironment env = ExecutionEnvironment.getExecutionEnvironment();

	// data
	DataSet<Tuple3<Integer, Long, String>> ds = CollectionDataSets.get3TupleDataSet(env);

	// partition and group data
	UnsortedGrouping<Tuple3<Integer, Long, String>> partitionedDS = ds.partitionByHash(0).groupBy(1);

	List<Tuple2<Long, Integer>> result = partitionedDS
			.combineGroup(
					new GroupCombineFunction<Tuple3<Integer, Long, String>, Tuple2<Long, Integer>>() {
		@Override
		public void combine(Iterable<Tuple3<Integer, Long, String>> values, Collector<Tuple2<Long, Integer>> out) throws Exception {
			int count = 0;
			long key = 0;
			for (Tuple3<Integer, Long, String> value : values) {
				key = value.f1;
				count++;
			}
			out.collect(new Tuple2<>(key, count));
		}
	}).collect();

	String[] localExpected = new String[] { "(6,6)", "(5,5)" + "(4,4)", "(3,3)", "(2,2)", "(1,1)" };

	String[] resultAsStringArray = new String[result.size()];
	for (int i = 0; i < resultAsStringArray.length; ++i) {
		resultAsStringArray[i] = result.get(i).toString();
	}
	Arrays.sort(resultAsStringArray);

	Assert.assertEquals("The two arrays were identical.", false, Arrays.equals(localExpected, resultAsStringArray));
}
 

@Test
// check if parallelism of 1 results in the same data like a shuffle
public void testCheckPartitionShuffleDOP1() throws Exception {

	final ExecutionEnvironment env = ExecutionEnvironment.getExecutionEnvironment();

	env.setParallelism(1);

	// data
	DataSet<Tuple3<Integer, Long, String>> ds = CollectionDataSets.get3TupleDataSet(env);

	// partition and group data
	UnsortedGrouping<Tuple3<Integer, Long, String>> partitionedDS = ds.partitionByHash(0).groupBy(1);

	List<Tuple2<Long, Integer>> result = partitionedDS
			.combineGroup(
			new GroupCombineFunction<Tuple3<Integer, Long, String>, Tuple2<Long, Integer>>() {
				@Override
				public void combine(Iterable<Tuple3<Integer, Long, String>> values, Collector<Tuple2<Long, Integer>> out) throws Exception {
					int count = 0;
					long key = 0;
					for (Tuple3<Integer, Long, String> value : values) {
						key = value.f1;
						count++;
					}
					out.collect(new Tuple2<>(key, count));
				}
			}).collect();

	String expected = "6,6\n" +
			"5,5\n" +
			"4,4\n" +
			"3,3\n" +
			"2,2\n" +
			"1,1\n";

	compareResultAsTuples(result, expected);
}
 

/**
 * Constructor for a grouped reduce.
 *
 * @param input The grouped input to be processed group-wise by the groupReduce function.
 * @param function The user-defined GroupReduce function.
 */
public GroupCombineOperator(Grouping<IN> input, TypeInformation<OUT> resultType, GroupCombineFunction<IN, OUT> function, String defaultName) {
	super(input != null ? input.getInputDataSet() : null, resultType);

	this.function = function;
	this.grouper = input;
	this.defaultName = defaultName;
}
 

/**
 * Marks the group reduce operation as combinable. Combinable operations may pre-reduce the
 * data before the actual group reduce operations. Combinable user-defined functions
 * must implement the interface {@link GroupCombineFunction}.
 * 
 * @param combinable Flag to mark the group reduce operation as combinable.
 */
public void setCombinable(boolean combinable) {
	// sanity check
	if (combinable && !GroupCombineFunction.class.isAssignableFrom(this.userFunction.getUserCodeClass())) {
		throw new IllegalArgumentException("Cannot set a UDF as combinable if it does not implement the interface " +
				GroupCombineFunction.class.getName());
	} else {
		this.combinable = combinable;
	}
}
 

@SuppressWarnings("unchecked")
@Override
public void combine(Iterable<Tuple2<K, IN>> values, Collector<Tuple2<K, IN>> out) throws Exception {

	iter.set(values.iterator());
	coll.set(out);
	((GroupCombineFunction<IN, IN>) this.wrappedFunction).combine(iter, coll);
}
 

private TupleUnwrappingGroupCombinableGroupReducer(GroupReduceFunction<IN, OUT> wrapped) {
	super(wrapped);

	if (!GroupCombineFunction.class.isAssignableFrom(wrappedFunction.getClass())) {
		throw new IllegalArgumentException("Wrapped reduce function does not implement the GroupCombineFunction interface.");
	}

	this.iter = new Tuple3UnwrappingIterator<>();
	this.coll = new Tuple3WrappingCollector<>(this.iter);
}
 

@SuppressWarnings("unchecked")
@Override
public void combine(Iterable<Tuple3<K1, K2, IN>> values, Collector<Tuple3<K1, K2, IN>> out) throws Exception {
	iter.set(values.iterator());
	coll.set(out);
	((GroupCombineFunction<IN, IN>) this.wrappedFunction).combine(iter, coll);
}