tysonjh commented on a change in pull request #13059: URL: https://github.com/apache/beam/pull/13059#discussion_r526611139
########## File path: sdks/java/extensions/sorter/src/main/java/org/apache/beam/sdk/extensions/sorter/MergeSortCombineFn.java ########## @@ -0,0 +1,243 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.beam.sdk.extensions.sorter; + +import com.google.common.primitives.UnsignedBytes; +import java.io.IOException; +import java.io.InputStream; +import java.io.OutputStream; +import java.util.ArrayList; +import java.util.Collections; +import java.util.Comparator; +import java.util.Iterator; +import java.util.List; +import org.apache.beam.sdk.coders.AtomicCoder; +import org.apache.beam.sdk.coders.BooleanCoder; +import org.apache.beam.sdk.coders.ByteArrayCoder; +import org.apache.beam.sdk.coders.Coder; +import org.apache.beam.sdk.coders.CoderException; +import org.apache.beam.sdk.coders.CoderRegistry; +import org.apache.beam.sdk.coders.KvCoder; +import org.apache.beam.sdk.coders.ListCoder; +import org.apache.beam.sdk.extensions.sorter.MergeSortCombineFn.SortingAccumulator; +import org.apache.beam.sdk.extensions.sorter.MergeSortCombineFn.SortingAccumulator.SortingAccumulatorCoder; +import org.apache.beam.sdk.transforms.Combine; +import org.apache.beam.sdk.transforms.Combine.CombineFn; +import org.apache.beam.sdk.transforms.SerializableFunction; +import org.apache.beam.sdk.util.CoderUtils; +import org.apache.beam.sdk.values.KV; +import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Iterators; +import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Lists; + +/** + * An implementation of a distributed sorter that works by accumulating inputs while maintaining + * their sorted order according to a sort key K. Accumulators are merged lazily to maintain sorted + * order every time the input is reduced. The sort key K is extracted from each input V when it is + * added to the accumulator. + * + * <p>Unlike {@link SortValues}, MergeSortCombineFn will NOT spill to disk if available memory is + * exceeded. + * + * @param <K> the class of the sort key + * @param <V> the class of the input PCollection elements + */ +public class MergeSortCombineFn<K, V> extends CombineFn<V, SortingAccumulator, Iterable<V>> { Review comment: One of the problems described in the Jira issue were that, to combat the OOM issues, spilling to disk when sorting was required. Because of a few hot-keys with lots of values and that worker VMs are homogeneous, the disk requirements across all the workers increased, but that resulted in poor disk utilization overall. Are the increased disk requirements because the records spilled to disk for sorting are essentially duplicates of those being read from the GroupByKey? It seems like this is how the SortValues implementation works. Could you explain how the disk requirements here are reduced? My best guess is that because partial sorting is done before the GroupByKey (i.e. a lifted combiner performing a partial GroupByKey and combine), there is no longer a need to spill to disk when sorting in the post-GroupByKey combine because they're already sorted on disk, ready for merge sorting. ########## File path: sdks/java/extensions/sorter/src/main/java/org/apache/beam/sdk/extensions/sorter/MergeSortCombineFn.java ########## @@ -0,0 +1,243 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.beam.sdk.extensions.sorter; + +import com.google.common.primitives.UnsignedBytes; +import java.io.IOException; +import java.io.InputStream; +import java.io.OutputStream; +import java.util.ArrayList; +import java.util.Collections; +import java.util.Comparator; +import java.util.Iterator; +import java.util.List; +import org.apache.beam.sdk.coders.AtomicCoder; +import org.apache.beam.sdk.coders.BooleanCoder; +import org.apache.beam.sdk.coders.ByteArrayCoder; +import org.apache.beam.sdk.coders.Coder; +import org.apache.beam.sdk.coders.CoderException; +import org.apache.beam.sdk.coders.CoderRegistry; +import org.apache.beam.sdk.coders.KvCoder; +import org.apache.beam.sdk.coders.ListCoder; +import org.apache.beam.sdk.extensions.sorter.MergeSortCombineFn.SortingAccumulator; +import org.apache.beam.sdk.extensions.sorter.MergeSortCombineFn.SortingAccumulator.SortingAccumulatorCoder; +import org.apache.beam.sdk.transforms.Combine; +import org.apache.beam.sdk.transforms.Combine.CombineFn; +import org.apache.beam.sdk.transforms.SerializableFunction; +import org.apache.beam.sdk.util.CoderUtils; +import org.apache.beam.sdk.values.KV; +import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Iterators; +import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Lists; + +/** + * An implementation of a distributed sorter that works by accumulating inputs while maintaining + * their sorted order according to a sort key K. Accumulators are merged lazily to maintain sorted + * order every time the input is reduced. The sort key K is extracted from each input V when it is + * added to the accumulator. + * + * <p>Unlike {@link SortValues}, MergeSortCombineFn will NOT spill to disk if available memory is + * exceeded. + * + * @param <K> the class of the sort key + * @param <V> the class of the input PCollection elements + */ +public class MergeSortCombineFn<K, V> extends CombineFn<V, SortingAccumulator, Iterable<V>> { + private final Coder<K> sortKeyCoder; + private final Coder<V> valueCoder; + private final SerializableFunction<V, K> sortKeyFn; + + public static <KeyT, ValueT> MergeSortCombineFn<KeyT, ValueT> of( + Coder<ValueT> valueCoder, + Coder<KeyT> sortKeyCoder, + SerializableFunction<ValueT, KeyT> sortKeyFn) { + return new MergeSortCombineFn<>(valueCoder, sortKeyCoder, sortKeyFn); + } + + private MergeSortCombineFn( + Coder<V> valueCoder, Coder<K> sortKeyCoder, SerializableFunction<V, K> sortKeyFn) { + this.valueCoder = valueCoder; + this.sortKeyCoder = sortKeyCoder; + this.sortKeyFn = sortKeyFn; + } + + public static <KeyT, SortingKeyT, ValueT> Combine.PerKey<KeyT, ValueT, Iterable<ValueT>> perKey( + Coder<ValueT> valueCoder, + Coder<SortingKeyT> sortKeyCoder, + SerializableFunction<ValueT, SortingKeyT> sortKeyFn) { + return Combine.perKey(MergeSortCombineFn.of(valueCoder, sortKeyCoder, sortKeyFn)); + } + + public static <KeyT, SortingKeyT, ValueT> + Combine.GroupedValues<KeyT, ValueT, Iterable<ValueT>> groupedValues( + Coder<ValueT> valueCoder, + Coder<SortingKeyT> sortKeyCoder, + SerializableFunction<ValueT, SortingKeyT> sortKeyFn) { + return Combine.groupedValues(MergeSortCombineFn.of(valueCoder, sortKeyCoder, sortKeyFn)); + } + + @Override + public Coder<SortingAccumulator> getAccumulatorCoder( + CoderRegistry registry, Coder<V> inputCoder) { + return SortingAccumulatorCoder.of(); + } + + @Override + public SortingAccumulator createAccumulator() { + return new SortingAccumulator(); + } + + @Override + public SortingAccumulator addInput(SortingAccumulator mutableAccumulator, V input) { + final K sortKey = sortKeyFn.apply(input); + try { + mutableAccumulator.add( + KV.of( + CoderUtils.encodeToByteArray(sortKeyCoder, sortKey), + CoderUtils.encodeToByteArray(valueCoder, input))); + } catch (CoderException e) { + throw new RuntimeException("Caught exception encoding input", e); + } + return mutableAccumulator; + } + + @Override + public SortingAccumulator mergeAccumulators(Iterable<SortingAccumulator> accumulators) { + return new SortingAccumulator(new MergeSortingIterable(accumulators)); + } + + // Lazily decodes each byte[] pair back to its original class V. + @Override + public Iterable<V> extractOutput(SortingAccumulator accumulator) { + return () -> + Iterators.transform( + accumulator.iterator(), + (KV<byte[], byte[]> item) -> { + if (item == null) { + throw new RuntimeException("Found unexpected null value in encoded output"); + } + + try { + return CoderUtils.decodeFromByteArray(valueCoder, item.getValue()); + } catch (CoderException e) { + throw new RuntimeException("Caught decoding exception", e); + } + }); + } + + /** + * An Iterable composed of several pre-sorted Sources. It merges its Sources maintaining their + * sorted lexicographic byte order. + */ + static class MergeSortingIterable implements Iterable<KV<byte[], byte[]>> { + private final List<Iterator<KV<byte[], byte[]>>> sources; + + MergeSortingIterable(Iterable<SortingAccumulator> accumulators) { + sources = new ArrayList<>(); + accumulators.forEach(a -> sources.add(a.materialized().iterator())); + } + + @Override + public Iterator<KV<byte[], byte[]>> iterator() { + return Iterators.mergeSorted(sources, SortingAccumulator.KV_COMPARATOR); + } + } + + /** + * Accumulator of `KV<byte[], byte[]>` pairs that maintains them in a lexicographically sorted + * order at all times. + */ + static class SortingAccumulator implements Iterable<KV<byte[], byte[]>> { + private Iterable<KV<byte[], byte[]>> sortedItems; + private boolean isMaterialized; + + static final Comparator<KV<byte[], byte[]>> KV_COMPARATOR = + (kv1, kv2) -> UnsignedBytes.lexicographicalComparator().compare(kv1.getKey(), kv2.getKey()); + + SortingAccumulator() { + this.sortedItems = new ArrayList<>(); + this.isMaterialized = true; + } + + // Items are always sorted during ser/de and when accumulators merge, so we know this + // Iterable is always sorted + SortingAccumulator(Iterable<KV<byte[], byte[]>> sortedItems) { + this.sortedItems = sortedItems; + this.isMaterialized = List.class.isAssignableFrom(sortedItems.getClass()); Review comment: Could you explain this to me please? ########## File path: sdks/java/extensions/sorter/src/main/java/org/apache/beam/sdk/extensions/sorter/MergeSortCombineFn.java ########## @@ -0,0 +1,243 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.beam.sdk.extensions.sorter; + +import com.google.common.primitives.UnsignedBytes; +import java.io.IOException; +import java.io.InputStream; +import java.io.OutputStream; +import java.util.ArrayList; +import java.util.Collections; +import java.util.Comparator; +import java.util.Iterator; +import java.util.List; +import org.apache.beam.sdk.coders.AtomicCoder; +import org.apache.beam.sdk.coders.BooleanCoder; +import org.apache.beam.sdk.coders.ByteArrayCoder; +import org.apache.beam.sdk.coders.Coder; +import org.apache.beam.sdk.coders.CoderException; +import org.apache.beam.sdk.coders.CoderRegistry; +import org.apache.beam.sdk.coders.KvCoder; +import org.apache.beam.sdk.coders.ListCoder; +import org.apache.beam.sdk.extensions.sorter.MergeSortCombineFn.SortingAccumulator; +import org.apache.beam.sdk.extensions.sorter.MergeSortCombineFn.SortingAccumulator.SortingAccumulatorCoder; +import org.apache.beam.sdk.transforms.Combine; +import org.apache.beam.sdk.transforms.Combine.CombineFn; +import org.apache.beam.sdk.transforms.SerializableFunction; +import org.apache.beam.sdk.util.CoderUtils; +import org.apache.beam.sdk.values.KV; +import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Iterators; +import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Lists; + +/** + * An implementation of a distributed sorter that works by accumulating inputs while maintaining + * their sorted order according to a sort key K. Accumulators are merged lazily to maintain sorted + * order every time the input is reduced. The sort key K is extracted from each input V when it is + * added to the accumulator. + * + * <p>Unlike {@link SortValues}, MergeSortCombineFn will NOT spill to disk if available memory is + * exceeded. + * + * @param <K> the class of the sort key + * @param <V> the class of the input PCollection elements + */ +public class MergeSortCombineFn<K, V> extends CombineFn<V, SortingAccumulator, Iterable<V>> { + private final Coder<K> sortKeyCoder; + private final Coder<V> valueCoder; + private final SerializableFunction<V, K> sortKeyFn; + + public static <KeyT, ValueT> MergeSortCombineFn<KeyT, ValueT> of( + Coder<ValueT> valueCoder, + Coder<KeyT> sortKeyCoder, + SerializableFunction<ValueT, KeyT> sortKeyFn) { + return new MergeSortCombineFn<>(valueCoder, sortKeyCoder, sortKeyFn); + } + + private MergeSortCombineFn( + Coder<V> valueCoder, Coder<K> sortKeyCoder, SerializableFunction<V, K> sortKeyFn) { + this.valueCoder = valueCoder; + this.sortKeyCoder = sortKeyCoder; + this.sortKeyFn = sortKeyFn; + } + + public static <KeyT, SortingKeyT, ValueT> Combine.PerKey<KeyT, ValueT, Iterable<ValueT>> perKey( + Coder<ValueT> valueCoder, + Coder<SortingKeyT> sortKeyCoder, + SerializableFunction<ValueT, SortingKeyT> sortKeyFn) { + return Combine.perKey(MergeSortCombineFn.of(valueCoder, sortKeyCoder, sortKeyFn)); + } + + public static <KeyT, SortingKeyT, ValueT> + Combine.GroupedValues<KeyT, ValueT, Iterable<ValueT>> groupedValues( + Coder<ValueT> valueCoder, + Coder<SortingKeyT> sortKeyCoder, + SerializableFunction<ValueT, SortingKeyT> sortKeyFn) { + return Combine.groupedValues(MergeSortCombineFn.of(valueCoder, sortKeyCoder, sortKeyFn)); + } + + @Override + public Coder<SortingAccumulator> getAccumulatorCoder( + CoderRegistry registry, Coder<V> inputCoder) { + return SortingAccumulatorCoder.of(); + } + + @Override + public SortingAccumulator createAccumulator() { + return new SortingAccumulator(); + } + + @Override + public SortingAccumulator addInput(SortingAccumulator mutableAccumulator, V input) { + final K sortKey = sortKeyFn.apply(input); + try { + mutableAccumulator.add( + KV.of( + CoderUtils.encodeToByteArray(sortKeyCoder, sortKey), + CoderUtils.encodeToByteArray(valueCoder, input))); + } catch (CoderException e) { + throw new RuntimeException("Caught exception encoding input", e); + } + return mutableAccumulator; + } + + @Override + public SortingAccumulator mergeAccumulators(Iterable<SortingAccumulator> accumulators) { + return new SortingAccumulator(new MergeSortingIterable(accumulators)); + } + + // Lazily decodes each byte[] pair back to its original class V. + @Override + public Iterable<V> extractOutput(SortingAccumulator accumulator) { + return () -> + Iterators.transform( + accumulator.iterator(), + (KV<byte[], byte[]> item) -> { + if (item == null) { + throw new RuntimeException("Found unexpected null value in encoded output"); + } + + try { + return CoderUtils.decodeFromByteArray(valueCoder, item.getValue()); + } catch (CoderException e) { + throw new RuntimeException("Caught decoding exception", e); + } + }); + } + + /** + * An Iterable composed of several pre-sorted Sources. It merges its Sources maintaining their + * sorted lexicographic byte order. + */ + static class MergeSortingIterable implements Iterable<KV<byte[], byte[]>> { + private final List<Iterator<KV<byte[], byte[]>>> sources; + + MergeSortingIterable(Iterable<SortingAccumulator> accumulators) { + sources = new ArrayList<>(); + accumulators.forEach(a -> sources.add(a.materialized().iterator())); + } + + @Override + public Iterator<KV<byte[], byte[]>> iterator() { + return Iterators.mergeSorted(sources, SortingAccumulator.KV_COMPARATOR); + } + } + + /** + * Accumulator of `KV<byte[], byte[]>` pairs that maintains them in a lexicographically sorted + * order at all times. + */ + static class SortingAccumulator implements Iterable<KV<byte[], byte[]>> { + private Iterable<KV<byte[], byte[]>> sortedItems; + private boolean isMaterialized; + + static final Comparator<KV<byte[], byte[]>> KV_COMPARATOR = + (kv1, kv2) -> UnsignedBytes.lexicographicalComparator().compare(kv1.getKey(), kv2.getKey()); + + SortingAccumulator() { + this.sortedItems = new ArrayList<>(); + this.isMaterialized = true; + } + + // Items are always sorted during ser/de and when accumulators merge, so we know this + // Iterable is always sorted + SortingAccumulator(Iterable<KV<byte[], byte[]>> sortedItems) { + this.sortedItems = sortedItems; + this.isMaterialized = List.class.isAssignableFrom(sortedItems.getClass()); + } + + // Returns a possibly lazily evaluating iterator. + @Override + public Iterator<KV<byte[], byte[]>> iterator() { + return sortedItems.iterator(); + } + + // Eagerly materializes items into a List. + @SuppressWarnings(("unchecked")) + List<KV<byte[], byte[]>> materialized() { + if (!isMaterialized) { + sortedItems = Lists.newArrayList(iterator()); + isMaterialized = true; + } + + return (List<KV<byte[], byte[]>>) sortedItems; + } + + // Merge a new item into the sorted list + void add(KV<byte[], byte[]> item) { + int index = Collections.binarySearch(materialized(), item, KV_COMPARATOR); + if (index < 0) { + index = -index - 1; + } + materialized().add(index, item); + } + + static class SortingAccumulatorCoder extends AtomicCoder<SortingAccumulator> { Review comment: This coder eagerly decodes the entire value set for a single SortingAccumulator. During execution, there will be multiple of these in memory, in preparation for `mergeAccumulators` being called repeatedly until there is a single SortingAccumulator left and all values are in memory. Could you explain how this helps avoid the problems you described in the Jira issue? They were described as OOMs resulting from too many elements in memory when using SortValues, which I think would still happen here. ---------------------------------------------------------------- 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: [email protected]
