Read batch size (#102)

Signed-off-by: sunby <sunbingyi1992@gmail.com>
milvus-io · Dec 26, 2023 · 5751940 · 5751940
1 parent 422d03d
commit 5751940
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 5 deletions.
diff --git 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
@@ -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
 	}