[Go][Parquet] pqarrow to support RowGroupTotalBytesWritten() for writer using WriteBuffered()

[DuanWeiFan]

Describe the enhancement requested

The method RowGroupTotalBytesWritten() is available but only make sense for the ones that use fileWriter.Write()

For the ones that use fileWriter.WriteBuffered() , it won't populate the bytes written until the writer is closed.

We are wondering if there is a way of populating or give an estimation of the size before closing the rowGroupWriter. This would be helpful to determine how large of a parquet file we want to build.

Component(s)

Parquet

[zeroshade]

Does RowGroupWriter.TotalCompressedBytes() provide the value you're looking for? As far as I can tell, it should be populated as data is written rather than waiting for data to get flushed.

[DuanWeiFan]

Hi @zeroshade thanks for getting back to me.
I tried to add a test in file_writer_test.go trying to log out the RowGroupTotalCompressedBytes().

It seems to be always returning zero.

It's only after the rowGroup writer is closed, then the RowGroupTotalBytesWritten() returned non-zero value.

	require.NoError(t, writer.WriteBuffered(record))

	t.Log(writer.RowGroupTotalCompressedBytes()) // output: 0
	t.Log(writer.RowGroupTotalBytesWritten()) // output: 0

	require.NoError(t, writer.Close())
	assert.Equal(t, 4, writer.NumRows())
	t.Log(writer.RowGroupTotalCompressedBytes()) // output: 0
	t.Log(writer.RowGroupTotalBytesWritten()) // output: 188

[DuanWeiFan]

After some investigation, I realize RowGroupTotalCompressedBytes() will get populated as soon as one of the DataPage got flushed.
Wondering if that makes sense to add a totalCompressedBytes & totalBytesWritten in FileWriter such that it can track the total number of bytes for every row group instead of just the current row group.

The two new methods: TotalBytesWritten() & TotalCompressedBytes() will then be able to report the total bytes written by the FileWriter.

type FileWriter struct {
  ...
  totalCompressedBytes int64
  totalBytesWritten    int64
}
// NewRowGroup does what it says on the tin, creates a new row group in the underlying file.
// Equivalent to `AppendRowGroup` on a file.Writer
func (fw *FileWriter) NewRowGroup() {
  if fw.rgw != nil {
    fw.totalCompressedBytes += fw.rgw.TotalCompressedBytes()
    fw.totalBytesWritten += fw.rgw.TotalBytesWritten()
    fw.rgw.Close()
  }
  fw.rgw = fw.wr.AppendRowGroup()
  fw.colIdx = 0
}

// NewBufferedRowGroup starts a new memory Buffered Row Group to allow writing columns / records
// without immediately flushing them to disk. This allows using WriteBuffered to write records
// and decide where to break your row group based on the TotalBytesWritten rather than on the max
// row group len. If using Records, this should be paired with WriteBuffered, while
// Write will always write a new record as a row group in and of itself.
func (fw *FileWriter) NewBufferedRowGroup() {
  if fw.rgw != nil {
    fw.totalCompressedBytes += fw.rgw.TotalCompressedBytes()
    fw.totalBytesWritten += fw.rgw.TotalBytesWritten()
    fw.rgw.Close()
  }
  fw.rgw = fw.wr.AppendBufferedRowGroup()
  fw.colIdx = 0
}

// TotalCompressedBytes returns the total number of bytes after compression
// that have been written to the file so far. It includes all the closed row groups
// and the current row group.
func (fw *FileWriter) TotalCompressedBytes() int64 {
  return fw.totalCompressedBytes + fw.RowGroupTotalCompressedBytes()
}

// TotalBytesWritten returns the total number of bytes
// that have been written to the file so far. It includes all the closed row groups
// and the current row group.
func (fw *FileWriter) TotalBytesWritten() int64 {
  return fw.totalBytesWritten + fw.RowGroupTotalBytesWritten()
}

[zeroshade]

Sorry for the delay in response here!

Your idea makes sense, would you be willing to contribute a PR for it?

[DuanWeiFan]

No worries. Yeah I’ll put together a PR! Thanks for checking