diff --git a/parquet/file/column_reader.go b/parquet/file/column_reader.go index b9c96f73..53cc284a 100644 --- a/parquet/file/column_reader.go +++ b/parquet/file/column_reader.go @@ -141,8 +141,10 @@ type ColumnChunkReader interface { // // This will clear any current error in the reader but does not // automatically read the first page of the page reader passed in until - // HasNext which will read in the next page. - setPageReader(PageReader) + // HasNext which will read in the next page. An error is returned if the + // current page reader cannot be closed; in that case the new reader is not + // installed and the error is retained by the column reader. + setPageReader(PageReader) error // Close releases the resources held by the column reader. Close() error } @@ -252,19 +254,31 @@ func (c *columnChunkReader) Descriptor() *schema.Column { return c.descr } func (c *columnChunkReader) consumeBufferedValues(n int64) { c.numDecoded += n } func (c *columnChunkReader) numAvailValues() int64 { return c.numBuffered - c.numDecoded } func (c *columnChunkReader) pager() PageReader { return c.rdr } -func (c *columnChunkReader) setPageReader(rdr PageReader) { - c.Close() + +func (c *columnChunkReader) setPageReader(rdr PageReader) error { + if err := c.Close(); err != nil { + c.err = err + return err + } c.rdr, c.err = rdr, nil c.decoders = make(map[format.Encoding]encoding.TypedDecoder) c.dictState = dictNotRead c.numBuffered, c.numDecoded = 0, 0 + return nil } -// Close closes the page raeder and the page if set. +// Close closes the page reader and the page if set. func (c *columnChunkReader) Close() error { - // c.rdr owns curPage (c.curPage == c.rdr.Page()) and releases it on Close. - if c.rdr != nil { - return c.rdr.Close() + page := c.curPage + rdr := c.rdr + c.curPage = nil + c.rdr = nil + + if rdr != nil { + return rdr.Close() + } + if page != nil { + page.Release() } return nil } diff --git a/parquet/file/column_reader_close_test.go b/parquet/file/column_reader_close_test.go new file mode 100644 index 00000000..86ee36db --- /dev/null +++ b/parquet/file/column_reader_close_test.go @@ -0,0 +1,108 @@ +// 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 file + +import ( + "errors" + "testing" + + "github.com/apache/arrow-go/v18/parquet" + "github.com/apache/arrow-go/v18/parquet/compress" + format "github.com/apache/arrow-go/v18/parquet/internal/gen-go/parquet" + "github.com/stretchr/testify/require" +) + +type closeTrackingPage struct { + releaseCalls int +} + +func (p *closeTrackingPage) Type() PageType { return PageTypeDataPage } +func (p *closeTrackingPage) Data() []byte { return nil } +func (p *closeTrackingPage) Encoding() format.Encoding { return format.Encoding_PLAIN } +func (p *closeTrackingPage) NumValues() int32 { return 0 } +func (p *closeTrackingPage) Release() { p.releaseCalls++ } + +type closeTrackingPageReader struct { + closeCalls int + closeErr error + page Page +} + +func (r *closeTrackingPageReader) SetMaxPageHeaderSize(int) {} +func (r *closeTrackingPageReader) Page() Page { return nil } +func (r *closeTrackingPageReader) Next() bool { return false } +func (r *closeTrackingPageReader) Err() error { return nil } +func (r *closeTrackingPageReader) Reset(parquet.BufferedReader, int64, compress.Compression, *CryptoContext) { +} +func (r *closeTrackingPageReader) GetDictionaryPage() (*DictionaryPage, error) { return nil, nil } +func (r *closeTrackingPageReader) SeekToPageWithRow(int64) error { return nil } +func (r *closeTrackingPageReader) Close() error { + r.closeCalls++ + if r.page != nil { + r.page.Release() + r.page = nil + } + return r.closeErr +} + +func TestColumnChunkReaderCloseIsIdempotent(t *testing.T) { + page := &closeTrackingPage{} + rdr := &closeTrackingPageReader{page: page} + reader := &columnChunkReader{curPage: page, rdr: rdr} + + require.NoError(t, reader.Close()) + require.NoError(t, reader.Close()) + require.Equal(t, 1, page.releaseCalls) + require.Equal(t, 1, rdr.closeCalls) + require.Nil(t, reader.curPage) + require.Nil(t, reader.rdr) +} + +func TestColumnChunkReaderSetPageReaderReleasesBeforeReplacement(t *testing.T) { + page := &closeTrackingPage{} + oldRdr := &closeTrackingPageReader{page: page} + newRdr := &closeTrackingPageReader{} + reader := &columnChunkReader{curPage: page, rdr: oldRdr} + + require.NoError(t, reader.setPageReader(newRdr)) + require.Equal(t, 1, page.releaseCalls) + require.Equal(t, 1, oldRdr.closeCalls) + require.Equal(t, 0, newRdr.closeCalls) + require.Same(t, newRdr, reader.rdr) + require.Nil(t, reader.curPage) + + require.NoError(t, reader.Close()) + require.Equal(t, 1, newRdr.closeCalls) +} + +func TestColumnChunkReaderSetPageReaderPropagatesCloseError(t *testing.T) { + closeErr := errors.New("page reader close failed") + page := &closeTrackingPage{} + oldRdr := &closeTrackingPageReader{closeErr: closeErr, page: page} + newRdr := &closeTrackingPageReader{} + reader := &columnChunkReader{curPage: page, rdr: oldRdr} + + require.ErrorIs(t, reader.setPageReader(newRdr), closeErr) + require.ErrorIs(t, reader.err, closeErr) + require.Equal(t, 1, page.releaseCalls) + require.Equal(t, 1, oldRdr.closeCalls) + require.Equal(t, 0, newRdr.closeCalls) + require.Nil(t, reader.curPage) + require.Nil(t, reader.rdr) + require.NoError(t, reader.Close()) + require.Equal(t, 1, oldRdr.closeCalls) +} diff --git a/parquet/file/record_reader.go b/parquet/file/record_reader.go index 40ff69d0..3b98cdf0 100644 --- a/parquet/file/record_reader.go +++ b/parquet/file/record_reader.go @@ -85,6 +85,14 @@ type RecordReader interface { SeekToRow(int64) error } +// RecordReaderWithError is implemented by record readers that expose failures +// while replacing their page reader. It supplements RecordReader without +// changing the existing no-error SetPageReader method. +type RecordReaderWithError interface { + RecordReader + SetPageReaderWithError(PageReader) error +} + // BinaryRecordReader provides an extra GetBuilderChunks function above and beyond // the plain RecordReader to allow for efficiently building chunked arrays. type BinaryRecordReader interface { @@ -112,6 +120,7 @@ type recordReaderImpl interface { NullCount() int64 Values() []byte SetPageReader(PageReader) + SetPageReaderWithError(PageReader) error Retain() Release() } @@ -173,7 +182,11 @@ func (pr *primitiveRecordReader) Release() { } func (pr *primitiveRecordReader) SetPageReader(rdr PageReader) { - pr.setPageReader(rdr) + _ = pr.setPageReader(rdr) +} + +func (pr *primitiveRecordReader) SetPageReaderWithError(rdr PageReader) error { + return pr.setPageReader(rdr) } func (pr *primitiveRecordReader) ReleaseValidBits() *memory.Buffer { @@ -393,8 +406,12 @@ func (rr *recordReader) HasMore() bool { } func (rr *recordReader) SetPageReader(pr PageReader) { + _ = rr.SetPageReaderWithError(pr) +} + +func (rr *recordReader) SetPageReaderWithError(pr PageReader) error { rr.atRecStart = true - rr.recordReaderImpl.SetPageReader(pr) + return rr.recordReaderImpl.SetPageReaderWithError(pr) } func (rr *recordReader) ValuesWritten() int { diff --git a/parquet/pqarrow/column_reader_close_test.go b/parquet/pqarrow/column_reader_close_test.go new file mode 100644 index 00000000..d6dcfef9 --- /dev/null +++ b/parquet/pqarrow/column_reader_close_test.go @@ -0,0 +1,110 @@ +// 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 pqarrow + +import ( + "bytes" + "context" + "errors" + "testing" + + "github.com/apache/arrow-go/v18/arrow" + "github.com/apache/arrow-go/v18/arrow/array" + "github.com/apache/arrow-go/v18/arrow/memory" + "github.com/apache/arrow-go/v18/parquet" + "github.com/apache/arrow-go/v18/parquet/file" + "github.com/stretchr/testify/require" +) + +type pageReaderReplacementErrorRecordReader struct { + file.RecordReader + err error + seekCalls int +} + +func (r *pageReaderReplacementErrorRecordReader) SetPageReaderWithError(pr file.PageReader) error { + return r.err +} + +func (r *pageReaderReplacementErrorRecordReader) SeekToRow(row int64) error { + r.seekCalls++ + return r.RecordReader.SeekToRow(row) +} + +func (r *pageReaderReplacementErrorRecordReader) Release() { + r.SetPageReader(nil) + r.RecordReader.Release() +} + +type trackingPageReader struct { + file.PageReader + closed bool +} + +func (r *trackingPageReader) Close() error { + r.closed = true + return nil +} + +func TestLeafReaderSeekToRowPropagatesPageReaderReplacementError(t *testing.T) { + mem := memory.NewCheckedAllocator(memory.DefaultAllocator) + defer mem.AssertSize(t, 0) + + builder := array.NewInt32Builder(mem) + builder.AppendValues([]int32{1, 2}, nil) + values := builder.NewInt32Array() + builder.Release() + defer values.Release() + + field := arrow.Field{Name: "value", Type: arrow.PrimitiveTypes.Int32} + schema := arrow.NewSchema([]arrow.Field{field}, nil) + column := arrow.NewColumnFromArr(field, values) + tbl := array.NewTable(schema, []arrow.Column{column}, -1) + column.Release() + defer tbl.Release() + + var buf bytes.Buffer + require.NoError(t, WriteTable(tbl, &buf, tbl.NumRows(), nil, NewArrowWriterProperties(WithAllocator(mem)))) + + pfile, err := file.NewParquetReader(bytes.NewReader(buf.Bytes()), file.WithReadProps(parquet.NewReaderProperties(mem))) + require.NoError(t, err) + defer pfile.Close() + + reader, err := NewFileReader(pfile, ArrowReadProperties{}, mem) + require.NoError(t, err) + + columnReader, err := reader.GetColumn(context.Background(), 0) + require.NoError(t, err) + defer columnReader.Release() + + leaf, ok := columnReader.colReaderImpl.(*leafReader) + require.True(t, ok) + + closeErr := errors.New("page reader close failed") + wrapped := &pageReaderReplacementErrorRecordReader{ + RecordReader: leaf.recordRdr, + err: closeErr, + } + leaf.recordRdr = wrapped + + require.ErrorIs(t, columnReader.SeekToRow(0), closeErr) + require.Zero(t, wrapped.seekCalls) + + replacement := &trackingPageReader{} + require.ErrorIs(t, leaf.setPageReader(replacement), closeErr) + require.True(t, replacement.closed) +} diff --git a/parquet/pqarrow/column_readers.go b/parquet/pqarrow/column_readers.go index 2d571038..5c216946 100644 --- a/parquet/pqarrow/column_readers.go +++ b/parquet/pqarrow/column_readers.go @@ -162,13 +162,32 @@ func (lr *leafReader) clearOut() (out *arrow.Chunked) { func (lr *leafReader) Field() *arrow.Field { return lr.field } +func (lr *leafReader) setPageReader(pr file.PageReader) error { + rr, ok := lr.recordRdr.(file.RecordReaderWithError) + if !ok { + if pr != nil { + _ = pr.Close() + } + return errors.New("record reader does not support error-returning page reader replacement") + } + if err := rr.SetPageReaderWithError(pr); err != nil { + if pr != nil { + _ = pr.Close() + } + return err + } + return nil +} + func (lr *leafReader) SeekToRow(rowIdx int64) error { pr, offset, err := lr.input.FindChunkForRow(rowIdx) if err != nil { return err } - lr.recordRdr.SetPageReader(pr) + if err := lr.setPageReader(pr); err != nil { + return err + } return lr.recordRdr.SeekToRow(offset) } @@ -180,7 +199,9 @@ func (lr *leafReader) nextRowGroup(remainingRows int64) error { if err != nil { return err } - lr.recordRdr.SetPageReader(pr) + if err := lr.setPageReader(pr); err != nil { + return err + } lr.curRGUncompressedBytes = uncompressedBytes lr.curRGNumRows = numRows // When called mid-batch, extend the builder's data buffer for the new row group.