parquet write performance fix

rilldata · Jun 24, 2024 · e234e04 · e234e04
1 parent 740889f
commit e234e04
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 8 deletions.
diff --git 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
@@ -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)
 }