-
Notifications
You must be signed in to change notification settings - Fork 8
/
s2_ids.go
121 lines (104 loc) · 3.04 KB
/
s2_ids.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
package segmenters
import (
"encoding/json"
"fmt"
"log"
"github.com/golang/geo/s2"
_segmenters "github.com/caraml-dev/xp/common/segmenters"
)
const (
// MinS2CellLevel is the permissible minimum value for S2 Cell level
MinS2CellLevel = 0
// MaxS2CellLevel is the permissible maximum value for S2 Cell level
MaxS2CellLevel = 30
)
type S2IDSegmenterConfig struct {
MinS2CellLevel int `json:"mins2celllevel"`
MaxS2CellLevel int `json:"maxs2celllevel"`
}
func NewS2IDSegmenter(configData json.RawMessage) (Segmenter, error) {
segmenterErrTpl := "failed to create segmenter (s2_ids): %s"
var config S2IDSegmenterConfig
err := json.Unmarshal(configData, &config)
if err != nil {
return nil, fmt.Errorf(segmenterErrTpl, err)
}
// Verify configured levels
if config.MinS2CellLevel < MinS2CellLevel ||
config.MinS2CellLevel > MaxS2CellLevel ||
config.MaxS2CellLevel < MinS2CellLevel ||
config.MaxS2CellLevel > MaxS2CellLevel {
return nil, fmt.Errorf(segmenterErrTpl,
fmt.Sprintf("S2 cell levels should be in the range %d - %d", MinS2CellLevel, MaxS2CellLevel))
}
if config.MinS2CellLevel > config.MaxS2CellLevel {
return nil, fmt.Errorf(segmenterErrTpl, "Min S2 cell level cannot be greater than max")
}
s2IDConfig := &_segmenters.SegmenterConfiguration{
Name: "s2_ids",
Type: _segmenters.SegmenterValueType_INTEGER,
Options: map[string]*_segmenters.SegmenterValue{},
MultiValued: true,
TreatmentRequestFields: &_segmenters.ListExperimentVariables{
Values: []*_segmenters.ExperimentVariables{
{
Value: []string{"s2id"},
},
{
Value: []string{"latitude", "longitude"},
},
},
},
Required: false,
Description: fmt.Sprintf("S2 Cell IDs between levels %d and %d are supported.", config.MinS2CellLevel, config.MaxS2CellLevel),
}
levels := []int{}
for i := config.MinS2CellLevel; i <= config.MaxS2CellLevel; i++ {
levels = append(levels, i)
}
return &s2ids{
NewBaseSegmenter(s2IDConfig),
levels,
}, nil
}
type s2ids struct {
Segmenter
AllowedLevels []int
}
func (s *s2ids) ValidateSegmenterAndConstraints(segment map[string]*_segmenters.ListSegmenterValue) error {
err := s.Segmenter.ValidateSegmenterAndConstraints(segment)
if err != nil {
return err
}
name := s.GetName()
// Additional check to see that s2id values are valid
listInputValues := segment[name]
for _, val := range listInputValues.GetValues() {
cellID := val.GetInteger()
cellIDObj := s2.CellID(uint64(cellID))
if cellID <= 0 || !cellIDObj.IsValid() {
return fmt.Errorf("One or more %s values is invalid", name)
} else {
cellLevel := cellIDObj.Level()
if !s.isValidLevel(cellLevel) {
return fmt.Errorf("One or more %s values is at level %d, only the following levels are allowed: %v",
name, cellLevel, s.AllowedLevels)
}
}
}
return nil
}
func (s *s2ids) isValidLevel(level int) bool {
for _, val := range s.AllowedLevels {
if val == level {
return true
}
}
return false
}
func init() {
err := Register("s2_ids", NewS2IDSegmenter)
if err != nil {
log.Fatal(err)
}
}