Repository: segmentio/parquet-go
Pull request #417 (2022-11-30 18:19:23)

Add support for synchronous page reading
This PR adds a File level config option that controls whether or not pages are read synchronously. It defaults to async to maintain backwards compatibility.

I'm not married to any of the names in this PR so feel free to suggest alternatives.

cc @asubiotto 


Patch:
diff --git a/buffer.go b/buffer.go
index 42248a4..c29a027 100644
--- a/buffer.go
+++ b/buffer.go
@@ -256,7 +256,7 @@ func (buf *Buffer) WriteRowGroup(rowGroup RowGroup) (int64, error) {
 //
 // The buffer and the returned reader share memory. Mutating the buffer
 // concurrently to reading rows may result in non-deterministic behavior.
-func (buf *Buffer) Rows() Rows { return &rowGroupRows{rowGroup: buf} }
+func (buf *Buffer) Rows() Rows { return newRowGroupRows(buf, ReadModeSync) }
 
 // bufferWriter is an adapter for Buffer which implements both RowWriter and
 // PageWriter to enable optimizations in CopyRows for types that support writing
diff --git a/config.go b/config.go
index ca3a12c..9885544 100644
--- a/config.go
+++ b/config.go
@@ -10,6 +10,14 @@ import (
 	"github.com/segmentio/parquet-go/compress"
 )
 
+// ReadMode is an enum that is used to configure the way that a File reads pages.
+type ReadMode int
+
+const (
+	ReadModeSync  ReadMode = iota // ReadModeSync reads pages synchronously on demand.
+	ReadModeAsync                 // ReadModeAsync reads pages asynchronously in the background.
+)
+
 const (
 	DefaultColumnIndexSizeLimit = 16
 	DefaultColumnBufferCapacity = 16 * 1024
@@ -20,6 +28,7 @@ const (
 	DefaultSkipPageIndex        = false
 	DefaultSkipBloomFilters     = false
 	DefaultMaxRowsPerRowGroup   = math.MaxInt64
+	DefaultReadMode             = ReadModeSync
 )
 
 const (
@@ -83,6 +92,7 @@ type FileConfig struct {
 	SkipPageIndex    bool
 	SkipBloomFilters bool
 	ReadBufferSize   int
+	ReadMode         ReadMode
 }
 
 // DefaultFileConfig returns a new FileConfig value initialized with the
@@ -92,6 +102,7 @@ func DefaultFileConfig() *FileConfig {
 		SkipPageIndex:    DefaultSkipPageIndex,
 		SkipBloomFilters: DefaultSkipBloomFilters,
 		ReadBufferSize:   defaultReadBufferSize,
+		ReadMode:         DefaultReadMode,
 	}
 }
 
@@ -118,6 +129,8 @@ func (c *FileConfig) ConfigureFile(config *FileConfig) {
 	*config = FileConfig{
 		SkipPageIndex:    config.SkipPageIndex,
 		SkipBloomFilters: config.SkipBloomFilters,
+		ReadBufferSize:   config.ReadBufferSize,
+		ReadMode:         config.ReadMode,
 	}
 }
 
@@ -369,6 +382,15 @@ func SkipBloomFilters(skip bool) FileOption {
 	return fileOption(func(config *FileConfig) { config.SkipBloomFilters = skip })
 }
 
+// FileReadMode is a file configuration option which controls the way pages
+// are read. Currently the only two options are PageReadModeAsync and PageReadModeSync
+// which control whether or not pages are loaded asynchronously.
+//
+// Defaults to ReadModeAsync.
+func FileReadMode(mode ReadMode) FileOption {
+	return fileOption(func(config *FileConfig) { config.ReadMode = mode })
+}
+
 // ReadBufferSize is a file configuration option which controls the default
 // buffer sizes for reads made to the provided io.Reader. The default of 4096
 // is appropriate for disk based access but if your reader is backed by something
diff --git a/file.go b/file.go
index 2322b73..8b6c078 100644
--- a/file.go
+++ b/file.go
@@ -363,11 +363,13 @@ type fileRowGroup struct {
 	rowGroup *format.RowGroup
 	columns  []ColumnChunk
 	sorting  []SortingColumn
+	config   *FileConfig
 }
 
 func (g *fileRowGroup) init(file *File, schema *Schema, columns []*Column, rowGroup *format.RowGroup) {
 	g.schema = schema
 	g.rowGroup = rowGroup
+	g.config = file.config
 	g.columns = make([]ColumnChunk, len(rowGroup.Columns))
 	g.sorting = make([]SortingColumn, len(rowGroup.SortingColumns))
 	fileColumnChunks := make([]fileColumnChunk, len(rowGroup.Columns))
@@ -402,7 +404,7 @@ func (g *fileRowGroup) Schema() *Schema                 { return g.schema }
 func (g *fileRowGroup) NumRows() int64                  { return g.rowGroup.NumRows }
 func (g *fileRowGroup) ColumnChunks() []ColumnChunk     { return g.columns }
 func (g *fileRowGroup) SortingColumns() []SortingColumn { return g.sorting }
-func (g *fileRowGroup) Rows() Rows                      { return &rowGroupRows{rowGroup: g} }
+func (g *fileRowGroup) Rows() Rows                      { return newRowGroupRows(g, g.config.ReadMode) }
 
 type fileSortingColumn struct {
 	column     *Column
diff --git a/page.go b/page.go
index 9f594d5..883e7c5 100644
--- a/page.go
+++ b/page.go
@@ -238,11 +238,6 @@ func readPages(pages Pages, read chan<- asyncPage, seek <-chan int64, done <-cha
 	}
 }
 
-func copyPagesAndClose(w PageWriter, r Pages) (int64, error) {
-	defer r.Close()
-	return CopyPages(w, r)
-}
-
 type singlePage struct {
 	page    Page
 	seek    int64
@@ -301,34 +296,6 @@ func CopyPages(dst PageWriter, src PageReader) (numValues int64, err error) {
 	}
 }
 
-func forEachPageSlice(page Page, wantSize int64, do func(Page) error) error {
-	numRows := page.NumRows()
-	if numRows == 0 {
-		return nil
-	}
-
-	pageSize := page.Size()
-	numPages := (pageSize + (wantSize - 1)) / wantSize
-	rowIndex := int64(0)
-	if numPages < 2 {
-		return do(page)
-	}
-
-	for numPages > 0 {
-		lastRowIndex := rowIndex + ((numRows - rowIndex) / numPages)
-		pageSlice := page.Slice(rowIndex, lastRowIndex)
-		err := do(pageSlice)
-		Release(pageSlice)
-		if err != nil {
-			return err
-		}
-		rowIndex = lastRowIndex
-		numPages--
-	}
-
-	return nil
-}
-
 // errorPage is an implementation of the Page interface which always errors when
 // attempting to read its values.
 //
diff --git a/row_group.go b/row_group.go
index 9fd7bac..b3a9158 100644
--- a/row_group.go
+++ b/row_group.go
@@ -245,20 +245,21 @@ func (r *rowGroup) NumRows() int64                  { return r.numRows }
 func (r *rowGroup) ColumnChunks() []ColumnChunk     { return r.columns }
 func (r *rowGroup) SortingColumns() []SortingColumn { return r.sorting }
 func (r *rowGroup) Schema() *Schema                 { return r.schema }
-func (r *rowGroup) Rows() Rows                      { return &rowGroupRows{rowGroup: r} }
+func (r *rowGroup) Rows() Rows                      { return newRowGroupRows(r, ReadModeSync) }
 
 func NewRowGroupRowReader(rowGroup RowGroup) Rows {
-	return &rowGroupRows{rowGroup: rowGroup}
+	return newRowGroupRows(rowGroup, ReadModeSync)
 }
 
 type rowGroupRows struct {
-	rowGroup RowGroup
-	buffers  []Value
-	readers  []asyncPages
-	columns  []columnChunkRows
-	inited   bool
-	closed   bool
-	done     chan<- struct{}
+	rowGroup     RowGroup
+	buffers      []Value
+	readers      []Pages
+	columns      []columnChunkRows
+	inited       bool
+	closed       bool
+	done         chan<- struct{}
+	pageReadMode ReadMode
 }
 
 type columnChunkRows struct {
@@ -277,18 +278,35 @@ func (r *rowGroupRows) buffer(i int) []Value {
 	return r.buffers[j:k:k]
 }
 
+func newRowGroupRows(rowGroup RowGroup, pageReadMode ReadMode) *rowGroupRows {
+	return &rowGroupRows{
+		rowGroup:     rowGroup,
+		pageReadMode: pageReadMode,
+	}
+}
+
 func (r *rowGroupRows) init() {
 	columns := r.rowGroup.ColumnChunks()
 
 	r.buffers = make([]Value, len(columns)*columnBufferSize)
-	r.readers = make([]asyncPages, len(columns))
+	r.readers = make([]Pages, len(columns))
 	r.columns = make([]columnChunkRows, len(columns))
 
-	done := make(chan struct{})
-	r.done = done
-
-	for i, column := range columns {
-		r.readers[i].init(column.Pages(), done)
+	switch r.pageReadMode {
+	case ReadModeAsync:
+		done := make(chan struct{})
+		r.done = done
+		readers := make([]asyncPages, len(columns))
+		for i, column := range columns {
+			readers[i].init(column.Pages(), done)
+			r.readers[i] = &readers[i]
+		}
+	case ReadModeSync:
+		for i, column := range columns {
+			r.readers[i] = column.Pages()
+		}
+	default:
+		panic(fmt.Sprintf("parquet: invalid page read mode: %d", r.pageReadMode))
 	}
 
 	r.inited = true
diff --git a/value.go b/value.go
index 5cf0053..4a1ae6a 100644
--- a/value.go
+++ b/value.go
@@ -204,6 +204,38 @@ func ValueOf(v interface{}) Value {
 	return makeValue(k, nil, reflect.ValueOf(v))
 }
 
+// BooleanValue constructs a BOOLEAN parquet value from the bool passed as
+// argument.
+func BooleanValue(value bool) Value { return makeValueBoolean(value) }
+
+// Int32Value constructs a INT32 parquet value from the int32 passed as
+// argument.
+func Int32Value(value int32) Value { return makeValueInt32(value) }
+
+// Int64Value constructs a INT64 parquet value from the int64 passed as
+// argument.
+func Int64Value(value int64) Value { return makeValueInt64(value) }
+
+// Int96Value constructs a INT96 parquet value from the deprecated.Int96 passed
+// as argument.
+func Int96Value(value deprecated.Int96) Value { return makeValueInt96(value) }
+
+// FloatValue constructs a FLOAT parquet value from the float32 passed as
+// argument.
+func FloatValue(value float32) Value { return makeValueFloat(value) }
+
+// DoubleValue constructs a DOUBLE parquet value from the float64 passed as
+// argument.
+func DoubleValue(value float64) Value { return makeValueDouble(value) }
+
+// ByteArrayValue constructs a BYTE_ARRAY parquet value from the byte slice
+// passed as argument.
+func ByteArrayValue(value []byte) Value { return makeValueBytes(ByteArray, value) }
+
+// FixedLenByteArrayValue constructs a BYTE_ARRAY parquet value from the byte
+// slice passed as argument.
+func FixedLenByteArrayValue(value []byte) Value { return makeValueBytes(FixedLenByteArray, value) }
+
 func makeValue(k Kind, lt *format.LogicalType, v reflect.Value) Value {
 	switch v.Type() {
 	case reflect.TypeOf(time.Time{}):
diff --git a/writer.go b/writer.go
index 00ac1de..0ee2e24 100644
--- a/writer.go
+++ b/writer.go
@@ -179,7 +179,9 @@ func (w *Writer) WriteRowGroup(rowGroup RowGroup) (int64, error) {
 		return 0, err
 	}
 	w.writer.configureBloomFilters(rowGroup.ColumnChunks())
-	n, err := CopyRows(w.writer, rowGroup.Rows())
+	rows := rowGroup.Rows()
+	defer rows.Close()
+	n, err := CopyRows(w.writer, rows)
 	if err != nil {
 		return n, err
 	}
@@ -747,14 +749,6 @@ func (w *writer) WriteValues(values []Value) (numValues int, err error) {
 	return w.columns[values[0].Column()].WriteValues(values)
 }
 
-// This WritePage method satisfies the PageWriter interface as a mechanism to
-// allow writing whole pages of values instead of individual rows. It is called
-// indirectly by readers that implement WriteRowsTo and are able to leverage
-// the method to optimize writes.
-func (w *writer) WritePage(page Page) (int64, error) {
-	return w.columns[page.Column()].WritePage(page)
-}
-
 // One writerBuffers is used by each writer instance, the memory buffers here
 // are shared by all columns of the writer because serialization is not done
 // concurrently, which helps keep memory utilization low, both in the total
@@ -945,10 +939,6 @@ func (c *writerColumn) totalRowCount() int64 {
 	return n
 }
 
-func (c *writerColumn) canFlush() bool {
-	return c.columnBuffer.Size() >= int64(c.bufferSize/2)
-}
-
 func (c *writerColumn) flush() (err error) {
 	if c.numValues != 0 {
 		c.numValues = 0
@@ -1046,57 +1036,6 @@ func (c *writerColumn) WriteValues(values []Value) (numValues int, err error) {
 	return numValues, err
 }
 
-func (c *writerColumn) WritePage(page Page) (numValues int64, err error) {
-	// Page write optimizations are only available the column is not reindexing
-	// the values. If a dictionary is present, the column needs to see each
-	// individual value in order to re-index them in the dictionary.
-	if c.dictionary == nil || c.dictionary == page.Dictionary() {
-		// If the column had buffered values, we continue writing values from
-		// the page into the column buffer if it would have caused producing a
-		// page less than half the size of the target; if there were enough
-		// buffered values already, we have to flush the buffered page instead
-		// otherwise values would get reordered.
-		if c.numValues > 0 && c.canFlush() {
-			if err := c.flush(); err != nil {
-				return 0, err
-			}
-		}
-
-		// If we were successful at flushing buffered values, we attempt to
-		// optimize the write path by copying whole pages without decoding them
-		// into a sequence of values.
-		if c.numValues == 0 {
-			// Buffered pages may be larger than the target page size on the
-			// column, in which case multiple pages get written by slicing
-			// the original page into sub-pages.
-			err = forEachPageSlice(page, int64(c.bufferSize), func(p Page) error {
-				n, err := c.writeDataPage(p)
-				numValues += n
-				return err
-			})
-			return numValues, err
-		}
-	}
-
-	// Pages that implement neither of those interfaces can still be
-	// written by copying their values into the column buffer and flush
-	// them to compressed page buffers as if the program had written
-	// rows individually.
-	return c.writePageValues(page.Values())
-}
-
-func (c *writerColumn) writePageValues(page ValueReader) (numValues int64, err error) {
-	numValues, err = CopyValues(c, page)
-	if err == nil && c.canFlush() {
-		// Always attempt to flush after writing a full page if we have enough
-		// buffered values; the intent is to leave the column clean so that
-		// subsequent calls to the WritePage method can use optimized write path
-		// to bypass buffering.
-		err = c.flush()
-	}
-	return numValues, err
-}
-
 func (c *writerColumn) writeBloomFilter(w io.Writer) error {
 	e := thrift.NewEncoder(c.header.protocol.NewWriter(w))
 	h := bloomFilterHeader(c.columnFilter)
@@ -1417,9 +1356,7 @@ var (
 	_ RowGroupWriter      = (*Writer)(nil)
 
 	_ RowWriter   = (*writer)(nil)
-	_ PageWriter  = (*writer)(nil)
 	_ ValueWriter = (*writer)(nil)
 
-	_ PageWriter  = (*writerColumn)(nil)
 	_ ValueWriter = (*writerColumn)(nil)
 )


Test patch:
diff --git a/row_group_test.go b/row_group_test.go
index 4474438..cb85a97 100644
--- a/row_group_test.go
+++ b/row_group_test.go
@@ -2,6 +2,7 @@ package parquet_test
 
 import (
 	"bytes"
+	"io"
 	"reflect"
 	"sort"
 	"testing"
@@ -421,3 +422,27 @@ func TestMergeRowGroups(t *testing.T) {
 		})
 	}
 }
+
+func TestWriteRowGroupClosesRows(t *testing.T) {
+	var rows []*wrappedRows
+	rg := wrappedRowGroup{
+		RowGroup: newPeopleFile([]Person{{}}),
+		rowsCallback: func(r parquet.Rows) parquet.Rows {
+			wrapped := &wrappedRows{Rows: r}
+			rows = append(rows, wrapped)
+			return wrapped
+		},
+	}
+	writer := parquet.NewWriter(io.Discard)
+	if _, err := writer.WriteRowGroup(rg); err != nil {
+		t.Fatal(err)
+	}
+	if err := writer.Close(); err != nil {
+		t.Fatal(err)
+	}
+	for _, r := range rows {
+		if !r.closed {
+			t.Fatal("rows not closed")
+		}
+	}
+}