From e234e04219ffc0ef0b0215e4c094d82505843a4f Mon Sep 17 00:00:00 2001 From: Anshul Khandelwal Date: Mon, 24 Jun 2024 09:27:03 +0530 Subject: [PATCH] parquet write performance fix --- go/arrow/compute/utils.go | 6 +++--- go/parquet/internal/encoding/types.go | 10 +++++----- 2 files changed, 8 insertions(+), 8 deletions(-) diff --git a/go/arrow/compute/utils.go b/go/arrow/compute/utils.go index cc4d6edc48ec8..b35cb5e517c2b 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/v14/arrow" @@ -30,6 +29,7 @@ import ( "github.com/apache/arrow/go/v14/arrow/compute/internal/kernels" "github.com/apache/arrow/go/v14/arrow/internal/debug" "github.com/apache/arrow/go/v14/arrow/memory" + "github.com/apache/arrow/go/v14/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.MaxInt(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 2d13a9c8f877a..6282f5e94e9f2 100644 --- a/go/parquet/internal/encoding/types.go +++ b/go/parquet/internal/encoding/types.go @@ -185,9 +185,9 @@ func (b *PooledBufferWriter) Reserve(nbytes int) { b.buf = bufferPool.Get().(*memory.Buffer) } - newCap := utils.MaxInt(b.buf.Cap()+b.offset, 256) + newCap := utils.MaxInt(b.buf.Cap(), 256) for newCap < b.pos+nbytes { - newCap = bitutil.NextPowerOf2(newCap) + newCap = bitutil.NextPowerOf2(b.pos + nbytes) } b.buf.Reserve(newCap) } @@ -375,9 +375,9 @@ func (b *BufferWriter) Reserve(nbytes int) { if b.buffer == nil { b.buffer = memory.NewResizableBuffer(b.mem) } - newCap := utils.MaxInt(b.buffer.Cap()+b.offset, 256) - for newCap < b.pos+nbytes+b.offset { - newCap = bitutil.NextPowerOf2(newCap) + newCap := utils.MaxInt(b.buffer.Cap(), 256) + for newCap < b.pos+nbytes { + newCap = bitutil.NextPowerOf2(b.pos + nbytes) } b.buffer.Reserve(newCap) }