I am trying to implement a ListFlatten function, I have implemented it using SimpleDoFn which is working fine but for parallelizing. I am converting the function to Splittable Do Function. I managed to get a unit test running in local with 5000 elements using DirectRunner while running the same in DataFlow, it is failing with below error.
Error Details:
java.lang.RuntimeException: org.apache.beam.sdk.util.UserCodeException: java.lang.RuntimeException: java.io.IOException: INVALID_ARGUMENT: Shuffle key too large:3749653 > 1572864
at org.apache.beam.runners.dataflow.worker.GroupAlsoByWindowsParDoFn$1.output (GroupAlsoByWindowsParDoFn.java:184)
at org.apache.beam.runners.dataflow.worker.GroupAlsoByWindowFnRunner$1.outputWindowedValue (GroupAlsoByWindowFnRunner.java:102)
at org.apache.beam.runners.dataflow.worker.util.BatchGroupAlsoByWindowViaIteratorsFn.processElement (BatchGroupAlsoByWindowViaIteratorsFn.java:126)
at org.apache.beam.runners.dataflow.worker.util.BatchGroupAlsoByWindowViaIteratorsFn.processElement (BatchGroupAlsoByWindowViaIteratorsFn.java:54)
at org.apache.beam.runners.dataflow.worker.GroupAlsoByWindowFnRunner.invokeProcessElement (GroupAlsoByWindowFnRunner.java:115)
at org.apache.beam.runners.dataflow.worker.GroupAlsoByWindowFnRunner.processElement (GroupAlsoByWindowFnRunner.java:73)
at org.apache.beam.runners.dataflow.worker.GroupAlsoByWindowsParDoFn.processElement (GroupAlsoByWindowsParDoFn.java:114)
at org.apache.beam.runners.dataflow.worker.util.common.worker.ParDoOperation.process (ParDoOperation.java:44)
at org.apache.beam.runners.dataflow.worker.util.common.worker.OutputReceiver.process (OutputReceiver.java:49)
at org.apache.beam.runners.dataflow.worker.util.common.worker.ReadOperation.runReadLoop (ReadOperation.java:201)
Caused by: org.apache.beam.sdk.util.UserCodeException: java.lang.RuntimeException: java.io.IOException: INVALID_ARGUMENT: Shuffle key too large:3749653 > 1572864
at com.abc.common.batch.functions.AbcListFlattenFn.splitRestriction (AbcListFlattenFn.java:68)
The differences in data between local DirectRunner and Cloud DataFlow runner are given below.
DirectRunner in local:
- It has 5000 abcs in the sample input PCollection element
DataflowRunner in cloud:
- There are varying sizes of abcs in 600 input PCollection elements
- Few input elements have 50000 abcs to flatten
public class AbcList implements Serializable {
private List<Abc> abcs;
private List<Xyz> xyzs;
}
public class AbcListFlattenFn extends DoFn<AbcList, KV<Abc, List<Xyz>> {
@ProcessElement
public void process(@Element AbcList input,
ProcessContext context, RestrictionTracker<OffsetRange, Long> tracker) {
try {
/* Below commented lines are without the Splittable DoFn
input.getAbcs().stream().forEach(abc -> {
context.output(KV.of(abc, input.getXyzs()));
}); */
for (long index = tracker.currentRestriction().getFrom(); tracker.tryClaim(index);
++index) {
context.output(KV.of(input.getAbcs().get(Math.toIntExact(index),input.getXyzs())));
}
} catch (Exception e) {
log.error("Flattening AbcList has failed ", e);
}
}
@GetInitialRestriction
public OffsetRange getInitialRestriction(AbcList input) {
return new OffsetRange(0, input.getAbcs().size());
}
@SplitRestriction
public void splitRestriction(final AbcList input,
final OffsetRange range, final OutputReceiver<OffsetRange> receiver) {
List<OffsetRange> ranges =
range.split(input.getAbcs().size() > 5000 ? 5000
: input.getAbcs().size(), 2000);
for (final OffsetRange p : ranges) {
receiver.output(p);
}
}
@NewTracker
public OffsetRangeTracker newTracker(OffsetRange range) {
return new OffsetRangeTracker(range);
}
}
Can someone suggest what is wrong with the ListFlatten function here? is splitRestriction causing the below issue? How can I fix this Shuffle key size problem?
