feat(parquet): utilize memory allocator in serializedPageReader

[joechenrh]

Rationale for this change

Previously, only data page v2 in the reader will use the memory allocated from allocator. And seems there exists some misuage of the allocator.

What changes are included in this PR?

Use memory allocator to allocate buffer for serializedPageReader, and introduce Close function to make sure all the memory will be freed after usage.

Are these changes tested?

Track the memory usage in existing test to ensure all the allocated memory are freed.

Are there any user-facing changes?

For user with no explicitly set allocators, they don't need to change anything. Otherwise, they should remember to call ColumnChunkReader.Close after reading after chunk.

[joechenrh]

The background is that we are using this library to parse parquet and convert to our row data format. I passed a simple allocator into the reader to track the memory usage. But it appreard that the allocated size was not zero after the program finished.

type testAllocator struct {
	allocated atomic.Int64
}

func (a *testAllocator) Allocate(n int) []byte {
	a.allocated.Add(int64(n))
	return make([]byte, n)
}

func (a *testAllocator) Free(b []byte) {
	a.allocated.Add(-int64(len(b)))
}

...

And I think there are some problems with the usage of Allocator in serializedPageReader

• We only retain the buffer for data page v2. But we are still using this memory after buf.Release() for v1 and dict.
• For data page v2, we may lose the chance to release the page, as well as the buffer it retained, if we only need to read a portion of rows. (The Release() is called in Next)
• We can refer to the usage of decompressBuffer to reuse the buffer for each page type.

And this is the purpose of this PR, no performance improvement, but some refactor. I'm not sure if I understand right, and if it's necessary. Could you give me some suggestions? Thanks!

[zeroshade]

why remove this?

[joechenrh]

Because this PR change it to memory.Buffer, it's redundant to call Reset

[zeroshade]

this can become a problem due to race conditions. If a new DictionaryPage is read in by the reader while this page is still in use, you'll end up overwriting data that is still being used. That's why we were creating a new buffer here, to avoid the potential race condition.

[joechenrh]

I understand your meaning, but I think there should be no concurrency, or, data race here.

• First, dictionary page is only read once for each column chunk reader, either reading it at first, or reading it after we skip some rows and find it's not initialized.
• Second, the page read by serializedPageReader is stored in the columnChunkReader, so there should be only one page in use at any time. And user shouldn't access one chunk reader in parallel, which I think is an incorrect usage.

arrow-go/parquet/file/column_reader.go

Lines 326 to 345 in 3211594

	func (c *columnChunkReader) readNewPage() bool {
	for c.rdr.Next() { // keep going until we get a data page
	c.curPage = c.rdr.Page()
	if c.curPage == nil {
	break
	}
	gotDataPage, err := c.processPage()
	if err != nil {
	c.err = err
	return false
	}
	if gotDataPage {
	return true
	}
	}
	c.err = c.rdr.Err()
	return false
	}

[zeroshade]

same comment as above. we should utilize a buffer pool instead most likely

[zeroshade]

same issue as above

[zeroshade]

But it appreard that the allocated size was not zero after the program finished.

This is likely due to the buffer pool that we utilize. The GC will eventually clean it up, but we utilize a buffer pool with the passed in allocator.

[joechenrh]

This is likely due to the buffer pool that we utilize.

No, the default allocator GoAllocator doesn't use buffer pool.
In page reader, we get buffer from allocator directly. For data page V2, this buffer is released until next time we call serializedPageReader.Next. So if we don't read to the end of the file, this buffer won't be freed. That's why I add a Close function to manually reclaim the resource.

Besides, I think the semantic of "allocator" is that we manually manage memory, rather than relying on GC.

[joechenrh]

So many failed tests, let me check them later 😢
Well, forgot to call Close in test.

[joechenrh]

I found another problem with the current serializedPageReader, either using memory allocator or buffer pool.

For example, suppose each data page contains 4 ByteArray, and we want to read 12 data. We first read the data in 1st and 2nd page, store them in values. And these values may point to the internal buffer in the reader (depends on the decoder we use). If the pool/allocator give us a previously used buffer when reading the 3rd page, old values might be overwritten.

So I add the wrapper here to make a copy of the output. Do you think it's acceptable? @zeroshade

[zeroshade]

The pool/allocator shouldn't be able to return a previously used buffer until it has been put back into the pool, which shouldn't happen until we release it. Did you find a situation where this was happening?

[joechenrh]

I didn't meet this situation, but I think it could happen. Just like the above example:

• Get buffer1 -> Read data from page 1 -> Release buffer1
• Get buffer2 -> Read data from page 2 -> Release buffer2
• Get buffer 1 again

Since the values might directly point to these buffers, they could be modified when reading subsequent pages

arrow-go/parquet/file/column_reader_types.gen.go

Lines 216 to 220 in c6ce2ef

	func (cr *ByteArrayColumnChunkReader) ReadBatch(batchSize int64, values []parquet.ByteArray, defLvls, repLvls []int16) (total int64, valuesRead int, err error) {
	return cr.readBatch(batchSize, defLvls, repLvls, func(start, len int64) (int, error) {
	return cr.curDecoder.(encoding.ByteArrayDecoder).Decode(values[start : start+len])
	})
	}

[zeroshade]

I think that's also part of why we were creating new buffers and copying the data rather than being able to reuse the buffers in the page reader (the thing you're changing in this PR)

[zeroshade]

I think it would be good if you could craft a test case that validates this issue so that we can prove to ourselves that the updates you're making here prevent it from happening.

[joechenrh]

I think that's also part of why we were creating new buffers and copying the data rather than being able to reuse the buffers in the page reader (the thing you're changing in this PR)

But memory.NewBufferBytes doesn't copy the data, just ref it. That's the reason I open this PR.

arrow-go/arrow/memory/buffer.go

Lines 50 to 53 in c6ce2ef

	// NewBufferBytes creates a fixed-size buffer from the specified data.
	func NewBufferBytes(data []byte) *Buffer {
	return &Buffer{buf: data, length: len(data)}
	}

[zeroshade]

Yes, but remember that the code you're changing called buf := memory.NewResizableBuffer(p.mem) which does allocate a new buffer that the data is copied into before the changes you're making. The changes in this PR are introducing buffer re-use that wasn't there before.

I'm not saying it's a bad thing, just pointing out part of the reason for the previous behavior that allocated new buffers was to prevent the problem you've identified. Now that we're reusing these buffers in this change, we have to do the copying you're proposing.

Thats all

[joechenrh]

From eec53776b4aba9af1011a5a419472768c7c4dd0e Mon Sep 17 00:00:00 2001
From: Ruihao Chen <joechenrh@gmail.com>
Date: Fri, 29 Aug 2025 12:41:17 -0400
Subject: [PATCH] Add simple test case

Signed-off-by: Ruihao Chen <joechenrh@gmail.com>
---
 parquet/file/column_writer_test.go | 36 ++++++++++++++++++++++++++++++
 1 file changed, 36 insertions(+)

diff --git a/parquet/file/column_writer_test.go b/parquet/file/column_writer_test.go
index 90b239e4..3670de55 100644
--- a/parquet/file/column_writer_test.go
+++ b/parquet/file/column_writer_test.go
@@ -56,6 +56,41 @@ const (
 	DictionaryPageSize = 1024 * 1024
 )
 
+type simpleAllocator struct {
+	q chan []byte
+}
+
+func (a *simpleAllocator) Allocate(n int) []byte {
+	l := len(a.q)
+	for range l {
+		b := <-a.q
+		if cap(b) >= n {
+			return b[:n]
+		}
+		a.q <- b
+	}
+	return make([]byte, n)
+}
+
+func (a *simpleAllocator) Free(b []byte) {
+	if b == nil {
+		return
+	}
+	select {
+	case a.q <- b:
+	default: // discard if full
+	}
+}
+
+func (a *simpleAllocator) Reallocate(n int, old []byte) []byte {
+	a.Free(old)
+	return a.Allocate(n)
+}
+
+func newSimpleAllocator() *simpleAllocator {
+	return &simpleAllocator{q: make(chan []byte, 64)}
+}
+
 type mockpagewriter struct {
 	mock.Mock
 }
@@ -258,6 +293,7 @@ func (p *PrimitiveWriterTestSuite) TearDownTest() {
 }
 
 func (p *PrimitiveWriterTestSuite) buildReader(nrows int64, compression compress.Compression) file.ColumnChunkReader {
+	mem := newSimpleAllocator()
 	p.readbuffer = p.sink.Finish()
 	pagereader, _ := file.NewPageReader(arrutils.NewByteReader(p.readbuffer.Bytes()), nrows, compression, mem, nil)
 	return file.NewColumnReader(p.descr, pagereader, mem, &p.bufferPool)
-- 
2.47.3

I've crafted a customized allocator, which could make the UT in column_writer_test on main branch fail.

[zeroshade]

it would be better for the cloning to be done by the ByteArrayDecoder itself inside of the Decode method rather than needing this wrapper.

[joechenrh]

Because there are several different encoder for ByteArray type, comparing to modify every implementation, I prefer add a wrapper for this work. Besides, by only applying the wrapper to page reader (which the target of this PR), we can ensure that other part of the code won't be affected by this additional overhead.

[joechenrh]

After inspect the code again, I guess the problem arises from ColumnChunkReader.ReadBatch function itself, which is used in test to read data.

For other place that read data from column chunk, like recordReaderImpl, seems it has special logic for ByteArray and FixedLenByteArray. That is, copy the value from the decoder to a new, separate buffer. Like the below code, bldr itself contains a buffer allocated from allocator.

arrow-go/parquet/file/record_reader.go

Lines 842 to 846 in c6ce2ef

	switch bldr := br.bldr.(type) {
	case *array.BinaryBuilder:
	for _, val := range values {
	bldr.Append(val)
	}

Since it's used in test, let me revert the change and try to rewrite this function separately.

[zeroshade]

Just one nitpick change, but otherwise looks great to me. thanks for all this!

[zeroshade]

We shouldn't export this.

Also the constraint should probably be ~[]byte, yea?

[joechenrh]

Do you think we need to add any additional test cases？

[zeroshade]

An additional test case would be fantastic! Thanks!

[zeroshade]

@joechenrh do you want me to hold this open for you to add a test? Or would you want to add the test as a separate PR given this is already a fairly large one?

[joechenrh]

Sorry, I may not have time to do this these days, I perfer hold this, I'll try to add test(maybe not too much) this week.

[zeroshade]

Sure thing, I'll hold this open for you until the test gets added.

[zeroshade]

Thanks for the test!