darshanj commented on a change in pull request #11682:
URL: https://github.com/apache/beam/pull/11682#discussion_r428399671
##########
File path:
sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/Partition.java
##########
@@ -76,6 +93,45 @@
int partitionFor(T elem, int numPartitions);
}
+ /**
+ * A function object that chooses an output partition for an element.
+ *
+ * @param <T> the type of the elements being partitioned
+ */
+ public interface PartitionWithSideInputsFn<T> extends Serializable {
+ /**
+ * Chooses the partition into which to put the given element.
+ *
+ * @param elem the element to be partitioned
+ * @param numPartitions the total number of partitions ({@code >= 1})
+ * @param c the {@link Contextful.Fn.Context} needed to access sideInputs.
+ * @return index of the selected partition (in the range {@code
[0..numPartitions-1]})
+ */
+ int partitionFor(T elem, int numPartitions, Contextful.Fn.Context c);
+ }
+
+ /**
+ * Returns a new {@code Partition} {@code PTransform} that divides its input
{@code PCollection}
+ * into the given number of partitions, using the given partitioning
function.
+ *
+ * @param numPartitions the number of partitions to divide the input {@code
PCollection} into
+ * @param partitionFn the function to invoke on each element to choose its
output partition
+ * @param requirements the {@link Requirements} needed to run it.
+ * @throws IllegalArgumentException if {@code numPartitions <= 0}
+ */
+ public static <T> Partition<T> of(
+ int numPartitions,
+ PartitionWithSideInputsFn<? super T> partitionFn,
+ Requirements requirements) {
+ Contextful ctfFn =
+ Contextful.fn(
+ (T element, Contextful.Fn.Context c) ->
+ partitionFn.partitionFor(element, numPartitions, c),
+ requirements);
+ Object aClass = partitionFn;
Review comment:
Hi, I don't get your suggestion here. I will need to wrap interface
function in ContextFul.Fn. Can you eloborate please?
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
us...@infra.apache.org
