TheNeuralBit commented on code in PR #21844:
URL: https://github.com/apache/beam/pull/21844#discussion_r897115826
##########
sdks/python/apache_beam/transforms/core.py:
##########
@@ -746,10 +762,18 @@ def get_input_batch_type(self) ->
typing.Optional[TypeConstraint]:
# TODO(BEAM-14340): Consider supporting an alternative (dynamic?)
approach
# for declaring input type
raise TypeError(
- f"{self.__class__.__name__}.process_batch() does not have a type "
- "annotation on its first parameter. This is required for "
- "process_batch implementations.")
- return typehints.native_type_compatibility.convert_to_beam_type(input_type)
+ f"Either {self.__class__.__name__}.process_batch() must have a type "
+ f"annotation on its first parameter, or {self.__class__.__name__} "
+ "must override get_input_batch_type.")
+ return input_type
+
+ def _get_input_batch_type_normalized(self, input_element_type):
+ return typehints.native_type_compatibility.convert_to_beam_type(
+ self.get_input_batch_type(input_element_type))
+
+ def _get_output_batch_type_normalized(self, input_element_type):
+ return typehints.native_type_compatibility.convert_to_beam_type(
+ self.get_output_batch_type(input_element_type))
Review Comment:
These are convenience functions I provided for our internal use, users
shouldn't call them. Users shouldn't call the others
(`get_{input,output}_batch_type`) either - but they are part of the public API
since users can override them if they need to.
Come to think of it I should probably mark some other convenience functions
we added as protected. I'll follow up with a PR for that.
##########
sdks/python/apache_beam/runners/portability/fn_api_runner/fn_runner_test.py:
##########
@@ -136,6 +136,37 @@ def test_batch_pardo(self):
assert_that(res, equal_to([6, 12, 18]))
+ def test_batch_pardo_override_type_inference(self):
+
+ class ArrayMultiplyTransposedDoFn(beam.DoFn):
+
Review Comment:
Done
--
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]
