diff --git a/go/parquet/internal/encoding/types.go b/go/parquet/internal/encoding/types.go index 147c1746c515a..2d7a5d6b1d166 100644 --- a/go/parquet/internal/encoding/types.go +++ b/go/parquet/internal/encoding/types.go @@ -361,11 +361,16 @@ func (b *BufferWriter) Truncate() { func (b *BufferWriter) Reset(initial int) { if b.buffer != nil { b.buffer.Release() + } else { + b.buffer = memory.NewResizableBuffer(b.mem) } b.pos = 0 b.offset = 0 - b.Reserve(initial) + + if initial > 0 { + b.Reserve(initial) + } } // Reserve ensures that there is at least enough capacity to write nbytes diff --git a/go/parquet/pqarrow/file_writer_test.go b/go/parquet/pqarrow/file_writer_test.go index 425e4479f6d5c..fc965279a928d 100644 --- a/go/parquet/pqarrow/file_writer_test.go +++ b/go/parquet/pqarrow/file_writer_test.go @@ -18,6 +18,7 @@ package pqarrow_test import ( "bytes" + "math" "strings" "testing" @@ -87,3 +88,44 @@ func TestFileWriterNumRows(t *testing.T) { require.NoError(t, writer.Close()) assert.Equal(t, 4, writer.NumRows()) } + +func TestFileWriterBuffered(t *testing.T) { + schema := arrow.NewSchema([]arrow.Field{ + {Name: "one", Nullable: true, Type: arrow.PrimitiveTypes.Float64}, + {Name: "two", Nullable: true, Type: arrow.PrimitiveTypes.Float64}, + }, nil) + + data := `[ + {"one": 1, "two": 2}, + {"one": 1, "two": null}, + {"one": null, "two": 2}, + {"one": null, "two": null} + ]` + + alloc := memory.NewCheckedAllocator(memory.DefaultAllocator) + defer alloc.AssertSize(t, 0) + + record, _, err := array.RecordFromJSON(alloc, schema, strings.NewReader(data)) + require.NoError(t, err) + defer record.Release() + + output := &bytes.Buffer{} + writer, err := pqarrow.NewFileWriter( + schema, + output, + parquet.NewWriterProperties( + parquet.WithAllocator(alloc), + // Ensure enough space so we can close the writer with rows still buffered + parquet.WithMaxRowGroupLength(math.MaxInt64), + ), + pqarrow.NewArrowWriterProperties( + pqarrow.WithAllocator(alloc), + ), + ) + require.NoError(t, err) + + require.NoError(t, writer.WriteBuffered(record)) + + require.NoError(t, writer.Close()) + assert.Equal(t, 4, writer.NumRows()) +}