From c4013482c9172be894fe9748b8156f31d021dd5f Mon Sep 17 00:00:00 2001 From: Nozomi Isozaki <63992859+nontan-pixiv@users.noreply.github.com> Date: Thu, 11 Jul 2024 03:57:47 +0900 Subject: [PATCH] GH-41541: [Go][Parquet] More fixes for writer performance regression (#42003) ### Rationale for this change This PR is complementary to #41638 . The prior PR reduces reallocations in `PooledBufferWriter`. However the problematic formula it addressed is still used in other functions. In addition to this, `(*PooledBufferWriter).Reserve()` simply doubles the capacity of buffers regardless of its argument `nbytes`. This may result in excessive allocations in some cases. ### What changes are included in this PR? - Applied the fixed formula to `(*BufferWriter).Reserve()`. - Updated the new capacity passed to `(*memory.Buffer).Reserve()`. - Now using `bitutil.NextPowerOf2(b.pos + nbytes)` to avoid reallocations when adding `nbytes`. - Replaced `math.Max` with `utils.Max` in `(*bufferWriteSeeker).Reserve()` to avoid unnecessary type conversions. ### Are these changes tested? Yes. The following commands pass. ``` $ export PARQUET_TEST_DATA=$PWD/cpp/submodules/parquet-testing/data $ (cd go && go test ./...) ``` ### Are there any user-facing changes? No, but it may reduce the number of allocations and improve the throughput. Before: ``` $ go test -test.run='^$' -test.bench='^BenchmarkWriteColumn$' -benchmem ./parquet/pqarrow/... goos: linux goarch: arm64 pkg: github.com/apache/arrow/go/v17/parquet/pqarrow BenchmarkWriteColumn/int32_not_nullable-10 1190 1016705 ns/op 4125.39 MB/s 5443579 B/op 240 allocs/op BenchmarkWriteColumn/int32_nullable-10 52 24780561 ns/op 169.26 MB/s 12048944 B/op 249 allocs/op BenchmarkWriteColumn/int64_not_nullable-10 632 1717090 ns/op 4885.36 MB/s 5445954 B/op 265 allocs/op BenchmarkWriteColumn/int64_nullable-10 51 22949770 ns/op 365.52 MB/s 12209860 B/op 262 allocs/op BenchmarkWriteColumn/float32_not_nullable-10 519 2234718 ns/op 1876.88 MB/s 5452627 B/op 1263 allocs/op BenchmarkWriteColumn/float32_nullable-10 56 23423793 ns/op 179.06 MB/s 12057540 B/op 1272 allocs/op BenchmarkWriteColumn/float64_not_nullable-10 416 2761247 ns/op 3037.98 MB/s 5507068 B/op 1292 allocs/op BenchmarkWriteColumn/float64_nullable-10 51 25767881 ns/op 325.55 MB/s 12059614 B/op 1285 allocs/op PASS ok github.com/apache/arrow/go/v17/parquet/pqarrow 10.592s ``` After: ``` $ go test -test.run='^$' -test.bench='^BenchmarkWriteColumn$' -benchmem ./parquet/pqarrow/... goos: linux goarch: arm64 pkg: github.com/apache/arrow/go/v17/parquet/pqarrow BenchmarkWriteColumn/int32_not_nullable-10 1196 959528 ns/op 4371.22 MB/s 5420349 B/op 238 allocs/op BenchmarkWriteColumn/int32_nullable-10 51 23017598 ns/op 182.22 MB/s 14138480 B/op 248 allocs/op BenchmarkWriteColumn/int64_not_nullable-10 690 1671710 ns/op 5017.98 MB/s 5419878 B/op 263 allocs/op BenchmarkWriteColumn/int64_nullable-10 50 23196051 ns/op 361.64 MB/s 13728465 B/op 261 allocs/op BenchmarkWriteColumn/float32_not_nullable-10 540 2185075 ns/op 1919.52 MB/s 5459392 B/op 1261 allocs/op BenchmarkWriteColumn/float32_nullable-10 54 21796783 ns/op 192.43 MB/s 14150622 B/op 1271 allocs/op BenchmarkWriteColumn/float64_not_nullable-10 418 2708292 ns/op 3097.38 MB/s 5455095 B/op 1290 allocs/op BenchmarkWriteColumn/float64_nullable-10 51 22174952 ns/op 378.29 MB/s 14142791 B/op 1283 allocs/op PASS ok github.com/apache/arrow/go/v17/parquet/pqarrow 10.210s ``` * GitHub Issue: #41541 --- go/arrow/compute/utils.go | 6 +++--- go/parquet/internal/encoding/types.go | 8 ++++---- 2 files changed, 7 insertions(+), 7 deletions(-) diff --git a/go/arrow/compute/utils.go b/go/arrow/compute/utils.go index b20688539f146..899fe4cfbf4cc 100644 --- a/go/arrow/compute/utils.go +++ b/go/arrow/compute/utils.go @@ -21,7 +21,6 @@ package compute import ( "fmt" "io" - "math" "time" "github.com/apache/arrow/go/v17/arrow" @@ -30,6 +29,7 @@ import ( "github.com/apache/arrow/go/v17/arrow/compute/internal/kernels" "github.com/apache/arrow/go/v17/arrow/internal/debug" "github.com/apache/arrow/go/v17/arrow/memory" + "github.com/apache/arrow/go/v17/internal/utils" "golang.org/x/xerrors" ) @@ -43,9 +43,9 @@ func (b *bufferWriteSeeker) Reserve(nbytes int) { if b.buf == nil { b.buf = memory.NewResizableBuffer(b.mem) } - newCap := int(math.Max(float64(b.buf.Cap()), 256)) + newCap := utils.Max(b.buf.Cap(), 256) for newCap < b.pos+nbytes { - newCap = bitutil.NextPowerOf2(newCap) + newCap = bitutil.NextPowerOf2(b.pos + nbytes) } b.buf.Reserve(newCap) } diff --git a/go/parquet/internal/encoding/types.go b/go/parquet/internal/encoding/types.go index 2d7a5d6b1d166..6962c95d4f818 100644 --- a/go/parquet/internal/encoding/types.go +++ b/go/parquet/internal/encoding/types.go @@ -187,7 +187,7 @@ func (b *PooledBufferWriter) Reserve(nbytes int) { newCap := utils.Max(b.buf.Cap(), 256) for newCap < b.pos+nbytes { - newCap = bitutil.NextPowerOf2(newCap) + newCap = bitutil.NextPowerOf2(b.pos + nbytes) } b.buf.Reserve(newCap) } @@ -380,9 +380,9 @@ func (b *BufferWriter) Reserve(nbytes int) { if b.buffer == nil { b.buffer = memory.NewResizableBuffer(b.mem) } - newCap := utils.Max(b.buffer.Cap()+b.offset, 256) - for newCap < b.pos+nbytes+b.offset { - newCap = bitutil.NextPowerOf2(newCap) + newCap := utils.Max(b.buffer.Cap(), 256) + for newCap < b.pos+nbytes { + newCap = bitutil.NextPowerOf2(b.pos + nbytes) } b.buffer.Reserve(newCap) }