damccorm commented on code in PR #41216:
URL: https://github.com/apache/beam/pull/41216#discussion_r4542033301
##########
sdks/python/apache_beam/coders/coder_impl.py:
##########
@@ -1158,19 +1166,27 @@ def _construct_from_components(self, components):
def encode_to_stream(self, value, out, nested):
# type: (Any, create_OutputStream, bool) -> None
values = self._extract_components(value)
- if len(self._coder_impls) != len(values):
+ if self._num_coders != len(values):
raise ValueError('Number of components does not match number of coders.')
- for i in range(0, len(self._coder_impls)):
+ if self._is_kv:
+ self._coder_0.encode_to_stream(values[0], out, True)
+ self._coder_1.encode_to_stream(values[1], out, nested)
+ return
+ last = self._num_coders - 1
+ for i in range(0, self._num_coders):
c = self._coder_impls[i] # type cast
- c.encode_to_stream(
- values[i], out, nested or i + 1 < len(self._coder_impls))
+ c.encode_to_stream(values[i], out, nested or i < last)
def decode_from_stream(self, in_stream, nested):
# type: (create_InputStream, bool) -> Any
+ if self._is_kv:
+ return self._construct_from_components((
+ self._coder_0.decode_from_stream(in_stream, True),
+ self._coder_1.decode_from_stream(in_stream, nested)))
+ last = self._num_coders - 1
return self._construct_from_components([
- c.decode_from_stream(
- in_stream, nested or i + 1 < len(self._coder_impls))
- for i, c in enumerate(self._coder_impls)
+ self._coder_impls[i].decode_from_stream(in_stream, nested or i < last)
Review Comment:
Fixed in 7e134ed9a91: kept `for i, c in enumerate(self._coder_impls)` so `c`
remains typed as `CoderImpl` via `@cython.locals(c=CoderImpl, i=int, last=int)`.
##########
sdks/python/apache_beam/coders/coder_impl.pxd:
##########
@@ -150,16 +150,20 @@ cdef class SingletonCoderImpl(CoderImpl):
cdef class AbstractComponentCoderImpl(StreamCoderImpl):
cdef tuple _coder_impls
+ cdef int _num_coders
+ cdef bint _is_kv
+ cdef CoderImpl _coder_0
+ cdef CoderImpl _coder_1
cpdef _extract_components(self, value)
cpdef _construct_from_components(self, components)
- @cython.locals(c=CoderImpl)
+ @cython.locals(c=CoderImpl, i=int, last=int)
cpdef encode_to_stream(self, value, OutputStream stream, bint nested)
- @cython.locals(c=CoderImpl)
+ @cython.locals(c=CoderImpl, i=int, last=int)
cpdef decode_from_stream(self, InputStream stream, bint nested)
- @cython.locals(c=CoderImpl)
+ @cython.locals(c=CoderImpl, i=int, last=int)
Review Comment:
Fixed in 7e134ed9a91: reverted `@cython.locals` on
`get_estimated_size_and_observables` back to `@cython.locals(c=CoderImpl)`.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
