Skip to content

Ensure all non-final multipart uploads in filesystem sink are the same size - #889

Merged
mwylde merged 6 commits into
masterfrom
multipart_sizes
May 16, 2025
Merged

Ensure all non-final multipart uploads in filesystem sink are the same size#889
mwylde merged 6 commits into
masterfrom
multipart_sizes

Conversation

@mwylde

@mwylde mwylde commented May 16, 2025

Copy link
Copy Markdown
Member

R2 requires that all parts of a multipart upload (aside from the last one) be the same size. This PR changes how the pieces of the FilesystemSink work together to meet that requirement, while also simplifying the logic around multipart handling.

The FilesystemSink has a few major traits, which each have several implementations to abstract over json/parquet and object store/local filesystem:

  • BatchBufferingWriter, which handles the construction and buffering of a particular format (parquet or json)
  • BatchMultipartWriter, which coordinates the multipart writing process
  • MultipartManager, which manages a single multipart upload

Previously, the BatchBufferingWriters were "multipart aware"—they were responsible for deciding when to break up a batch into a new multipart write. For parquet, this would happen in the following process:

  1. A batch gets inserted into the ArrowWriter (which converts Arrow to parquet)
  2. If the batch exceeds the maximum number of rows for a row group, the ArrowWriter would flush the bytes for the row group into a buffer
  3. If the buffer exceeds the max part size, the BatchBufferingWriter would return all of those bytes to the BatchMultipartWriter to write

This PR changes the responsibilities such that it's more cleanly layered. The BatchBufferingWriter is now responsible only for converting its format and buffering, and has no awareness of multipart uploads. The BatchMultipartWriter decides when and how much to upload for a particular part. The new process is:

  1. A batch gets inserted into the ArrowWriter (which converts Arrow to parquet)
  2. If the row group is now larger than the configured row_group_size_bytes, we flush it to our buffer
  3. The BatchMultipartManager checks the size of the buffer
    a. If we don't have an active multipart upload, we check if the buffer is larger than our desired multipart upload size
    b. If we do have one, we check if it's larger than the actual multipart upload size for this upload
  4. If the appropriate condition is matched, we instruct the BatchBufferingWriter to split its buffer at the desired point, and upload that using the MultipartManager

A bit of additional complexity comes in once we're ready to finish the multipart upload, as we need to ensure that the final part is smaller than our part size. Before sending the last part we check if that's the case, and if not we split it into two parts.

In addition to the functional changes, there are also some performance improvements from moving from Vec to Bytes/BytesMut, which allows us to avoid some copies.

Breaking changes

This PR also includes some breaking changes to how we configure filesystem options; these are aimed at making the configuration more useful and easier to understand:

  • parquet_row_batch_size is removed (this was not actually be used anywhere)
  • parquet_row_group_size is replaced with parquet.row_group_size_bytes, as I believe size-based configuration is more useful than row-count base configuration
  • parquet_compression is renamed parquet.compression to match our standard option style

Testing

These changes were tested against R2 and Minio, with parquet, json, and deltalake. To ensure that checkpointing correctness is unaffected, I used the following query:

create table impulse with (
    connector = 'impulse',
    event_rate = 500000
);

create table sink with (
    connector = 'filesystem',
    type = 'sink',
    format = 'parquet',
    target_file_size = 268435456,
    rollover_seconds = 120,
    target_part_size = 5242880,
    parquet_compression = 'zstd',
    path = 's3::http://127.0.2.3:9000/test/outputs/consistency3'

);

insert into sink
select counter, sha512(counter), sha256(counter), md5(counter) from impulse;

which produces an incrementing counter (along with junk data to run up the file size). While that was running, I stopped/started and randomly kill -9'd the worker and controller processes and allowed it to recover. Then, I verified that the data was consistent with this duckdb query:

D create table events as SELECT counter FROM read_parquet('s3://test/outputs/consistency/*.parquet');
D select count(*), count(distinct counter), max(counter)+1 from events;
┌─────────────────┬─────────────────────────┬─────────────────┐
│  count_star()   │ count(DISTINCT counter) │  max(counter)   │
│      int64      │          int64          │     uint64      │
├─────────────────┼─────────────────────────┼─────────────────┤
│    628440756284407562844075     │
│ (62.84 million) │     (62.84 million)     │ (62.84 million) │
└─────────────────┴─────────────────────────┴─────────────────┘

which should produce the same number for each (the +1 on the max is due to 0-based indexing on the counter)

@mwylde
mwylde requested a review from cmackenzie1 May 16, 2025 17:50

@cmackenzie1 cmackenzie1 left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

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

Overall LGTM but I'm still getting up to speed on the rest of the codebase.
Once other quirk about R2 is the minimum part size is 5 MiB. Does the multipart stuff need to be aware of that limitation or at least enforce target_part_size is >= 5 MiB when destination is R2?

Comment thread crates/arroyo-connectors/src/filesystem/sink/parquet.rs
@mwylde

mwylde commented May 16, 2025

Copy link
Copy Markdown
Member Author

S3 actually has the same limitation, so I've modified the configuration to enforce that target_part_size is at least 5MB.

@mwylde
mwylde merged commit 13324f9 into master May 16, 2025
@mwylde
mwylde deleted the multipart_sizes branch May 16, 2025 20:29
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants