[Go][Parquet] Use custom allocator in ReaderProperties.GetStream

[daniel-adam-tfs]

Describe the enhancement requested

We're testing reading parquet files from S3 and I'm looking at the memory profile. We're using a 2GB file for testing and in the profile I see that the make call in

arrow-go/parquet/reader_properties.go

Line 80 in c8d2478

data := make([]byte, nbytes)

allocates about 2GB of data, which is about a third of the total allocations that our test that reads the whole file does.

I see that the ReaderProperties has an allocator field, but the allocation there is made using the standard Go call. Is there a reason why the custom allocator is not used there? I wanted to test using a custom allocator, so maybe we can switch the make call to r.alloc.Allocate ?

Component(s)

Parquet

[zeroshade]

Seems like a reasonable approach. Would you be up for making a PR for this?

[daniel-adam-tfs]

I was looking at version 18.4.1, since that is the latest release and what we have in our repository. And wondering how to free the pool allocated buffer stored in the serializedPagedReader struct, but I see that in main f0b6fd9 did most of the job for me. I should be able to utilize that commit and make it work. Let me take a stab at it.

[daniel-adam-tfs]

So I switched to main and redid the memory profile. If I understand f0b6fd9 correctly, the point was for serializedPageReader.decompress to allocate memory using the memory package for the decompression buffer.

However, the call to io.CopyN might trigger a rellocation, in my memory profile it did so 2/3 of the time, because io.CopyN wraps the source buffer in LimitReader struct

https://github.com/golang/go/blob/39ed968832ad8923a4bd1fb6bc3d9090ddd98401/src/io/io.go#L364C27-L364C38

and continues to

https://github.com/golang/go/blob/39ed968832ad8923a4bd1fb6bc3d9090ddd98401/src/io/io.go#L415

which in our case goes through

https://github.com/golang/go/blob/39ed968832ad8923a4bd1fb6bc3d9090ddd98401/src/bytes/buffer.go#L215

And the loop here is causing the issue, because the loop exists only when you hit EOF, however LimitReader never returns EOF on the first read, so even when you read everything in one go. So the iteration in ReadFrom continues for one more run, but that means that you must have at least MinRead (=512) bytes still available in your buffer, which we don't have, since we resize to match exactly the size of the source.

TL;DR you need to resize the decompress buffer to desired size + MinRead to avoid the reallocation. 🏆

[joechenrh]

However, the call to io.CopyN might trigger a rellocation

So i think we can change to io.ReadFull? Then there is no reallocation:

b := bytes.NewBuffer(p.decompressBuffer.Bytes()[:lenCompressed])
if n, err := io.ReadFull(rd, b); err != nil || n != lenCompressed {
	return nil, err
}

[daniel-adam-tfs]

However, the call to io.CopyN might trigger a rellocation

So i think we can change to io.ReadFull? Then there is no reallocation:

b := bytes.NewBuffer(p.decompressBuffer.Bytes()[:lenCompressed])
if n, err := io.ReadFull(rd, b); err != nil || n != lenCompressed {
return nil, err
}

Yeah, that version would work, without the weird extra allocation size. Though the behavior between io.ReadFull and io.CopyN is different - one single read vs. multiple reads in a loop. @zeroshade was there a reason why io.CopyN was used instead of io.ReadFull?

Also, another idea I just had - with no compression I could read directly to the output buffer, right? For the nocodec decoder the Decode method is just a copy:

arrow-go/parquet/compress/compress.go

Lines 131 to 136 in 6e2e50b

	func (nocodec) Decode(dst, src []byte) []byte {
	if dst != nil {
	copy(dst, src)
	}
	return dst
	}

So I instead of reading to p.decompressBuffer and then copying to buf, I could just read directly to buf in this case and avoid the intermediate buffer totally.

EDIT: Which is what happens for PageType_DATA_PAGE_V2, but couldn't the same be done for format.PageType_DICTIONARY_PAGE and format.PageType_DATA_PAGE in serializedPageReader.Next?

[joechenrh]

one single read vs. multiple reads

Maybe using io.ReadFull also needs to call r.Read multiple times when reading from S3 🧐 (We are also reading files from S3).

[zeroshade]

Maybe using io.ReadFull also needs to call r.Read multiple times when reading from S3 🧐 (We are also reading files from S3).

That's precisely the case, both io.ReadFull and io.CopyN will perform multiple reads in a loop if the first read doesn't return enough (see https://cs.opensource.google/go/go/+/refs/tags/go1.25.3:src/io/io.go;l=333)

So I instead of reading to p.decompressBuffer and then copying to buf, I could just read directly to buf in this case and avoid the intermediate buffer totally.

EDIT: Which is what happens for PageType_DATA_PAGE_V2, but couldn't the same be done for format.PageType_DICTIONARY_PAGE and format.PageType_DATA_PAGE in serializedPageReader.Next?

That's definitely an oversight I think and a valid solution here. When there's no compression you could read directly to buf, that should be fine.

@zeroshade was there a reason why io.CopyN was used instead of io.ReadFull?

I don't remember offhand, the biggest difference between ReadFull and CopyN in this case would be that ReadFull can take the byte slice directly (instead of needing the intermediate bytes.NewBuffer).