diff --git a/go/common/arrow_util/arrow_util.go b/go/common/arrow_util/arrow_util.go index ee21b0c..032ed5a 100644 --- a/go/common/arrow_util/arrow_util.go +++ b/go/common/arrow_util/arrow_util.go @@ -21,6 +21,7 @@ import ( "github.com/apache/arrow/go/v12/arrow/memory" "github.com/apache/arrow/go/v12/parquet/file" "github.com/apache/arrow/go/v12/parquet/pqarrow" + "github.com/milvus-io/milvus-storage/go/common/constant" "github.com/milvus-io/milvus-storage/go/io/fs" "github.com/milvus-io/milvus-storage/go/storage/options" ) @@ -34,7 +35,7 @@ func MakeArrowFileReader(fs fs.Fs, filePath string) (*pqarrow.FileReader, error) if err != nil { return nil, err } - return pqarrow.NewFileReader(parquetReader, pqarrow.ArrowReadProperties{}, memory.DefaultAllocator) + return pqarrow.NewFileReader(parquetReader, pqarrow.ArrowReadProperties{BatchSize: constant.ReadBatchSize}, memory.DefaultAllocator) } func MakeArrowRecordReader(reader *pqarrow.FileReader, opts *options.ReadOptions) (array.RecordReader, error) { diff --git a/go/io/format/parquet/file_reader.go b/go/io/format/parquet/file_reader.go index b41df24..d305102 100644 --- a/go/io/format/parquet/file_reader.go +++ b/go/io/format/parquet/file_reader.go @@ -1,11 +1,11 @@ // Copyright 2023 Zilliz -// +// // Licensed 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. @@ -24,6 +24,7 @@ import ( "github.com/apache/arrow/go/v12/parquet/metadata" "github.com/apache/arrow/go/v12/parquet/pqarrow" "github.com/bits-and-blooms/bitset" + "github.com/milvus-io/milvus-storage/go/common/constant" "github.com/milvus-io/milvus-storage/go/filter" "github.com/milvus-io/milvus-storage/go/io/fs" "github.com/milvus-io/milvus-storage/go/storage/options" @@ -210,7 +211,7 @@ func NewFileReader(fs fs.Fs, filePath string, options *options.ReadOptions) (*Fi return nil, err } - reader, err := pqarrow.NewFileReader(parquetReader, pqarrow.ArrowReadProperties{BatchSize: 1}, memory.DefaultAllocator) + reader, err := pqarrow.NewFileReader(parquetReader, pqarrow.ArrowReadProperties{BatchSize: constant.ReadBatchSize}, memory.DefaultAllocator) if err != nil { return nil, err }