This is an automated email from the ASF dual-hosted git repository.
zeroshade pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/arrow-go.git
The following commit(s) were added to refs/heads/main by this push:
new 243e1f7a fix(parquet/encoding): honor BufferWriter offsets (#1144)
243e1f7a is described below
commit 243e1f7a40b3fb4d935fa896b19174929cae2b5c
Author: Minh Vu <[email protected]>
AuthorDate: Wed Aug 12 20:47:22 2026 +0200
fix(parquet/encoding): honor BufferWriter offsets (#1144)
### Rationale for this change
`BufferWriter` mixed logical positions with physical buffer offsets.
This made reservations too small, placed `UnsafeWriteCopy` at the wrong
location, truncated buffers after `UnsafeWrite`, and applied the offset
twice during `Seek`.
### What changes are included in this PR?
Keep reservations, writes, length reporting, and seeking consistent with
the configured offset. Add focused tests for regular writes, unsafe
pattern writes, and seeking.
### Are these changes tested?
- `go test ./parquet/internal/encoding`
- `go test -race ./parquet/internal/encoding -run TestBufferWriter`
### Are there any user-facing changes?
Parquet buffer writers now preserve the expected logical position and
length when an offset is used.
---
parquet/internal/encoding/types.go | 29 ++++++++----
parquet/internal/encoding/types_test.go | 80 +++++++++++++++++++++++++++++++++
2 files changed, 101 insertions(+), 8 deletions(-)
diff --git a/parquet/internal/encoding/types.go
b/parquet/internal/encoding/types.go
index 7c4015df..afbda0d6 100644
--- a/parquet/internal/encoding/types.go
+++ b/parquet/internal/encoding/types.go
@@ -340,18 +340,27 @@ func (b *BufferWriter) SetOffset(offset int) {
b.offset = offset
}
+func (b *BufferWriter) ensureOffset() {
+ if b.buffer.Len() < b.offset {
+ b.buffer.ResizeNoShrink(b.offset)
+ }
+}
+
// Bytes returns the current bytes slice of slice Len
func (b *BufferWriter) Bytes() []byte {
+ b.ensureOffset()
return b.buffer.Bytes()[b.offset:]
}
// Len provides the current Length of the byte slice
func (b *BufferWriter) Len() int {
+ b.ensureOffset()
return b.buffer.Len() - b.offset
}
// Cap returns the current capacity of the underlying buffer
func (b *BufferWriter) Cap() int {
+ b.ensureOffset()
return b.buffer.Cap() - b.offset
}
@@ -406,8 +415,8 @@ func (b *BufferWriter) Reserve(nbytes int) {
b.buffer = memory.NewResizableBuffer(b.mem)
}
newCap := utils.Max(b.buffer.Cap(), 256)
- for newCap < b.pos+nbytes {
- newCap = bitutil.NextPowerOf2(b.pos + nbytes)
+ for newCap < b.offset+b.pos+nbytes {
+ newCap = bitutil.NextPowerOf2(b.offset + b.pos + nbytes)
}
b.buffer.Reserve(newCap)
}
@@ -423,7 +432,7 @@ func (b *BufferWriter) WriteAt(p []byte, offset int64) (n
int, err error) {
need := int(offset) + len(p)
if need >= b.buffer.Cap() {
- b.Reserve(need - b.pos)
+ b.Reserve(need - b.offset - b.pos)
}
copy(b.buffer.Buf()[offset:], p)
@@ -449,13 +458,16 @@ func (b *BufferWriter) Write(buf []byte) (int, error) {
func (b *BufferWriter) UnsafeWriteCopy(ncopies int, pattern []byte) (int,
error) {
nbytes := len(pattern) * ncopies
- slc := b.buffer.Buf()[b.pos : b.pos+nbytes]
+ start := b.pos + b.offset
+ slc := b.buffer.Buf()[start : start+nbytes]
copy(slc, pattern)
for j := len(pattern); j < len(slc); j *= 2 {
copy(slc[j:], slc[:j])
}
b.pos += nbytes
- b.buffer.ResizeNoShrink(b.pos)
+ if b.buffer.Len() < b.pos+b.offset {
+ b.buffer.ResizeNoShrink(b.pos + b.offset)
+ }
return nbytes, nil
}
@@ -463,7 +475,9 @@ func (b *BufferWriter) UnsafeWriteCopy(ncopies int, pattern
[]byte) (int, error)
func (b *BufferWriter) UnsafeWrite(buf []byte) (int, error) {
copy(b.buffer.Buf()[b.pos+b.offset:], buf)
b.pos += len(buf)
- b.buffer.ResizeNoShrink(b.pos)
+ if b.buffer.Len() < b.pos+b.offset {
+ b.buffer.ResizeNoShrink(b.pos + b.offset)
+ }
return len(buf), nil
}
@@ -471,14 +485,13 @@ func (b *BufferWriter) UnsafeWrite(buf []byte) (int,
error) {
// whence must be io.SeekStart, io.SeekCurrent or io.SeekEnd or it will be
ignored.
func (b *BufferWriter) Seek(offset int64, whence int) (int64, error) {
newPos, offs := 0, int(offset)
- offs += b.offset
switch whence {
case io.SeekStart:
newPos = offs
case io.SeekCurrent:
newPos = b.pos + offs
case io.SeekEnd:
- newPos = b.buffer.Len() + offs
+ newPos = b.Len() + offs
}
if newPos < 0 {
return 0, errors.New("negative result pos")
diff --git a/parquet/internal/encoding/types_test.go
b/parquet/internal/encoding/types_test.go
new file mode 100644
index 00000000..0afc515b
--- /dev/null
+++ b/parquet/internal/encoding/types_test.go
@@ -0,0 +1,80 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package encoding
+
+import (
+ "io"
+ "testing"
+
+ "github.com/apache/arrow-go/v18/arrow/memory"
+ "github.com/stretchr/testify/require"
+)
+
+func TestBufferWriterOffset(t *testing.T) {
+ writer := NewBufferWriter(0, memory.DefaultAllocator)
+ defer writer.Release()
+
+ writer.SetOffset(1024)
+ require.Zero(t, writer.Cap())
+ require.Zero(t, writer.Len())
+
+ n, err := writer.Write([]byte("hello"))
+ require.NoError(t, err)
+ require.Equal(t, 5, n)
+ require.Equal(t, []byte("hello"), writer.Bytes())
+ require.Equal(t, 5, writer.Len())
+ require.Equal(t, int64(5), writer.Tell())
+}
+
+func TestBufferWriterUnsafeWriteCopyWithOffset(t *testing.T) {
+ writer := NewBufferWriter(0, memory.DefaultAllocator)
+ defer writer.Release()
+
+ writer.SetOffset(4)
+ writer.Reserve(4)
+
+ n, err := writer.UnsafeWriteCopy(2, []byte("ab"))
+ require.NoError(t, err)
+ require.Equal(t, 4, n)
+ require.Equal(t, []byte("abab"), writer.Bytes())
+ require.Equal(t, 4, writer.Len())
+}
+
+func TestBufferWriterSeekWithOffset(t *testing.T) {
+ writer := NewBufferWriter(0, memory.DefaultAllocator)
+ defer writer.Release()
+
+ writer.SetOffset(4)
+ _, err := writer.Write([]byte("abc"))
+ require.NoError(t, err)
+
+ pos, err := writer.Seek(0, io.SeekStart)
+ require.NoError(t, err)
+ require.Equal(t, int64(0), pos)
+
+ _, err = writer.Write([]byte("x"))
+ require.NoError(t, err)
+ require.Equal(t, []byte("xbc"), writer.Bytes())
+
+ pos, err = writer.Seek(-1, io.SeekEnd)
+ require.NoError(t, err)
+ require.Equal(t, int64(2), pos)
+
+ _, err = writer.Write([]byte("y"))
+ require.NoError(t, err)
+ require.Equal(t, []byte("xby"), writer.Bytes())
+}
