perf(parquet/pqarrow): avoid goroutines for serial struct reads

parquet/pqarrow/column_readers.go

@@ -295,12 +295,19 @@ func (sr *structReader) GetRepLevels() ([]int16, error) {
 }
 
 func (sr *structReader) SeekToRow(rowIdx int64) error {
-	var g errgroup.Group
 	if !sr.props.Parallel {
-		g.SetLimit(1)
+		var firstErr error
+		for _, rdr := range sr.children {
+			if err := rdr.SeekToRow(rowIdx); err != nil && firstErr == nil {
+				firstErr = err
+			}
+		}
+		return firstErr
 	}
 
+	var g errgroup.Group
 	for _, rdr := range sr.children {
+		rdr := rdr
 		g.Go(func() error {
 			return rdr.SeekToRow(rowIdx)
 		})
@@ -310,14 +317,21 @@ func (sr *structReader) SeekToRow(rowIdx int64) error {
 }
 
 func (sr *structReader) LoadBatch(nrecords int64) error {
+	if !sr.props.Parallel {
+		var firstErr error
+		for _, rdr := range sr.children {
+			if err := rdr.LoadBatch(nrecords); err != nil && firstErr == nil {
+				firstErr = err
+			}
+		}
+		return firstErr
+	}
+
 	// Load batches in parallel
 	// When reading structs with large numbers of columns, the serial load is very slow.
 	// This is especially true when reading Cloud Storage. Loading concurrently
 	// greatly improves performance.
 	g := new(errgroup.Group)
-	if !sr.props.Parallel {
-		g.SetLimit(1)
-	}
 	for _, rdr := range sr.children {
 		rdr := rdr
 		g.Go(func() error {

parquet/pqarrow/struct_reader_bench_test.go

@@ -0,0 +1,131 @@
+// 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_test
+
+import (
+	"bytes"
+	"context"
+	"fmt"
+	"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/compress"
+	"github.com/apache/arrow-go/v18/parquet/file"
+	"github.com/apache/arrow-go/v18/parquet/pqarrow"
+)
+
+func BenchmarkReadNestedStructSerial(b *testing.B) {
+	const (
+		nrows        = 1024
+		rowGroupSize = 64
+	)
+
+	for _, nchildren := range []int{1, 8, 32, 128} {
+		b.Run(fmt.Sprintf("children=%d", nchildren), func(b *testing.B) {
+			mem := memory.DefaultAllocator
+			tbl := makeWideNestedInt32Table(mem, nchildren, nrows)
+			defer tbl.Release()
+
+			var buf bytes.Buffer
+			writerProps := parquet.NewWriterProperties(parquet.WithCompression(compress.Codecs.Uncompressed))
+			if err := pqarrow.WriteTable(tbl, &buf, rowGroupSize, writerProps, pqarrow.DefaultWriterProps()); err != nil {
+				b.Fatal(err)
+			}
+			parquetData := buf.Bytes()
+
+			pf, err := file.NewParquetReader(bytes.NewReader(parquetData))
+			if err != nil {
+				b.Fatal(err)
+			}
+			defer pf.Close()
+
+			reader, err := pqarrow.NewFileReader(pf, pqarrow.ArrowReadProperties{
+				BatchSize: nrows,
+				Parallel:  false,
+			}, mem)
+			if err != nil {
+				b.Fatal(err)
+			}
+
+			rowGroups := make([]int, nrows/rowGroupSize)
+			for i := range rowGroups {
+				rowGroups[i] = i
+			}
+			includedLeaves := make(map[int]bool, nchildren)
+			for i := 0; i < nchildren; i++ {
+				includedLeaves[i] = true
+			}
+			fieldReader, err := reader.GetFieldReader(context.Background(), 0, includedLeaves, rowGroups)
+			if err != nil {
+				b.Fatal(err)
+			}
+			defer fieldReader.Release()
+
+			b.ReportAllocs()
+			b.SetBytes(int64(len(parquetData)))
+			b.ResetTimer()
+			for range b.N {
+				if err := fieldReader.SeekToRow(0); err != nil {
+					b.Fatal(err)
+				}
+				out, err := fieldReader.NextBatch(nrows)
+				if err != nil {
+					b.Fatal(err)
+				}
+				out.Release()
+			}
+		})
+	}
+}
+
+func makeWideNestedInt32Table(mem memory.Allocator, nchildren, nrows int) arrow.Table {
+	childFields := make([]arrow.Field, nchildren)
+	for i := range childFields {
+		childFields[i] = arrow.Field{
+			Name: fmt.Sprintf("child_%d", i),
+			Type: arrow.PrimitiveTypes.Int32,
+		}
+	}
+
+	structType := arrow.StructOf(childFields...)
+	schema := arrow.NewSchema([]arrow.Field{{Name: "nested", Type: structType}}, nil)
+	builder := array.NewStructBuilder(mem, structType)
+	defer builder.Release()
+
+	valid := make([]bool, nrows)
+	values := make([]int32, nrows)
+	for i := range valid {
+		valid[i] = true
+		values[i] = int32(i)
+	}
+	builder.AppendValues(valid)
+	for i := 0; i < nchildren; i++ {
+		builder.FieldBuilder(i).(*array.Int32Builder).AppendValues(values, nil)
+	}
+
+	arr := builder.NewStructArray()
+	chunked := arrow.NewChunked(structType, []arrow.Array{arr})
+	column := arrow.NewColumn(schema.Field(0), chunked)
+	table := array.NewTable(schema, []arrow.Column{*column}, int64(nrows))
+	column.Release()
+	chunked.Release()
+	arr.Release()
+	return table
+}

parquet/pqarrow/struct_reader_test.go

@@ -0,0 +1,120 @@
+// 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 (
+	"errors"
+	"testing"
+
+	"github.com/apache/arrow-go/v18/arrow"
+	"github.com/stretchr/testify/require"
+)
+
+type structReaderTestChild struct {
+	name      string
+	events    *[]string
+	seekRow   int64
+	batchSize int64
+	seekErr   error
+	loadErr   error
+}
+
+func (r *structReaderTestChild) LoadBatch(nrecords int64) error {
+	*r.events = append(*r.events, r.name)
+	r.batchSize = nrecords
+	return r.loadErr
+}
+
+func (r *structReaderTestChild) BuildArray(int64) (*arrow.Chunked, error) { return nil, nil }
+
+func (r *structReaderTestChild) GetDefLevels() ([]int16, error) { return nil, nil }
+
+func (r *structReaderTestChild) GetRepLevels() ([]int16, error) { return nil, nil }
+
+func (r *structReaderTestChild) Field() *arrow.Field {
+	return &arrow.Field{Name: r.name, Type: arrow.PrimitiveTypes.Int32}
+}
+
+func (r *structReaderTestChild) SeekToRow(row int64) error {
+	*r.events = append(*r.events, r.name)
+	r.seekRow = row
+	return r.seekErr
+}
+
+func (r *structReaderTestChild) IsOrHasRepeatedChild() bool { return false }
+
+func (r *structReaderTestChild) Retain() {}
+
+func (r *structReaderTestChild) Release() {}
+
+func TestStructReaderSerialOperationsVisitEveryChild(t *testing.T) {
+	seekErr := errors.New("seek failed")
+	loadErr := errors.New("load failed")
+
+	tests := []struct {
+		name        string
+		call        func(*structReader) error
+		expectedErr error
+		check       func(*testing.T, []*structReaderTestChild)
+	}{
+		{
+			name:        "seek to row",
+			call:        func(reader *structReader) error { return reader.SeekToRow(42) },
+			expectedErr: seekErr,
+			check: func(t *testing.T, children []*structReaderTestChild) {
+				for _, child := range children {
+					require.Equal(t, int64(42), child.seekRow)
+				}
+			},
+		},
+		{
+			name:        "load batch",
+			call:        func(reader *structReader) error { return reader.LoadBatch(128) },
+			expectedErr: loadErr,
+			check: func(t *testing.T, children []*structReaderTestChild) {
+				for _, child := range children {
+					require.Equal(t, int64(128), child.batchSize)
+				}
+			},
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			events := make([]string, 0, 3)
+			children := []*structReaderTestChild{
+				{name: "first", events: &events},
+				{name: "second", events: &events, seekErr: seekErr, loadErr: loadErr},
+				{name: "third", events: &events, seekErr: errors.New("later seek failed"), loadErr: errors.New("later load failed")},
+			}
+
+			readers := make([]*ColumnReader, len(children))
+			for i, child := range children {
+				readers[i] = &ColumnReader{colReaderImpl: child}
+			}
+
+			reader := &structReader{children: readers}
+			err := tt.call(reader)
+
+			require.ErrorIs(t, err, tt.expectedErr, "the first child error should be returned")
+			require.Equal(t, []string{"first", "second", "third"}, events)
+			tt.check(t, children)
+		})
+	}
+}
+
+var _ colReaderImpl = (*structReaderTestChild)(nil)