zeroshade commented on code in PR #1506:
URL: https://github.com/apache/arrow-go/pull/1506#discussion_r4528398292
##########
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 here: under `-race` the `parquet/file` package takes 24.5 s, mostly
this test. Sampling offsets, as suggested on the pqarrow test, would cut that
down.
##########
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:
Failing at every byte offset is slow under `-race`, which CI uses on macOS.
Those jobs are already close to their timeout (see the comment at
`ci/scripts/test.sh:149`). Locally on arm64 this test alone takes 113.8 s with
`-race` (51 s unbuffered + 62 s buffered).
Could this test a sample of offsets instead? For example every ~37th byte,
plus the bytes just before and after each page and column chunk boundary, taken
from the successful file's metadata. That covers the same failure points in a
fraction of the time. Writing fewer rows would also help.
--
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]
