fix(parquet): make column reader cleanup idempotent

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
 }

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)
+}

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 {

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)
+}

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.