diff --git a/parquet/pqarrow/file_reader.go b/parquet/pqarrow/file_reader.go index 34992f8a..b4e46008 100644 --- a/parquet/pqarrow/file_reader.go +++ b/parquet/pqarrow/file_reader.go @@ -28,6 +28,7 @@ import ( "github.com/apache/arrow-go/v18/arrow" "github.com/apache/arrow-go/v18/arrow/array" "github.com/apache/arrow-go/v18/arrow/arrio" + "github.com/apache/arrow-go/v18/arrow/bitutil" "github.com/apache/arrow-go/v18/arrow/memory" "github.com/apache/arrow-go/v18/internal/utils" "github.com/apache/arrow-go/v18/parquet" @@ -516,9 +517,11 @@ func (fr *FileReader) GetRecordReader(ctx context.Context, colIndices, rowGroups nrows += fr.rdr.MetaData().RowGroup(rg).NumRows() } - batchSize := fr.Props.BatchSize + var batchSize int64 if fr.Props.BatchSize <= 0 { batchSize = nrows + } else { + batchSize = min(fr.Props.BatchSize, int64(bitutil.NextPowerOf2(int(nrows)))) } rr := &recordReader{ numRows: nrows,