beam-commits mailing list archives

Site index · List index
Message view « Date » · « Thread »
Top « Date » · « Thread »
From rober...@apache.org
Subject [4/4] incubator-beam git commit: Implement coder optimized for coding primitives.
Date Tue, 19 Jul 2016 00:50:21 GMT
Implement coder optimized for coding primitives.

This coder falls back to another coder (e.g. PickleCoder)
for other types, but is much faster for ints, strings, etc.


Project: http://git-wip-us.apache.org/repos/asf/incubator-beam/repo
Commit: http://git-wip-us.apache.org/repos/asf/incubator-beam/commit/f58c2bda
Tree: http://git-wip-us.apache.org/repos/asf/incubator-beam/tree/f58c2bda
Diff: http://git-wip-us.apache.org/repos/asf/incubator-beam/diff/f58c2bda

Branch: refs/heads/python-sdk
Commit: f58c2bdaae241718e0b4e250a90304317445613f
Parents: 69f895a
Author: Robert Bradshaw <robertwb@google.com>
Authored: Tue Jul 12 15:58:56 2016 -0700
Committer: Robert Bradshaw <robertwb@google.com>
Committed: Mon Jul 18 17:48:33 2016 -0700

----------------------------------------------------------------------
 sdks/python/apache_beam/coders/coder_impl.pxd   |  8 +++
 sdks/python/apache_beam/coders/coder_impl.py    | 64 ++++++++++++++++++++
 sdks/python/apache_beam/coders/coders.py        | 43 +++++++++++++
 .../apache_beam/coders/coders_test_common.py    |  8 +++
 4 files changed, 123 insertions(+)
----------------------------------------------------------------------


http://git-wip-us.apache.org/repos/asf/incubator-beam/blob/f58c2bda/sdks/python/apache_beam/coders/coder_impl.pxd
----------------------------------------------------------------------
diff --git a/sdks/python/apache_beam/coders/coder_impl.pxd b/sdks/python/apache_beam/coders/coder_impl.pxd
index 0c92c5b..ee1590d 100644
--- a/sdks/python/apache_beam/coders/coder_impl.pxd
+++ b/sdks/python/apache_beam/coders/coder_impl.pxd
@@ -59,6 +59,14 @@ cdef class DeterministicPickleCoderImpl(CoderImpl):
   cdef bint _check_safe(self, value) except -1
 
 
+cdef object NoneType
+cdef char UNKNOWN_TYPE, NONE_TYPE, INT_TYPE, FLOAT_TYPE
+cdef char STR_TYPE, UNICODE_TYPE, LIST_TYPE, TUPLE_TYPE
+
+cdef class FastPrimitivesCoderImpl(StreamCoderImpl):
+  cdef CoderImpl fallback_coder_impl
+
+
 cdef class BytesCoderImpl(CoderImpl):
   pass
 

http://git-wip-us.apache.org/repos/asf/incubator-beam/blob/f58c2bda/sdks/python/apache_beam/coders/coder_impl.py
----------------------------------------------------------------------
diff --git a/sdks/python/apache_beam/coders/coder_impl.py b/sdks/python/apache_beam/coders/coder_impl.py
index ca64d9c..473cff5 100644
--- a/sdks/python/apache_beam/coders/coder_impl.py
+++ b/sdks/python/apache_beam/coders/coder_impl.py
@@ -26,6 +26,7 @@ coder_impl.pxd file for type hints.
 """
 
 import collections
+from types import NoneType
 
 
 # pylint: disable=wrong-import-order, wrong-import-position, ungrouped-imports
@@ -147,6 +148,69 @@ class DeterministicPickleCoderImpl(CoderImpl):
     return self._pickle_coder.decode(encoded)
 
 
+UNKNOWN_TYPE = 0xFF
+NONE_TYPE = 0
+INT_TYPE = 1
+FLOAT_TYPE = 2
+STR_TYPE = 3
+UNICODE_TYPE = 4
+LIST_TYPE = 5
+TUPLE_TYPE = 6
+
+
+class FastPrimitivesCoderImpl(StreamCoderImpl):
+
+  def __init__(self, fallback_coder_impl):
+    self.fallback_coder_impl = fallback_coder_impl
+
+  def encode_to_stream(self, value, stream, nested):
+    t = type(value)
+    if t is NoneType:
+      stream.write_byte(NONE_TYPE)
+    elif t is int:
+      stream.write_byte(INT_TYPE)
+      stream.write_var_int64(value)
+    elif t is float:
+      stream.write_byte(FLOAT_TYPE)
+      stream.write_bigendian_double(value)
+    elif t is str:
+      stream.write_byte(STR_TYPE)
+      stream.write(value, nested)
+    elif t is unicode:
+      stream.write_byte(UNICODE_TYPE)
+      stream.write(value.encode('utf-8'), nested)
+    elif t is list or t is tuple:
+      stream.write_byte(LIST_TYPE if t is list else TUPLE_TYPE)
+      stream.write_var_int64(len(value))
+      for e in value:
+        self.encode_to_stream(e, stream, True)
+    else:
+      stream.write_byte(UNKNOWN_TYPE)
+      self.fallback_coder_impl.encode_to_stream(value, stream, nested)
+
+  def decode_from_stream(self, stream, nested):
+    t = stream.read_byte()
+    if t == NONE_TYPE:
+      return None
+    elif t == INT_TYPE:
+      return stream.read_var_int64()
+    elif t == FLOAT_TYPE:
+      return stream.read_bigendian_double()
+    elif t == STR_TYPE:
+      return stream.read_all(nested)
+    elif t == UNICODE_TYPE:
+      return stream.read_all(nested).decode('utf-8')
+    elif t == LIST_TYPE or t == TUPLE_TYPE:
+      vlen = stream.read_var_int64()
+      vlist = [self.decode_from_stream(stream, True) for _ in range(vlen)]
+      if t == LIST_TYPE:
+        return vlist
+      else:
+        return tuple(vlist)
+    else:
+      return self.fallback_coder_impl.decode_from_stream(stream, nested)
+
+
 class BytesCoderImpl(CoderImpl):
   """A coder for bytes/str objects."""
 

http://git-wip-us.apache.org/repos/asf/incubator-beam/blob/f58c2bda/sdks/python/apache_beam/coders/coders.py
----------------------------------------------------------------------
diff --git a/sdks/python/apache_beam/coders/coders.py b/sdks/python/apache_beam/coders/coders.py
index 619586f..cf5ca6d 100644
--- a/sdks/python/apache_beam/coders/coders.py
+++ b/sdks/python/apache_beam/coders/coders.py
@@ -354,6 +354,49 @@ class DeterministicPickleCoder(FastCoder):
     return self
 
 
+class FastPrimitivesCoder(FastCoder):
+  """Encodes simple primitives (e.g. str, int) efficiently.
+
+  For unknown types, falls back to another coder (e.g. PickleCoder).
+  """
+  def __init__(self, fallback_coder):
+    self._fallback_coder = fallback_coder
+
+  def _create_impl(self):
+    return coder_impl.FastPrimitivesCoderImpl(
+        self._fallback_coder.get_impl())
+
+  def is_deterministic(self):
+    return self._fallback_coder.is_deterministic()
+
+  def as_cloud_object(self, is_pair_like=True):
+    value = super(FastCoder, self).as_cloud_object()
+    # We currently use this coder in places where we cannot infer the coder to
+    # use for the value type in a more granular way.  In places where the
+    # service expects a pair, it checks for the "is_pair_like" key, in which
+    # case we would fail without the hack below.
+    if is_pair_like:
+      value['is_pair_like'] = True
+      value['component_encodings'] = [
+          self.as_cloud_object(is_pair_like=False),
+          self.as_cloud_object(is_pair_like=False)
+      ]
+
+    return value
+
+  # We allow .key_coder() and .value_coder() to be called on PickleCoder since
+  # we can't always infer the return values of lambdas in ParDo operations, the
+  # result of which may be used in a GroupBykey.
+  def is_kv_coder(self):
+    return True
+
+  def key_coder(self):
+    return self
+
+  def value_coder(self):
+    return self
+
+
 class Base64PickleCoder(Coder):
   """Coder of objects by Python pickle, then base64 encoding."""
   # TODO(robertwb): Do base64 encoding where it's needed (e.g. in json) rather

http://git-wip-us.apache.org/repos/asf/incubator-beam/blob/f58c2bda/sdks/python/apache_beam/coders/coders_test_common.py
----------------------------------------------------------------------
diff --git a/sdks/python/apache_beam/coders/coders_test_common.py b/sdks/python/apache_beam/coders/coders_test_common.py
index 0266fdc..b084947 100644
--- a/sdks/python/apache_beam/coders/coders_test_common.py
+++ b/sdks/python/apache_beam/coders/coders_test_common.py
@@ -111,6 +111,14 @@ class CodersTest(unittest.TestCase):
         coders.TupleCoder((coders.VarIntCoder(), coders.DillCoder())),
         (1, cell_value))
 
+  def test_fast_primitives_coder(self):
+    coder = coders.FastPrimitivesCoder(coders.SingletonCoder(len))
+    self.check_coder(coder, None, 1, -1, 1.5, 'str\0str', u'unicode\0\u0101')
+    self.check_coder(coder, (), (1, 2, 3))
+    self.check_coder(coder, [], [1, 2, 3])
+    self.check_coder(coder, len)
+    self.check_coder(coders.TupleCoder((coder,)), ('a',), (1,))
+
   def test_bytes_coder(self):
     self.check_coder(coders.BytesCoder(), 'a', '\0', 'z' * 1000)
 


Mime
View raw message