diff --git a/sdks/python/apache_beam/io/gcp/gcsio.py b/sdks/python/apache_beam/io/gcp/gcsio.py index fc9a0cc1fd42..a673845811da 100644 --- a/sdks/python/apache_beam/io/gcp/gcsio.py +++ b/sdks/python/apache_beam/io/gcp/gcsio.py @@ -67,6 +67,11 @@ _LOGGER = logging.getLogger(__name__) DEFAULT_READ_BUFFER_SIZE = 16 * 1024 * 1024 +DEFAULT_WRITE_BUFFER_SIZE = 16 * 1024 * 1024 + +# Writes are performed as resumable uploads, which require the chunk size to be +# a multiple of 256 KiB. +WRITE_BUFFER_SIZE_MULTIPLE = 256 * 1024 # Maximum number of operations permitted in GcsIO.copy_batch() and # GcsIO.delete_batch(). @@ -240,6 +245,14 @@ def __init__( gcsio_retry.get_retry(pipeline_options) if GCS_INSTALLED else None) self._use_blob_generation = getattr( google_cloud_options, 'enable_gcsio_blob_generation', False) + self._read_buffer_size = getattr( + google_cloud_options, 'gcs_read_buffer_size_bytes', None) + if self._read_buffer_size is None: + self._read_buffer_size = DEFAULT_READ_BUFFER_SIZE + self._write_buffer_size = getattr( + google_cloud_options, 'gcs_write_buffer_size_bytes', None) + if self._write_buffer_size is None: + self._write_buffer_size = DEFAULT_WRITE_BUFFER_SIZE def get_project_number(self, bucket): if bucket not in self.bucket_to_project_number: @@ -286,15 +299,23 @@ def open( self, filename, mode='r', - read_buffer_size=DEFAULT_READ_BUFFER_SIZE, - mime_type='application/octet-stream'): + read_buffer_size=None, + mime_type='application/octet-stream', + write_buffer_size=None): """Open a GCS file path for reading or writing. Args: filename (str): GCS file path in the form ``gs:///``. mode (str): ``'r'`` for reading or ``'w'`` for writing. read_buffer_size (int): Buffer size to use during read operations. + Defaults to the value of the ``--gcs_read_buffer_size_bytes`` + pipeline option, or ``DEFAULT_READ_BUFFER_SIZE`` when that option is + not set. mime_type (str): Mime type to set for write operations. + write_buffer_size (int): Buffer size to use during write operations. + Must be a multiple of 256 KiB. Defaults to the value of the + ``--gcs_write_buffer_size_bytes`` pipeline option, or + ``DEFAULT_WRITE_BUFFER_SIZE`` when that option is not set. Returns: GCS file object. @@ -302,6 +323,11 @@ def open( Raises: ValueError: Invalid open file mode. """ + if read_buffer_size is None: + read_buffer_size = self._read_buffer_size + if write_buffer_size is None: + write_buffer_size = self._write_buffer_size + bucket_name, blob_name = parse_gcs_path(filename) bucket = self.client.bucket(bucket_name) @@ -317,6 +343,7 @@ def open( return BeamBlobWriter( blob, mime_type, + chunk_size=write_buffer_size, enable_write_bucket_metric=self.enable_write_bucket_metric, retry=self._storage_client_retry) else: @@ -705,6 +732,19 @@ def is_soft_delete_enabled(self, gcs_path): class BeamBlobReader(BlobReader): + """A reader for GCS blobs. + + Note that constructing this reader does not issue any request to GCS. Object + metadata is fetched lazily by the underlying ``BlobReader``, on the first + read or seek. + + Known limitation: doubly compressed objects, i.e. those stored with both + "content-encoding=gzip" and "content-type=application/gzip" (or + "application/x-gzip"), are not supported. Detecting this up front would + require an extra metadata request on every open, which is too costly to do + unconditionally. See + https://github.com/googleapis/google-cloud-python/issues/18423. + """ def __init__( self, blob, @@ -717,21 +757,6 @@ def __init__( # (https://cloud.google.com/storage/docs/transcoding). super().__init__( blob, chunk_size=chunk_size, retry=retry, raw_download=raw_download) - # TODO: Remove this after - # https://github.com/googleapis/python-storage/issues/1406 is fixed. - # As a workaround, we manually trigger a reload here. Otherwise, an internal - # call of reader.seek() will cause an exception if raw_download is set - # when initializing BlobReader(), - blob.reload() - - # TODO: Currently there is a bug in GCS server side when a client requests - # a file with "content-encoding=gzip" and "content-type=application/gzip" or - # "content-type=application/x-gzip", which will lead to infinite loop. - # We skip the support of this type of files until the GCS bug is fixed. - # Internal bug id: 203845981. - if (blob.content_encoding == "gzip" and - blob.content_type in ["application/gzip", "application/x-gzip"]): - raise NotImplementedError("Doubly compressed files not supported.") self.enable_read_bucket_metric = enable_read_bucket_metric self.mode = "r" @@ -751,7 +776,7 @@ def __init__( self, blob, content_type, - chunk_size=16 * 1024 * 1024, + chunk_size=DEFAULT_WRITE_BUFFER_SIZE, ignore_flush=True, enable_write_bucket_metric=False, retry=DEFAULT_RETRY): diff --git a/sdks/python/apache_beam/io/gcp/gcsio_integration_test.py b/sdks/python/apache_beam/io/gcp/gcsio_integration_test.py index b86fb4e85455..1ab8e74809f1 100644 --- a/sdks/python/apache_beam/io/gcp/gcsio_integration_test.py +++ b/sdks/python/apache_beam/io/gcp/gcsio_integration_test.py @@ -281,9 +281,12 @@ class GcsIOReadGzipTest(unittest.TestCase): ] @parameterized.expand([ - (gzip_test_files[0], CompressionTypes.UNCOMPRESSED, NotImplementedError), - (gzip_test_files[0], CompressionTypes.GZIP, NotImplementedError), - (gzip_test_files[0], CompressionTypes.AUTO, NotImplementedError), + # TODO(https://github.com/googleapis/google-cloud-python/issues/18423): + # Re-enable these once doubly compressed objects ("content-encoding=gzip" + # + "content-type=application/gzip") are handled again. + # (gzip_test_files[0], CompressionTypes.UNCOMPRESSED, NotImplementedError), + # (gzip_test_files[0], CompressionTypes.GZIP, NotImplementedError), + # (gzip_test_files[0], CompressionTypes.AUTO, NotImplementedError), (gzip_test_files[1], CompressionTypes.UNCOMPRESSED, UnicodeDecodeError), (gzip_test_files[1], CompressionTypes.GZIP, None), (gzip_test_files[1], CompressionTypes.AUTO, None), diff --git a/sdks/python/apache_beam/options/pipeline_options.py b/sdks/python/apache_beam/options/pipeline_options.py index 98e3dea38d2c..6443094fda12 100644 --- a/sdks/python/apache_beam/options/pipeline_options.py +++ b/sdks/python/apache_beam/options/pipeline_options.py @@ -1205,6 +1205,23 @@ def _add_argparse_args(cls, parser): 'Entries are key value pairs separated by = ' '(e.g. --gcs_custom_audit_entry key=value) or a JSON string ' '(e.g. --gcs_custom_audit_entries=\'{ "user": "test", "id": "12" }\').') + parser.add_argument( + '--gcs_read_buffer_size_bytes', + type=int, + default=None, + help='Size in bytes of the buffer used when reading from GCS. A ' + 'larger buffer reduces the number of requests sent to GCS at the ' + 'cost of more memory per reader. When unset, the GCS client in Beam ' + 'uses its default buffer size (16 MiB).') + parser.add_argument( + '--gcs_write_buffer_size_bytes', + type=int, + default=None, + help='Size in bytes of the buffer used when writing to GCS. Must be ' + 'a multiple of 256 KiB, since writes are performed as resumable ' + 'uploads. A larger buffer reduces the number of requests sent to GCS ' + 'at the cost of more memory per writer. When unset, the GCS client ' + 'in Beam uses its default buffer size (16 MiB).') def _create_default_gcs_bucket(self): try: @@ -1314,6 +1331,24 @@ def validate(self, validator): validator.validate_repeatable_argument_passed_as_list( self, 'dataflow_service_options')) + if (self.gcs_read_buffer_size_bytes is not None and + self.gcs_read_buffer_size_bytes <= 0): + errors.append( + '--gcs_read_buffer_size_bytes must be a positive number of bytes, ' + 'got %s.' % self.gcs_read_buffer_size_bytes) + + if self.gcs_write_buffer_size_bytes is not None: + # GCS resumable uploads require the chunk size to be a multiple of + # 256 KiB. Checking here avoids a failure deep inside the GCS client + # on the first flush. + write_buffer_size_multiple = 256 * 1024 + if (self.gcs_write_buffer_size_bytes <= 0 or + self.gcs_write_buffer_size_bytes % write_buffer_size_multiple != 0): + errors.append( + '--gcs_write_buffer_size_bytes must be a positive multiple of ' + '%d bytes, got %s.' % + (write_buffer_size_multiple, self.gcs_write_buffer_size_bytes)) + return errors def get_cloud_profiler_service_name(self): diff --git a/sdks/python/setup.py b/sdks/python/setup.py index 5ea8257a4b4b..26a68d7c9fe3 100644 --- a/sdks/python/setup.py +++ b/sdks/python/setup.py @@ -525,7 +525,7 @@ def get_portability_package_data(): 'google-auth-httplib2>=0.1.0,<0.3.0', 'google-cloud-datastore>=2.0.0,<3', 'google-cloud-pubsub>=2.1.0,<3', - 'google-cloud-storage>=2.18.2,<4', + 'google-cloud-storage>=3.0.0,<4', 'google-cloud-resource-manager>=1.12.0,<2', 'google-cloud-dataflow-client>=0.13.0,<0.14.0', # GCP packages required by tests