diff --git a/parquet/file/record_reader.go b/parquet/file/record_reader.go index 87ae20f3..62090443 100644 --- a/parquet/file/record_reader.go +++ b/parquet/file/record_reader.go @@ -522,7 +522,7 @@ func (rr *recordReader) Reset() { // process no more levels than necessary to delimit the indicated // number of logical records. updates internal state of recordreader // returns number of records delimited -func (rr *recordReader) delimitRecords(numRecords int64) (recordsRead, valsToRead int64) { +func (rr *recordReader) delimitRecords(numRecords int64) (recordsRead, valsToRead int64, err error) { var ( curRep int16 curDef int16 @@ -533,6 +533,9 @@ func (rr *recordReader) delimitRecords(numRecords int64) (recordsRead, valsToRea for rr.levelsPos < rr.levelsWritten { curRep, repLevels = repLevels[0], repLevels[1:] + if rr.atRecStart && curRep != 0 { + return 0, 0, errors.New("parquet: record starts with a nonzero repetition level") + } if curRep == 0 { // if at record start, we are seeing the start of a record // for the second time, such as after repeated calls to delimitrecords. @@ -576,7 +579,10 @@ func (rr *recordReader) ReadRecordData(numRecords int64) (int64, error) { ) if rr.Descriptor().MaxRepetitionLevel() > 0 { - recordsRead, valuesToRead = rr.delimitRecords(numRecords) + recordsRead, valuesToRead, err = rr.delimitRecords(numRecords) + if err != nil { + return 0, err + } } else if rr.Descriptor().MaxDefinitionLevel() > 0 { // no repetition levels, skip delimiting logic. each level // represents null or not null entry diff --git a/parquet/file/record_reader_internal_test.go b/parquet/file/record_reader_internal_test.go new file mode 100644 index 00000000..e9285edd --- /dev/null +++ b/parquet/file/record_reader_internal_test.go @@ -0,0 +1,46 @@ +// 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 ( + "testing" + + "github.com/apache/arrow-go/v18/arrow" + "github.com/apache/arrow-go/v18/arrow/memory" + "github.com/apache/arrow-go/v18/parquet" + "github.com/apache/arrow-go/v18/parquet/schema" + "github.com/stretchr/testify/require" +) + +func TestDelimitRecordsRejectsNonzeroInitialRepetitionLevel(t *testing.T) { + descr := schema.NewColumn(schema.NewInt32Node("values", parquet.Repetitions.Repeated, -1), 1, 1) + rr := newRecordReader(descr, LevelInfo{DefLevel: 1}, memory.DefaultAllocator, nil).(*recordReader) + defer rr.Release() + + rr.repLevels.ResizeNoShrink(arrow.Int16SizeBytes) + rr.defLevels.ResizeNoShrink(arrow.Int16SizeBytes) + rr.RepLevels()[0] = 1 + rr.DefLevels()[0] = 1 + rr.levelsWritten = 1 + rr.atRecStart = true + + records, values, err := rr.delimitRecords(1) + require.Error(t, err) + require.Zero(t, records) + require.Zero(t, values) + require.Zero(t, rr.levelsPos) +}