diff --git a/CHANGES.md b/CHANGES.md index b9b04c32fc13..48e492d36a45 100644 --- a/CHANGES.md +++ b/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 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