diff --git a/parquet/metadata/row_group.go b/parquet/metadata/row_group.go index b578e343a..1c698fc88 100644 --- a/parquet/metadata/row_group.go +++ b/parquet/metadata/row_group.go @@ -20,6 +20,7 @@ import ( "fmt" "reflect" + "github.com/apache/arrow-go/v18/arrow" "github.com/apache/arrow-go/v18/parquet" "github.com/apache/arrow-go/v18/parquet/internal/encryption" format "github.com/apache/arrow-go/v18/parquet/internal/gen-go/parquet" @@ -79,8 +80,8 @@ func (r *RowGroupMetaData) Ordinal() int16 { return r.rowGroup.GetOrdinal() } // ColumnChunk returns the metadata for the requested (0-based) chunk index func (r *RowGroupMetaData) ColumnChunk(i int) (*ColumnChunkMetaData, error) { - if i >= r.NumColumns() { - panic(fmt.Errorf("parquet: the file only has %d columns, requested metadata for column: %d", r.NumColumns(), i)) + if i < 0 || i >= r.NumColumns() { + return nil, fmt.Errorf("%w: parquet file only has %d columns, requested metadata for column %d", arrow.ErrIndex, r.NumColumns(), i) } return NewColumnChunkMetaData(r.rowGroup.Columns[i], r.Schema.Column(i), r.version, r.rowGroup.GetOrdinal(), int16(i), r.fileDecryptor) @@ -95,6 +96,9 @@ func (r *RowGroupMetaData) SortingColumns() []parquet.SortingColumn { // directly from the underlying thrift struct, avoiding the overhead of // constructing a full ColumnChunkMetaData. func (r *RowGroupMetaData) ColumnIndexLocation(i int) (IndexLocation, bool) { + if i < 0 || i >= r.NumColumns() { + return IndexLocation{}, false + } col := r.rowGroup.Columns[i] if col.IsSetColumnIndexOffset() { return IndexLocation{ @@ -109,6 +113,9 @@ func (r *RowGroupMetaData) ColumnIndexLocation(i int) (IndexLocation, bool) { // directly from the underlying thrift struct, avoiding the overhead of // constructing a full ColumnChunkMetaData. func (r *RowGroupMetaData) OffsetIndexLocation(i int) (IndexLocation, bool) { + if i < 0 || i >= r.NumColumns() { + return IndexLocation{}, false + } col := r.rowGroup.Columns[i] if col.IsSetOffsetIndexOffset() { return IndexLocation{ diff --git a/parquet/metadata/row_group_bounds_test.go b/parquet/metadata/row_group_bounds_test.go new file mode 100644 index 000000000..489777e3b --- /dev/null +++ b/parquet/metadata/row_group_bounds_test.go @@ -0,0 +1,79 @@ +// 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 metadata_test + +import ( + "testing" + + "github.com/apache/arrow-go/v18/arrow" + "github.com/apache/arrow-go/v18/parquet" + format "github.com/apache/arrow-go/v18/parquet/internal/gen-go/parquet" + "github.com/apache/arrow-go/v18/parquet/metadata" + "github.com/apache/arrow-go/v18/parquet/schema" + "github.com/stretchr/testify/require" +) + +func TestRowGroupColumnIndexBounds(t *testing.T) { + columnIndexOffset, columnIndexLength := int64(100), int32(10) + offsetIndexOffset, offsetIndexLength := int64(200), int32(20) + primitive := schema.Must(schema.NewPrimitiveNode( + "value", parquet.Repetitions.Required, parquet.Types.Int32, -1, -1, + )) + parquetSchema := schema.MustGroup(schema.NewGroupNode( + "schema", parquet.Repetitions.Required, schema.FieldList{primitive}, -1, + )) + rowGroup := metadata.NewRowGroupMetaData( + &format.RowGroup{Columns: []*format.ColumnChunk{{ + MetaData: &format.ColumnMetaData{ + Type: format.Type_INT32, + Encodings: []format.Encoding{format.Encoding_PLAIN}, + PathInSchema: []string{"value"}, + }, + ColumnIndexOffset: &columnIndexOffset, + ColumnIndexLength: &columnIndexLength, + OffsetIndexOffset: &offsetIndexOffset, + OffsetIndexLength: &offsetIndexLength, + }}}, + schema.NewSchema(parquetSchema), + nil, + nil, + ) + + _, err := rowGroup.ColumnChunk(-1) + require.ErrorIs(t, err, arrow.ErrIndex) + _, err = rowGroup.ColumnChunk(1) + require.ErrorIs(t, err, arrow.ErrIndex) + column, err := rowGroup.ColumnChunk(0) + require.NoError(t, err) + require.NotNil(t, column) + + _, ok := rowGroup.ColumnIndexLocation(-1) + require.False(t, ok) + _, ok = rowGroup.ColumnIndexLocation(1) + require.False(t, ok) + location, ok := rowGroup.ColumnIndexLocation(0) + require.True(t, ok) + require.Equal(t, metadata.IndexLocation{Offset: columnIndexOffset, Length: columnIndexLength}, location) + + _, ok = rowGroup.OffsetIndexLocation(-1) + require.False(t, ok) + _, ok = rowGroup.OffsetIndexLocation(1) + require.False(t, ok) + location, ok = rowGroup.OffsetIndexLocation(0) + require.True(t, ok) + require.Equal(t, metadata.IndexLocation{Offset: offsetIndexOffset, Length: offsetIndexLength}, location) +}