-
Notifications
You must be signed in to change notification settings - Fork 15
/
record_reader.go
94 lines (83 loc) · 3.05 KB
/
record_reader.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
// 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.
// See the License for the specific language governing permissions and
// limitations under the License.
package record_reader
import (
"github.com/apache/arrow/go/v12/arrow/array"
"github.com/milvus-io/milvus-storage/go/file/fragment"
"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/manifest"
"github.com/milvus-io/milvus-storage/go/storage/options"
"github.com/milvus-io/milvus-storage/go/storage/schema"
)
func MakeRecordReader(
m *manifest.Manifest,
s *schema.Schema,
f fs.Fs,
deleteFragments fragment.DeleteFragmentVector,
options *options.ReadOptions,
) array.RecordReader {
relatedColumns := make([]string, 0)
for _, column := range options.Columns {
relatedColumns = append(relatedColumns, column)
}
for _, filter := range options.Filters {
relatedColumns = append(relatedColumns, filter.GetColumnName())
}
scalarData := m.GetScalarFragments()
vectorData := m.GetVectorFragments()
onlyScalar := onlyContainScalarColumns(s, relatedColumns)
onlyVector := onlyContainVectorColumns(s, relatedColumns)
if onlyScalar || onlyVector {
var dataFragments fragment.FragmentVector
if onlyScalar {
dataFragments = scalarData
} else {
dataFragments = vectorData
}
return NewScanRecordReader(s, options, f, dataFragments, deleteFragments)
}
if len(options.Filters) > 0 && filtersOnlyContainPKAndVersion(s, options.FiltersV2) {
return NewMergeRecordReader(s, options, f, scalarData, vectorData, deleteFragments)
}
return NewFilterQueryReader(s, options, f, scalarData, vectorData, deleteFragments)
}
func onlyContainVectorColumns(schema *schema.Schema, relatedColumns []string) bool {
for _, column := range relatedColumns {
if schema.Options().VectorColumn != column && schema.Options().PrimaryColumn != column && schema.Options().VersionColumn != column {
return false
}
}
return true
}
func onlyContainScalarColumns(schema *schema.Schema, relatedColumns []string) bool {
for _, column := range relatedColumns {
if schema.Options().VectorColumn == column {
return false
}
}
return true
}
func filtersOnlyContainPKAndVersion(s *schema.Schema, filters []filter.Filter) bool {
for _, f := range filters {
if f.GetColumnName() != s.Options().PrimaryColumn &&
f.GetColumnName() != s.Options().VersionColumn {
return false
}
}
return true
}
func MakeScanDeleteReader(manifest *manifest.Manifest, fs fs.Fs) array.RecordReader {
return NewMultiFilesSequentialReader(fs, manifest.GetDeleteFragments(), manifest.GetSchema().DeleteSchema(), options.NewReadOptions())
}