[Python] Fix WriteToFiles crashing when the sink writes compressed files

CHANGES.md

@@ -130,6 +130,7 @@
 * (Python/YAML) Fixed `WindowInto` `'d'` (day) duration unit parsing (`24h` instead of `12h`) and fixed `ExtractWindowingInfo` duplicate field name validation ([#40902](https://github.com/apache/beam/issues/40902)).
 * (Python) Fixed `GeminiModelHandler` returning only one result per batch: each prompt is now sent as its own request (up to 32 at a time), so every input gets its own result. A custom `request_fn` must now return one response per input ([#40913](https://github.com/apache/beam/issues/40913)).
 * (Java) BigQuery Storage Write API successful-row output now propagates fused downstream exceptions instead of allowing the bundle to succeed with missing downstream results ([#40004](https://github.com/apache/beam/pull/40004)).
+* (Python) Fixed `fileio.WriteToFiles` failing for sinks that ask for a compressed file in `create_metadata`. Flushing a compressed file ended the compressed stream, so closing the file afterwards raised ([#31863](https://github.com/apache/beam/issues/31863)).
 * Fixed X (Java/Python) ([#X](https://github.com/apache/beam/issues/X)).
 
 ## Security Fixes

sdks/python/apache_beam/io/fileio_test.py

@@ -810,6 +810,15 @@ def write(self, record):
       self._fh.write(json.dumps(record).encode('utf8'))
       self._fh.write('\n'.encode('utf8'))
 
+  class CompressedTextSink(fileio.TextSink):
+    def __init__(self, compression_type):
+      self._compression_type = compression_type
+
+    def create_metadata(self, destination, full_file_name):
+      return fileio.FileMetadata(
+          mime_type='application/octet-stream',
+          compression_type=self._compression_type)
+
   def test_write_to_single_file_batch(self):
 
     dir = self._new_tempdir()
@@ -831,6 +840,32 @@ def test_write_to_single_file_batch(self):
 
       assert_that(result, equal_to([row for row in self.SIMPLE_COLLECTION]))
 
+  def test_write_compressed_files(self):
+    # The sink asks for compression in create_metadata. WriteToFiles flushes
+    # the sink and then closes the file.
+    sink_class = WriteFilesTest.CompressedTextSink
+    for compression_type, suffix in [(CompressionTypes.BZIP2, '.bz2'),
+                                     (CompressionTypes.GZIP, '.gz')]:
+      dir = self._new_tempdir()
+
+      with TestPipeline() as p:
+        _ = (
+            p
+            | beam.Create(WriteFilesTest.LARGER_COLLECTION)
+            | fileio.WriteToFiles(
+                path=dir,
+                sink=lambda _, c=compression_type: sink_class(c),
+                file_naming=fileio.default_file_naming('output', suffix)))
+
+      with TestPipeline() as p:
+        result = (
+            p
+            | fileio.MatchFiles(FileSystems.join(dir, '*' + suffix))
+            | fileio.ReadMatches()
+            | beam.FlatMap(lambda f: f.read_utf8().strip().split('\n')))
+
+        assert_that(result, equal_to(WriteFilesTest.LARGER_COLLECTION))
+
   def test_write_to_dynamic_destination(self):
 
     sink_params = [

sdks/python/apache_beam/io/filesystem.py

@@ -331,9 +331,21 @@ def close(self) -> None:
     self._file.close()
 
   def flush(self) -> None:
+    """Writes buffered data to the file without finishing the compressed stream.
+
+    Only close() finishes the stream, so more data can be written after a flush.
+    """
     if self.writeable():
       assert self._compressor
-      self._file.write(self._compressor.flush())
+      if self._compression_type in (CompressionTypes.GZIP,
+                                    CompressionTypes.DEFLATE):
+        self._file.write(self._compressor.flush(zlib.Z_SYNC_FLUSH))
+      elif self._compression_type == CompressionTypes.ZSTD:
+        self._file.write(
+            self._compressor.flush(zstandard.COMPRESSOBJ_FLUSH_BLOCK))
+      # BZIP2 and LZMA compressors cannot flush without finishing the stream,
+      # so their buffered data is written by close(), as in bz2.BZ2File and
+      # lzma.LZMAFile.
     self._file.flush()
 
   @property

sdks/python/apache_beam/io/filesystem_test.py

@@ -505,6 +505,46 @@ def test_tell(self):
         if not line:
           break
 
+  def test_close_after_flush(self):
+    # fileio.WriteToFiles flushes the sink and then closes the file.
+    for compression_type in [CompressionTypes.BZIP2,
+                             CompressionTypes.DEFLATE,
+                             CompressionTypes.GZIP,
+                             CompressionTypes.ZSTD,
+                             CompressionTypes.LZMA]:
+      with self.subTest(compression_type=compression_type):
+        file_name = self._create_temp_file()
+        with open(file_name, 'wb') as f:
+          writeable = CompressedFile(f, compression_type)
+          writeable.write(self.content)
+          writeable.flush()
+          writeable.close()
+
+        with open(file_name, 'rb') as f:
+          readable = CompressedFile(
+              f, compression_type, read_size=self.read_block_size)
+          self.assertEqual(readable.read(), self.content)
+
+  def test_write_after_flush(self):
+    for compression_type in [CompressionTypes.BZIP2,
+                             CompressionTypes.DEFLATE,
+                             CompressionTypes.GZIP,
+                             CompressionTypes.ZSTD,
+                             CompressionTypes.LZMA]:
+      with self.subTest(compression_type=compression_type):
+        file_name = self._create_temp_file()
+        with open(file_name, 'wb') as f:
+          writeable = CompressedFile(f, compression_type)
+          writeable.write(self.content[:10])
+          writeable.flush()
+          writeable.write(self.content[10:])
+          writeable.close()
+
+        with open(file_name, 'rb') as f:
+          readable = CompressedFile(
+              f, compression_type, read_size=self.read_block_size)
+          self.assertEqual(readable.read(), self.content)
+
   def test_concatenated_compressed_file(self):
     # The test apache_beam.io.textio_test.test_read_gzip_concat
     # does not encounter the problem in the Beam 2.13 and earlier