[
https://issues.apache.org/jira/browse/BEAM-11048?focusedWorklogId=513948&page=com.atlassian.jira.plugin.system.issuetabpanels:worklog-tabpanel#worklog-513948
]
ASF GitHub Bot logged work on BEAM-11048:
-----------------------------------------
Author: ASF GitHub Bot
Created on: 19/Nov/20 06:07
Start Date: 19/Nov/20 06:07
Worklog Time Spent: 10m
Work Description: 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]
Issue Time Tracking
-------------------
Worklog Id: (was: 513948)
Time Spent: 0.5h (was: 20m)
> Add alternate Sorting transform as an implementation of CombineFn
> -----------------------------------------------------------------
>
> Key: BEAM-11048
> URL: https://issues.apache.org/jira/browse/BEAM-11048
> Project: Beam
> Issue Type: Improvement
> Components: extensions-java-sorter
> Reporter: Claire McGinty
> Priority: P2
> Labels: Clarified
> Time Spent: 0.5h
> Remaining Estimate: 0h
>
> My team has been using the
> [SortValues|https://github.com/apache/beam/blob/master/sdks/java/extensions/sorter/src/main/java/org/apache/beam/sdk/extensions/sorter/SortValues.java]
> transform in `extensions-java-sorter` to sort pre-grouped values by a
> secondary sorter key. However, for large key groups, we've run into many OOM
> issues and have to increase disk size quite a bit to accommodate the larger
> key groups spilling to disk, even if there are only a few large key groups
> and most fit in memory.
> I drafted a new iteration of a Sorter that's a distributed merge-sort
> implemented as a `CombineFn`: each Accumulator maintains an always-sorted
> list of elements, and those Accumulators can be merged simply by zipping
> their lists together. This has the extra advantage that `extractOutput` can
> be lazily evaluated as a merging Iterator rather than as a fully materialized
> list. I also observed that this implementation is able to scale more
> effectively than the old SortValues, and for several use cases where
> `SortValues` ran OOM, the CombineFn-based implementation was able to complete
> using only the default Dataflow disk specs.
> Finally, from an API perspective, I think it's a little easier to use,
> because the user doesn't have to extract the sortKey out into the PCollection
> itself, but instead provide a function mapping each element type T to its
> sort key K, which will be evaluated inside the combiner. So I think in that
> sense it's more intuitive and similar to a Comparator-style sort.
--
This message was sent by Atlassian Jira
(v8.3.4#803005)
