-
Notifications
You must be signed in to change notification settings - Fork 8
/
s2idx.go
160 lines (133 loc) · 3.79 KB
/
s2idx.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
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
package index
import (
"bytes"
"encoding/binary"
"github.com/akhenakh/oureadb/index/geodata"
"github.com/akhenakh/oureadb/store"
"github.com/golang/geo/s2"
"github.com/pkg/errors"
)
// S2FlatIdx a flat S2 region cover index
type S2FlatIdx struct {
// s2 level to index
level int
// prefix for the keys
prefix []byte
store.KVStore
}
// NewS2FlatIdx returns a new indexer
func NewS2FlatIdx(s store.KVStore, prefix []byte, level int) *S2FlatIdx {
return &S2FlatIdx{
KVStore: s,
prefix: prefix,
level: level,
}
}
// GeoIndex is geo indexing the geo data
// it's not storing GeoData itself but only the geo index of the cover
// id is the key referring to the GeoData stored somewhere else
func (idx *S2FlatIdx) GeoIndex(gd *geodata.GeoData, id GeoID) error {
cu, err := idx.Covering(gd)
if err != nil {
return errors.Wrap(err, "generating cover failed")
}
// no cover for this geo object this is probably an error
if len(cu) == 0 {
return errors.New("geo object can't be indexed, empty cover")
}
kv, err := idx.KVStore.Writer()
if err != nil {
return err
}
batch := kv.NewBatch()
defer batch.Close()
// For each cell we store
// a key prefix+cellid+id -> no values
for _, c := range cu {
k := make([]byte, len(idx.prefix))
copy(k, idx.prefix)
k = append(k, itob(uint64(c))...)
k = append(k, []byte(id)...)
batch.Set(k, nil)
}
return kv.ExecuteBatch(batch)
}
// GeoIdsAtCell returns all GeoData keys contained in the cell
func (idx *S2FlatIdx) GeoIdsAtCell(c s2.CellID) ([]GeoID, error) {
if c.Level() != idx.level {
return nil, errors.New("requested a cellID with a different level than the index")
}
// add the prefix to the queried key
k := make([]byte, len(idx.prefix))
copy(k, idx.prefix)
k = append(k, itob(uint64(c))...)
var res []GeoID
kv, err := idx.Reader()
defer kv.Close()
if err != nil {
return nil, err
}
// iterate entries with cell's prefix
iter := kv.PrefixIterator(k)
defer iter.Close()
for {
kid, _, ok := iter.Current()
if !ok {
break
}
_, id, err := idx.keyToValues(kid)
if err != nil {
return nil, errors.Wrap(err, "read back failed key from db")
}
res = append(res, id)
iter.Next()
}
return res, nil
}
// GeoIdsAtCells returns all GeoData keys contained in the cells, without duplicates
func (idx *S2FlatIdx) GeoIdsAtCells(cells []s2.CellID) ([]GeoID, error) {
m := make(map[string]struct{})
for _, c := range cells {
ids, err := idx.GeoIdsAtCell(c)
if err != nil {
return nil, errors.Wrap(err, "fetching geo ids from cells failed")
}
for _, id := range ids {
m[string(id)] = struct{}{}
}
}
res := make([]GeoID, len(m))
var i int
for k := range m {
res[i] = []byte(k)
i++
}
return res, nil
}
// GeoIdsRadiusQuery returns the GeoID found in the index inside radius
// note you should check the returned GeoData is really inside/intersects the cap
func (idx *S2FlatIdx) GeoIdsRadiusQuery(lat, lng, radius float64) ([]GeoID, error) {
center := s2.PointFromLatLng(s2.LatLngFromDegrees(lat, lng))
cap := s2.CapFromCenterArea(center, s2RadialAreaMeters(radius))
coverer := &s2.RegionCoverer{MinLevel: idx.level, MaxLevel: idx.level}
cu := coverer.Covering(cap)
return idx.GeoIdsAtCells(cu)
}
// Covering is generating the cover of a GeoData
func (idx *S2FlatIdx) Covering(gd *geodata.GeoData) (s2.CellUnion, error) {
coverer := &s2.RegionCoverer{MinLevel: idx.level, MaxLevel: idx.level}
return gd.Cover(coverer)
}
func (idx *S2FlatIdx) keyToValues(k []byte) (c s2.CellID, id GeoID, err error) {
// prefix+cellid+id
if len(k) <= len(idx.prefix)+8 {
return c, id, errors.New("invalid key")
}
buf := bytes.NewBuffer(k[len(idx.prefix):])
err = binary.Read(buf, binary.BigEndian, &c)
if err != nil {
return c, id, errors.Wrap(err, "read back cell key failed")
}
id = k[len(idx.prefix)+8:]
return
}