runners/gearpump/src/main/java/org/apache/beam/runners/gearpump/translators/GroupByKeyTranslator.java - beam - Git at Google

 /*
  * 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.runners.gearpump.translators;

 import io.gearpump.streaming.dsl.api.functions.FoldFunction;
 import io.gearpump.streaming.dsl.api.functions.MapFunction;
 import io.gearpump.streaming.dsl.javaapi.JavaStream;
 import io.gearpump.streaming.dsl.javaapi.functions.GroupByFunction;
 import io.gearpump.streaming.dsl.window.api.Discarding$;
 import io.gearpump.streaming.dsl.window.api.EventTimeTrigger$;
 import io.gearpump.streaming.dsl.window.api.WindowFunction;
 import io.gearpump.streaming.dsl.window.api.Windows;
 import io.gearpump.streaming.dsl.window.impl.Window;
 import java.io.Serializable;
 import java.nio.ByteBuffer;
 import java.util.ArrayList;
 import java.util.Collection;
 import java.util.List;
 import org.apache.beam.runners.gearpump.translators.utils.TranslatorUtils;
 import org.apache.beam.sdk.coders.Coder;
 import org.apache.beam.sdk.coders.CoderException;
 import org.apache.beam.sdk.coders.KvCoder;
 import org.apache.beam.sdk.transforms.GroupByKey;
 import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
 import org.apache.beam.sdk.transforms.windowing.TimestampCombiner;
 import org.apache.beam.sdk.transforms.windowing.WindowFn;
 import org.apache.beam.sdk.util.CoderUtils;
 import org.apache.beam.sdk.util.WindowedValue;
 import org.apache.beam.sdk.values.KV;
 import org.apache.beam.sdk.values.PCollection;
 import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Iterables;
 import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Lists;
 import org.joda.time.Instant;

 /** {@link GroupByKey} is translated to Gearpump groupBy function. */
 @SuppressWarnings({"rawtypes", "unchecked"})
 public class GroupByKeyTranslator<K, V> implements TransformTranslator<GroupByKey<K, V>> {

   private static final long serialVersionUID = -8742202583992787659L;

   @Override
   public void translate(GroupByKey<K, V> transform, TranslationContext context) {
     PCollection<KV<K, V>> input = (PCollection<KV<K, V>>) context.getInput();
     Coder<K> inputKeyCoder = ((KvCoder<K, V>) input.getCoder()).getKeyCoder();
     JavaStream<WindowedValue<KV<K, V>>> inputStream = context.getInputStream(input);
     int parallelism = context.getPipelineOptions().getParallelism();
     TimestampCombiner timestampCombiner = input.getWindowingStrategy().getTimestampCombiner();
     WindowFn<KV<K, V>, BoundedWindow> windowFn =
         (WindowFn<KV<K, V>, BoundedWindow>) input.getWindowingStrategy().getWindowFn();
     JavaStream<WindowedValue<KV<K, List<V>>>> outputStream =
         inputStream
             .window(
                 Windows.apply(
                     new GearpumpWindowFn(windowFn.isNonMerging()),
                     EventTimeTrigger$.MODULE$,
                     Discarding$.MODULE$,
                     windowFn.toString()))
             .groupBy(new GroupByFn<>(inputKeyCoder), parallelism, "group_by_Key_and_Window")
             .map(new KeyedByTimestamp<>(windowFn, timestampCombiner), "keyed_by_timestamp")
             .fold(new Merge<>(windowFn, timestampCombiner), "merge")
             .map(new Values<>(), "values");

     context.setOutputStream(context.getOutput(), outputStream);
   }

   /** A transform used internally to translate Beam's Window to Gearpump's Window. */
   protected static class GearpumpWindowFn<T, W extends BoundedWindow>
       implements WindowFunction, Serializable {

     private final boolean isNonMerging;

     public GearpumpWindowFn(boolean isNonMerging) {
       this.isNonMerging = isNonMerging;
     }

     @Override
     public <T2> Window[] apply(Context<T2> context) {
       try {
         Object element = context.element();
         if (element instanceof TranslatorUtils.RawUnionValue) {
           element = ((TranslatorUtils.RawUnionValue) element).getValue();
         }
         return toGearpumpWindows(
             ((WindowedValue<T>) element).getWindows().toArray(new BoundedWindow[0]));
       } catch (Exception e) {
         throw new RuntimeException(e);
       }
     }

     @Override
     public boolean isNonMerging() {
       return isNonMerging;
     }

     private Window[] toGearpumpWindows(BoundedWindow[] windows) {
       Window[] gwins = new Window[windows.length];
       for (int i = 0; i < windows.length; i++) {
         gwins[i] = TranslatorUtils.boundedWindowToGearpumpWindow(windows[i]);
       }
       return gwins;
     }
   }

   /** A transform used internally to group KV message by its key. */
   protected static class GroupByFn<K, V>
       extends GroupByFunction<WindowedValue<KV<K, V>>, ByteBuffer> {

     private static final long serialVersionUID = -807905402490735530L;
     private final Coder<K> keyCoder;

     GroupByFn(Coder<K> keyCoder) {
       this.keyCoder = keyCoder;
     }

     @Override
     public ByteBuffer groupBy(WindowedValue<KV<K, V>> wv) {
       try {
         return ByteBuffer.wrap(CoderUtils.encodeToByteArray(keyCoder, wv.getValue().getKey()));
       } catch (CoderException e) {
         throw new RuntimeException(e);
       }
     }
   }

   /** A transform used internally to transform WindowedValue to KV. */
   protected static class KeyedByTimestamp<K, V>
       extends MapFunction<WindowedValue<KV<K, V>>, KV<Instant, WindowedValue<KV<K, V>>>> {

     private final WindowFn<KV<K, V>, BoundedWindow> windowFn;
     private final TimestampCombiner timestampCombiner;

     public KeyedByTimestamp(
         WindowFn<KV<K, V>, BoundedWindow> windowFn, TimestampCombiner timestampCombiner) {
       this.windowFn = windowFn;
       this.timestampCombiner = timestampCombiner;
     }

     @Override
     public KV<Instant, WindowedValue<KV<K, V>>> map(WindowedValue<KV<K, V>> wv) {
       BoundedWindow window = Iterables.getOnlyElement(wv.getWindows());
       Instant timestamp =
           timestampCombiner.assign(window, windowFn.getOutputTime(wv.getTimestamp(), window));
       return KV.of(timestamp, wv);
     }
   }

   /** A transform used internally by Gearpump which encapsulates the merge logic. */
   protected static class Merge<K, V>
       extends FoldFunction<
           KV<Instant, WindowedValue<KV<K, V>>>, KV<Instant, WindowedValue<KV<K, List<V>>>>> {

     private final WindowFn<KV<K, V>, BoundedWindow> windowFn;
     private final TimestampCombiner timestampCombiner;

     Merge(WindowFn<KV<K, V>, BoundedWindow> windowFn, TimestampCombiner timestampCombiner) {
       this.windowFn = windowFn;
       this.timestampCombiner = timestampCombiner;
     }

     @Override
     public KV<Instant, WindowedValue<KV<K, List<V>>>> init() {
       return KV.of(null, null);
     }

     @Override
     public KV<Instant, WindowedValue<KV<K, List<V>>>> fold(
         KV<Instant, WindowedValue<KV<K, List<V>>>> accum,
         KV<Instant, WindowedValue<KV<K, V>>> iter) {
       if (accum.getKey() == null) {
         WindowedValue<KV<K, V>> wv = iter.getValue();
         KV<K, V> kv = wv.getValue();
         V v = kv.getValue();
         List<V> nv = Lists.newArrayList(v);
         return KV.of(iter.getKey(), wv.withValue(KV.of(kv.getKey(), nv)));
       }

       Instant t1 = accum.getKey();
       Instant t2 = iter.getKey();

       final WindowedValue<KV<K, List<V>>> wv1 = accum.getValue();
       final WindowedValue<KV<K, V>> wv2 = iter.getValue();
       wv1.getValue().getValue().add(wv2.getValue().getValue());

       final List<BoundedWindow> mergedWindows = new ArrayList<>();
       if (!windowFn.isNonMerging()) {
         try {
           windowFn.mergeWindows(
               windowFn.new MergeContext() {

                 @Override
                 public Collection<BoundedWindow> windows() {
                   ArrayList<BoundedWindow> windows = new ArrayList<>();
                   windows.addAll(wv1.getWindows());
                   windows.addAll(wv2.getWindows());
                   return windows;
                 }

                 @Override
                 public void merge(Collection<BoundedWindow> toBeMerged, BoundedWindow mergeResult)
                     throws Exception {
                   mergedWindows.add(mergeResult);
                 }
               });
         } catch (Exception e) {
           throw new RuntimeException(e);
         }
       } else {
         mergedWindows.addAll(wv1.getWindows());
       }

       Instant timestamp = timestampCombiner.combine(t1, t2);
       return KV.of(
           timestamp, WindowedValue.of(wv1.getValue(), timestamp, mergedWindows, wv1.getPane()));
     }
   }

   private static class Values<K, V>
       extends MapFunction<
           KV<Instant, WindowedValue<KV<K, List<V>>>>, WindowedValue<KV<K, List<V>>>> {

     @Override
     public WindowedValue<KV<K, List<V>>> map(KV<Instant, WindowedValue<KV<K, List<V>>>> kv) {
       Instant timestamp = kv.getKey();
       WindowedValue<KV<K, List<V>>> wv = kv.getValue();
       return WindowedValue.of(wv.getValue(), timestamp, wv.getWindows(), wv.getPane());
     }
   }
 }
	/*
	* 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.runners.gearpump.translators;

	import io.gearpump.streaming.dsl.api.functions.FoldFunction;
	import io.gearpump.streaming.dsl.api.functions.MapFunction;
	import io.gearpump.streaming.dsl.javaapi.JavaStream;
	import io.gearpump.streaming.dsl.javaapi.functions.GroupByFunction;
	import io.gearpump.streaming.dsl.window.api.Discarding$;
	import io.gearpump.streaming.dsl.window.api.EventTimeTrigger$;
	import io.gearpump.streaming.dsl.window.api.WindowFunction;
	import io.gearpump.streaming.dsl.window.api.Windows;
	import io.gearpump.streaming.dsl.window.impl.Window;
	import java.io.Serializable;
	import java.nio.ByteBuffer;
	import java.util.ArrayList;
	import java.util.Collection;
	import java.util.List;
	import org.apache.beam.runners.gearpump.translators.utils.TranslatorUtils;
	import org.apache.beam.sdk.coders.Coder;
	import org.apache.beam.sdk.coders.CoderException;
	import org.apache.beam.sdk.coders.KvCoder;
	import org.apache.beam.sdk.transforms.GroupByKey;
	import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
	import org.apache.beam.sdk.transforms.windowing.TimestampCombiner;
	import org.apache.beam.sdk.transforms.windowing.WindowFn;
	import org.apache.beam.sdk.util.CoderUtils;
	import org.apache.beam.sdk.util.WindowedValue;
	import org.apache.beam.sdk.values.KV;
	import org.apache.beam.sdk.values.PCollection;
	import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Iterables;
	import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Lists;
	import org.joda.time.Instant;

	/** {@link GroupByKey} is translated to Gearpump groupBy function. */
	@SuppressWarnings({"rawtypes", "unchecked"})
	public class GroupByKeyTranslator<K, V> implements TransformTranslator<GroupByKey<K, V>> {

	private static final long serialVersionUID = -8742202583992787659L;

	@Override
	public void translate(GroupByKey<K, V> transform, TranslationContext context) {
	PCollection<KV<K, V>> input = (PCollection<KV<K, V>>) context.getInput();
	Coder<K> inputKeyCoder = ((KvCoder<K, V>) input.getCoder()).getKeyCoder();
	JavaStream<WindowedValue<KV<K, V>>> inputStream = context.getInputStream(input);
	int parallelism = context.getPipelineOptions().getParallelism();
	TimestampCombiner timestampCombiner = input.getWindowingStrategy().getTimestampCombiner();
	WindowFn<KV<K, V>, BoundedWindow> windowFn =
	(WindowFn<KV<K, V>, BoundedWindow>) input.getWindowingStrategy().getWindowFn();
	JavaStream<WindowedValue<KV<K, List<V>>>> outputStream =
	inputStream
	.window(
	Windows.apply(
	new GearpumpWindowFn(windowFn.isNonMerging()),
	EventTimeTrigger$.MODULE$,
	Discarding$.MODULE$,
	windowFn.toString()))
	.groupBy(new GroupByFn<>(inputKeyCoder), parallelism, "group_by_Key_and_Window")
	.map(new KeyedByTimestamp<>(windowFn, timestampCombiner), "keyed_by_timestamp")
	.fold(new Merge<>(windowFn, timestampCombiner), "merge")
	.map(new Values<>(), "values");

	context.setOutputStream(context.getOutput(), outputStream);
	}

	/** A transform used internally to translate Beam's Window to Gearpump's Window. */
	protected static class GearpumpWindowFn<T, W extends BoundedWindow>
	implements WindowFunction, Serializable {

	private final boolean isNonMerging;

	public GearpumpWindowFn(boolean isNonMerging) {
	this.isNonMerging = isNonMerging;
	}

	@Override
	public <T2> Window[] apply(Context<T2> context) {
	try {
	Object element = context.element();
	if (element instanceof TranslatorUtils.RawUnionValue) {
	element = ((TranslatorUtils.RawUnionValue) element).getValue();
	}
	return toGearpumpWindows(
	((WindowedValue<T>) element).getWindows().toArray(new BoundedWindow[0]));
	} catch (Exception e) {
	throw new RuntimeException(e);
	}
	}

	@Override
	public boolean isNonMerging() {
	return isNonMerging;
	}

	private Window[] toGearpumpWindows(BoundedWindow[] windows) {
	Window[] gwins = new Window[windows.length];
	for (int i = 0; i < windows.length; i++) {
	gwins[i] = TranslatorUtils.boundedWindowToGearpumpWindow(windows[i]);
	}
	return gwins;
	}
	}

	/** A transform used internally to group KV message by its key. */
	protected static class GroupByFn<K, V>
	extends GroupByFunction<WindowedValue<KV<K, V>>, ByteBuffer> {

	private static final long serialVersionUID = -807905402490735530L;
	private final Coder<K> keyCoder;

	GroupByFn(Coder<K> keyCoder) {
	this.keyCoder = keyCoder;
	}

	@Override
	public ByteBuffer groupBy(WindowedValue<KV<K, V>> wv) {
	try {
	return ByteBuffer.wrap(CoderUtils.encodeToByteArray(keyCoder, wv.getValue().getKey()));
	} catch (CoderException e) {
	throw new RuntimeException(e);
	}
	}
	}

	/** A transform used internally to transform WindowedValue to KV. */
	protected static class KeyedByTimestamp<K, V>
	extends MapFunction<WindowedValue<KV<K, V>>, KV<Instant, WindowedValue<KV<K, V>>>> {

	private final WindowFn<KV<K, V>, BoundedWindow> windowFn;
	private final TimestampCombiner timestampCombiner;

	public KeyedByTimestamp(
	WindowFn<KV<K, V>, BoundedWindow> windowFn, TimestampCombiner timestampCombiner) {
	this.windowFn = windowFn;
	this.timestampCombiner = timestampCombiner;
	}

	@Override
	public KV<Instant, WindowedValue<KV<K, V>>> map(WindowedValue<KV<K, V>> wv) {
	BoundedWindow window = Iterables.getOnlyElement(wv.getWindows());
	Instant timestamp =
	timestampCombiner.assign(window, windowFn.getOutputTime(wv.getTimestamp(), window));
	return KV.of(timestamp, wv);
	}
	}

	/** A transform used internally by Gearpump which encapsulates the merge logic. */
	protected static class Merge<K, V>
	extends FoldFunction<
	KV<Instant, WindowedValue<KV<K, V>>>, KV<Instant, WindowedValue<KV<K, List<V>>>>> {

	private final WindowFn<KV<K, V>, BoundedWindow> windowFn;
	private final TimestampCombiner timestampCombiner;

	Merge(WindowFn<KV<K, V>, BoundedWindow> windowFn, TimestampCombiner timestampCombiner) {
	this.windowFn = windowFn;
	this.timestampCombiner = timestampCombiner;
	}

	@Override
	public KV<Instant, WindowedValue<KV<K, List<V>>>> init() {
	return KV.of(null, null);
	}

	@Override
	public KV<Instant, WindowedValue<KV<K, List<V>>>> fold(
	KV<Instant, WindowedValue<KV<K, List<V>>>> accum,
	KV<Instant, WindowedValue<KV<K, V>>> iter) {
	if (accum.getKey() == null) {
	WindowedValue<KV<K, V>> wv = iter.getValue();
	KV<K, V> kv = wv.getValue();
	V v = kv.getValue();
	List<V> nv = Lists.newArrayList(v);
	return KV.of(iter.getKey(), wv.withValue(KV.of(kv.getKey(), nv)));
	}

	Instant t1 = accum.getKey();
	Instant t2 = iter.getKey();

	final WindowedValue<KV<K, List<V>>> wv1 = accum.getValue();
	final WindowedValue<KV<K, V>> wv2 = iter.getValue();
	wv1.getValue().getValue().add(wv2.getValue().getValue());

	final List<BoundedWindow> mergedWindows = new ArrayList<>();
	if (!windowFn.isNonMerging()) {
	try {
	windowFn.mergeWindows(
	windowFn.new MergeContext() {

	@Override
	public Collection<BoundedWindow> windows() {
	ArrayList<BoundedWindow> windows = new ArrayList<>();
	windows.addAll(wv1.getWindows());
	windows.addAll(wv2.getWindows());
	return windows;
	}

	@Override
	public void merge(Collection<BoundedWindow> toBeMerged, BoundedWindow mergeResult)
	throws Exception {
	mergedWindows.add(mergeResult);
	}
	});
	} catch (Exception e) {
	throw new RuntimeException(e);
	}
	} else {
	mergedWindows.addAll(wv1.getWindows());
	}

	Instant timestamp = timestampCombiner.combine(t1, t2);
	return KV.of(
	timestamp, WindowedValue.of(wv1.getValue(), timestamp, mergedWindows, wv1.getPane()));
	}
	}

	private static class Values<K, V>
	extends MapFunction<
	KV<Instant, WindowedValue<KV<K, List<V>>>>, WindowedValue<KV<K, List<V>>>> {

	@Override
	public WindowedValue<KV<K, List<V>>> map(KV<Instant, WindowedValue<KV<K, List<V>>>> kv) {
	Instant timestamp = kv.getKey();
	WindowedValue<KV<K, List<V>>> wv = kv.getValue();
	return WindowedValue.of(wv.getValue(), timestamp, wv.getWindows(), wv.getPane());
	}
	}
	}