-
Notifications
You must be signed in to change notification settings - Fork 342
/
orc_reader.go
128 lines (116 loc) · 3.5 KB
/
orc_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
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
package s3inventory
import (
"context"
"time"
"github.com/hashicorp/go-multierror"
"github.com/scritchley/orc"
)
type OrcInventoryFileReader struct {
reader *orc.Reader
cursor *orc.Cursor
ctx context.Context
orcSelect *OrcSelect
orcFile *OrcFile
}
type OrcField struct {
IndexInFile int
IndexInSelect int
}
type OrcSelect struct {
SelectFields []string // the list of fields to select from the file
IndexInSelect map[string]int // for each field, its index in the select query
IndexInFile map[string]int // for each field, its index in the original file
}
func getOrcSelect(typeDescription *orc.TypeDescription) *OrcSelect {
res := &OrcSelect{
SelectFields: nil,
IndexInFile: make(map[string]int),
IndexInSelect: make(map[string]int),
}
for i, field := range typeDescription.Columns() {
res.IndexInFile[field] = i
}
j := 0
for _, field := range inventoryFields {
if _, ok := res.IndexInFile[field]; ok {
res.SelectFields = append(res.SelectFields, field)
res.IndexInSelect[field] = j
j++
}
}
return res
}
func (r *OrcInventoryFileReader) inventoryObjectFromRow(rowData []interface{}) *InventoryObject {
obj := NewInventoryObject()
obj.Bucket = rowData[r.orcSelect.IndexInSelect[bucketFieldName]].(string)
obj.Key = rowData[r.orcSelect.IndexInSelect[keyFieldName]].(string)
if sizeIdx, ok := r.orcSelect.IndexInSelect[sizeFieldName]; ok && rowData[sizeIdx] != nil {
obj.Size = rowData[sizeIdx].(int64)
}
if lastModifiedIdx, ok := r.orcSelect.IndexInSelect[lastModifiedDateFieldName]; ok && rowData[lastModifiedIdx] != nil {
tm := rowData[lastModifiedIdx].(time.Time)
obj.LastModified = &tm
}
if eTagIdx, ok := r.orcSelect.IndexInSelect[eTagFieldName]; ok && rowData[eTagIdx] != nil {
obj.Checksum = rowData[eTagIdx].(string)
}
if isLatestIdx, ok := r.orcSelect.IndexInSelect[isLatestFieldName]; ok && rowData[isLatestIdx] != nil {
obj.IsLatest = rowData[isLatestIdx].(bool)
}
if isDeleteMarkerIdx, ok := r.orcSelect.IndexInSelect[isDeleteMarkerFieldName]; ok && rowData[isDeleteMarkerIdx] != nil {
obj.IsDeleteMarker = rowData[isDeleteMarkerIdx].(bool)
}
return obj
}
func (r *OrcInventoryFileReader) Read(n int) ([]*InventoryObject, error) {
res := make([]*InventoryObject, 0, n)
for {
select {
case <-r.ctx.Done():
return nil, r.ctx.Err()
default:
}
if !r.cursor.Next() {
if !r.cursor.Stripes() {
break
}
if !r.cursor.Next() {
break
}
}
obj := r.inventoryObjectFromRow(r.cursor.Row())
res = append(res, obj)
if len(res) == n {
break
}
}
return res, nil
}
func (r *OrcInventoryFileReader) GetNumRows() int64 {
return int64(r.reader.NumRows())
}
func (r *OrcInventoryFileReader) Close() error {
var combinedErr error
if err := r.cursor.Close(); err != nil {
combinedErr = multierror.Append(combinedErr, err)
}
if err := r.reader.Close(); err != nil {
combinedErr = multierror.Append(combinedErr, err)
}
if err := r.orcFile.Close(); err != nil {
combinedErr = multierror.Append(combinedErr, err)
}
return combinedErr
}
func (r *OrcInventoryFileReader) FirstObjectKey() string {
if len(r.reader.Metadata().StripeStats) == 0 {
return ""
}
return *r.reader.Metadata().StripeStats[0].GetColStats()[r.orcSelect.IndexInFile[keyFieldName]+1].StringStatistics.Minimum
}
func (r *OrcInventoryFileReader) LastObjectKey() string {
if len(r.reader.Metadata().StripeStats) == 0 {
return ""
}
return *r.reader.Metadata().StripeStats[0].GetColStats()[r.orcSelect.IndexInFile[keyFieldName]+1].StringStatistics.Maximum
}