sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/windowing/WindowingTest.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.sdk.transforms.windowing;

 import java.io.File;
 import java.io.FileOutputStream;
 import java.io.PrintStream;
 import java.io.Serializable;
 import java.util.List;
 import org.apache.beam.sdk.coders.StringUtf8Coder;
 import org.apache.beam.sdk.io.TextIO;
 import org.apache.beam.sdk.testing.NeedsRunner;
 import org.apache.beam.sdk.testing.PAssert;
 import org.apache.beam.sdk.testing.TestPipeline;
 import org.apache.beam.sdk.testing.ValidatesRunner;
 import org.apache.beam.sdk.transforms.Count;
 import org.apache.beam.sdk.transforms.Create;
 import org.apache.beam.sdk.transforms.DoFn;
 import org.apache.beam.sdk.transforms.Flatten;
 import org.apache.beam.sdk.transforms.PTransform;
 import org.apache.beam.sdk.transforms.ParDo;
 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.TimestampedValue;
 import org.apache.beam.vendor.guava.v20_0.com.google.common.base.Splitter;
 import org.joda.time.Duration;
 import org.joda.time.Instant;
 import org.junit.Rule;
 import org.junit.Test;
 import org.junit.experimental.categories.Category;
 import org.junit.rules.TemporaryFolder;
 import org.junit.runner.RunWith;
 import org.junit.runners.JUnit4;

 /** Unit tests for bucketing. */
 @RunWith(JUnit4.class)
 public class WindowingTest implements Serializable {

   @Rule public final transient TestPipeline p = TestPipeline.create();

   @Rule public transient TemporaryFolder tmpFolder = new TemporaryFolder();

   private static class WindowedCount extends PTransform<PCollection<String>, PCollection<String>> {

     private static final class FormatCountsDoFn extends DoFn<KV<String, Long>, String> {
       @ProcessElement
       public void processElement(ProcessContext c, BoundedWindow window) {
         c.output(
             c.element().getKey()
                 + ":"
                 + c.element().getValue()
                 + ":"
                 + c.timestamp().getMillis()
                 + ":"
                 + window);
       }
     }

     private WindowFn<? super String, ?> windowFn;

     public WindowedCount(WindowFn<? super String, ?> windowFn) {
       this.windowFn = windowFn;
     }

     @Override
     public PCollection<String> expand(PCollection<String> in) {
       return in.apply(
               "Window",
               Window.<String>into(windowFn).withTimestampCombiner(TimestampCombiner.EARLIEST))
           .apply(Count.perElement())
           .apply("FormatCounts", ParDo.of(new FormatCountsDoFn()))
           .setCoder(StringUtf8Coder.of());
     }
   }

   private String output(String value, int count, int timestamp, int windowStart, int windowEnd) {
     return value
         + ":"
         + count
         + ":"
         + timestamp
         + ":["
         + new Instant(windowStart)
         + ".."
         + new Instant(windowEnd)
         + ")";
   }

   @Test
   @Category(ValidatesRunner.class)
   public void testPartitioningWindowing() {
     PCollection<String> input =
         p.apply(
             Create.timestamped(
                 TimestampedValue.of("a", new Instant(1)),
                 TimestampedValue.of("b", new Instant(2)),
                 TimestampedValue.of("b", new Instant(3)),
                 TimestampedValue.of("c", new Instant(11)),
                 TimestampedValue.of("d", new Instant(11))));

     PCollection<String> output = input.apply(new WindowedCount(FixedWindows.of(new Duration(10))));

     PAssert.that(output)
         .containsInAnyOrder(
             output("a", 1, 1, 0, 10),
             output("b", 2, 2, 0, 10),
             output("c", 1, 11, 10, 20),
             output("d", 1, 11, 10, 20));

     p.run();
   }

   @Test
   @Category(ValidatesRunner.class)
   public void testNonPartitioningWindowing() {
     PCollection<String> input =
         p.apply(
             Create.timestamped(
                 TimestampedValue.of("a", new Instant(1)),
                 TimestampedValue.of("a", new Instant(7)),
                 TimestampedValue.of("b", new Instant(8))));

     PCollection<String> output =
         input.apply(new WindowedCount(SlidingWindows.of(new Duration(10)).every(new Duration(5))));

     PAssert.that(output)
         .containsInAnyOrder(
             output("a", 1, 1, -5, 5),
             output("a", 2, 5, 0, 10),
             output("a", 1, 10, 5, 15),
             output("b", 1, 8, 0, 10),
             output("b", 1, 10, 5, 15));

     p.run();
   }

   @Test
   @Category(ValidatesRunner.class)
   public void testMergingWindowing() {
     PCollection<String> input =
         p.apply(
             Create.timestamped(
                 TimestampedValue.of("a", new Instant(1)),
                 TimestampedValue.of("a", new Instant(5)),
                 TimestampedValue.of("a", new Instant(20))));

     PCollection<String> output =
         input.apply(new WindowedCount(Sessions.withGapDuration(new Duration(10))));

     PAssert.that(output).containsInAnyOrder(output("a", 2, 1, 1, 15), output("a", 1, 20, 20, 30));

     p.run();
   }

   @Test
   @Category(ValidatesRunner.class)
   public void testWindowPreservation() {
     PCollection<String> input1 =
         p.apply(
             "Create12",
             Create.timestamped(
                 TimestampedValue.of("a", new Instant(1)),
                 TimestampedValue.of("b", new Instant(2))));

     PCollection<String> input2 =
         p.apply(
             "Create34",
             Create.timestamped(
                 TimestampedValue.of("a", new Instant(3)),
                 TimestampedValue.of("b", new Instant(4))));

     PCollectionList<String> input = PCollectionList.of(input1).and(input2);

     PCollection<String> output =
         input
             .apply(Flatten.pCollections())
             .apply(new WindowedCount(FixedWindows.of(new Duration(5))));

     PAssert.that(output).containsInAnyOrder(output("a", 2, 1, 0, 5), output("b", 2, 2, 0, 5));

     p.run();
   }

   @Test
   @Category(NeedsRunner.class)
   public void testEmptyInput() {
     PCollection<String> input = p.apply(Create.empty(StringUtf8Coder.of()));

     PCollection<String> output = input.apply(new WindowedCount(FixedWindows.of(new Duration(10))));

     PAssert.that(output).empty();

     p.run();
   }

   @Test
   @Category(NeedsRunner.class)
   public void testTextIoInput() throws Exception {
     File tmpFile = tmpFolder.newFile("file.txt");
     String filename = tmpFile.getPath();

     try (PrintStream writer = new PrintStream(new FileOutputStream(tmpFile))) {
       writer.println("a 1");
       writer.println("b 2");
       writer.println("b 3");
       writer.println("c 11");
       writer.println("d 11");
     }

     PCollection<String> output =
         p.begin()
             .apply("ReadLines", TextIO.read().from(filename))
             .apply(ParDo.of(new ExtractWordsWithTimestampsFn()))
             .apply(new WindowedCount(FixedWindows.of(Duration.millis(10))));

     PAssert.that(output)
         .containsInAnyOrder(
             output("a", 1, 1, 0, 10),
             output("b", 2, 2, 0, 10),
             output("c", 1, 11, 10, 20),
             output("d", 1, 11, 10, 20));

     p.run();
   }

   /** A DoFn that tokenizes lines of text into individual words. */
   static class ExtractWordsWithTimestampsFn extends DoFn<String, String> {
     @ProcessElement
     public void processElement(ProcessContext c) {
       List<String> words = Splitter.onPattern("[^a-zA-Z0-9']+").splitToList(c.element());
       if (words.size() == 2) {
         c.outputWithTimestamp(words.get(0), new Instant(Long.parseLong(words.get(1))));
       }
     }
   }
 }
	/*
	* 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.transforms.windowing;

	import java.io.File;
	import java.io.FileOutputStream;
	import java.io.PrintStream;
	import java.io.Serializable;
	import java.util.List;
	import org.apache.beam.sdk.coders.StringUtf8Coder;
	import org.apache.beam.sdk.io.TextIO;
	import org.apache.beam.sdk.testing.NeedsRunner;
	import org.apache.beam.sdk.testing.PAssert;
	import org.apache.beam.sdk.testing.TestPipeline;
	import org.apache.beam.sdk.testing.ValidatesRunner;
	import org.apache.beam.sdk.transforms.Count;
	import org.apache.beam.sdk.transforms.Create;
	import org.apache.beam.sdk.transforms.DoFn;
	import org.apache.beam.sdk.transforms.Flatten;
	import org.apache.beam.sdk.transforms.PTransform;
	import org.apache.beam.sdk.transforms.ParDo;
	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.TimestampedValue;
	import org.apache.beam.vendor.guava.v20_0.com.google.common.base.Splitter;
	import org.joda.time.Duration;
	import org.joda.time.Instant;
	import org.junit.Rule;
	import org.junit.Test;
	import org.junit.experimental.categories.Category;
	import org.junit.rules.TemporaryFolder;
	import org.junit.runner.RunWith;
	import org.junit.runners.JUnit4;

	/** Unit tests for bucketing. */
	@RunWith(JUnit4.class)
	public class WindowingTest implements Serializable {

	@Rule public final transient TestPipeline p = TestPipeline.create();

	@Rule public transient TemporaryFolder tmpFolder = new TemporaryFolder();

	private static class WindowedCount extends PTransform<PCollection<String>, PCollection<String>> {

	private static final class FormatCountsDoFn extends DoFn<KV<String, Long>, String> {
	@ProcessElement
	public void processElement(ProcessContext c, BoundedWindow window) {
	c.output(
	c.element().getKey()
	+ ":"
	+ c.element().getValue()
	+ ":"
	+ c.timestamp().getMillis()
	+ ":"
	+ window);
	}
	}

	private WindowFn<? super String, ?> windowFn;

	public WindowedCount(WindowFn<? super String, ?> windowFn) {
	this.windowFn = windowFn;
	}

	@Override
	public PCollection<String> expand(PCollection<String> in) {
	return in.apply(
	"Window",
	Window.<String>into(windowFn).withTimestampCombiner(TimestampCombiner.EARLIEST))
	.apply(Count.perElement())
	.apply("FormatCounts", ParDo.of(new FormatCountsDoFn()))
	.setCoder(StringUtf8Coder.of());
	}
	}

	private String output(String value, int count, int timestamp, int windowStart, int windowEnd) {
	return value
	+ ":"
	+ count
	+ ":"
	+ timestamp
	+ ":["
	+ new Instant(windowStart)
	+ ".."
	+ new Instant(windowEnd)
	+ ")";
	}

	@Test
	@Category(ValidatesRunner.class)
	public void testPartitioningWindowing() {
	PCollection<String> input =
	p.apply(
	Create.timestamped(
	TimestampedValue.of("a", new Instant(1)),
	TimestampedValue.of("b", new Instant(2)),
	TimestampedValue.of("b", new Instant(3)),
	TimestampedValue.of("c", new Instant(11)),
	TimestampedValue.of("d", new Instant(11))));

	PCollection<String> output = input.apply(new WindowedCount(FixedWindows.of(new Duration(10))));

	PAssert.that(output)
	.containsInAnyOrder(
	output("a", 1, 1, 0, 10),
	output("b", 2, 2, 0, 10),
	output("c", 1, 11, 10, 20),
	output("d", 1, 11, 10, 20));

	p.run();
	}

	@Test
	@Category(ValidatesRunner.class)
	public void testNonPartitioningWindowing() {
	PCollection<String> input =
	p.apply(
	Create.timestamped(
	TimestampedValue.of("a", new Instant(1)),
	TimestampedValue.of("a", new Instant(7)),
	TimestampedValue.of("b", new Instant(8))));

	PCollection<String> output =
	input.apply(new WindowedCount(SlidingWindows.of(new Duration(10)).every(new Duration(5))));

	PAssert.that(output)
	.containsInAnyOrder(
	output("a", 1, 1, -5, 5),
	output("a", 2, 5, 0, 10),
	output("a", 1, 10, 5, 15),
	output("b", 1, 8, 0, 10),
	output("b", 1, 10, 5, 15));

	p.run();
	}

	@Test
	@Category(ValidatesRunner.class)
	public void testMergingWindowing() {
	PCollection<String> input =
	p.apply(
	Create.timestamped(
	TimestampedValue.of("a", new Instant(1)),
	TimestampedValue.of("a", new Instant(5)),
	TimestampedValue.of("a", new Instant(20))));

	PCollection<String> output =
	input.apply(new WindowedCount(Sessions.withGapDuration(new Duration(10))));

	PAssert.that(output).containsInAnyOrder(output("a", 2, 1, 1, 15), output("a", 1, 20, 20, 30));

	p.run();
	}

	@Test
	@Category(ValidatesRunner.class)
	public void testWindowPreservation() {
	PCollection<String> input1 =
	p.apply(
	"Create12",
	Create.timestamped(
	TimestampedValue.of("a", new Instant(1)),
	TimestampedValue.of("b", new Instant(2))));

	PCollection<String> input2 =
	p.apply(
	"Create34",
	Create.timestamped(
	TimestampedValue.of("a", new Instant(3)),
	TimestampedValue.of("b", new Instant(4))));

	PCollectionList<String> input = PCollectionList.of(input1).and(input2);

	PCollection<String> output =
	input
	.apply(Flatten.pCollections())
	.apply(new WindowedCount(FixedWindows.of(new Duration(5))));

	PAssert.that(output).containsInAnyOrder(output("a", 2, 1, 0, 5), output("b", 2, 2, 0, 5));

	p.run();
	}

	@Test
	@Category(NeedsRunner.class)
	public void testEmptyInput() {
	PCollection<String> input = p.apply(Create.empty(StringUtf8Coder.of()));

	PCollection<String> output = input.apply(new WindowedCount(FixedWindows.of(new Duration(10))));

	PAssert.that(output).empty();

	p.run();
	}

	@Test
	@Category(NeedsRunner.class)
	public void testTextIoInput() throws Exception {
	File tmpFile = tmpFolder.newFile("file.txt");
	String filename = tmpFile.getPath();

	try (PrintStream writer = new PrintStream(new FileOutputStream(tmpFile))) {
	writer.println("a 1");
	writer.println("b 2");
	writer.println("b 3");
	writer.println("c 11");
	writer.println("d 11");
	}

	PCollection<String> output =
	p.begin()
	.apply("ReadLines", TextIO.read().from(filename))
	.apply(ParDo.of(new ExtractWordsWithTimestampsFn()))
	.apply(new WindowedCount(FixedWindows.of(Duration.millis(10))));

	PAssert.that(output)
	.containsInAnyOrder(
	output("a", 1, 1, 0, 10),
	output("b", 2, 2, 0, 10),
	output("c", 1, 11, 10, 20),
	output("d", 1, 11, 10, 20));

	p.run();
	}

	/** A DoFn that tokenizes lines of text into individual words. */
	static class ExtractWordsWithTimestampsFn extends DoFn<String, String> {
	@ProcessElement
	public void processElement(ProcessContext c) {
	List<String> words = Splitter.onPattern("[^a-zA-Z0-9']+").splitToList(c.element());
	if (words.size() == 2) {
	c.outputWithTimestamp(words.get(0), new Instant(Long.parseLong(words.get(1))));
	}
	}
	}
	}