y1chi commented on a change in pull request #14460:
URL: https://github.com/apache/beam/pull/14460#discussion_r637716952
##########
File path: sdks/python/apache_beam/io/mongodbio.py
##########
@@ -236,60 +325,109 @@ def split(self, desired_bundle_size,
start_position=None, stop_position=None):
weight=weight,
source=self,
start_position=bundle_start,
- stop_position=stop_position)
+ stop_position=stop_position,
+ )
- def get_range_tracker(self, start_position, stop_position):
+ def get_range_tracker(
+ self,
+ start_position: Union[int, str, ObjectId] = None,
+ stop_position: Union[int, str, ObjectId] = None,
+ ) -> Union[
+ _ObjectIdRangeTracker, OffsetRangeTracker, LexicographicKeyRangeTracker]:
+ """Returns a RangeTracker for a given position range depending on type.
+
+ Args:
+ start_position: starting position of the range. If 'None' default start
+ position of the source must be used.
+ stop_position: ending position of the range. If 'None' default stop
+ position of the source must be used.
+ Returns:
+ a ``_ObjectIdRangeTracker``, ``OffsetRangeTracker``
+ or ``LexicographicKeyRangeTracker`` depending on the given position
range.
+ """
start_position, stop_position = self._replace_none_positions(
- start_position, stop_position)
- return _ObjectIdRangeTracker(start_position, stop_position)
+ start_position, stop_position
+ )
+
+ if isinstance(start_position, ObjectId):
+ return _ObjectIdRangeTracker(start_position, stop_position)
+
+ if isinstance(start_position, int):
+ return OffsetRangeTracker(start_position, stop_position)
+
+ if isinstance(start_position, str):
+ return LexicographicKeyRangeTracker(start_position, stop_position)
Review comment:
I think LexicographicKeyRangeTracker will deal with bytes, so you need
to encode the start_position and stop_position when feeding it into the range
tracker, and decode the position back into the str once read from the range
tracker.
--
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.
For queries about this service, please contact Infrastructure at:
[email protected]
