[Python] Make read and write buffer size configurable in gcsio

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,22 +299,35 @@ 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://<bucket>/<object>``.
       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.
 
     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):

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),

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):

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