You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@beam.apache.org by ro...@apache.org on 2017/09/19 22:23:52 UTC
[3/3] beam git commit: Allow subclasses of tuple, list,
and dict as pvaluish inputs/outputs.
Allow subclasses of tuple, list, and dict as pvaluish inputs/outputs.
Project: http://git-wip-us.apache.org/repos/asf/beam/repo
Commit: http://git-wip-us.apache.org/repos/asf/beam/commit/72960b31
Tree: http://git-wip-us.apache.org/repos/asf/beam/tree/72960b31
Diff: http://git-wip-us.apache.org/repos/asf/beam/diff/72960b31
Branch: refs/heads/master
Commit: 72960b31843d1dcdf2b43a55db0797a15f48ef18
Parents: 28d4f09
Author: Robert Bradshaw <ro...@gmail.com>
Authored: Fri Sep 8 17:53:09 2017 -0700
Committer: Robert Bradshaw <ro...@gmail.com>
Committed: Tue Sep 19 15:23:28 2017 -0700
----------------------------------------------------------------------
sdks/python/apache_beam/pipeline.py | 2 +-
.../python/apache_beam/transforms/ptransform.py | 62 ++++++++++----------
.../apache_beam/transforms/ptransform_test.py | 16 +++++
3 files changed, 48 insertions(+), 32 deletions(-)
----------------------------------------------------------------------
http://git-wip-us.apache.org/repos/asf/beam/blob/72960b31/sdks/python/apache_beam/pipeline.py
----------------------------------------------------------------------
diff --git a/sdks/python/apache_beam/pipeline.py b/sdks/python/apache_beam/pipeline.py
index 1ebd099..c670978 100644
--- a/sdks/python/apache_beam/pipeline.py
+++ b/sdks/python/apache_beam/pipeline.py
@@ -438,7 +438,7 @@ class Pipeline(object):
if type_options is not None and type_options.pipeline_type_check:
transform.type_check_outputs(pvalueish_result)
- for result in ptransform.GetPValues().visit(pvalueish_result):
+ for result in ptransform.get_nested_pvalues(pvalueish_result):
assert isinstance(result, (pvalue.PValue, pvalue.DoOutputsTuple))
# Make sure we set the producer only for a leaf node in the transform DAG.
http://git-wip-us.apache.org/repos/asf/beam/blob/72960b31/sdks/python/apache_beam/transforms/ptransform.py
----------------------------------------------------------------------
diff --git a/sdks/python/apache_beam/transforms/ptransform.py b/sdks/python/apache_beam/transforms/ptransform.py
index eccaccd..f630977 100644
--- a/sdks/python/apache_beam/transforms/ptransform.py
+++ b/sdks/python/apache_beam/transforms/ptransform.py
@@ -73,27 +73,25 @@ class _PValueishTransform(object):
This visits a PValueish, contstructing a (possibly mutated) copy.
"""
- def visit(self, node, *args):
- return getattr(
- self,
- 'visit_' + node.__class__.__name__,
- lambda x, *args: x)(node, *args)
-
- def visit_list(self, node, *args):
- return [self.visit(x, *args) for x in node]
-
- def visit_tuple(self, node, *args):
- return tuple(self.visit(x, *args) for x in node)
-
- def visit_dict(self, node, *args):
- return {key: self.visit(value, *args) for (key, value) in node.items()}
+ def visit_nested(self, node, *args):
+ if isinstance(node, (tuple, list)):
+ # namedtuples require unpacked arguments in their constructor,
+ # but do have a _make method that takes a sequence.
+ return getattr(node.__class__, '_make', node.__class__)(
+ [self.visit(x, *args) for x in node])
+ elif isinstance(node, dict):
+ return node.__class__(
+ {key: self.visit(value, *args) for (key, value) in node.items()})
+ else:
+ return node
class _SetInputPValues(_PValueishTransform):
def visit(self, node, replacements):
if id(node) in replacements:
return replacements[id(node)]
- return super(_SetInputPValues, self).visit(node, replacements)
+ else:
+ return self.visit_nested(node, replacements)
class _MaterializedDoOutputsTuple(pvalue.DoOutputsTuple):
@@ -116,22 +114,25 @@ class _MaterializePValues(_PValueishTransform):
return self._pvalue_cache.get_unwindowed_pvalue(node)
elif isinstance(node, pvalue.DoOutputsTuple):
return _MaterializedDoOutputsTuple(node, self._pvalue_cache)
- return super(_MaterializePValues, self).visit(node)
+ else:
+ return self.visit_nested(node)
-class GetPValues(_PValueishTransform):
- def visit(self, node, pvalues=None):
- if pvalues is None:
- pvalues = []
- self.visit(node, pvalues)
- return pvalues
- elif isinstance(node, (pvalue.PValue, pvalue.DoOutputsTuple)):
+class _GetPValues(_PValueishTransform):
+ def visit(self, node, pvalues):
+ if isinstance(node, (pvalue.PValue, pvalue.DoOutputsTuple)):
pvalues.append(node)
else:
- super(GetPValues, self).visit(node, pvalues)
+ self.visit_nested(node, pvalues)
+
+
+def get_nested_pvalues(pvalueish):
+ pvalues = []
+ _GetPValues().visit(pvalueish, pvalues)
+ return pvalues
-class _ZipPValues(_PValueishTransform):
+class _ZipPValues(object):
"""Pairs each PValue in a pvalueish with a value in a parallel out sibling.
Sibling should have the same nested structure as pvalueish. Leaves in
@@ -153,10 +154,12 @@ class _ZipPValues(_PValueishTransform):
return pairs
elif isinstance(pvalueish, (pvalue.PValue, pvalue.DoOutputsTuple)):
pairs.append((context, pvalueish, sibling))
- else:
- super(_ZipPValues, self).visit(pvalueish, sibling, pairs, context)
+ elif isinstance(pvalueish, (list, tuple)):
+ self.visit_sequence(pvalueish, sibling, pairs, context)
+ elif isinstance(pvalueish, dict):
+ self.visit_dict(pvalueish, sibling, pairs, context)
- def visit_list(self, pvalueish, sibling, pairs, context):
+ def visit_sequence(self, pvalueish, sibling, pairs, context):
if isinstance(sibling, (list, tuple)):
for ix, (p, s) in enumerate(zip(
pvalueish, list(sibling) + [None] * len(pvalueish))):
@@ -165,9 +168,6 @@ class _ZipPValues(_PValueishTransform):
for p in pvalueish:
self.visit(p, sibling, pairs, context)
- def visit_tuple(self, pvalueish, sibling, pairs, context):
- self.visit_list(pvalueish, sibling, pairs, context)
-
def visit_dict(self, pvalueish, sibling, pairs, context):
if isinstance(sibling, dict):
for key, p in pvalueish.items():
http://git-wip-us.apache.org/repos/asf/beam/blob/72960b31/sdks/python/apache_beam/transforms/ptransform_test.py
----------------------------------------------------------------------
diff --git a/sdks/python/apache_beam/transforms/ptransform_test.py b/sdks/python/apache_beam/transforms/ptransform_test.py
index 435270e..c328cb1 100644
--- a/sdks/python/apache_beam/transforms/ptransform_test.py
+++ b/sdks/python/apache_beam/transforms/ptransform_test.py
@@ -19,6 +19,7 @@
from __future__ import absolute_import
+import collections
import operator
import re
import unittest
@@ -670,6 +671,21 @@ class PTransformTest(unittest.TestCase):
self.assertEqual(['x', 'x', 'y', 'y', 'z'], sorted(res['b']))
self.assertEqual([], sorted(res['c']))
+ def test_named_tuple(self):
+ MinMax = collections.namedtuple('MinMax', ['min', 'max'])
+
+ class MinMaxTransform(PTransform):
+ def expand(self, pcoll):
+ return MinMax(
+ min=pcoll | beam.CombineGlobally(min).without_defaults(),
+ max=pcoll | beam.CombineGlobally(max).without_defaults())
+ res = [1, 2, 4, 8] | MinMaxTransform()
+ self.assertIsInstance(res, MinMax)
+ self.assertEqual(res, MinMax(min=[1], max=[8]))
+
+ flat = res | beam.Flatten()
+ self.assertEqual(sorted(flat), [1, 8])
+
@beam.ptransform_fn
def SamplePTransform(pcoll):