Skip to content

feat(file): stream multipart uploads for File writes with abort on rollback#7216

Open
BABTUNA wants to merge 5 commits into
Eventual-Inc:mainfrom
BABTUNA:feat/file-write-multipart
Open

feat(file): stream multipart uploads for File writes with abort on rollback#7216
BABTUNA wants to merge 5 commits into
Eventual-Inc:mainfrom
BABTUNA:feat/file-write-multipart

Conversation

@BABTUNA

@BABTUNA BABTUNA commented Jul 7, 2026

Copy link
Copy Markdown
Contributor

Summary

Stacked on #7206. Adds multipart streaming to daft.File writes: 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 the with block) 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 MultipartWriter machinery. It also addresses the review finding on the earlier revision of #7206: the MultipartWriter trait had no abort operation, so abandoned writes would leave incomplete multipart uploads on S3 indefinitely. This PR adds abort to the trait and wires it into every non-commit path.

Changes Made

  • src/daft-io/src/multipart.rs: add abort to the MultipartWriter trait
  • src/daft-io/src/s3_like.rs: add S3LikeSource::abort_multipart_upload (S3 AbortMultipartUpload call with a new UnableToAbortMultipartUpload error variant); S3MultipartWriter::abort cancels in-flight part tasks via JoinSet::shutdown and then aborts the upload
  • src/daft-io/src/opendal_source.rs: OpenDALMultipartWriter::abort delegates to opendal::Writer::abort
  • src/daft-file/src/file.rs: restore the FileWriter::Multipart variant 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 failures
  • daft/file/file.py: docstring notes streamed parts are aborted on rollback
  • tests/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 a ListMultipartUploads assertion that no incomplete upload lingers, and text-mode writes

Behavior

  • Large writes stream part-sized chunks (default 8MB) in the background; memory is bounded by the part size instead of the object size
  • Commit semantics are unchanged from feat(file): add native write support to daft.File.open #7206: nothing is visible at the destination until a clean close
  • New: abandoning a write (exception in the with block) or a failed part upload now aborts the multipart upload on the store; previously buffered-only, so there was nothing to leak
  • Local files and stores without a multipart writer keep the buffered single-put path from feat(file): add native write support to daft.File.open #7206

Test Plan

  • pytest -m integration tests/integration/io/test_file_write_s3_minio.py against a local minio container: 4 passed
  • Negative control verified separately: an incomplete multipart upload created directly via boto3 does show up in ListMultipartUploads on minio, so the empty-list assertion in the rollback test genuinely proves the abort ran
  • DAFT_RUNNER=native pytest tests/file/test_file_basics.py tests/file/test_hdf5.py: 79 passed
  • cargo check -p daft-file -p daft-io --features python clean
  • cargo fmt -p daft-file -p daft-io --check clean
  • cargo clippy -p daft-file -p daft-io --features python --no-deps clean

Related Issues

Stacked on #7206; the diff includes its commits until it merges. Part of the follow-up promised there. Related to #6730.

BABTUNA added 5 commits July 4, 2026 14:07
…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.
@BABTUNA BABTUNA requested a review from a team as a code owner July 7, 2026 02:49
@github-actions github-actions Bot added the feat label Jul 7, 2026
@greptile-apps

greptile-apps Bot commented Jul 7, 2026

Copy link
Copy Markdown
Contributor

Greptile Summary

This PR upgrades daft.File writes from a single in-memory buffer/PUT to streaming multipart uploads on S3-compatible stores, bounding memory to the part size rather than the full object size. It also adds abort to the MultipartWriter trait and wires it into every non-commit exit path so that abandoned or failed uploads do not leave incomplete multipart uploads on the store.

  • Streaming multipart path (FileWriter::Multipart): a background task drains part-sized chunks from a bounded channel and either complete()s or abort()s the upload depending on whether close_writer was called with commit=true or the with block exited with an exception.
  • Rollback semantics: __exit__ passes exc_type.is_none() as the commit flag, with the atomic should_commit flag coordinating between the Python thread and the IO runtime task.
  • MultipartWriter::abort trait method: added to both S3MultipartWriter and OpenDALMultipartWriter.

Confidence Score: 3/5

Safe 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

Filename Overview
src/daft-file/src/file.rs Adds FileWriter enum with Multipart and Buffered variants, create_writer, write, and close_writer. Empty writes to S3 multipart path fail on close; assert! for part_size > 0 should be a graceful error.
src/daft-file/src/python.rs Adds write/flush/readable/writable/seekable/isatty to PyDaftFile, context-manager rollback via exc_type check, and _create_writer static method. Deref/DerefMut to DaftFile are used throughout.
src/daft-io/src/s3_like.rs Adds abort_multipart_upload on S3LikeSource and S3MultipartWriter::abort that cancels in-flight JoinSet tasks then calls AbortMultipartUpload. Implementation is correct.
src/daft-io/src/multipart.rs Extends the MultipartWriter trait with an abort method. Clean addition with appropriate documentation.
src/daft-io/src/opendal_source.rs Implements abort for OpenDALMultipartWriter by delegating to opendal::Writer::abort. Straightforward and correct.
daft/file/file.py Updates File.open() to accept mode parameter with read/write variants and correct writable() delegation. Mode validation and error messages are correct.
tests/integration/io/test_file_write_s3_minio.py New integration tests covering multipart commit/read-back, visibility before close, exception rollback with ListMultipartUploads assertion, and text-mode writes. Does not cover the 0-byte write case.
tests/file/test_file_basics.py Comprehensive local write tests covering binary/text modes, context exception rollback, mode validation, and UDF usage.

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
Loading
%%{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
Loading

Comments Outside Diff (1)

  1. src/daft-file/src/file.rs, line 277-347 (link)

    P1 Empty write to S3 raises an S3 API error on close

    When write() is never called (or only called with a total of 0 bytes), part_buffer remains empty. In close_writer(commit=true) the early-return condition !part_buffer.is_empty() is false, so the final part is never sent through the channel. The upload task then sees the channel close, checks should_commit = true, and calls complete() via S3MultipartWriter::shutdown(). That function calls CompleteMultipartUpload with an empty parts list — AWS S3 (and MinIO) reject this with an InvalidRequest/MalformedXML error.

    The multipart upload is subsequently aborted (the task aborts on result.is_err()), so there is no resource leak, but a user who opens a file for writing and immediately closes it without writing any data will see an unexpected error. The Buffered path (local files) handles this correctly by issuing a put with an empty buffer.

    A minimal fix would be to track whether any parts have been sent and fall back to abort() rather than complete() when no data was written, or route through a single-PUT path for the zero-byte case.

Reviews (1): Last reviewed commit: "feat(file): stream multipart uploads for..." | Re-trigger Greptile

Comment thread src/daft-file/src/file.rs
{
Some(mut multipart_writer) => {
let part_size = multipart_writer.part_size();
assert!(part_size > 0, "Object multipart part size must be non-zero");

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 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.

Suggested change
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!

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant