twuebi commented on code in PR #1506:
URL: https://github.com/apache/arrow-go/pull/1506#discussion_r4529127840
##########
parquet/pqarrow/file_writer_test.go:
##########
@@ -352,3 +356,82 @@ func TestWriteOnClosedFileWriter(t *testing.T) {
require.ErrorContains(t, writer.WriteColumnChunked(nil, 0, 0), "already
closed")
require.ErrorContains(t, writer.WriteColumnData(nil), "already closed")
}
+
+type byteLimitSink struct {
+ buf bytes.Buffer
+ limit int
+ closed bool
+}
+
+func (s *byteLimitSink) Close() error {
+ s.closed = true
+ return nil
+}
+
+var errSinkFull = errors.New("sink full")
+
+func (s *byteLimitSink) Write(p []byte) (int, error) {
+ if room := s.limit - s.buf.Len(); len(p) > room {
+ s.buf.Write(p[:max(room, 0)])
+ return max(room, 0), errSinkFull
+ }
+ return s.buf.Write(p)
+}
+
+func TestFileWriterCloseReleasesBuffersWhenSinkFails(t *testing.T) {
+ schema := arrow.NewSchema([]arrow.Field{
+ {Name: "num", Nullable: true, Type: arrow.PrimitiveTypes.Int8},
+ {Name: "dict", Nullable: true, Type:
&arrow.DictionaryType{IndexType: arrow.PrimitiveTypes.Int32, ValueType:
arrow.BinaryTypes.String}},
+ {Name: "str", Nullable: true, Type: arrow.BinaryTypes.String},
+ }, nil)
+
+ buildRecord := func(mem memory.Allocator) arrow.RecordBatch {
+ bldr := array.NewRecordBuilder(mem, schema)
+ defer bldr.Release()
+ num := bldr.Field(0).(*array.Int8Builder)
+ dict := bldr.Field(1).(*array.BinaryDictionaryBuilder)
+ str := bldr.Field(2).(*array.StringBuilder)
+ for i := 0; i < 1000; i++ {
+ require.NoError(t,
dict.AppendString("value-"+strconv.Itoa(i%10)))
+ str.Append("row-" + strconv.Itoa(i))
+ num.Append(int8(i))
+ }
+ return bldr.NewRecordBatch()
+ }
+
+ writeFile := func(t *testing.T, mem memory.Allocator, sink io.Writer,
buffered bool) error {
+ rec := buildRecord(mem)
+ defer rec.Release()
+
+ writer, err := pqarrow.NewFileWriter(schema, sink,
+ parquet.NewWriterProperties(parquet.WithAllocator(mem),
parquet.WithDictionaryFor("str", false)),
+
pqarrow.NewArrowWriterProperties(pqarrow.WithAllocator(mem),
pqarrow.WithStoreSchema()))
+ require.NoError(t, err)
+
+ if buffered {
+ err = writer.WriteBuffered(rec)
+ } else {
+ err = writer.Write(rec)
+ }
+ return errors.Join(err, writer.Close())
+ }
+
+ for _, buffered := range []bool{false, true} {
+ t.Run(fmt.Sprintf("buffered=%t", buffered), func(t *testing.T) {
+ var full bytes.Buffer
+ require.NoError(t, writeFile(t,
memory.DefaultAllocator, &full, buffered))
+
+ for limit := 4; limit < full.Len(); limit++ {
Review Comment:
Thanks, done in 62acf2e0. The test now writes the file once successfully and
records where every sink write starts, which covers every page header, page
body and footer piece. From the file's metadata it adds each column chunk's
dictionary page offset, data page offset and end (start plus total compressed
size), then a stride sample every 37 bytes. It fails one byte before, at and
after each of those offsets, keeping only offsets past the magic bytes. I also
cut the rows from 1000 to 300. The assertions are unchanged: the error wraps
the sink error, the sink is closed and `mem.AssertSize(t, 0)` holds.
Under `-race -tags assert,test` on macOS arm64 the test now takes 1.34 s
(0.66 s unbuffered, 0.69 s buffered), down from 109 s. That includes the small
data pages added in 41c33ed2 for the row-count mismatch path. Reverting any one
of the fixes still fails at least one subtest, and a full every-byte sweep at
the new row count passes with the fix, so the sampling hides nothing at this
size.
##########
parquet/file/file_writer_test.go:
##########
@@ -1674,3 +1675,97 @@ func TestPqarrowNewFileWriter_PropagatesShortWrite(t
*testing.T) {
require.Nil(t, writer)
require.Contains(t, err.Error(), "short write of magic number")
}
+
+type byteLimitSink struct {
+ buf bytes.Buffer
+ limit int
+}
+
+var errSinkFull = errors.New("sink full")
+
+func (s *byteLimitSink) Write(p []byte) (int, error) {
+ if room := s.limit - s.buf.Len(); len(p) > room {
+ s.buf.Write(p[:max(room, 0)])
+ return max(room, 0), errSinkFull
+ }
+ return s.buf.Write(p)
+}
+
+func newPlainAndDictStringSchema(t *testing.T) *schema.GroupNode {
+ t.Helper()
+ fields := schema.FieldList{
+ schema.Must(schema.NewPrimitiveNodeLogical("plain",
parquet.Repetitions.Optional,
+ schema.StringLogicalType{}, parquet.Types.ByteArray,
-1, -1)),
+ schema.Must(schema.NewPrimitiveNodeLogical("dict",
parquet.Repetitions.Optional,
+ schema.StringLogicalType{}, parquet.Types.ByteArray,
-1, -1)),
+ }
+ sc, err := schema.NewGroupNode("schema", parquet.Repetitions.Required,
fields, 0)
+ require.NoError(t, err)
+ return sc
+}
+
+func writeStringColumn(cw file.ColumnChunkWriter) error {
+ const n = 1000
+ values := make([]parquet.ByteArray, n)
+ defLevels := make([]int16, n)
+ for i := range values {
+ values[i] = parquet.ByteArray("value-" + strconv.Itoa(i%10))
+ defLevels[i] = 1
+ }
+ _, err := cw.(*file.ByteArrayColumnChunkWriter).WriteBatch(values,
defLevels, nil)
+ return err
+}
+
+func TestWriterReleasesBuffersWhenSinkFails(t *testing.T) {
+ writeFile := func(t *testing.T, mem memory.Allocator, sink io.Writer,
buffered bool) error {
+ props := parquet.NewWriterProperties(
+ parquet.WithAllocator(mem),
+ parquet.WithDictionaryDefault(true),
+ parquet.WithDictionaryFor("plain", false),
+ parquet.WithDataPageSize(1024),
+ parquet.WithBatchSize(100),
+ )
+ writer, err := file.NewParquetWriterWithError(sink,
newPlainAndDictStringSchema(t), file.WithWriterProps(props))
+ require.NoError(t, err)
+
+ if buffered {
+ rgw, err := writer.AppendBufferedRowGroupChecked()
+ require.NoError(t, err)
+ for i := 0; i < rgw.NumColumns(); i++ {
+ cw, err := rgw.Column(i)
+ require.NoError(t, err)
+ require.NoError(t, writeStringColumn(cw))
+ }
+ } else {
+ rgw, err := writer.AppendRowGroupChecked()
+ require.NoError(t, err)
+ for i := 0; i < rgw.NumColumns(); i++ {
+ cw, err := rgw.NextColumn()
+ if err == nil {
+ err = writeStringColumn(cw)
+ }
+ if err != nil {
+ return errors.Join(err, writer.Close())
+ }
+ }
+ }
+ return writer.Close()
+ }
+
+ for _, buffered := range []bool{false, true} {
+ t.Run(fmt.Sprintf("buffered=%t", buffered), func(t *testing.T) {
+ var full bytes.Buffer
+ require.NoError(t, writeFile(t,
memory.DefaultAllocator, &full, buffered))
+
+ for limit := 4; limit < full.Len(); limit++ {
Review Comment:
Same change here in 62acf2e0, with 300 rows instead of 1000. The test takes
0.40 s under `-race -tags assert,test` (0.14 s serial, 0.26 s buffered), down
from 20.9 s, and the `parquet/file` package now runs in 2.1 s. The same revert
check and every-byte cross-check pass.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
