Repository: incubator-beam Updated Branches: refs/heads/master 2a30f5236 -> 84573e0b4
Add Pane matchers to PAssert This allows PCollections that have upstream triggers that fire multiple times to assert the contents of panes with specific timings, rather than restricting the trigger to fire only once. Project: http://git-wip-us.apache.org/repos/asf/incubator-beam/repo Commit: http://git-wip-us.apache.org/repos/asf/incubator-beam/commit/ae0293d0 Tree: http://git-wip-us.apache.org/repos/asf/incubator-beam/tree/ae0293d0 Diff: http://git-wip-us.apache.org/repos/asf/incubator-beam/diff/ae0293d0 Branch: refs/heads/master Commit: ae0293d0c84e1e49031f62e1026ab7a9b362bbf7 Parents: ff72e27 Author: Thomas Groh <[email protected]> Authored: Mon Jun 20 14:38:11 2016 -0700 Committer: Kenneth Knowles <[email protected]> Committed: Thu Jul 14 14:30:49 2016 -0700 ---------------------------------------------------------------------- .../org/apache/beam/sdk/testing/PAssert.java | 222 ++++++++++--- .../apache/beam/sdk/testing/PaneExtractors.java | 140 ++++++++ .../apache/beam/sdk/testing/PAssertTest.java | 4 +- .../beam/sdk/testing/PaneExtractorsTest.java | 323 +++++++++++++++++++ 4 files changed, 649 insertions(+), 40 deletions(-) ---------------------------------------------------------------------- http://git-wip-us.apache.org/repos/asf/incubator-beam/blob/ae0293d0/sdks/java/core/src/main/java/org/apache/beam/sdk/testing/PAssert.java ---------------------------------------------------------------------- diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/testing/PAssert.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/testing/PAssert.java index 883b2b3..b8bd431 100644 --- a/sdks/java/core/src/main/java/org/apache/beam/sdk/testing/PAssert.java +++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/testing/PAssert.java @@ -17,7 +17,6 @@ */ package org.apache.beam.sdk.testing; -import static com.google.common.base.Preconditions.checkState; import static org.hamcrest.Matchers.containsInAnyOrder; import static org.hamcrest.Matchers.equalTo; import static org.hamcrest.Matchers.not; @@ -38,9 +37,11 @@ import org.apache.beam.sdk.transforms.DoFn; import org.apache.beam.sdk.transforms.DoFn.RequiresWindowAccess; import org.apache.beam.sdk.transforms.Flatten; import org.apache.beam.sdk.transforms.GroupByKey; +import org.apache.beam.sdk.transforms.MapElements; import org.apache.beam.sdk.transforms.PTransform; 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.transforms.Sum; import org.apache.beam.sdk.transforms.Values; import org.apache.beam.sdk.transforms.View; @@ -49,10 +50,13 @@ import org.apache.beam.sdk.transforms.windowing.BoundedWindow; import org.apache.beam.sdk.transforms.windowing.GlobalWindow; import org.apache.beam.sdk.transforms.windowing.GlobalWindows; import org.apache.beam.sdk.transforms.windowing.Never; +import org.apache.beam.sdk.transforms.windowing.Trigger; import org.apache.beam.sdk.transforms.windowing.Window; +import org.apache.beam.sdk.transforms.windowing.Window.ClosingBehavior; import org.apache.beam.sdk.util.CoderUtils; import org.apache.beam.sdk.util.GatherAllPanes; import org.apache.beam.sdk.util.WindowedValue; +import org.apache.beam.sdk.util.WindowingStrategy; import org.apache.beam.sdk.values.KV; import org.apache.beam.sdk.values.PBegin; import org.apache.beam.sdk.values.PCollection; @@ -63,16 +67,15 @@ import org.apache.beam.sdk.values.PDone; import com.google.common.collect.Iterables; import com.google.common.collect.Lists; +import org.joda.time.Duration; import org.slf4j.Logger; import org.slf4j.LoggerFactory; import java.io.IOException; import java.io.Serializable; -import java.util.ArrayList; import java.util.Arrays; import java.util.Collection; import java.util.Collections; -import java.util.List; import java.util.Map; import java.util.NoSuchElementException; @@ -127,12 +130,51 @@ public class PAssert { * Creates a new {@link IterableAssert} like this one, but with the assertion restricted to only * run on the provided window. * + * <p>The assertion will concatenate all panes present in the provided window if the + * {@link Trigger} produces multiple panes. If the windowing strategy accumulates fired panes + * and triggers fire multple times, consider using instead {@link #inFinalPane(BoundedWindow)} + * or {@link #inOnTimePane(BoundedWindow)}. + * * @return a new {@link IterableAssert} like this one but with the assertion only applied to the * specified window. */ IterableAssert<T> inWindow(BoundedWindow window); /** + * Creates a new {@link IterableAssert} like this one, but with the assertion restricted to only + * run on the provided window, running the checker only on the final pane for each key. + * + * <p>If the input {@link WindowingStrategy} does not always produce final panes, the assertion + * may be executed over an empty input even if the trigger has fired previously. To ensure that + * a final pane is always produced, set the {@link ClosingBehavior} of the windowing strategy + * (via {@link Window.Bound#withAllowedLateness(Duration, ClosingBehavior)} setting + * {@link ClosingBehavior} to {@link ClosingBehavior#FIRE_ALWAYS}). + * + * @return a new {@link IterableAssert} like this one but with the assertion only applied to the + * specified window. + */ + IterableAssert<T> inFinalPane(BoundedWindow window); + + /** + * Creates a new {@link IterableAssert} like this one, but with the assertion restricted to only + * run on the provided window. + * + * @return a new {@link IterableAssert} like this one but with the assertion only applied to the + * specified window. + */ + IterableAssert<T> inOnTimePane(BoundedWindow window); + + /** + * Creates a new {@link IterableAssert} like this one, but with the assertion restricted to only + * run on the provided window across all panes that were not produced by the arrival of late + * data. + * + * @return a new {@link IterableAssert} like this one but with the assertion only applied to the + * specified window. + */ + IterableAssert<T> inCombinedNonLatePanes(BoundedWindow window); + + /** * Asserts that the iterable in question contains the provided elements. * * @return the same {@link IterableAssert} builder for further assertions @@ -170,10 +212,38 @@ public class PAssert { * Creates a new {@link SingletonAssert} like this one, but with the assertion restricted to * only run on the provided window. * + * <p>The assertion will expect outputs to be produced to the provided window exactly once. If + * the upstream {@link Trigger} may produce output multiple times, consider instead using + * {@link #inFinalPane(BoundedWindow)} or {@link #inOnTimePane(BoundedWindow)}. + * + * @return a new {@link SingletonAssert} like this one but with the assertion only applied to + * the specified window. + */ + SingletonAssert<T> inOnlyPane(BoundedWindow window); + + /** + * Creates a new {@link SingletonAssert} like this one, but with the assertion restricted to + * only run on the provided window, running the checker only on the final pane for each key. + * + * <p>If the input {@link WindowingStrategy} does not always produce final panes, the assertion + * may be executed over an empty input even if the trigger has fired previously. To ensure that + * a final pane is always produced, set the {@link ClosingBehavior} of the windowing strategy + * (via {@link Window.Bound#withAllowedLateness(Duration, ClosingBehavior)} setting + * {@link ClosingBehavior} to {@link ClosingBehavior#FIRE_ALWAYS}). + * + * @return a new {@link SingletonAssert} like this one but with the assertion only applied to + * the specified window. + */ + SingletonAssert<T> inFinalPane(BoundedWindow window); + + /** + * Creates a new {@link SingletonAssert} like this one, but with the assertion restricted to + * only run on the provided window, running the checker only on the on-time pane for each key. + * * @return a new {@link SingletonAssert} like this one but with the assertion only applied to * the specified window. */ - SingletonAssert<T> inWindow(BoundedWindow window); + SingletonAssert<T> inOnTimePane(BoundedWindow window); /** * Asserts that the value in question is equal to the provided value, according to @@ -275,22 +345,49 @@ public class PAssert { private static class PCollectionContentsAssert<T> implements IterableAssert<T> { private final PCollection<T> actual; private final AssertionWindows rewindowingStrategy; + private final SimpleFunction<Iterable<WindowedValue<T>>, Iterable<T>> paneExtractor; public PCollectionContentsAssert(PCollection<T> actual) { - this(actual, IntoGlobalWindow.<T>of()); + this(actual, IntoGlobalWindow.<T>of(), PaneExtractors.<T>onlyPane()); } - public PCollectionContentsAssert(PCollection<T> actual, AssertionWindows rewindowingStrategy) { + public PCollectionContentsAssert( + PCollection<T> actual, + AssertionWindows rewindowingStrategy, + SimpleFunction<Iterable<WindowedValue<T>>, Iterable<T>> paneExtractor) { this.actual = actual; this.rewindowingStrategy = rewindowingStrategy; + this.paneExtractor = paneExtractor; } @Override public PCollectionContentsAssert<T> inWindow(BoundedWindow window) { + return withPane(window, PaneExtractors.<T>allPanes()); + } + + @Override + public PCollectionContentsAssert<T> inFinalPane(BoundedWindow window) { + return withPane(window, PaneExtractors.<T>finalPane()); + } + + @Override + public PCollectionContentsAssert<T> inOnTimePane(BoundedWindow window) { + return withPane(window, PaneExtractors.<T>onTimePane()); + } + + @Override + public PCollectionContentsAssert<T> inCombinedNonLatePanes(BoundedWindow window) { + return withPane(window, PaneExtractors.<T>nonLatePanes()); + } + + private PCollectionContentsAssert<T> withPane( + BoundedWindow window, + SimpleFunction<Iterable<WindowedValue<T>>, Iterable<T>> paneExtractor) { @SuppressWarnings({"unchecked", "rawtypes"}) Coder<BoundedWindow> windowCoder = (Coder) actual.getWindowingStrategy().getWindowFn().windowCoder(); - return new PCollectionContentsAssert<>(actual, IntoStaticWindows.<T>of(windowCoder, window)); + return new PCollectionContentsAssert<>( + actual, IntoStaticWindows.<T>of(windowCoder, window), paneExtractor); } /** @@ -323,7 +420,9 @@ public class PAssert { @Override public PCollectionContentsAssert<T> satisfies( SerializableFunction<Iterable<T>, Void> checkerFn) { - actual.apply(nextAssertionName(), new GroupThenAssert<>(checkerFn, rewindowingStrategy)); + actual.apply( + nextAssertionName(), + new GroupThenAssert<>(checkerFn, rewindowingStrategy, paneExtractor)); return this; } @@ -364,7 +463,8 @@ public class PAssert { SerializableFunction<Iterable<T>, Void> checkerFn = (SerializableFunction) new MatcherCheckerFn<>(matcher); actual.apply( - "PAssert$" + (assertCount++), new GroupThenAssert<>(checkerFn, rewindowingStrategy)); + "PAssert$" + (assertCount++), + new GroupThenAssert<>(checkerFn, rewindowingStrategy, paneExtractor)); return this; } @@ -414,13 +514,17 @@ public class PAssert { private final PCollection<Iterable<T>> actual; private final Coder<T> elementCoder; private final AssertionWindows rewindowingStrategy; + private final SimpleFunction<Iterable<WindowedValue<Iterable<T>>>, Iterable<Iterable<T>>> + paneExtractor; public PCollectionSingletonIterableAssert(PCollection<Iterable<T>> actual) { - this(actual, IntoGlobalWindow.<Iterable<T>>of()); + this(actual, IntoGlobalWindow.<Iterable<T>>of(), PaneExtractors.<Iterable<T>>onlyPane()); } public PCollectionSingletonIterableAssert( - PCollection<Iterable<T>> actual, AssertionWindows rewindowingStrategy) { + PCollection<Iterable<T>> actual, + AssertionWindows rewindowingStrategy, + SimpleFunction<Iterable<WindowedValue<Iterable<T>>>, Iterable<Iterable<T>>> paneExtractor) { this.actual = actual; @SuppressWarnings("unchecked") @@ -428,15 +532,37 @@ public class PAssert { this.elementCoder = typedCoder; this.rewindowingStrategy = rewindowingStrategy; + this.paneExtractor = paneExtractor; } @Override public PCollectionSingletonIterableAssert<T> inWindow(BoundedWindow window) { + return withPanes(window, PaneExtractors.<Iterable<T>>allPanes()); + } + + @Override + public PCollectionSingletonIterableAssert<T> inFinalPane(BoundedWindow window) { + return withPanes(window, PaneExtractors.<Iterable<T>>finalPane()); + } + + @Override + public PCollectionSingletonIterableAssert<T> inOnTimePane(BoundedWindow window) { + return withPanes(window, PaneExtractors.<Iterable<T>>onTimePane()); + } + + @Override + public PCollectionSingletonIterableAssert<T> inCombinedNonLatePanes(BoundedWindow window) { + return withPanes(window, PaneExtractors.<Iterable<T>>nonLatePanes()); + } + + private PCollectionSingletonIterableAssert<T> withPanes( + BoundedWindow window, + SimpleFunction<Iterable<WindowedValue<Iterable<T>>>, Iterable<Iterable<T>>> paneExtractor) { @SuppressWarnings({"unchecked", "rawtypes"}) Coder<BoundedWindow> windowCoder = (Coder) actual.getWindowingStrategy().getWindowFn().windowCoder(); return new PCollectionSingletonIterableAssert<>( - actual, IntoStaticWindows.<Iterable<T>>of(windowCoder, window)); + actual, IntoStaticWindows.<Iterable<T>>of(windowCoder, window), paneExtractor); } @Override @@ -460,7 +586,7 @@ public class PAssert { SerializableFunction<Iterable<T>, Void> checkerFn) { actual.apply( "PAssert$" + (assertCount++), - new GroupThenAssertForSingleton<>(checkerFn, rewindowingStrategy)); + new GroupThenAssertForSingleton<>(checkerFn, rewindowingStrategy, paneExtractor)); return this; } @@ -480,33 +606,53 @@ public class PAssert { private final PCollection<ElemT> actual; private final PTransform<PCollection<ElemT>, PCollectionView<ViewT>> view; private final AssertionWindows rewindowActuals; + private final SimpleFunction<Iterable<WindowedValue<ElemT>>, Iterable<ElemT>> paneExtractor; private final Coder<ViewT> coder; protected PCollectionViewAssert( PCollection<ElemT> actual, PTransform<PCollection<ElemT>, PCollectionView<ViewT>> view, Coder<ViewT> coder) { - this(actual, view, IntoGlobalWindow.<ElemT>of(), coder); + this(actual, view, IntoGlobalWindow.<ElemT>of(), PaneExtractors.<ElemT>onlyPane(), coder); } private PCollectionViewAssert( PCollection<ElemT> actual, PTransform<PCollection<ElemT>, PCollectionView<ViewT>> view, AssertionWindows rewindowActuals, + SimpleFunction<Iterable<WindowedValue<ElemT>>, Iterable<ElemT>> paneExtractor, Coder<ViewT> coder) { this.actual = actual; this.view = view; this.rewindowActuals = rewindowActuals; + this.paneExtractor = paneExtractor; this.coder = coder; } @Override - public PCollectionViewAssert<ElemT, ViewT> inWindow(BoundedWindow window) { + public PCollectionViewAssert<ElemT, ViewT> inOnlyPane(BoundedWindow window) { + return inPane(window, PaneExtractors.<ElemT>onlyPane()); + } + + @Override + public PCollectionViewAssert<ElemT, ViewT> inFinalPane(BoundedWindow window) { + return inPane(window, PaneExtractors.<ElemT>finalPane()); + } + + @Override + public PCollectionViewAssert<ElemT, ViewT> inOnTimePane(BoundedWindow window) { + return inPane(window, PaneExtractors.<ElemT>onTimePane()); + } + + private PCollectionViewAssert<ElemT, ViewT> inPane( + BoundedWindow window, + SimpleFunction<Iterable<WindowedValue<ElemT>>, Iterable<ElemT>> paneExtractor) { return new PCollectionViewAssert<>( actual, view, IntoStaticWindows.of( (Coder) actual.getWindowingStrategy().getWindowFn().windowCoder(), window), + paneExtractor, coder); } @@ -528,7 +674,7 @@ public class PAssert { .apply( "PAssert$" + (assertCount++), new OneSideInputAssert<ViewT>( - CreateActual.from(actual, rewindowActuals, view), + CreateActual.from(actual, rewindowActuals, paneExtractor, view), rewindowActuals.<Integer>windowDummy(), checkerFn)); return this; @@ -578,21 +724,25 @@ public class PAssert { private final transient PCollection<T> actual; private final transient AssertionWindows rewindowActuals; + private final transient SimpleFunction<Iterable<WindowedValue<T>>, Iterable<T>> extractPane; private final transient PTransform<PCollection<T>, PCollectionView<ActualT>> actualView; public static <T, ActualT> CreateActual<T, ActualT> from( PCollection<T> actual, AssertionWindows rewindowActuals, + SimpleFunction<Iterable<WindowedValue<T>>, Iterable<T>> extractPane, PTransform<PCollection<T>, PCollectionView<ActualT>> actualView) { - return new CreateActual<>(actual, rewindowActuals, actualView); + return new CreateActual<>(actual, rewindowActuals, extractPane, actualView); } private CreateActual( PCollection<T> actual, AssertionWindows rewindowActuals, + SimpleFunction<Iterable<WindowedValue<T>>, Iterable<T>> extractPane, PTransform<PCollection<T>, PCollectionView<ActualT>> actualView) { this.actual = actual; this.rewindowActuals = rewindowActuals; + this.extractPane = extractPane; this.actualView = actualView; } @@ -601,6 +751,10 @@ public class PAssert { final Coder<T> coder = actual.getCoder(); return actual .apply("FilterActuals", rewindowActuals.<T>prepareActuals()) + .apply("GatherPanes", GatherAllPanes.<T>globally()) + .apply("ExtractPane", MapElements.via(extractPane)) + .setCoder(IterableCoder.of(actual.getCoder())) + .apply(Flatten.<T>iterables()) .apply("RewindowActuals", rewindowActuals.<T>windowActuals()) .apply( ParDo.of( @@ -741,18 +895,22 @@ public class PAssert { implements Serializable { private final SerializableFunction<Iterable<T>, Void> checkerFn; private final AssertionWindows rewindowingStrategy; + private final SimpleFunction<Iterable<WindowedValue<T>>, Iterable<T>> paneExtractor; private GroupThenAssert( - SerializableFunction<Iterable<T>, Void> checkerFn, AssertionWindows rewindowingStrategy) { + SerializableFunction<Iterable<T>, Void> checkerFn, + AssertionWindows rewindowingStrategy, + SimpleFunction<Iterable<WindowedValue<T>>, Iterable<T>> paneExtractor) { this.checkerFn = checkerFn; this.rewindowingStrategy = rewindowingStrategy; + this.paneExtractor = paneExtractor; } @Override public PDone apply(PCollection<T> input) { input .apply("GroupGlobally", new GroupGlobally<T>(rewindowingStrategy)) - .apply("GetOnlyPane", ParDo.of(new ExtractOnlyPane<T>())) + .apply("GetPane", MapElements.via(paneExtractor)) .setCoder(IterableCoder.of(input.getCoder())) .apply("RunChecks", ParDo.of(new GroupedValuesCheckerDoFn<>(checkerFn))); @@ -768,18 +926,23 @@ public class PAssert { extends PTransform<PCollection<Iterable<T>>, PDone> implements Serializable { private final SerializableFunction<Iterable<T>, Void> checkerFn; private final AssertionWindows rewindowingStrategy; + private final SimpleFunction<Iterable<WindowedValue<Iterable<T>>>, Iterable<Iterable<T>>> + paneExtractor; private GroupThenAssertForSingleton( - SerializableFunction<Iterable<T>, Void> checkerFn, AssertionWindows rewindowingStrategy) { + SerializableFunction<Iterable<T>, Void> checkerFn, + AssertionWindows rewindowingStrategy, + SimpleFunction<Iterable<WindowedValue<Iterable<T>>>, Iterable<Iterable<T>>> paneExtractor) { this.checkerFn = checkerFn; this.rewindowingStrategy = rewindowingStrategy; + this.paneExtractor = paneExtractor; } @Override public PDone apply(PCollection<Iterable<T>> input) { input .apply("GroupGlobally", new GroupGlobally<Iterable<T>>(rewindowingStrategy)) - .apply("GetOnlyPane", ParDo.of(new ExtractOnlyPane<Iterable<T>>())) + .apply("GetPane", MapElements.via(paneExtractor)) .setCoder(IterableCoder.of(input.getCoder())) .apply("RunChecks", ParDo.of(new SingletonCheckerDoFn<>(checkerFn))); @@ -862,23 +1025,6 @@ public class PAssert { } } - private static class ExtractOnlyPane<T> extends DoFn<Iterable<WindowedValue<T>>, Iterable<T>> { - @Override - public void processElement(ProcessContext c) throws Exception { - List<T> outputs = new ArrayList<>(); - for (WindowedValue<T> value : c.element()) { - checkState( - value.getPane().isFirst() && value.getPane().isLast(), - "Expected elements to be produced by a trigger that fires at most once, but got" - + "a value in a pane that is %s. Actual Pane Info: %s", - value.getPane().isFirst() ? "not the last pane" : "not the first pane", - value.getPane()); - outputs.add(value.getValue()); - } - c.output(outputs); - } - } - /** * A {@link DoFn} that runs a checking {@link SerializableFunction} on the contents of * the single iterable element of the input {@link PCollection} and adjusts counters and http://git-wip-us.apache.org/repos/asf/incubator-beam/blob/ae0293d0/sdks/java/core/src/main/java/org/apache/beam/sdk/testing/PaneExtractors.java ---------------------------------------------------------------------- diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/testing/PaneExtractors.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/testing/PaneExtractors.java new file mode 100644 index 0000000..f699bfc --- /dev/null +++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/testing/PaneExtractors.java @@ -0,0 +1,140 @@ +/* + * 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.testing; + +import static com.google.common.base.Preconditions.checkState; + +import org.apache.beam.sdk.transforms.PTransform; +import org.apache.beam.sdk.transforms.SimpleFunction; +import org.apache.beam.sdk.transforms.windowing.PaneInfo; +import org.apache.beam.sdk.transforms.windowing.PaneInfo.Timing; +import org.apache.beam.sdk.util.WindowedValue; +import org.apache.beam.sdk.values.PCollection; +import org.apache.beam.sdk.values.TypeDescriptor; + +import java.util.ArrayList; +import java.util.List; + +/** + * {@link PTransform PTransforms} which take an {@link Iterable} of {@link WindowedValue + * WindowedValues} and outputs an {@link Iterable} of all values in the specified pane, dropping the + * {@link WindowedValue} metadata. + * + * <p>Although all of the method signatures return SimpleFunction, users should ensure to set the + * coder of any output {@link PCollection}, as appropriate {@link TypeDescriptor TypeDescriptors} + * cannot be obtained when the extractor is created. + */ +final class PaneExtractors { + private PaneExtractors() { + } + + static <T> SimpleFunction<Iterable<WindowedValue<T>>, Iterable<T>> onlyPane() { + return new ExtractOnlyPane<>(); + } + + static <T> SimpleFunction<Iterable<WindowedValue<T>>, Iterable<T>> onTimePane() { + return new ExtractOnTimePane<>(); + } + + static <T> SimpleFunction<Iterable<WindowedValue<T>>, Iterable<T>> finalPane() { + return new ExtractFinalPane<>(); + } + + static <T> SimpleFunction<Iterable<WindowedValue<T>>, Iterable<T>> nonLatePanes() { + return new ExtractNonLatePanes<>(); + } + + static <T> SimpleFunction<Iterable<WindowedValue<T>>, Iterable<T>> allPanes() { + return new ExtractAllPanes<>(); + } + + private static class ExtractOnlyPane<T> + extends SimpleFunction<Iterable<WindowedValue<T>>, Iterable<T>> { + @Override + public Iterable<T> apply(Iterable<WindowedValue<T>> input) { + List<T> outputs = new ArrayList<>(); + for (WindowedValue<T> value : input) { + checkState(value.getPane().isFirst() && value.getPane().isLast(), + "Expected elements to be produced by a trigger that fires at most once, but got" + + "a value in a pane that is %s. Actual Pane Info: %s", + value.getPane().isFirst() ? "not the last pane" : "not the first pane", + value.getPane()); + outputs.add(value.getValue()); + } + return outputs; + } + } + + + private static class ExtractOnTimePane<T> + extends SimpleFunction<Iterable<WindowedValue<T>>, Iterable<T>> { + @Override + public Iterable<T> apply(Iterable<WindowedValue<T>> input) { + List<T> outputs = new ArrayList<>(); + for (WindowedValue<T> value : input) { + if (value.getPane().getTiming().equals(Timing.ON_TIME)) { + outputs.add(value.getValue()); + } + } + return outputs; + } + } + + + private static class ExtractFinalPane<T> + extends SimpleFunction<Iterable<WindowedValue<T>>, Iterable<T>> { + @Override + public Iterable<T> apply(Iterable<WindowedValue<T>> input) { + List<T> outputs = new ArrayList<>(); + for (WindowedValue<T> value : input) { + if (value.getPane().isLast()) { + outputs.add(value.getValue()); + } + } + return outputs; + } + } + + + private static class ExtractAllPanes<T> + extends SimpleFunction<Iterable<WindowedValue<T>>, Iterable<T>> { + @Override + public Iterable<T> apply(Iterable<WindowedValue<T>> input) { + List<T> outputs = new ArrayList<>(); + for (WindowedValue<T> value : input) { + outputs.add(value.getValue()); + } + return outputs; + } + } + + + private static class ExtractNonLatePanes<T> + extends SimpleFunction<Iterable<WindowedValue<T>>, Iterable<T>> { + @Override + public Iterable<T> apply(Iterable<WindowedValue<T>> input) { + List<T> outputs = new ArrayList<>(); + for (WindowedValue<T> value : input) { + if (value.getPane().getTiming() != PaneInfo.Timing.LATE) { + outputs.add(value.getValue()); + } + } + return outputs; + } + } +} http://git-wip-us.apache.org/repos/asf/incubator-beam/blob/ae0293d0/sdks/java/core/src/test/java/org/apache/beam/sdk/testing/PAssertTest.java ---------------------------------------------------------------------- diff --git a/sdks/java/core/src/test/java/org/apache/beam/sdk/testing/PAssertTest.java b/sdks/java/core/src/test/java/org/apache/beam/sdk/testing/PAssertTest.java index bafd897..1070dab 100644 --- a/sdks/java/core/src/test/java/org/apache/beam/sdk/testing/PAssertTest.java +++ b/sdks/java/core/src/test/java/org/apache/beam/sdk/testing/PAssertTest.java @@ -282,10 +282,10 @@ public class PAssertTest implements Serializable { TimestampedValue.of(22, new Instant(-250L)))) .apply(Window.<Integer>into(FixedWindows.of(Duration.millis(500L)))); PAssert.thatSingleton(pcollection) - .inWindow(new IntervalWindow(new Instant(0L), new Instant(500L))) + .inOnlyPane(new IntervalWindow(new Instant(0L), new Instant(500L))) .isEqualTo(43); PAssert.thatSingleton(pcollection) - .inWindow(new IntervalWindow(new Instant(-500L), new Instant(0L))) + .inOnlyPane(new IntervalWindow(new Instant(-500L), new Instant(0L))) .isEqualTo(22); pipeline.run(); } http://git-wip-us.apache.org/repos/asf/incubator-beam/blob/ae0293d0/sdks/java/core/src/test/java/org/apache/beam/sdk/testing/PaneExtractorsTest.java ---------------------------------------------------------------------- diff --git a/sdks/java/core/src/test/java/org/apache/beam/sdk/testing/PaneExtractorsTest.java b/sdks/java/core/src/test/java/org/apache/beam/sdk/testing/PaneExtractorsTest.java new file mode 100644 index 0000000..1487eea --- /dev/null +++ b/sdks/java/core/src/test/java/org/apache/beam/sdk/testing/PaneExtractorsTest.java @@ -0,0 +1,323 @@ +/* + * 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.testing; + +import static org.hamcrest.Matchers.containsInAnyOrder; +import static org.hamcrest.Matchers.emptyIterable; +import static org.junit.Assert.assertThat; + +import org.apache.beam.sdk.transforms.SerializableFunction; +import org.apache.beam.sdk.transforms.windowing.GlobalWindow; +import org.apache.beam.sdk.transforms.windowing.PaneInfo; +import org.apache.beam.sdk.transforms.windowing.PaneInfo.Timing; +import org.apache.beam.sdk.util.WindowedValue; + +import com.google.common.collect.ImmutableList; + +import org.joda.time.Instant; +import org.junit.Rule; +import org.junit.Test; +import org.junit.rules.ExpectedException; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +/** + * Tests for {@link PaneExtractors}. + */ +@RunWith(JUnit4.class) +public class PaneExtractorsTest { + @Rule public ExpectedException thrown = ExpectedException.none(); + + @Test + public void onlyPaneNoFiring() { + SerializableFunction<Iterable<WindowedValue<Integer>>, Iterable<Integer>> extractor = + PaneExtractors.onlyPane(); + Iterable<WindowedValue<Integer>> noFiring = + ImmutableList.of( + WindowedValue.valueInGlobalWindow(9), WindowedValue.valueInEmptyWindows(19)); + assertThat(extractor.apply(noFiring), containsInAnyOrder(9, 19)); + } + + @Test + public void onlyPaneOnlyOneFiring() { + SerializableFunction<Iterable<WindowedValue<Integer>>, Iterable<Integer>> extractor = + PaneExtractors.onlyPane(); + Iterable<WindowedValue<Integer>> onlyFiring = + ImmutableList.of( + WindowedValue.of( + 2, new Instant(0L), GlobalWindow.INSTANCE, PaneInfo.ON_TIME_AND_ONLY_FIRING), + WindowedValue.of( + 1, new Instant(0L), GlobalWindow.INSTANCE, PaneInfo.ON_TIME_AND_ONLY_FIRING)); + + assertThat(extractor.apply(onlyFiring), containsInAnyOrder(2, 1)); + } + + @Test + public void onlyPaneMultiplePanesFails() { + SerializableFunction<Iterable<WindowedValue<Integer>>, Iterable<Integer>> extractor = + PaneExtractors.onlyPane(); + Iterable<WindowedValue<Integer>> multipleFiring = + ImmutableList.of( + WindowedValue.of( + 4, + new Instant(0L), + GlobalWindow.INSTANCE, + PaneInfo.createPane(true, false, Timing.EARLY)), + WindowedValue.of( + 2, + new Instant(0L), + GlobalWindow.INSTANCE, + PaneInfo.createPane(false, false, Timing.ON_TIME, 1L, 0L)), + WindowedValue.of( + 1, + new Instant(0L), + GlobalWindow.INSTANCE, + PaneInfo.createPane(false, false, Timing.LATE, 2L, 1L))); + + thrown.expect(IllegalStateException.class); + thrown.expectMessage("trigger that fires at most once"); + extractor.apply(multipleFiring); + } + + @Test + public void onTimePane() { + SerializableFunction<Iterable<WindowedValue<Integer>>, Iterable<Integer>> extractor = + PaneExtractors.onTimePane(); + Iterable<WindowedValue<Integer>> onlyOnTime = + ImmutableList.of( + WindowedValue.of( + 4, + new Instant(0L), + GlobalWindow.INSTANCE, + PaneInfo.createPane(false, false, Timing.ON_TIME, 1L, 0L)), + WindowedValue.of( + 2, + new Instant(0L), + GlobalWindow.INSTANCE, + PaneInfo.createPane(false, false, Timing.ON_TIME, 1L, 0L))); + + assertThat(extractor.apply(onlyOnTime), containsInAnyOrder(2, 4)); + } + + @Test + public void onTimePaneOnlyEarlyAndLate() { + SerializableFunction<Iterable<WindowedValue<Integer>>, Iterable<Integer>> extractor = + PaneExtractors.onTimePane(); + Iterable<WindowedValue<Integer>> onlyOnTime = + ImmutableList.of( + WindowedValue.of( + 8, + new Instant(0L), + GlobalWindow.INSTANCE, + PaneInfo.createPane(false, false, Timing.LATE, 2L, 1L)), + WindowedValue.of( + 4, + new Instant(0L), + GlobalWindow.INSTANCE, + PaneInfo.createPane(false, false, Timing.ON_TIME, 1L, 0L)), + WindowedValue.of( + 2, + new Instant(0L), + GlobalWindow.INSTANCE, + PaneInfo.createPane(false, false, Timing.ON_TIME, 1L, 0L)), + WindowedValue.of( + 1, + new Instant(0L), + GlobalWindow.INSTANCE, + PaneInfo.createPane(true, false, Timing.EARLY))); + + assertThat(extractor.apply(onlyOnTime), containsInAnyOrder(2, 4)); + } + + @Test + public void finalPane() { + SerializableFunction<Iterable<WindowedValue<Integer>>, Iterable<Integer>> extractor = + PaneExtractors.finalPane(); + Iterable<WindowedValue<Integer>> onlyOnTime = + ImmutableList.of( + WindowedValue.of( + 8, + new Instant(0L), + GlobalWindow.INSTANCE, + PaneInfo.createPane(false, true, Timing.LATE, 2L, 1L)), + WindowedValue.of( + 4, + new Instant(0L), + GlobalWindow.INSTANCE, + PaneInfo.createPane(false, false, Timing.ON_TIME, 1L, 0L)), + WindowedValue.of( + 1, + new Instant(0L), + GlobalWindow.INSTANCE, + PaneInfo.createPane(true, false, Timing.EARLY))); + + assertThat(extractor.apply(onlyOnTime), containsInAnyOrder(8)); + } + + @Test + public void finalPaneNoExplicitFinalEmpty() { + SerializableFunction<Iterable<WindowedValue<Integer>>, Iterable<Integer>> extractor = + PaneExtractors.finalPane(); + Iterable<WindowedValue<Integer>> onlyOnTime = + ImmutableList.of( + WindowedValue.of( + 8, + new Instant(0L), + GlobalWindow.INSTANCE, + PaneInfo.createPane(false, false, Timing.LATE, 2L, 1L)), + WindowedValue.of( + 4, + new Instant(0L), + GlobalWindow.INSTANCE, + PaneInfo.createPane(false, false, Timing.ON_TIME, 1L, 0L)), + WindowedValue.of( + 1, + new Instant(0L), + GlobalWindow.INSTANCE, + PaneInfo.createPane(true, false, Timing.EARLY))); + + assertThat(extractor.apply(onlyOnTime), emptyIterable()); + } + + @Test + public void nonLatePanesSingleOnTime() { + SerializableFunction<Iterable<WindowedValue<Integer>>, Iterable<Integer>> extractor = + PaneExtractors.nonLatePanes(); + Iterable<WindowedValue<Integer>> onlyOnTime = + ImmutableList.of( + WindowedValue.of( + 8, new Instant(0L), GlobalWindow.INSTANCE, PaneInfo.ON_TIME_AND_ONLY_FIRING), + WindowedValue.of( + 4, new Instant(0L), GlobalWindow.INSTANCE, PaneInfo.ON_TIME_AND_ONLY_FIRING), + WindowedValue.of( + 2, new Instant(0L), GlobalWindow.INSTANCE, PaneInfo.ON_TIME_AND_ONLY_FIRING)); + + assertThat(extractor.apply(onlyOnTime), containsInAnyOrder(2, 4, 8)); + } + + @Test + public void nonLatePanesSingleEarly() { + SerializableFunction<Iterable<WindowedValue<Integer>>, Iterable<Integer>> extractor = + PaneExtractors.nonLatePanes(); + Iterable<WindowedValue<Integer>> onlyOnTime = + ImmutableList.of( + WindowedValue.of( + 8, + new Instant(0L), + GlobalWindow.INSTANCE, + PaneInfo.createPane(true, false, Timing.EARLY)), + WindowedValue.of( + 4, + new Instant(0L), + GlobalWindow.INSTANCE, + PaneInfo.createPane(true, false, Timing.EARLY))); + + assertThat(extractor.apply(onlyOnTime), containsInAnyOrder(4, 8)); + } + + @Test + public void allPanesSingleLate() { + SerializableFunction<Iterable<WindowedValue<Integer>>, Iterable<Integer>> extractor = + PaneExtractors.nonLatePanes(); + Iterable<WindowedValue<Integer>> onlyOnTime = + ImmutableList.of( + WindowedValue.of( + 8, + new Instant(0L), + GlobalWindow.INSTANCE, + PaneInfo.createPane(false, false, Timing.LATE, 0L, 0L))); + + assertThat(extractor.apply(onlyOnTime), emptyIterable()); + } + + @Test + public void nonLatePanesMultiplePanes() { + SerializableFunction<Iterable<WindowedValue<Integer>>, Iterable<Integer>> extractor = + PaneExtractors.nonLatePanes(); + Iterable<WindowedValue<Integer>> onlyOnTime = + ImmutableList.of( + WindowedValue.of( + 8, + new Instant(0L), + GlobalWindow.INSTANCE, + PaneInfo.createPane(false, false, Timing.LATE, 2L, 1L)), + WindowedValue.of(7, new Instant(0L), GlobalWindow.INSTANCE, PaneInfo.NO_FIRING), + WindowedValue.of( + 4, + new Instant(0L), + GlobalWindow.INSTANCE, + PaneInfo.createPane(false, false, Timing.ON_TIME, 1L, 0L)), + WindowedValue.of( + 1, + new Instant(0L), + GlobalWindow.INSTANCE, + PaneInfo.createPane(true, false, Timing.EARLY))); + + assertThat(extractor.apply(onlyOnTime), containsInAnyOrder(4, 1, 7)); + } + + @Test + public void allPanesSinglePane() { + SerializableFunction<Iterable<WindowedValue<Integer>>, Iterable<Integer>> extractor = + PaneExtractors.allPanes(); + Iterable<WindowedValue<Integer>> onlyOnTime = + ImmutableList.of( + WindowedValue.of( + 8, new Instant(0L), GlobalWindow.INSTANCE, PaneInfo.ON_TIME_AND_ONLY_FIRING), + WindowedValue.of( + 4, new Instant(0L), GlobalWindow.INSTANCE, PaneInfo.ON_TIME_AND_ONLY_FIRING), + WindowedValue.of( + 2, new Instant(0L), GlobalWindow.INSTANCE, PaneInfo.ON_TIME_AND_ONLY_FIRING)); + + assertThat(extractor.apply(onlyOnTime), containsInAnyOrder(2, 4, 8)); + } + + @Test + public void allPanesMultiplePanes() { + SerializableFunction<Iterable<WindowedValue<Integer>>, Iterable<Integer>> extractor = + PaneExtractors.allPanes(); + Iterable<WindowedValue<Integer>> onlyOnTime = + ImmutableList.of( + WindowedValue.of( + 8, + new Instant(0L), + GlobalWindow.INSTANCE, + PaneInfo.createPane(false, false, Timing.LATE, 2L, 1L)), + WindowedValue.of( + 4, + new Instant(0L), + GlobalWindow.INSTANCE, + PaneInfo.createPane(false, false, Timing.ON_TIME, 1L, 0L)), + WindowedValue.of( + 1, + new Instant(0L), + GlobalWindow.INSTANCE, + PaneInfo.createPane(true, false, Timing.EARLY))); + + assertThat(extractor.apply(onlyOnTime), containsInAnyOrder(4, 8, 1)); + } + + @Test + public void allPanesEmpty() { + SerializableFunction<Iterable<WindowedValue<Integer>>, Iterable<Integer>> extractor = + PaneExtractors.allPanes(); + Iterable<WindowedValue<Integer>> noPanes = ImmutableList.of(); + + assertThat(extractor.apply(noPanes), emptyIterable()); + } +}
