[PDGGK]

Fixes #31863. #35363 hit the same failure with gzip.

fileio.WriteToFiles fails when a sink asks for a compressed file through create_metadata, which is how the FileSink docstring says to choose the metadata passed to FileSystems.create:

class GzipTextSink(fileio.TextSink):
  def create_metadata(self, destination, full_file_name):
    return fileio.FileMetadata(
        mime_type='application/octet-stream',
        compression_type=CompressionTypes.GZIP)

WriteToFiles calls sink.flush() and then writer.close(). TextSink.flush() calls CompressedFile.flush(), which called the compressor's flush() with its default mode. That mode finishes the compressed stream. close() then finished it a second time, and every compression type raises at that point:

compression	error from CompressedFile.close()
gzip, deflate	zlib.error: Error -2 while flushing: inconsistent stream state
bzip2, lzma	ValueError: Repeated call to flush()
zstd	ZstdError: compressor object already finished

A write after a flush also raises, because the compressor is already finished.

Change

CompressedFile.flush() now writes buffered data to the file without finishing the stream. Only close() finishes it.

• gzip and deflate: flush(zlib.Z_SYNC_FLUSH), the same mode gzip.GzipFile.flush() uses.
• zstd: flush(zstandard.COMPRESSOBJ_FLUSH_BLOCK).
• bzip2 and lzma: the compressor is not flushed, because these compressors cannot flush without finishing. bz2.BZ2File.flush() and lzma.LZMAFile.flush() do not flush the compressor either. The data is written by close().

The underlying file is still flushed in every case. close() is unchanged.

Before and after

A pipeline that writes five lines through WriteToFiles with the sink above, once per compression type, on the FnApiRunner. Each written file is then decompressed with the standard library or zstandard.

Beam 2.77.0 (current master is the same):

gzip     FAILED: zlib.error: Error -2 while flushing: inconsistent stream state [while running 'WriteToFiles/ParDo(_WriteUnshardedRecordsFn)/ParDo(_WriteUnshardedRecordsFn)']
bzip2    FAILED: ValueError: Repeated call to flush() [while running 'WriteToFiles/ParDo(_WriteUnshardedRecordsFn)/ParDo(_WriteUnshardedRecordsFn)']
deflate  FAILED: zlib.error: Error -2 while flushing: inconsistent stream state [while running 'WriteToFiles/ParDo(_WriteUnshardedRecordsFn)/ParDo(_WriteUnshardedRecordsFn)']
zstd     FAILED: zstandard.backend_c.ZstdError: compressor object already finished [while running 'WriteToFiles/ParDo(_WriteUnshardedRecordsFn)/ParDo(_WriteUnshardedRecordsFn)']
lzma     FAILED: ValueError: Repeated call to flush() [while running 'WriteToFiles/ParDo(_WriteUnshardedRecordsFn)/ParDo(_WriteUnshardedRecordsFn)']

With this change:

gzip     OK: read back 5 lines ['line-0', 'line-1', 'line-2', 'line-3', 'line-4']
bzip2    OK: read back 5 lines ['line-0', 'line-1', 'line-2', 'line-3', 'line-4']
deflate  OK: read back 5 lines ['line-0', 'line-1', 'line-2', 'line-3', 'line-4']
zstd     OK: read back 5 lines ['line-0', 'line-1', 'line-2', 'line-3', 'line-4']
lzma     OK: read back 5 lines ['line-0', 'line-1', 'line-2', 'line-3', 'line-4']

Tests

• filesystem_test.py: test_close_after_flush (write, flush, close) and test_write_after_flush (write, flush, write, close). Each runs for all five compression types and reads the file back through CompressedFile.
• fileio_test.py: test_write_compressed_files runs WriteToFiles with a sink that asks for bzip2 and then gzip, and reads the output back with ReadMatches. The compression is inferred from the .bz2 and .gz suffixes.

With the filesystem.py change reverted, all 10 CompressedFile subtests fail. The close-after-flush cases fail in close() with the errors in the table above, and the write-after-flush cases fail on the second write(). test_write_compressed_files fails with ValueError: Repeated call to flush() [while running 'WriteToFiles/...']. With the change, they all pass.

Run locally on Python 3.12, macOS, one process:

• filesystem_test.py: 41 passed, 10 subtests passed.
• fileio_test.py, WriteFilesTest and ReadTest: 13 passed and 1 skipped on the FnApiRunner. The two streaming tests (test_streaming_complex_timing, test_streaming_different_file_types) passed on the BundleBasedDirectRunner, because they use TestStream.
• textio_test.py, the compressed TextSinkTest cases: 10 passed. WriteToText closes without flushing, so this checks that path is unchanged.
• yapf --diff (0.43.0), ruff check (0.15.22) and pyrefly check (1.1.1) on the three changed Python files: clean.

I did not run the tests on the default Prism runner, on other Python versions, or the full Python suite. That is left to CI.

CHANGES.md has a Bugfixes entry.

[github-actions]

📋 Review Summary

This pull request properly resolves a bug where CompressedFile.flush() prematurely finalized the compression stream, causing subsequent write() or close() calls (such as in fileio.WriteToFiles) to fail across all supported compression types. The changes align CompressedFile.flush() behavior with standard library implementations (gzip.GzipFile, bz2.BZ2File, lzma.LZMAFile) and zstandard, and include thorough unit and pipeline tests.

🔍 General Feedback

• The use of zlib.Z_SYNC_FLUSH for GZIP/DEFLATE, zstandard.COMPRESSOBJ_FLUSH_BLOCK for ZSTD, and omitting compressor flushes until close() for BZIP2/LZMA is accurate and matches standard library semantics.
• The unit tests in filesystem_test.py and end-to-end pipeline test in fileio_test.py provide comprehensive coverage for both flush() -> close() and flush() -> write() -> close() workflows across compression types.

---

💡 Tip: Comment /gemini review on this pull request to trigger a new Gemini code review.

[github-actions]

Assigning reviewers:

R: @shunping for label python.
Backup reviewer: damccorm (if the main reviewer is unresponsive, you can escalate to them by tagging their username)

This pull request likely touches a core component ("core" label). Please review with scrutiny.

Note: If you would like to opt out of this review, comment assign to next reviewer.

Available commands:

• stop reviewer notifications - opt out of the automated review tooling
• remind me after tests pass - tag the comment author after tests pass
• waiting on author - shift the attention set back to the author (any comment or push by the author will return the attention set to the reviewers)

The PR bot will only process comments in the main thread (not review comments).

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 57.42%. Comparing base (8f5a492) to head (eefc258).
⚠️ Report is 2 commits behind head on master.

Additional details and impacted files

@@             Coverage Diff              @@
##             master   #40966      +/-   ##
============================================
- Coverage     58.87%   57.42%   -1.45%     
+ Complexity     4768     2492    -2276     
============================================
  Files          1461     1202     -259     
  Lines        185377   170995   -14382     
  Branches       3585     2103    -1482     
============================================
- Hits         109140    98201   -10939     
+ Misses        72892    70145    -2747     
+ Partials       3345     2649     -696     

Flag	Coverage Δ
python	80.17% <100.00%> (-0.01%)	⬇️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
• 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

[shunping]

LGTM! Please resolve the conflict and we should be good to go.

[PDGGK]

Thanks! I merged master and resolved the conflict in CHANGES.md, keeping both Bugfixes entries. The PR is mergeable again.