Struct mz_storage_operators::s3_oneshot_sink::parquet::ParquetUploader

pub(super) struct ParquetUploader {
    desc: Arc<RelationDesc>,
    next_file_index: usize,
    key_manager: S3KeyManager,
    batch: u64,
    max_file_size: u64,
    sdk_config: Arc<SdkConfig>,
    row_group_size_bytes: u64,
    arrow_builder_buffer_bytes: u64,
    active_file: Option<ParquetFile>,
    params: CopyToParameters,
}

A ParquetUploader that writes rows to parquet files and uploads them to S3.

Spawns all S3 operations in tokio tasks to avoid blocking the surrounding timely context.

Buffering

There are several layers of buffering in this uploader:

• The uploader will hold a ParquetFile object after the first row is added. This ParquetFile holds an ArrowBuilder and an ArrowWriter.

• The ArrowBuilder builds a structure of in-memory mz_arrow_util::builder::ColBuilders from incoming mz_repr::Rows. Each mz_arrow_util::builder::ColBuilder holds a specific arrow::array::builder type for constructing a column of the given type. The entire ArrowBuilder is flushed to the ParquetFile’s ArrowWriter by converting it into a arrow::record_batch::RecordBatch once we’ve given it more than the configured arrow_builder_buffer_bytes.

• The ParquetFile holds a ArrowWriter that buffers until it has enough data to write a parquet ‘row group’. The ‘row group’ size is usually based on the number of rows (in the ArrowWriter), but we also force it to flush based on data-size (see below for more details).

• When a row group is written out, the active ParquetFile provides a reference to the row group buffer to its S3MultiPartUploader which will copy the data to its own buffer. If this upload buffer exceeds the configured part size limit, the S3MultiPartUploader will upload parts to S3 until the upload buffer is below the limit.

• When the ParquetUploader is finished, it will flush the active ParquetFile which will flush its ArrowBuilder and any open row groups to the S3MultiPartUploader and upload the remaining parts to S3.

      ┌───────────────┐
      │ mz_repr::Rows │
      └───────┬───────┘
┌─────────────│───────────────────────────────────────────┐
│             │         ParquetFile                       │
│ ┌───────────▼─────────────┐                             │
│ │       ArrowBuilder      │                             │
│ │                         │    ┌──────────────────┐     │
│ │     Vec<ArrowColumn>    │    │    ArrowWriter   │     │
│ │ ┌─────────┐ ┌─────────┐ │    │                  │     │
│ │ │         │ │         │ │    │   ┌──────────┐   │     │
│ │ │ColBuildr│ │ColBuildr│ ├────┼──►│  buffer  │   │     │
│ │ │         │ │         │ │    │   └─────┬────┘   │     │
│ │ └─────────┘ └─────────┘ │    │         │        │     │
│ │                         │    │   ┌─────▼────┐   │     │
│ └─────────────────────────┘    │   │ row group│   │     │
│                                │   └─┬────────┘   │     │
│                                │     │            │     │
│                                └─────┼────────────┘     │
│                               ┌──────┼────────────────┐ │
│                               │      │     S3MultiPart│ │
│                               │ ┌────▼─────┐ Uploader │ │
│                               │ │  buffer  │          │ │
│    ┌─────────┐                │ └───┬─────┬┘          │ │
│    │ S3 API  │◄───────────────┤     │     │           │ │
│    └─────────┘                │ ┌───▼──┐ ┌▼─────┐     │ │
│                               │ │ part │ │ part │     │ │
│                               │ └──────┘ └──────┘     │ │
│                               │                       │ │
│                               └───────────────────────┘ │
│                                                         │
└─────────────────────────────────────────────────────────┘

File Size & Buffer Sizes

We expose a ‘MAX FILE SIZE’ parameter to the user, but this is difficult to enforce exactly since we don’t know the exact size of the data we’re writing before a parquet row-group is flushed. This is because the encoded size of the data is different than the in-memory representation and because the data pages within each column in a row-group are compressed. We also don’t know the exact size of the parquet metadata that will be written to the file.

Therefore we don’t use the S3MultiPartUploader’s hard file size limit since it’s difficult to handle those errors after we’ve already flushed data to the ArrowWriter. Instead we implement a crude check ourselves.

This check aims to hit the max-size limit but may exceed it by some amount. To ensure that amount is small, we set the max row-group size to a configurable ratio (e.g. 20%) of the max_file_size. This determines how often we’ll flush a row-group, but is only an approximation since the actual size of the row-group is not known until it’s written. After each row-group is flushed, the size of the file is checked and if it’s exceeded max-file-size a new file is started.

We also set the max ArrowBuilder buffer size to a ratio (e.g. 150%) of the row-group size to avoid the ArrowWriter buffering too much data itself before flushing a row-group. We’re aiming for the encoded & compressed size of the ArrowBuilder data to be roughly equal to the row-group size, but this is only an approximation.

TODO: We may want to consider adding additional limits to the buffer sizes to avoid memory issues if the user sets the max file size to be very large.

Fields

desc: Arc<RelationDesc>

The output description.

next_file_index: usize

The index of the next file to upload within the batch.

key_manager: S3KeyManager

Provides the appropriate bucket and object keys to use for uploads.

batch: u64

Identifies the batch that files uploaded by this uploader belong to.

max_file_size: u64

The desired file size. A new file upload will be started when the size exceeds this amount.

sdk_config: Arc<SdkConfig>

The aws sdk config.

row_group_size_bytes: u64

arrow_builder_buffer_bytes: u64

active_file: Option<ParquetFile>

The active parquet file being written to, stored in an option since it won’t be initialized until the builder is first flushed, and to make it easier to take ownership when calling in spawned tokio tasks (to avoid doing I/O in the surrounding timely context).

params: CopyToParameters

Upload and buffer params

Implementations

impl ParquetUploader

async fn start_new_file(&mut self) -> Result<&mut ParquetFile, Error>

Start a new parquet file for upload. Will finish the current file if one is active.

Trait Implementations

impl CopyToS3Uploader for ParquetUploader

fn new( sdk_config: SdkConfig, connection_details: S3UploadInfo, sink_id: &GlobalId, batch: u64, params: CopyToParameters, ) -> Result<ParquetUploader, Error>

async fn append_row(&mut self, row: &Row) -> Result<(), Error>

Append a row to the internal buffer, and optionally flush the buffer to S3.

async fn finish(&mut self) -> Result<(), Error>

Flush the full remaining internal buffer to S3, and close all open resources. This will be called when the input stream is finished.