feat(file): stream multipart uploads for File writes with abort on rollback#7216
feat(file): stream multipart uploads for File writes with abort on rollback#7216BABTUNA wants to merge 5 commits into
Conversation
…ord-only Drops the multipart streaming path in favor of a single put on close; multipart streaming moves to a stacked follow-up PR with abort support and integration tests. Also fixes writable() to mirror the byte-range write guard.
Greptile SummaryThis PR upgrades
Confidence Score: 3/5Safe to merge after addressing the zero-byte S3 write path, which produces an unexpected API error on close. The streaming and abort logic is architecturally sound and tests are thorough for the common path. When a user opens an S3 file for writing and closes it without writing any data, CompleteMultipartUpload is issued with an empty parts list, which AWS S3 and MinIO both reject. The error surfaces on close rather than corrupting data or leaking a multipart upload, but it makes writing zero bytes to S3 reliably broken. This case is not covered by the minio integration tests. src/daft-file/src/file.rs — specifically the close_writer path for the Multipart variant when no parts were ever sent Important Files Changed
Sequence Diagram%%{init: {'theme': 'neutral'}}%%
sequenceDiagram
participant PY as Python (GIL thread)
participant FW as FileWriter::Multipart
participant CH as mpsc channel (cap=1)
participant UT as Upload Task (IO runtime)
participant S3 as S3 / Object Store
PY->>FW: open(wb) create_writer
FW->>S3: CreateMultipartUpload
S3-->>FW: upload_id
FW->>UT: spawn upload task
FW->>PY: return writer
loop write calls
PY->>FW: write(data)
FW->>FW: accumulate in part_buffer
alt "part_buffer >= part_size"
FW->>CH: send(part) blocking
UT->>CH: recv(part)
UT->>S3: UploadPart spawned in JoinSet
end
end
alt Clean close
PY->>FW: close_writer commit true
FW->>FW: commit_flag store true
FW->>CH: send remaining part
FW->>CH: drop tx channel closed
UT->>S3: CompleteMultipartUpload
S3-->>UT: ETag
FW->>PY: Ok
else Exception rollback
PY->>FW: close_writer commit false
FW->>CH: drop tx channel closed
UT->>UT: should_commit false skip complete
UT->>S3: AbortMultipartUpload
S3-->>UT: Ok
FW->>PY: Ok
end
%%{init: {'theme': 'base', 'themeVariables': {"darkMode": true, "background": "#0d1117", "primaryColor": "#21262d", "primaryTextColor": "#e6edf3", "primaryBorderColor": "#8b949e", "lineColor": "#8b949e", "textColor": "#e6edf3", "edgeLabelBackground": "#161b22", "actorBkg": "#21262d", "actorBorder": "#8b949e", "actorTextColor": "#e6edf3", "actorLineColor": "#8b949e", "signalColor": "#8b949e", "signalTextColor": "#e6edf3", "noteBkgColor": "#373320", "noteBorderColor": "#d4a72c", "noteTextColor": "#f0e6c0", "labelBoxBkgColor": "#21262d", "labelBoxBorderColor": "#8b949e", "labelTextColor": "#e6edf3", "loopTextColor": "#e6edf3", "activationBkgColor": "#30363d", "activationBorderColor": "#8b949e"}}}%%
sequenceDiagram
participant PY as Python (GIL thread)
participant FW as FileWriter::Multipart
participant CH as mpsc channel (cap=1)
participant UT as Upload Task (IO runtime)
participant S3 as S3 / Object Store
PY->>FW: open(wb) create_writer
FW->>S3: CreateMultipartUpload
S3-->>FW: upload_id
FW->>UT: spawn upload task
FW->>PY: return writer
loop write calls
PY->>FW: write(data)
FW->>FW: accumulate in part_buffer
alt "part_buffer >= part_size"
FW->>CH: send(part) blocking
UT->>CH: recv(part)
UT->>S3: UploadPart spawned in JoinSet
end
end
alt Clean close
PY->>FW: close_writer commit true
FW->>FW: commit_flag store true
FW->>CH: send remaining part
FW->>CH: drop tx channel closed
UT->>S3: CompleteMultipartUpload
S3-->>UT: ETag
FW->>PY: Ok
else Exception rollback
PY->>FW: close_writer commit false
FW->>CH: drop tx channel closed
UT->>UT: should_commit false skip complete
UT->>S3: AbortMultipartUpload
S3-->>UT: Ok
FW->>PY: Ok
end
|
| { | ||
| Some(mut multipart_writer) => { | ||
| let part_size = multipart_writer.part_size(); | ||
| assert!(part_size > 0, "Object multipart part size must be non-zero"); |
There was a problem hiding this comment.
Using
assert! here will panic at runtime if any MultipartWriter implementation incorrectly returns 0 for part_size(). A graceful DaftResult::Err is more appropriate for a production-facing code path.
| assert!(part_size > 0, "Object multipart part size must be non-zero"); | |
| if part_size == 0 { | |
| return Err(DaftError::InternalError( | |
| "Object multipart part size must be non-zero".to_string(), | |
| )); | |
| } |
Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!
Summary
Stacked on #7206. Adds multipart streaming to
daft.Filewrites: on sources that support multipart uploads (S3-compatible stores), part-sized chunks are streamed to the destination in the background as they fill, instead of buffering the whole object in memory. If the write is abandoned (exception in thewithblock) or a part upload fails, the multipart upload is explicitly aborted so the store does not retain, and bill for, parts of an incomplete upload.Why
#7206 lands write support with an in-memory buffer and a single put on close, which holds the entire object in memory. This follow-up bounds memory for large writes by streaming parts through the existing
MultipartWritermachinery. It also addresses the review finding on the earlier revision of #7206: theMultipartWritertrait had no abort operation, so abandoned writes would leave incomplete multipart uploads on S3 indefinitely. This PR addsabortto the trait and wires it into every non-commit path.Changes Made
src/daft-io/src/multipart.rs: addabortto theMultipartWritertraitsrc/daft-io/src/s3_like.rs: addS3LikeSource::abort_multipart_upload(S3AbortMultipartUploadcall with a newUnableToAbortMultipartUploaderror variant);S3MultipartWriter::abortcancels in-flight part tasks viaJoinSet::shutdownand then aborts the uploadsrc/daft-io/src/opendal_source.rs:OpenDALMultipartWriter::abortdelegates toopendal::Writer::abortsrc/daft-file/src/file.rs: restore theFileWriter::Multipartvariant with the background upload task; the task now aborts the upload whenever it does not commit, covering both the exception-rollback path and mid-stream part failuresdaft/file/file.py: docstring notes streamed parts are aborted on rollbacktests/integration/io/test_file_write_s3_minio.py: 4 minio integration tests covering multi-part commit and read-back, nothing-visible-before-close, exception rollback with aListMultipartUploadsassertion that no incomplete upload lingers, and text-mode writesBehavior
withblock) or a failed part upload now aborts the multipart upload on the store; previously buffered-only, so there was nothing to leakTest Plan
pytest -m integration tests/integration/io/test_file_write_s3_minio.pyagainst a local minio container: 4 passedListMultipartUploadson minio, so the empty-list assertion in the rollback test genuinely proves the abort ranDAFT_RUNNER=native pytest tests/file/test_file_basics.py tests/file/test_hdf5.py: 79 passedcargo check -p daft-file -p daft-io --features pythoncleancargo fmt -p daft-file -p daft-io --checkcleancargo clippy -p daft-file -p daft-io --features python --no-depscleanRelated Issues
Stacked on #7206; the diff includes its commits until it merges. Part of the follow-up promised there. Related to #6730.