diff --git a/parquet/pqarrow/column_readers.go b/parquet/pqarrow/column_readers.go index 3456fccb..b90072ea 100644 --- a/parquet/pqarrow/column_readers.go +++ b/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 { diff --git a/parquet/pqarrow/struct_reader_bench_test.go b/parquet/pqarrow/struct_reader_bench_test.go new file mode 100644 index 00000000..682028a7 --- /dev/null +++ b/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 +} diff --git a/parquet/pqarrow/struct_reader_test.go b/parquet/pqarrow/struct_reader_test.go new file mode 100644 index 00000000..65762191 --- /dev/null +++ b/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)