From 7fcaee0f2274eeebddc7972c9a6392cb7e85eeb8 Mon Sep 17 00:00:00 2001 From: Zihan Dai Date: Tue, 3 Nov 2026 19:41:22 +1100 Subject: [PATCH] [Python] Fix WriteToFiles crashing when the sink writes compressed files CompressedFile.flush() called the compressor's flush() with its default mode, which finishes the compressed stream. WriteToFiles flushes the sink and then closes the file, and close() finishes the stream a second time, so it raised for every compression type: zlib.error for gzip and deflate, ValueError for bzip2 and lzma, and ZstdError for zstd. flush() now writes out buffered data without finishing the stream: a sync flush for gzip and deflate, a block flush for zstd, and nothing for bzip2 and lzma, which cannot flush without finishing. bz2.BZ2File and lzma.LZMAFile behave the same way. close() still finishes the stream. --- CHANGES.md | 1 + sdks/python/apache_beam/io/fileio_test.py | 35 ++++++++++++++++ sdks/python/apache_beam/io/filesystem.py | 14 ++++++- sdks/python/apache_beam/io/filesystem_test.py | 40 +++++++++++++++++++ 4 files changed, 89 insertions(+), 1 deletion(-) diff --git a/CHANGES.md b/CHANGES.md index 24cbe75e6c71..983f1054a390 100644 --- a/CHANGES.md +++ b/CHANGES.md @@ -84,6 +84,7 @@ * (Python) RunInference now measures model memory (the `model_byte_size` metric) when `torch` is not installed. It was always 0, because a failed `ModelManager` import also disabled the `resource` module ([#40854](https://github.com/apache/beam/issues/40854)). * (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)). +* (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 diff --git a/sdks/python/apache_beam/io/fileio_test.py b/sdks/python/apache_beam/io/fileio_test.py index 3bb250a34bd2..4920cec1ae8e 100644 --- a/sdks/python/apache_beam/io/fileio_test.py +++ b/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 = [ diff --git a/sdks/python/apache_beam/io/filesystem.py b/sdks/python/apache_beam/io/filesystem.py index bfddb42dd912..2d93ae94f415 100644 --- a/sdks/python/apache_beam/io/filesystem.py +++ b/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 diff --git a/sdks/python/apache_beam/io/filesystem_test.py b/sdks/python/apache_beam/io/filesystem_test.py index fa68bc5ef6df..fb2f0ed71fb6 100644 --- a/sdks/python/apache_beam/io/filesystem_test.py +++ b/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