| /* |
| * 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.euphoria.core.translate; |
| |
| import static java.util.Objects.requireNonNull; |
| import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkState; |
| |
| import java.util.stream.StreamSupport; |
| import javax.annotation.Nullable; |
| import org.apache.beam.sdk.coders.CannotProvideCoderException; |
| import org.apache.beam.sdk.coders.Coder; |
| import org.apache.beam.sdk.coders.CoderRegistry; |
| import org.apache.beam.sdk.extensions.euphoria.core.client.accumulators.AccumulatorProvider; |
| import org.apache.beam.sdk.extensions.euphoria.core.client.functional.BinaryFunction; |
| import org.apache.beam.sdk.extensions.euphoria.core.client.functional.CombinableBinaryFunction; |
| import org.apache.beam.sdk.extensions.euphoria.core.client.functional.ReduceFunctor; |
| import org.apache.beam.sdk.extensions.euphoria.core.client.functional.UnaryFunction; |
| import org.apache.beam.sdk.extensions.euphoria.core.client.functional.VoidFunction; |
| import org.apache.beam.sdk.extensions.euphoria.core.client.operator.ReduceByKey; |
| import org.apache.beam.sdk.extensions.euphoria.core.client.type.TypeAwareness; |
| import org.apache.beam.sdk.extensions.euphoria.core.client.util.PCollectionLists; |
| import org.apache.beam.sdk.extensions.euphoria.core.translate.collector.AdaptableCollector; |
| import org.apache.beam.sdk.extensions.euphoria.core.translate.collector.SingleValueCollector; |
| import org.apache.beam.sdk.transforms.Combine; |
| import org.apache.beam.sdk.transforms.DoFn; |
| import org.apache.beam.sdk.transforms.GroupByKey; |
| import org.apache.beam.sdk.transforms.MapElements; |
| import org.apache.beam.sdk.transforms.ParDo; |
| import org.apache.beam.sdk.transforms.SerializableFunction; |
| import org.apache.beam.sdk.transforms.SimpleFunction; |
| import org.apache.beam.sdk.values.KV; |
| import org.apache.beam.sdk.values.PCollection; |
| import org.apache.beam.sdk.values.PCollectionList; |
| import org.apache.beam.sdk.values.TypeDescriptor; |
| import org.apache.beam.sdk.values.TypeDescriptors; |
| |
| /** Translator for {@code ReduceByKey} operator. */ |
| public class ReduceByKeyTranslator<InputT, KeyT, ValueT, OutputT> |
| implements OperatorTranslator< |
| InputT, KV<KeyT, OutputT>, ReduceByKey<InputT, KeyT, ValueT, ?, OutputT>> { |
| |
| @Override |
| public PCollection<KV<KeyT, OutputT>> translate( |
| ReduceByKey<InputT, KeyT, ValueT, ?, OutputT> operator, PCollectionList<InputT> inputs) { |
| |
| // todo Could we even do values sorting in Beam ? And do we want it? |
| checkState(!operator.getValueComparator().isPresent(), "Values sorting is not supported."); |
| |
| final UnaryFunction<InputT, KeyT> keyExtractor = operator.getKeyExtractor(); |
| final UnaryFunction<InputT, ValueT> valueExtractor = operator.getValueExtractor(); |
| |
| final PCollection<InputT> input = |
| operator |
| .getWindow() |
| .map(window -> PCollectionLists.getOnlyElement(inputs).apply(window)) |
| .orElseGet(() -> PCollectionLists.getOnlyElement(inputs)); |
| |
| // ~ create key & value extractor |
| final MapElements<InputT, KV<KeyT, ValueT>> extractor = |
| MapElements.via(new KeyValueExtractor<>(keyExtractor, valueExtractor)); |
| |
| final PCollection<KV<KeyT, ValueT>> extracted = |
| input |
| .apply("extract-keys", extractor) |
| .setTypeDescriptor( |
| TypeDescriptors.kvs( |
| TypeAwareness.orObjects(operator.getKeyType()), |
| TypeAwareness.orObjects(operator.getValueType()))); |
| |
| final AccumulatorProvider accumulators = |
| new LazyAccumulatorProvider(AccumulatorProvider.of(inputs.getPipeline())); |
| |
| if (operator.isCombinable()) { |
| // if operator is combinable we can process it in more efficient way |
| @SuppressWarnings("unchecked") |
| final PCollection combined; |
| if (operator.isCombineFnStyle()) { |
| combined = extracted.apply("combine", Combine.perKey(asCombineFn(operator))); |
| } else { |
| combined = |
| extracted.apply( |
| "combine", |
| Combine.perKey( |
| asCombiner( |
| operator.getReducer(), accumulators, operator.getName().orElse(null)))); |
| } |
| @SuppressWarnings("unchecked") |
| final PCollection<KV<KeyT, OutputT>> cast = (PCollection) combined; |
| return cast.setTypeDescriptor( |
| operator |
| .getOutputType() |
| .orElseThrow( |
| () -> new IllegalStateException("Unable to infer output type descriptor."))); |
| } |
| |
| return extracted |
| .apply("group", GroupByKey.create()) |
| .setTypeDescriptor( |
| TypeDescriptors.kvs( |
| TypeAwareness.orObjects(operator.getKeyType()), |
| TypeDescriptors.iterables(TypeAwareness.orObjects(operator.getValueType())))) |
| .apply( |
| "reduce", |
| ParDo.of( |
| new ReduceDoFn<>( |
| operator.getReducer(), accumulators, operator.getName().orElse(null)))) |
| .setTypeDescriptor( |
| operator |
| .getOutputType() |
| .orElseThrow( |
| () -> new IllegalStateException("Unable to infer output type descriptor."))); |
| } |
| |
| @Override |
| public boolean canTranslate(ReduceByKey operator) { |
| // translation of sorted values is not supported yet |
| return !operator.getValueComparator().isPresent(); |
| } |
| |
| private static <InputT, KeyT, ValueT, AccT, OutputT> |
| Combine.CombineFn<ValueT, AccT, OutputT> asCombineFn( |
| ReduceByKey<InputT, KeyT, ValueT, AccT, OutputT> operator) { |
| |
| @SuppressWarnings("unchecked") |
| ReduceByKey<InputT, KeyT, ValueT, AccT, OutputT> cast = (ReduceByKey) operator; |
| |
| VoidFunction<AccT> accumulatorFactory = cast.getAccumulatorFactory(); |
| BinaryFunction<AccT, ValueT, AccT> accumulate = cast.getAccumulate(); |
| CombinableBinaryFunction<AccT> mergeAccumulators = cast.getMergeAccumulators(); |
| UnaryFunction<AccT, OutputT> outputFn = cast.getOutputFn(); |
| TypeDescriptor<AccT> accumulatorType = cast.getAccumulatorType(); |
| |
| return new Combine.CombineFn<ValueT, AccT, OutputT>() { |
| |
| @Override |
| public AccT createAccumulator() { |
| return accumulatorFactory.apply(); |
| } |
| |
| @Override |
| public Coder<AccT> getAccumulatorCoder(CoderRegistry registry, Coder<ValueT> inputCoder) |
| throws CannotProvideCoderException { |
| return registry.getCoder(accumulatorType); |
| } |
| |
| @Override |
| public AccT addInput(AccT mutableAccumulator, ValueT input) { |
| return accumulate.apply(mutableAccumulator, input); |
| } |
| |
| @Override |
| public AccT mergeAccumulators(Iterable<AccT> accumulators) { |
| AccT accumulated = null; |
| for (AccT o : accumulators) { |
| if (accumulated == null) { |
| accumulated = o; |
| } else { |
| accumulated = mergeAccumulators.apply(accumulated, o); |
| } |
| } |
| return accumulated; |
| } |
| |
| @Override |
| public OutputT extractOutput(AccT accumulator) { |
| return outputFn.apply(accumulator); |
| } |
| }; |
| } |
| |
| private static <InputT, OutputT> SerializableFunction<Iterable<InputT>, InputT> asCombiner( |
| ReduceFunctor<InputT, OutputT> reducer, |
| AccumulatorProvider accumulatorProvider, |
| @Nullable String operatorName) { |
| |
| @SuppressWarnings("unchecked") |
| final ReduceFunctor<InputT, InputT> combiner = (ReduceFunctor<InputT, InputT>) reducer; |
| |
| return (Iterable<InputT> input) -> { |
| SingleValueCollector<InputT> collector = |
| new SingleValueCollector<>(accumulatorProvider, operatorName); |
| combiner.apply(StreamSupport.stream(input.spliterator(), false), collector); |
| return collector.get(); |
| }; |
| } |
| |
| /** |
| * Extract key and values from input data set. |
| * |
| * @param <InputT> type of input |
| * @param <KeyT> type of key |
| * @param <ValueT> type of value |
| */ |
| private static class KeyValueExtractor<InputT, KeyT, ValueT> |
| extends SimpleFunction<InputT, KV<KeyT, ValueT>> { |
| |
| private final UnaryFunction<InputT, KeyT> keyExtractor; |
| private final UnaryFunction<InputT, ValueT> valueExtractor; |
| |
| KeyValueExtractor( |
| UnaryFunction<InputT, KeyT> keyExtractor, UnaryFunction<InputT, ValueT> valueExtractor) { |
| this.keyExtractor = keyExtractor; |
| this.valueExtractor = valueExtractor; |
| } |
| |
| @Override |
| public KV<KeyT, ValueT> apply(InputT in) { |
| return KV.of(keyExtractor.apply(in), valueExtractor.apply(in)); |
| } |
| } |
| |
| /** |
| * Perform reduction of given elements. |
| * |
| * @param <KeyT> type of key |
| * @param <ValueT> type of value |
| * @param <OutputT> type of output |
| */ |
| private static class ReduceDoFn<KeyT, ValueT, OutputT> |
| extends DoFn<KV<KeyT, Iterable<ValueT>>, KV<KeyT, OutputT>> { |
| |
| private final ReduceFunctor<ValueT, OutputT> reducer; |
| private final AdaptableCollector<KV<KeyT, Iterable<ValueT>>, KV<KeyT, OutputT>, OutputT> |
| collector; |
| |
| ReduceDoFn( |
| ReduceFunctor<ValueT, OutputT> reducer, |
| AccumulatorProvider accumulators, |
| @Nullable String operatorName) { |
| this.reducer = reducer; |
| this.collector = |
| new AdaptableCollector<>( |
| accumulators, |
| operatorName, |
| (DoFn<KV<KeyT, Iterable<ValueT>>, KV<KeyT, OutputT>>.ProcessContext ctx, |
| OutputT out) -> ctx.output(KV.of(ctx.element().getKey(), out))); |
| } |
| |
| @ProcessElement |
| @SuppressWarnings("unused") |
| public void processElement(ProcessContext ctx) { |
| collector.setProcessContext(ctx); |
| reducer.apply( |
| StreamSupport.stream(requireNonNull(ctx.element().getValue()).spliterator(), false), |
| collector); |
| } |
| } |
| } |
