github-actions[bot] commented on code in PR #41216:
URL: https://github.com/apache/beam/pull/41216#discussion_r4541522011


##########
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:
   **🟡 Medium**: Because `_coder_impls` is typed as `cdef tuple`, indexing 
`self._coder_impls[i]` returns an untyped Python `object` rather than 
`CoderImpl`. Previously, `for i, c in enumerate(self._coder_impls)` bound `c` 
to `@cython.locals(c=CoderImpl)` so `c.decode_from_stream(...)` used Cython's 
C-level `cpdef` vtable dispatch. Casting with `cython.cast(CoderImpl, 
self._coder_impls[i])` preserves the typed C call for non-2-component tuples.
   
   ```suggestion
           cython.cast(CoderImpl, 
self._coder_impls[i]).decode_from_stream(in_stream, nested or i < last)
   ```



##########
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:
   **🟢 Low**: `get_estimated_size_and_observables` in `coder_impl.py` was not 
updated to define `last`, so `last=int` in `@cython.locals` here is unused (and 
`c=CoderImpl` on line 163 is also unused if `decode_from_stream` does not bind 
a local variable `c`).
   
   ```suggestion
     @cython.locals(c=CoderImpl, i=int)
   ```



-- 
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]

Reply via email to