forked from parsyl/parquet
-
Notifications
You must be signed in to change notification settings - Fork 0
/
structs.go
81 lines (70 loc) · 1.67 KB
/
structs.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
package structs
import (
"fmt"
"strings"
sch "github.com/inigolabs/parquet/schema"
)
// Struct generates a struct definition based on the
// parquet schema.
func Struct(structName string, schema []*sch.SchemaElement) string {
if len(schema) == 0 {
return ""
}
schema[0].Name = structName
_, out := getStruct(schema[0], schema[1:])
if strings.Contains(out, "%s") {
out = fmt.Sprintf(out, "")
}
return out
}
func getStruct(parent *sch.SchemaElement, children []*sch.SchemaElement) (int, string) {
str := fmt.Sprintf(`type %s struct {
%%s
}`, strings.Title(parent.Name))
var i, j int
var fields string
for i < int(*parent.NumChildren) {
ch := children[i+j]
fields = fmt.Sprintf("%s\n%s", fields, field(ch))
if ch.NumChildren != nil && int(*ch.NumChildren) > 0 {
n, s := getStruct(ch, children[i+j+1:])
j += n
str += fmt.Sprintf("\n\n%s", s)
}
i++
}
return i + j, fmt.Sprintf(str, fields)
}
func field(elem *sch.SchemaElement) string {
n := strings.Title(elem.Name)
t := n
if elem.Type != nil {
t = getType(elem.Type.String())
}
var ptr string
if elem.RepetitionType != nil && *elem.RepetitionType == sch.FieldRepetitionType_OPTIONAL {
ptr = "*"
}
return fmt.Sprintf("%s %s%s `parquet:\"%s\"`", n, ptr, t, elem.Name)
}
func getType(t string) string {
return parquetTypes[t]
}
var parquetTypes = map[string]string{
"BOOLEAN": "bool",
"INT32": "int32",
"INT64": "int64",
"FLOAT": "float32",
"DOUBLE": "float64",
"BYTE_ARRAY": "string",
}
var primitiveTypes = map[string]bool{
"bool": true,
"int32": true,
"uint32": true,
"int64": true,
"uint64": true,
"float32": true,
"float64": true,
"string": true,
}