| #!/usr/bin/env python |
| """A simple reader for file segs produced by GCS output writer.""" |
| |
| from mapreduce import output_writers |
| |
| # pylint: disable=protected-access |
| # pylint: disable=invalid-name |
| |
| # pylint: disable=g-import-not-at-top |
| # TODO(user): Cleanup imports if/when cloudstorage becomes part of runtime. |
| try: |
| # Check if the full cloudstorage package exists. The stub part is in runtime. |
| import cloudstorage |
| if hasattr(cloudstorage, "_STUB"): |
| cloudstorage = None |
| except ImportError: |
| pass # CloudStorage library not available |
| |
| |
| class _GCSFileSegReader(object): |
| """A simple reader for file segs produced by GCS output writer. |
| |
| Internal use only. |
| |
| This reader conforms to Python stream interface. |
| """ |
| |
| def __init__(self, seg_prefix, last_seg_index): |
| """Init. |
| |
| Instances are pickle safe. |
| |
| Args: |
| seg_prefix: filename prefix for all segs. It is expected |
| seg_prefix + index = seg filename. |
| last_seg_index: the last index of all segs. int. |
| """ |
| self._EOF = False |
| self._offset = 0 |
| |
| # fields related to seg. |
| self._seg_prefix = seg_prefix |
| self._last_seg_index = last_seg_index |
| self._seg_index = -1 |
| self._seg_valid_length = None |
| self._seg = None |
| self._next_seg() |
| |
| def read(self, n): |
| """Read data from file segs. |
| |
| Args: |
| n: max bytes to read. Must be positive. |
| |
| Returns: |
| some bytes. May be smaller than n bytes. "" when no more data is left. |
| """ |
| if self._EOF: |
| return "" |
| |
| while self._seg_index <= self._last_seg_index: |
| result = self._read_from_seg(n) |
| if result != "": |
| return result |
| else: |
| self._next_seg() |
| |
| self._EOF = True |
| return "" |
| |
| def close(self): |
| if self._seg: |
| self._seg.close() |
| |
| def tell(self): |
| """Returns the next offset to read.""" |
| return self._offset |
| |
| def _next_seg(self): |
| """Get next seg.""" |
| if self._seg: |
| self._seg.close() |
| self._seg_index += 1 |
| if self._seg_index > self._last_seg_index: |
| self._seg = None |
| return |
| |
| filename = self._seg_prefix + str(self._seg_index) |
| stat = cloudstorage.stat(filename) |
| writer = output_writers._GoogleCloudStorageOutputWriter |
| if writer._VALID_LENGTH not in stat.metadata: |
| raise ValueError( |
| "Expect %s in metadata for file %s." % |
| (writer._VALID_LENGTH, filename)) |
| self._seg_valid_length = int(stat.metadata[writer._VALID_LENGTH]) |
| if self._seg_valid_length > stat.st_size: |
| raise ValueError( |
| "Valid length %s is too big for file %s of length %s" % |
| (self._seg_valid_length, filename, stat.st_size)) |
| self._seg = cloudstorage.open(filename) |
| |
| def _read_from_seg(self, n): |
| """Read from current seg. |
| |
| Args: |
| n: max number of bytes to read. |
| |
| Returns: |
| valid bytes from the current seg. "" if no more is left. |
| """ |
| result = self._seg.read(size=n) |
| if result == "": |
| return result |
| offset = self._seg.tell() |
| if offset > self._seg_valid_length: |
| extra = offset - self._seg_valid_length |
| result = result[:-1*extra] |
| self._offset += len(result) |
| return result |
