Source:
protected <T> Iterable<List<T>> partitionByCountAndBytes(final Iterable<T> iterable, final int maximumCountPerPartition, final long maximumBytesPerPartition, final String tableName, final Function<T, Long> sizingFunction) { return new Iterable<List<T>>() { @Override public Iterator<List<T>> iterator() { return new UnmodifiableIterator<List<T>>() { PeekingIterator<T> pi = Iterators.peekingIterator(iterable.iterator()); private int remainingEntries = Iterables.size(iterable); @Override public boolean hasNext() { return pi.hasNext(); } @Override public List<T> next() { if (!hasNext()) { throw new NoSuchElementException(); } List<T> entries = Lists.newArrayListWithCapacity(Math.min(maximumCountPerPartition, remainingEntries)); long runningSize = 0; T firstEntry = pi.next(); runningSize += sizingFunction.apply(firstEntry); entries.add(firstEntry); if (runningSize > maximumBytesPerPartition && log.isWarnEnabled()) { final String message = "Encountered an entry of approximate size bytes, " + "larger than maximum size of defined per entire batch, " + "while doing a write to {}. Attempting to batch anyways."; if (AtlasDbConstants.TABLES_KNOWN_TO_BE_POORLY_DESIGNED.contains( TableReference.createWithEmptyNamespace(tableName))) { log.warn(message, sizingFunction.apply(firstEntry), maximumBytesPerPartition, tableName); } else { final String longerMessage = message + " This can potentially cause out-of-memory errors."; log.warn(longerMessage, sizingFunction.apply(firstEntry), maximumBytesPerPartition, tableName); } } while (pi.hasNext() && entries.size() < maximumCountPerPartition) { runningSize += sizingFunction.apply(pi.peek()); if (runningSize > maximumBytesPerPartition) { break; } entries.add(pi.next()); } remainingEntries -= entries.size(); return entries; } }; } }; }

Target:
protected <T> Iterable<List<T>> partitionByCountAndBytes(final Iterable<T> iterable, final int maximumCountPerPartition, final long maximumBytesPerPartition, final String tableName, final Function<T, Long> sizingFunction) { return new Iterable<List<T>>() { @Override public Iterator<List<T>> iterator() { return new UnmodifiableIterator<List<T>>() { PeekingIterator<T> pi = Iterators.peekingIterator(iterable.iterator()); private int remainingEntries = Iterables.size(iterable); @Override public boolean hasNext() { return pi.hasNext(); } @Override public List<T> next() { if (!hasNext()) { throw new NoSuchElementException(); } List<T> entries = Lists.newArrayListWithCapacity(Math.min(maximumCountPerPartition, remainingEntries)); long runningSize = 0; T firstEntry = pi.next(); runningSize += sizingFunction.apply(firstEntry); entries.add(firstEntry); if (runningSize > maximumBytesPerPartition && log.isWarnEnabled()) { if (AtlasDbConstants.TABLES_KNOWN_TO_BE_POORLY_DESIGNED.contains( TableReference.createWithEmptyNamespace(tableName))) { log.warn(ENTRY_TOO_BIG_MESSAGE, sizingFunction.apply(firstEntry), maximumBytesPerPartition, tableName); } else { final String longerMessage = ENTRY_TOO_BIG_MESSAGE + " This can potentially cause out-of-memory errors."; log.warn(longerMessage, sizingFunction.apply(firstEntry), maximumBytesPerPartition, tableName); } } while (pi.hasNext() && entries.size() < maximumCountPerPartition) { runningSize += sizingFunction.apply(pi.peek()); if (runningSize > maximumBytesPerPartition) { break; } entries.add(pi.next()); } remainingEntries -= entries.size(); return entries; } }; } }; }