[BEAM-1872] Add IdentityWindowFn for use in Reshuffle #4040

udim · 2017-10-26T02:05:08Z

BEAM-1872 IdentityWindowFn is intended for internal use in a future
implementation of Reshuffle.

Add and pass current window to WindowFn.AssignContext, for
IdentityWindowFn implementation.

Add WindowedValueMatcher in testing/util.py, for use in IdentityWindowFn
unit tests.

Hi @robertwb. Can you please take a look?

chamikaramj · 2017-10-26T18:33:41Z

robertwb

I'm wondering if IdentityWindowFn is the right thing to use for Reshuffle. Instead, it seems it'd be better to reify the windows into the data, letting the windows be single-timestamp-windows, then after Reshuffle re-aligning the windows with what existed before (thinking about streaming in particular).

Is there a usecase for this other than Reshuffle? While I generally favor breaking changes up, it might make more sense to review this in the context where it's used.

robertwb · 2017-10-26T21:14:44Z

sdks/python/apache_beam/testing/util.py

+    super(_WindowedValueMatcherDoFn, self).__init__()
+    self.expected_elements = copy.copy(expected_elements)
+
+  def process(self, element, timestamp=DoFn.TimestampParam,


This will not work if there are multiple bundles.

I would propose that WindowedValueMatcher instead be implemented on top of assert_that, where the windows and timestamps are reified. Similar to how reify_windows is used in https://github.com/apache/beam/blob/master/sdks/python/apache_beam/transforms/window_test.py . If we want to make this a full util, we should probably provide a similar API to assert_that.

I believe this code does work with multiple bundles. It "Verifies that all processed elements are in the list of expected
elements."

That being said, I could try to rewrite this to use assert_that.

I should clarify, it's use in WindowedValueMatcher is not correct in the presence of multiple bundles. E.g. suppose the expected values were [a, b] and the actual values were [a, a] split into two bundles.

udim

PTAL. This is by no means complete, but I'd like to get some feedback from you Robert before I go further.
This implementation is based on https://github.com/apache/beam/blob/master/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/Reshuffle.java

robertwb · 2017-11-15T18:40:05Z

sdks/python/apache_beam/pvalue.py

@@ -124,6 +124,11 @@ def windowing(self):
          self.producer.inputs)
    return self._windowing

+  # TODO(ehudm): Make this internal.
+  @windowing.setter


As discussed, remove this setter.

robertwb · 2017-11-15T18:47:18Z

sdks/python/apache_beam/transforms/util.py

+    super(IdentityWindowFn, self).__init__()
+    if coder is None:
+      raise ValueError('coder should not be None')
+    self._coder = coder


Rename this window_coder?

robertwb · 2017-11-15T18:47:52Z

sdks/python/apache_beam/transforms/util.py

@@ -423,3 +431,115 @@ def expand(self, pcoll):
          self._batch_size_estimator))
    else:
      return pcoll | ParDo(_WindowAwareBatchingDoFn(self._batch_size_estimator))
+
+
+class IdentityWindowFn(NonMergingWindowFn):


Make this private (with an underscore)?

robertwb · 2017-11-15T18:48:09Z

sdks/python/apache_beam/transforms/util.py

+  def get_window_coder(self):
+    return self._coder
+
+class TriggerForEveryElement(TriggerFn):


Just use AfterCount(1).

robertwb · 2017-11-15T18:54:57Z

sdks/python/apache_beam/transforms/util.py

+# TODO: are these typehints necessary?
+@typehints.with_input_types(typehints.KV[K, V])
+@typehints.with_output_types(typehints.KV[K, V])
+class Reshuffle(PTransform):


As mentioned, rename this ReshufflePerKey, and add a Reshuffle that appends then strips a random key (e.g. random.getrandbits(32))

robertwb · 2017-11-15T18:55:59Z

sdks/python/apache_beam/transforms/util.py

+
+    # TODO: is it safe to reapply this value?
+    windowing_saved = pcoll.windowing
+    # TODO: add .with_input_types, .with_output_types to PTransforms below?


It should be able to infer.

robertwb · 2017-11-15T18:58:55Z

sdks/python/apache_beam/transforms/util.py

+  """
+
+  def expand(self, pcoll):
+    class ExpandIterableDoFn(DoFn):


A DoFn with nothing but a process method can be more simply implemented via beam.Map or beam.FlatMap.

robertwb · 2017-11-15T18:59:50Z

sdks/python/apache_beam/transforms/util.py

+                )
+            | 'GroupByKey' >> GroupByKey()
+            | 'ExpandIterable' >> ParDo(ExpandIterableDoFn()))
+    pcoll_intermediate.windowing = windowing_saved


Nit: I'd probably apply assign windowing as the very last thing rather than on this intermediate.

robertwb · 2017-11-15T19:03:37Z

sdks/python/apache_beam/testing/util.py

+      1. A matcher function taking as argument the actual value of a
+      materialized PCollection. The matcher validates this actual value
+      against expectations and raises BeamAssertException if they are not met.
+      2. An instance of WindowedValueMatcher.


Rather than branching here an introducing a new WindowedValueMatcher, perhaps introduce a windowed_equal_to sibling to equal_to.

I thought of doing that, but how does assert_that know if matcher accepts values or TestWindowedValues?

robertwb · 2017-11-15T19:05:12Z

sdks/python/apache_beam/testing/util.py

@@ -46,6 +56,26 @@ class BeamAssertException(Exception):
  pass


+# Used for reifying timestamps and windows for assert_that matchers.


I think we can get around the need for introducing this new class by pairing values with None internally on reification (or at least making it an internal implementation detail).

The idea was to use this class for user-implemented matchers (inheriting from WindowedValueMatcher). Am I wrong to assume that there are matchers defined outside this repo?

Ideally one should be able to use any callable, e.g. hamcrest matchers, rather than have to implement a windowed variant of each. I misspoke about having a windowed_equals_to, we should have a assert_that_windowed(pcoll, equal_to([WindowedValue(...), ...])), or assert_that(pcoll, equal_to([WindowedValue(...), ...], reify_windows=True).

udim · 2017-11-28T23:22:33Z

sdks/python/apache_beam/transforms/util.py

+  """
+
+  def expand(self, pcoll):
+    class ExpandIterableDoFn(DoFn):


robertwb wrote:
A DoFn with nothing but a process method can be more simply implemented via beam.Map or beam.FlatMap.

Done.

udim · 2017-11-28T23:22:33Z

sdks/python/apache_beam/transforms/util.py

@@ -423,3 +431,115 @@ def expand(self, pcoll):
          self._batch_size_estimator))
    else:
      return pcoll | ParDo(_WindowAwareBatchingDoFn(self._batch_size_estimator))
+
+
+class IdentityWindowFn(NonMergingWindowFn):


robertwb wrote:
Make this private (with an underscore)?

Done.

udim · 2017-11-28T23:22:33Z

sdks/python/apache_beam/testing/util.py

@@ -46,6 +56,26 @@ class BeamAssertException(Exception):
  pass


+# Used for reifying timestamps and windows for assert_that matchers.


robertwb wrote:
Ideally one should be able to use any callable, e.g. hamcrest matchers, rather than have to implement a windowed variant of each. I misspoke about having a windowed_equals_to, we should have a assert_that_windowed(pcoll, equal_to([WindowedValue(...), ...])), or assert_that(pcoll, equal_to([WindowedValue(...), ...], reify_windows=True).

Done.

udim · 2017-11-28T23:22:33Z

sdks/python/apache_beam/transforms/util.py

+    super(IdentityWindowFn, self).__init__()
+    if coder is None:
+      raise ValueError('coder should not be None')
+    self._coder = coder


robertwb wrote:
Rename this window_coder?

Done.

udim · 2017-11-28T23:22:33Z

sdks/python/apache_beam/pvalue.py

@@ -124,6 +124,11 @@ def windowing(self):
          self.producer.inputs)
    return self._windowing

+  # TODO(ehudm): Make this internal.
+  @windowing.setter


robertwb wrote:
As discussed, remove this setter.

Done.

udim · 2017-11-28T23:22:33Z

sdks/python/apache_beam/testing/util.py

+      1. A matcher function taking as argument the actual value of a
+      materialized PCollection. The matcher validates this actual value
+      against expectations and raises BeamAssertException if they are not met.
+      2. An instance of WindowedValueMatcher.


udim wrote:
I thought of doing that, but how does assert_that know if matcher accepts values or TestWindowedValues?

Acknowledged.

udim · 2017-11-28T23:22:33Z

sdks/python/apache_beam/transforms/util.py

+  def get_window_coder(self):
+    return self._coder
+
+class TriggerForEveryElement(TriggerFn):


robertwb wrote:
Just use AfterCount(1).

Done.

udim · 2017-11-28T23:22:33Z

sdks/python/apache_beam/transforms/util.py

+
+    # TODO: is it safe to reapply this value?
+    windowing_saved = pcoll.windowing
+    # TODO: add .with_input_types, .with_output_types to PTransforms below?


robertwb wrote:
It should be able to infer.

Done.

udim · 2017-11-28T23:22:33Z

sdks/python/apache_beam/transforms/util.py

+# TODO: are these typehints necessary?
+@typehints.with_input_types(typehints.KV[K, V])
+@typehints.with_output_types(typehints.KV[K, V])
+class Reshuffle(PTransform):


robertwb wrote:
As mentioned, rename this ReshufflePerKey, and add a Reshuffle that appends then strips a random key (e.g. random.getrandbits(32))

Done.

udim · 2017-11-28T23:22:34Z

sdks/python/apache_beam/transforms/util.py

+                )
+            | 'GroupByKey' >> GroupByKey()
+            | 'ExpandIterable' >> ParDo(ExpandIterableDoFn()))
+    pcoll_intermediate.windowing = windowing_saved


robertwb wrote:
Nit: I'd probably apply assign windowing as the very last thing rather than on this intermediate.

Done.

udim · 2017-11-29T18:57:37Z

retest this please

robertwb

Looking good, just a couple more comments.

robertwb · 2017-11-29T22:52:16Z

sdks/python/apache_beam/testing/util_test.py

  def test_assert_that_fails(self):
-    with self.assertRaises(Exception):
+    with self.assertRaises(BeamAssertException):


Not all runners preserve exceptions from remote workers. Please change back. (Same elsewhere.)

:/ I changed this because my code had a bug that raised another exception.

robertwb · 2017-11-29T22:56:06Z

sdks/python/apache_beam/transforms/util.py

+  """
+
+  def expand(self, pcoll):
+    class ReifyTimestampsIn(DoFn):


Reify means to make explicit. Should just be ReifyTimestamps.

robertwb · 2017-11-29T22:56:41Z

sdks/python/apache_beam/transforms/util.py

+      def process(self, element, timestamp=DoFn.TimestampParam):
+        if (isinstance(timestamp, type(DoFn.TimestampParam)) and
+            timestamp == DoFn.TimestampParam):
+          raise ValueError('timestamp was unset for element: %r' % element)


We shouldn't have to check for this.

robertwb · 2017-11-29T22:58:33Z

sdks/python/apache_beam/transforms/util.py

+
+    class ReifyTimestampsExtract(DoFn):
+      def process(self, element, window=DoFn.WindowParam):
+        # Return a WindowedValue so that IdentityWindowFn can reuse the window


I'm not following this comment...

Shouldn't things already be in the correct window? We just need to emit a TimestampedValue here.

robertwb · 2017-11-29T22:58:45Z

sdks/python/apache_beam/transforms/util.py

+          raise ValueError('timestamp was unset for element: %r' % element)
+        yield element[0], TimestampedValue(element[1], timestamp)
+
+    class ReifyTimestampsExtract(DoFn):


ApplyReifiedTimestamps?

robertwb · 2017-11-29T23:34:39Z

sdks/python/apache_beam/transforms/util.py

+
+    windowing_saved = pcoll.windowing
+    result = (pcoll
+              | 'ReifyTimestampsIn' >> ParDo(ReifyTimestampsIn())


Most of these explicit stage names are redundant with what the default would already be.

robertwb · 2017-11-29T23:35:10Z

sdks/python/apache_beam/transforms/util.py

+    return (pcoll
+            | 'AddRandomKeys' >> Map(lambda t: (random.getrandbits(32), t))
+            | ReshufflePerKey()
+            | 'RemoveTempKeys' >> Map(lambda t: t[1]))


RemoveRandomKeys (for symmetry).

Two flavors of Reshuffle: ReshufflePerKey operates on key-value pairs, while Reshuffle adds a random key to each element (key-value or other). Add _IdentityWindowFn, for internal use in Reshuffle. Add and pass current window to WindowFn.AssignContext, for IdentityWindowFn implementation. testing/util.py: - Extend assert_that with reify_windows keyword, allowing verification of timestamp values and windowing functions. - Add contains_in_any_order matcher.

udim · 2017-12-01T04:02:23Z

sdks/python/apache_beam/transforms/util.py

+
+    class ReifyTimestampsExtract(DoFn):
+      def process(self, element, window=DoFn.WindowParam):
+        # Return a WindowedValue so that IdentityWindowFn can reuse the window


robertwb wrote:
I'm not following this comment...

Shouldn't things already be in the correct window? We just need to emit a TimestampedValue here.

Passing a TimestampedValue would make the output processor for DoFn call the windowing function.
https://github.com/apache/beam/blob/master/sdks/python/apache_beam/runners/common.py#L451

Good point. Maybe call this RestoreTimestamps or RestoreWindows?

udim · 2017-12-01T04:02:23Z

sdks/python/apache_beam/transforms/util.py

+
+    windowing_saved = pcoll.windowing
+    result = (pcoll
+              | 'ReifyTimestampsIn' >> ParDo(ReifyTimestampsIn())


robertwb wrote:
Most of these explicit stage names are redundant with what the default would already be.

Done.

udim · 2017-12-01T04:02:23Z

sdks/python/apache_beam/transforms/util.py

+    return (pcoll
+            | 'AddRandomKeys' >> Map(lambda t: (random.getrandbits(32), t))
+            | ReshufflePerKey()
+            | 'RemoveTempKeys' >> Map(lambda t: t[1]))


robertwb wrote:
RemoveRandomKeys (for symmetry).

Done.

udim · 2017-12-01T04:02:23Z

sdks/python/apache_beam/transforms/util.py

+          raise ValueError('timestamp was unset for element: %r' % element)
+        yield element[0], TimestampedValue(element[1], timestamp)
+
+    class ReifyTimestampsExtract(DoFn):


robertwb wrote:
ApplyReifiedTimestamps?

Done.

udim · 2017-12-01T04:02:23Z

sdks/python/apache_beam/transforms/util.py

+  """
+
+  def expand(self, pcoll):
+    class ReifyTimestampsIn(DoFn):


robertwb wrote:
Reify means to make explicit. Should just be ReifyTimestamps.

Done.

udim · 2017-12-01T04:02:23Z

sdks/python/apache_beam/transforms/util.py

+      def process(self, element, timestamp=DoFn.TimestampParam):
+        if (isinstance(timestamp, type(DoFn.TimestampParam)) and
+            timestamp == DoFn.TimestampParam):
+          raise ValueError('timestamp was unset for element: %r' % element)


robertwb wrote:
We shouldn't have to check for this.

Done.

robertwb · 2017-12-02T00:57:36Z

Looks good. Jenkins is complaining.

Jenkins: retest this please.

udim · 2017-12-04T20:59:53Z

sdks/python/apache_beam/transforms/util.py

+
+    class ReifyTimestampsExtract(DoFn):
+      def process(self, element, window=DoFn.WindowParam):
+        # Return a WindowedValue so that IdentityWindowFn can reuse the window


robertwb wrote:
Good point. Maybe call this RestoreTimestamps or RestoreWindows?

Done.

robertwb · 2017-12-06T21:30:40Z

LGTM, thanks.

* Implement Reshuffle for Python SDK. Two flavors of Reshuffle: ReshufflePerKey operates on key-value pairs, while Reshuffle adds a random key to each element (key-value or other). Add _IdentityWindowFn, for internal use in Reshuffle. Add and pass current window to WindowFn.AssignContext, for IdentityWindowFn implementation. testing/util.py: - Extend assert_that with reify_windows keyword, allowing verification of timestamp values and windowing functions. - Add contains_in_any_order matcher.

robertwb reviewed Oct 26, 2017

View reviewed changes

udim force-pushed the reshuffle_py branch from cb5bd6d to 4fa4caa Compare October 31, 2017 23:55

udim commented Oct 31, 2017

View reviewed changes

robertwb reviewed Nov 15, 2017

View reviewed changes

udim force-pushed the reshuffle_py branch from 4fa4caa to f35e212 Compare November 28, 2017 23:20

udim commented Nov 28, 2017

View reviewed changes

robertwb reviewed Nov 29, 2017

View reviewed changes

udim added 2 commits November 30, 2017 16:45

Addressed review comments.

a7602e1

udim force-pushed the reshuffle_py branch from f35e212 to a7602e1 Compare December 1, 2017 04:02

udim commented Dec 1, 2017

View reviewed changes

Address review comment.

cae7c6d

udim commented Dec 4, 2017

View reviewed changes

robertwb merged commit 9c1eeb2 into apache:master Dec 6, 2017

udim deleted the reshuffle_py branch December 8, 2017 00:00

		@@ -46,6 +56,26 @@ class BeamAssertException(Exception):
		pass


		# Used for reifying timestamps and windows for assert_that matchers.

[BEAM-1872] Add IdentityWindowFn for use in Reshuffle #4040

[BEAM-1872] Add IdentityWindowFn for use in Reshuffle #4040

Conversation

udim commented Oct 26, 2017

chamikaramj commented Oct 26, 2017

robertwb left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

udim left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

udim commented Nov 29, 2017

robertwb left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

robertwb commented Dec 2, 2017

Choose a reason for hiding this comment

robertwb commented Dec 6, 2017