diff --git a/parquet/encryption_write_config_test.go b/parquet/encryption_write_config_test.go index 8bb7b547..9063ca7a 100644 --- a/parquet/encryption_write_config_test.go +++ b/parquet/encryption_write_config_test.go @@ -102,10 +102,11 @@ func (en *EncryptionConfigTestSuite) encryptFile(configs *parquet.FileEncryption ) if bufferedMode { - rgr = writer.AppendBufferedRowGroup() + rgr, err = writer.AppendBufferedRowGroupChecked() } else { - rgr = writer.AppendRowGroup() + rgr, err = writer.AppendRowGroupChecked() } + en.Require().NoError(err) nextColumn := func() file.ColumnChunkWriter { defer func() { colIndex++ }() diff --git a/parquet/file/file_writer.go b/parquet/file/file_writer.go index 291f2335..7221ada6 100644 --- a/parquet/file/file_writer.go +++ b/parquet/file/file_writer.go @@ -18,6 +18,7 @@ package file import ( "encoding/binary" + "errors" "fmt" "io" @@ -36,6 +37,7 @@ type Writer struct { props *parquet.WriterProperties rowGroups int nrows int + writeErr error metadata metadata.FileMetaDataBuilder fileEncryptor encryption.FileEncryptor rowGroupWriter *rowGroupWriter @@ -165,8 +167,26 @@ func (fw *Writer) TotalCompressedBytes() int64 { // is closed. // // When calling Close, all columns must have the same number of rows written. +// +// Deprecated: AppendBufferedRowGroup panics if closing the previous row group +// or starting the new row group fails. Use AppendBufferedRowGroupChecked to +// handle those errors explicitly. func (fw *Writer) AppendBufferedRowGroup() BufferedRowGroupWriter { - return fw.appendRowGroup(true) + rgw, err := fw.AppendBufferedRowGroupChecked() + if err != nil { + panic(err) + } + return rgw +} + +// AppendBufferedRowGroupChecked is the error-returning form of +// AppendBufferedRowGroup. +func (fw *Writer) AppendBufferedRowGroupChecked() (BufferedRowGroupWriter, error) { + rgw, err := fw.appendRowGroup(true) + if err != nil { + return nil, fw.recordError(err) + } + return rgw, nil } // AppendRowGroup appends a row group to the file and returns a writer @@ -174,26 +194,57 @@ func (fw *Writer) AppendBufferedRowGroup() BufferedRowGroupWriter { // // When calling NextColumn, the same number of rows need to have been written // to each column before moving on. Otherwise the rowgroup writer will panic. +// +// Deprecated: AppendRowGroup panics if closing the previous row group or +// starting the new row group fails. Use AppendRowGroupChecked to handle those +// errors explicitly. func (fw *Writer) AppendRowGroup() SerialRowGroupWriter { - return fw.appendRowGroup(false) + rgw, err := fw.AppendRowGroupChecked() + if err != nil { + panic(err) + } + return rgw } -func (fw *Writer) appendRowGroup(buffered bool) *rowGroupWriter { +// AppendRowGroupChecked is the error-returning form of AppendRowGroup. +func (fw *Writer) AppendRowGroupChecked() (SerialRowGroupWriter, error) { + rgw, err := fw.appendRowGroup(false) + if err != nil { + return nil, fw.recordError(err) + } + return rgw, nil +} + +func (fw *Writer) appendRowGroup(buffered bool) (*rowGroupWriter, error) { + if fw.writeErr != nil { + return nil, fw.writeErr + } + if fw.rowGroupWriter != nil { + if err := fw.rowGroupWriter.Close(); err != nil { + return nil, err + } fw.nrows += fw.rowGroupWriter.nrows - fw.rowGroupWriter.Close() } - fw.rowGroups++ - fw.footerFlushed = false - rgMeta := fw.metadata.AppendRowGroup() + if fw.pageIndexBuilder != nil { - fw.pageIndexBuilder.AppendRowGroup() + if err := fw.pageIndexBuilder.AppendRowGroup(); err != nil { + return nil, err + } } - fw.rowGroupWriter = newRowGroupWriter(fw.sink, rgMeta, int16(fw.rowGroups)-1, + rgMeta := fw.metadata.AppendRowGroup() + rgw, err := newRowGroupWriter(fw.sink, rgMeta, int16(fw.rowGroups), fw.props, buffered, fw.fileEncryptor, fw.pageIndexBuilder) - fw.bloomFilters.AppendRowGroup(rgMeta, fw.rowGroupWriter.bloomFilters) - return fw.rowGroupWriter + if err != nil { + return nil, err + } + + fw.rowGroups++ + fw.footerFlushed = false + fw.rowGroupWriter = rgw + fw.bloomFilters.AppendRowGroup(rgMeta, rgw.bloomFilters) + return rgw, nil } func (fw *Writer) startFile() error { @@ -228,12 +279,15 @@ func (fw *Writer) startFile() error { } n, err := fw.sink.Write(magic) - if err != nil { + if err != nil && !errors.Is(err, io.ErrShortWrite) { return fmt.Errorf("parquet: failed to write magic number: %w", err) } if n != len(magic) { return fmt.Errorf("parquet: short write of magic number: wrote %d of %d bytes", n, len(magic)) } + if err != nil { + return fmt.Errorf("parquet: failed to write magic number: %w", err) + } if fw.props.PageIndexEnabled() { fw.pageIndexBuilder = &metadata.PageIndexBuilder{ @@ -244,22 +298,35 @@ func (fw *Writer) startFile() error { return nil } -func (fw *Writer) writePageIndex() { +func (fw *Writer) writePageIndex() error { if fw.pageIndexBuilder != nil { // serialize page index after all row groups have been written and report // location to the file metadata fw.pageIndexBuilder.Finish() var pageIndexLocation metadata.PageIndexLocation - fw.pageIndexBuilder.WriteTo(fw.sink, &pageIndexLocation) + if err := fw.pageIndexBuilder.WriteTo(fw.sink, &pageIndexLocation); err != nil { + return err + } fw.metadata.SetPageIndexLocation(pageIndexLocation) } + return nil } // AppendKeyValueMetadata appends a key/value pair to the existing key/value metadata func (fw *Writer) AppendKeyValueMetadata(key string, value string) error { + if fw.writeErr != nil { + return fw.writeErr + } return fw.metadata.AppendKeyValueMetadata(key, value) } +func (fw *Writer) recordError(err error) error { + if fw.writeErr == nil { + fw.writeErr = err + } + return fw.writeErr +} + // Close closes any open row group writer and writes the file footer. Subsequent // calls to close will have no effect. If the underlying writer implements [io.Closer], // then this will also call Close on the underlying writer. @@ -276,15 +343,22 @@ func (fw *Writer) Close() (err error) { if ierr != nil { err = fmt.Errorf("error on close:%w, %s", err, ierr) } + fw.writeErr = err return } err = ierr + if err != nil { + fw.writeErr = err + } }() err = fw.FlushWithFooter() } - return + if err != nil { + return err + } + return fw.writeErr } // FileMetadata returns the current state of the FileMetadata that would be written @@ -299,31 +373,39 @@ func (fw *Writer) FileMetadata() (*metadata.FileMetaData, error) { // calls to FlushWithFooter or Close will be cumulative, so that only the last footer // written need ever be read by a reader. func (fw *Writer) FlushWithFooter() error { + if fw.writeErr != nil { + return fw.writeErr + } + if !fw.footerFlushed { if fw.rowGroupWriter != nil { + if err := fw.rowGroupWriter.Close(); err != nil { + return fw.recordError(err) + } fw.nrows += fw.rowGroupWriter.nrows - fw.rowGroupWriter.Close() + fw.rowGroupWriter = nil } - fw.rowGroupWriter = nil if err := fw.bloomFilters.WriteTo(fw.sink); err != nil { - return err + return fw.recordError(err) } - fw.writePageIndex() + if err := fw.writePageIndex(); err != nil { + return fw.recordError(err) + } fileMetadata, err := fw.metadata.Snapshot() if err != nil { - return err + return fw.recordError(err) } fileEncryptProps := fw.props.FileEncryptionProperties() if fileEncryptProps == nil { // non encrypted file if _, err = writeFileMetadata(fileMetadata, fw.sink); err != nil { - return err + return fw.recordError(err) } } else { if err := fw.flushEncryptedFile(fileMetadata, fileEncryptProps); err != nil { - return err + return fw.recordError(err) } } diff --git a/parquet/file/file_writer_test.go b/parquet/file/file_writer_test.go index 88d82d73..4633bcdf 100644 --- a/parquet/file/file_writer_test.go +++ b/parquet/file/file_writer_test.go @@ -1352,6 +1352,26 @@ func (f *flakyMagicSink) Write(p []byte) (int, error) { return f.buf.Write(p) } +type controlledSink struct { + buf bytes.Buffer + failAfter int + failErr error + short bool +} + +func (s *controlledSink) Write(p []byte) (int, error) { + if s.failAfter >= 0 { + if s.failAfter == 0 { + if s.short { + return len(p) - 1, nil + } + return 0, s.failErr + } + s.failAfter-- + } + return s.buf.Write(p) +} + func newSingleColumnSchema(t *testing.T) *schema.GroupNode { t.Helper() fields := schema.FieldList{schema.NewInt32Node("col", parquet.Repetitions.Required, 1)} @@ -1360,6 +1380,91 @@ func newSingleColumnSchema(t *testing.T) *schema.GroupNode { return sc } +func newTwoColumnSchema(t *testing.T) *schema.GroupNode { + t.Helper() + fields := schema.FieldList{ + schema.NewInt32Node("col0", parquet.Repetitions.Required, 1), + schema.NewInt32Node("col1", parquet.Repetitions.Required, 1), + } + sc, err := schema.NewGroupNode("schema", parquet.Repetitions.Required, fields, 0) + require.NoError(t, err) + return sc +} + +func TestAppendRowGroupCheckedPropagatesCloseError(t *testing.T) { + closeErr := errors.New("row group close failed") + sink := &controlledSink{failAfter: -1} + writer, err := file.NewParquetWriterWithError(sink, newSingleColumnSchema(t), + file.WithWriterProps(parquet.NewWriterProperties(parquet.WithDictionaryDefault(false)))) + require.NoError(t, err) + + rgw, err := writer.AppendRowGroupChecked() + require.NoError(t, err) + cw, err := rgw.NextColumn() + require.NoError(t, err) + _, err = cw.(*file.Int32ColumnChunkWriter).WriteBatch([]int32{1}, []int16{1}, nil) + require.NoError(t, err) + + sink.failErr = closeErr + sink.failAfter = 0 + _, err = writer.AppendRowGroupChecked() + require.ErrorIs(t, err, closeErr) + require.Zero(t, writer.NumRows()) + require.ErrorIs(t, writer.Close(), closeErr) + _, err = file.NewParquetReader(bytes.NewReader(sink.buf.Bytes())) + require.Error(t, err) +} + +func TestBufferedRowGroupPropagatesShortWrite(t *testing.T) { + sink := &controlledSink{failAfter: -1} + writer, err := file.NewParquetWriterWithError(sink, newSingleColumnSchema(t), + file.WithWriterProps(parquet.NewWriterProperties(parquet.WithDictionaryDefault(false)))) + require.NoError(t, err) + + rgw, err := writer.AppendBufferedRowGroupChecked() + require.NoError(t, err) + cw, err := rgw.Column(0) + require.NoError(t, err) + _, err = cw.(*file.Int32ColumnChunkWriter).WriteBatch([]int32{1}, []int16{1}, nil) + require.NoError(t, err) + + sink.short = true + sink.failAfter = 0 + require.ErrorIs(t, rgw.Close(), io.ErrShortWrite) + require.ErrorIs(t, writer.Close(), io.ErrShortWrite) + _, err = file.NewParquetReader(bytes.NewReader(sink.buf.Bytes())) + require.Error(t, err) +} + +func TestFlushWithFooterPropagatesPageIndexWriteError(t *testing.T) { + pageIndexErr := errors.New("page index write failed") + sink := &controlledSink{failAfter: -1} + props := parquet.NewWriterProperties( + parquet.WithDictionaryDefault(false), + parquet.WithPageIndexEnabled(true), + ) + writer, err := file.NewParquetWriterWithError(sink, newTwoColumnSchema(t), file.WithWriterProps(props)) + require.NoError(t, err) + + rgw, err := writer.AppendRowGroupChecked() + require.NoError(t, err) + for i := 0; i < 2; i++ { + cw, err := rgw.NextColumn() + require.NoError(t, err) + _, err = cw.(*file.Int32ColumnChunkWriter).WriteBatch([]int32{int32(i)}, []int16{1}, nil) + require.NoError(t, err) + require.NoError(t, cw.Close()) + } + require.NoError(t, rgw.Close()) + + sink.failErr = pageIndexErr + sink.failAfter = 1 + require.ErrorIs(t, writer.FlushWithFooter(), pageIndexErr) + require.ErrorIs(t, writer.Close(), pageIndexErr) + _, err = file.NewParquetReader(bytes.NewReader(sink.buf.Bytes())) + require.Error(t, err) +} + func TestNewParquetWriterWithError_Success(t *testing.T) { var buf bytes.Buffer writer, err := file.NewParquetWriterWithError(&buf, newSingleColumnSchema(t)) diff --git a/parquet/file/page_writer.go b/parquet/file/page_writer.go index e3b30bc7..7ab2b024 100644 --- a/parquet/file/page_writer.go +++ b/parquet/file/page_writer.go @@ -20,6 +20,7 @@ import ( "bytes" "errors" "fmt" + "io" "math" "sync" @@ -505,15 +506,26 @@ func (bw *bufferedPageWriter) Close(hasDict, fallback bool) error { DictEncodingStats: bw.pager.dictEncodingStats, DataEncodingStats: bw.pager.dataEncodingStats, } - bw.metadata.Finish(chunkInfo, hasDict, fallback, encodingStats) + if err := bw.metadata.Finish(chunkInfo, hasDict, fallback, encodingStats); err != nil { + return err + } bw.pager.FinishPageIndexes(position) - bw.metadata.WriteTo(bw.inMemSink) + if _, err := bw.metadata.WriteTo(bw.inMemSink); err != nil { + return err + } buf := bw.inMemSink.Finish() defer buf.Release() - _, err := bw.finalSink.Write(buf.Bytes()) - return err + data := buf.Bytes() + n, err := bw.finalSink.Write(data) + if err != nil { + return err + } + if n != len(data) { + return io.ErrShortWrite + } + return nil } func (bw *bufferedPageWriter) WriteDataPage(page DataPage) (int64, error) { diff --git a/parquet/file/row_group_writer.go b/parquet/file/row_group_writer.go index ff3361b5..8811c0ae 100644 --- a/parquet/file/row_group_writer.go +++ b/parquet/file/row_group_writer.go @@ -73,6 +73,7 @@ type rowGroupWriter struct { bytesWritten int64 compressedBytesWritten int64 closed bool + closeErr error ordinal int16 nextColumnIdx int nrows int @@ -86,7 +87,7 @@ type rowGroupWriter struct { bloomFilters map[string]metadata.BloomFilterBuilder } -func newRowGroupWriter(sink utils.WriterTell, rgMeta *metadata.RowGroupMetaDataBuilder, ordinal int16, props *parquet.WriterProperties, buffered bool, fileEncryptor encryption.FileEncryptor, pageIdxBldr *metadata.PageIndexBuilder) *rowGroupWriter { +func newRowGroupWriter(sink utils.WriterTell, rgMeta *metadata.RowGroupMetaDataBuilder, ordinal int16, props *parquet.WriterProperties, buffered bool, fileEncryptor encryption.FileEncryptor, pageIdxBldr *metadata.PageIndexBuilder) (*rowGroupWriter, error) { ret := &rowGroupWriter{ sink: sink, metadata: rgMeta, @@ -99,11 +100,13 @@ func newRowGroupWriter(sink utils.WriterTell, rgMeta *metadata.RowGroupMetaDataB } if buffered { - ret.initColumns() + if err := ret.initColumns(); err != nil { + return nil, err + } } else { ret.columnWriters = []ColumnChunkWriter{nil} } - return ret + return ret, nil } func (rg *rowGroupWriter) Buffered() bool { return rg.buffered } @@ -234,12 +237,14 @@ func (rg *rowGroupWriter) Close() error { if !rg.closed { rg.closed = true if err := rg.checkRowsWritten(); err != nil { + rg.closeErr = err return err } for _, wr := range rg.columnWriters { if wr != nil { if err := wr.Close(); err != nil { + rg.closeErr = err return err } rg.bytesWritten += wr.TotalBytesWritten() @@ -251,7 +256,7 @@ func (rg *rowGroupWriter) Close() error { rg.metadata.SetNumRows(rg.nrows) rg.metadata.Finish(rg.bytesWritten, rg.ordinal) } - return nil + return rg.closeErr } func (rg *rowGroupWriter) initColumns() error { diff --git a/parquet/internal/utils/write_utils.go b/parquet/internal/utils/write_utils.go index 28662c64..46113c56 100644 --- a/parquet/internal/utils/write_utils.go +++ b/parquet/internal/utils/write_utils.go @@ -53,5 +53,8 @@ func (w *TellWrapper) Tell() int64 { return w.pos } func (w *TellWrapper) Write(p []byte) (n int, err error) { n, err = w.Writer.Write(p) w.pos += int64(n) + if err == nil && n != len(p) { + err = io.ErrShortWrite + } return } diff --git a/parquet/pqarrow/encode_arrow_test.go b/parquet/pqarrow/encode_arrow_test.go index 197f6dcb..6c9e979d 100644 --- a/parquet/pqarrow/encode_arrow_test.go +++ b/parquet/pqarrow/encode_arrow_test.go @@ -839,7 +839,8 @@ func (ps *ParquetIOTestSuite) makeTestFile(mem memory.Allocator, typ arrow.DataT rowGroupSize := arr.Len() / numChunks for i := 0; i < numChunks; i++ { - rgw := writer.AppendRowGroup() + rgw, err := writer.AppendRowGroupChecked() + ps.NoError(err) cw, err := rgw.NextColumn() ps.NoError(err) @@ -1163,7 +1164,8 @@ func (ps *ParquetIOTestSuite) TestReadDecimals() { defer sink.Release() writer := file.NewParquetWriter(sink, sc) - rgw := writer.AppendRowGroup() + rgw, err := writer.AppendRowGroupChecked() + ps.NoError(err) cw, _ := rgw.NextColumn() cw.(*file.ByteArrayColumnChunkWriter).WriteBatch(bigEndian, nil, nil) cw.Close() @@ -1213,7 +1215,8 @@ func (ps *ParquetIOTestSuite) TestReadDecimal256() { defer sink.Release() writer := file.NewParquetWriter(sink, sc) - rgw := writer.AppendRowGroup() + rgw, err := writer.AppendRowGroupChecked() + ps.NoError(err) cw, _ := rgw.NextColumn() cw.(*file.ByteArrayColumnChunkWriter).WriteBatch(bigEndian, nil, nil) cw.Close() @@ -1298,7 +1301,8 @@ func (ps *ParquetIOTestSuite) TestReadDecimal256PartialWord() { sink := encoding.NewBufferWriter(0, mem) defer sink.Release() writer := file.NewParquetWriter(sink, sc) - rgw := writer.AppendRowGroup() + rgw, err := writer.AppendRowGroupChecked() + ps.NoError(err) cw, _ := rgw.NextColumn() cw.(*file.ByteArrayColumnChunkWriter).WriteBatch(bigEndian, nil, nil) cw.Close() diff --git a/parquet/pqarrow/file_writer.go b/parquet/pqarrow/file_writer.go index 827befaa..2bbb3c14 100644 --- a/parquet/pqarrow/file_writer.go +++ b/parquet/pqarrow/file_writer.go @@ -185,12 +185,19 @@ func NewFileWriter(arrschema *arrow.Schema, w io.Writer, props *parquet.WriterPr // NewRowGroup does what it says on the tin, creates a new row group in the underlying file. // Equivalent to `AppendRowGroup` on a file.Writer -func (fw *FileWriter) NewRowGroup() { +func (fw *FileWriter) NewRowGroup() error { if fw.rgw != nil { - fw.rgw.Close() + if err := fw.rgw.Close(); err != nil { + return err + } + } + rgw, err := fw.wr.AppendRowGroupChecked() + if err != nil { + return err } - fw.rgw = fw.wr.AppendRowGroup() + fw.rgw = rgw fw.colIdx = 0 + return nil } // NewBufferedRowGroup starts a new memory Buffered Row Group to allow writing columns / records @@ -198,12 +205,19 @@ func (fw *FileWriter) NewRowGroup() { // and decide where to break your row group based on the TotalBytesWritten rather than on the max // row group len. If using Records, this should be paired with WriteBuffered, while // Write will always write a new record as a row group in and of itself. -func (fw *FileWriter) NewBufferedRowGroup() { +func (fw *FileWriter) NewBufferedRowGroup() error { if fw.rgw != nil { - fw.rgw.Close() + if err := fw.rgw.Close(); err != nil { + return err + } + } + rgw, err := fw.wr.AppendBufferedRowGroupChecked() + if err != nil { + return err } - fw.rgw = fw.wr.AppendBufferedRowGroup() + fw.rgw = rgw fw.colIdx = 0 + return nil } // TotalCompressedBytes returns the total number of bytes after compression @@ -293,7 +307,9 @@ func (fw *FileWriter) WriteBuffered(rec arrow.RecordBatch) error { return err } } else { - fw.NewBufferedRowGroup() + if err := fw.NewBufferedRowGroup(); err != nil { + return err + } } if int64(curRows)+rec.NumRows() <= maxRows { @@ -310,7 +326,9 @@ func (fw *FileWriter) WriteBuffered(rec arrow.RecordBatch) error { for idx, r := range recList { if idx > 0 { - fw.NewBufferedRowGroup() + if err := fw.NewBufferedRowGroup(); err != nil { + return err + } } for i := 0; i < int(r.NumCols()); i++ { if err := fw.WriteColumnData(r.Column(i)); err != nil { @@ -353,7 +371,9 @@ func (fw *FileWriter) Write(rec arrow.RecordBatch) error { } for _, r := range recList { - fw.NewRowGroup() + if err := fw.NewRowGroup(); err != nil { + return err + } for i := 0; i < int(r.NumCols()); i++ { if err := fw.WriteColumnData(r.Column(i)); err != nil { fw.Close() @@ -382,7 +402,9 @@ func (fw *FileWriter) WriteTable(tbl arrow.Table, chunkSize int64) error { } writeRowGroup := func(offset, size int64) error { - fw.NewRowGroup() + if err := fw.NewRowGroup(); err != nil { + return err + } for i := 0; i < int(tbl.NumCols()); i++ { if err := fw.WriteColumnChunked(tbl.Column(i).Data(), offset, size); err != nil { return err