Initial Implementation of the Parquet Writer

[sharmrj]

This is the initial minimal implementation of the Parquet writer.

A Parquet file is a series of row groups followed by the file metadata (which contains the schema and the metadata for all the rowgroups, which, in turn, contain the metadata for each column chunk). Inside each rowgroup is a series of column chunks. Column chunks consist of a series of pages. Pages are the PageHeader followed by RLE encoded definition levels (if they exist), RLE encoded repetition levels (if they exist), and finally the encoded and then compressed data. I forgot about magic bytes.

For a parquet file to acheive efficient compression we tend to desirc row groupts of a specific size and for each of our column chunks to have pages that are of a specific size. So we must expose these fiddle factors to the user so they can tune the writer to have the behavior they want. (there are subtleties to this that are discussed further below)

We'll set the default Page size to 1 MiB and the default rowGroupSize to 128MiB (of course users will be able to adjust these numbers through write options). We need to hold the entire RowGroup in memory as we build is as the ColumnChunks need to be contiguous when written to disk. So we need to hold buffers for each individual columnChunk as we go row by row and build them; the columnChunks cannot be interleaved.

Since dataframe is columnar to begin with, we could, in theory, go golumn by column by estimating the size of a certain slice of a column, but I don't yet see a good way of doing this given we must run the gamut of encodings and compressions applied to each of those ColumnChunks (and those compression libraries have their own multifarious strategies for various kinds of data)

Each row group has to be a certain size, but each column in a row group must contain the same number of rows, even though each column may very well fit the same number of rows in very different amounts of space. So how do we ensure that we both hit our page size target, our record size target, and have the same number of rows in each column?

We must consider the page size and row group sizes to be best effort. They could be slightly above or below the target. The characteristics of the parquet file will depend on both the write options and the specific data being encoded. Arrow-rs runs batches of rows through the writer, flushing when they see that a page/rowgroup has met or exceeded its limit.

So a row group is flushed specifically only on batch boundaries and we get the same number of rows in every row group except the last which will be smaller than the rest. They also use sub batching. so as to not overshoot page size egregiously if the user sets a large batch size. Note: arrow-rs had an issue where some columns had really large values. We may need to implement batching and sub batching eventually.

If larger row groups are required (up to a gigabyte in size if not more), we should provide users who need to minimize memory usage an alternate two pass strategy where we first write to temporary files (one per columnChunk) until the temporary files have grown to the size of what a rowgroup should actually be and pipe the temporary files into the output. Essentially our rowgroup buffer is on disk instead of in memory. This is slower but should use less memory. In cases where there is extra RAM available but the user chooses the two pass strategy anyway, the temp files will tend to be held in the OS Page Cache (RAM) anyway.

Niceties like statistics and bloom filters and so on have not yet been implemented. We may need some extra machinery to keep track of row ranges so we can use them with the dataframe to generate our statistics.

We haven't yet implemented all the encodings and compressions possible. The writer should first be brought to parity with the reader, and then we should implement encodings and compressions in both together so neither lags behind the other.

We also don't yet support a way to have different compressions/endodings per page, and I imagine we would use some kind of heuristid to select these things, if we should want such a thing at all

Repetition levels and Definitions levlels above 1 are also not yet supported, but that may come hand in hand with bigger work where we work out the best way to support arbitraritly nested rows in dataframe in a general way (as opposed to what we have today)

The writer, as far as possible, attempts do almost all of its allocations up front and run in more or less constant memory. The current implementation is single threaded so in total we have 1 buffer per column, 1 page buffer + a scratch buffer for processing pages, a buffer for definition levels. New memory is also allocated for compression and accumulated throughout the write for the FileMetadata.

Note: we start each Column Chunk Buffer at the size of a PageBuffer and all buffers grow by 1.5 when full which is O(1) amortized. In the benchmark I did this doesn't seem to be an issue and the buffers grow to their ideal size relatively quickly.

I have a criterion benchmark and a stress test, both of which build a 10gb parquet file and write it to a temporary directory. The stress test specifically also reads it back using the parquet writer and confirms that the read back dataframe is equivalent to the initial dataframe.

Memory profile result

The wall time on this test does not reflect accurate performance as it is affected by GHC having to profile large amounts of memory (Refer to the criterion result below for the wall clock time for writing a 10gb file). we see that the majority of the memory is occupied by the dataframe itelf (10.4gb). Our buffers grow over the first few seconds and remain static until we get to the big spike from the parquet reader. The peak heap is roughly the size of our two dataframes + a little extra

Criterion Result:

Benchmark dataframe-parquet-writer-10gb: RUNNING...
benchmarking write 10 GiB dataframe
time                 23.27 s    (23.04 s .. 23.45 s)
                     1.000 R²   (1.000 R² .. 1.000 R²)
mean                 23.29 s    (23.25 s .. 23.32 s)
std dev              41.82 ms   (19.60 ms .. 57.89 ms)
variance introduced by outliers: 19% (moderately inflated)

A note on tests.

I put the tests inside the dataframe-parquet directory so I only had to compile relevant dependencies when working on just parquet. In the future we should move the remaining parquet tests to the parquet folder as well.

[mchav]

Will this evaluate it enough that this benchmark is meaningful?

[sharmrj]

I think whnfIO does enforce the IO action as it's just (). But I ran it withnfIO anyway just to make sure and got a similar result:

benchmarking write 10 GiB dataframe
time                 23.10 s    (22.94 s .. 23.30 s)
                     1.000 R²   (1.000 R² .. 1.000 R²)
mean                 23.25 s    (23.19 s .. 23.37 s)
std dev              107.8 ms   (22.15 ms .. 142.2 ms)
variance introduced by outliers: 19% (moderately inflated)

[mchav]

I think mixing implementation documentation like this with the code risks drift plus isn't a good look. Maybe attach in PR description or have an associated doc with state of implementation.

[sharmrj]

I've moved the comment into the PR description.