func addGeoData(t *testing.T, ps *store.Store, uid uint64, p geom.T, name string) {
g := types.Geo{p}
value, err := g.MarshalBinary()
require.NoError(t, err)
addEdgeToTypedValue(t, ps, "geometry", uid, types.GeoID, value)
addEdgeToTypedValue(t, ps, "name", uid, types.StringID, []byte(name))
}
func formData(t *testing.T, str string) string {
p, err := loadPolygon(str)
require.NoError(t, err)
d, err := wkb.Marshal(p, binary.LittleEndian)
require.NoError(t, err)
var g types.Geo
require.NoError(t, g.UnmarshalBinary(d))
gb, err := g.MarshalText()
require.NoError(t, err)
return string(gb)
}
func TestKeyGeneratorPolygon(t *testing.T) {
p, err := loadPolygon("zip.json")
require.NoError(t, err)
data, err := wkb.Marshal(p, binary.LittleEndian)
require.NoError(t, err)
var g types.Geo
err = g.UnmarshalBinary(data)
require.NoError(t, err)
keys, err := IndexTokens(&g)
require.NoError(t, err)
require.Len(t, keys, 65)
}
// FilterUids filters the uids based on the corresponding values and QueryData.
func FilterUids(uids *task.List, values []*task.Value, q *QueryData) *task.List {
x.AssertTruef(len(values) == len(uids.Uids), "lengths not matching")
rv := &task.List{}
for i := 0; i < len(values); i++ {
valBytes := values[i].Val
if bytes.Equal(valBytes, nil) {
continue
}
vType := values[i].ValType
if types.TypeID(vType) != types.GeoID {
continue
}
var g types.Geo
if err := g.UnmarshalBinary(valBytes); err != nil {
continue
}
if !q.MatchesFilter(g) {
continue
}
// we matched the geo filter, add the uid to the list
rv.Uids = append(rv.Uids, uids.Uids[i])
}
return rv
}
func formDataPolygon(t *testing.T, p *geom.Polygon) string {
d, err := wkb.Marshal(p, binary.LittleEndian)
require.NoError(t, err)
var g types.Geo
require.NoError(t, g.UnmarshalBinary(d))
gb, err := g.MarshalText()
require.NoError(t, err)
return string(gb)
}
func BenchmarkKeyGeneratorPoint(b *testing.B) {
p := geom.NewPoint(geom.XY).MustSetCoords(geom.Coord{-122.082506, 37.4249518})
data, err := wkb.Marshal(p, binary.LittleEndian)
if err != nil {
b.Error(err)
}
var g types.Geo
g.UnmarshalBinary(data)
b.ResetTimer()
for n := 0; n < b.N; n++ {
IndexTokens(&g)
}
}
func TestKeyGeneratorPoint(t *testing.T) {
p := geom.NewPoint(geom.XY).MustSetCoords(geom.Coord{-122.082506, 37.4249518})
data, err := wkb.Marshal(p, binary.LittleEndian)
require.NoError(t, err)
var g types.Geo
err = g.UnmarshalBinary(data)
require.NoError(t, err)
keys, err := IndexTokens(&g)
require.NoError(t, err)
require.Len(t, keys, MaxCellLevel-MinCellLevel+1+1) // +1 for the cover
}
func BenchmarkKeyGeneratorPolygon(b *testing.B) {
p, err := loadPolygon("zip.json")
if err != nil {
b.Error(err)
}
data, err := wkb.Marshal(p, binary.LittleEndian)
if err != nil {
b.Error(err)
}
var g types.Geo
g.UnmarshalBinary(data)
b.ResetTimer()
for n := 0; n < b.N; n++ {
IndexTokens(&g)
}
}
// IndexCells returns two cellunions. The first is a list of parents, which are all the cells upto
// the min level that contain this geometry. The second is the cover, which are the smallest
// possible cells required to cover the region. This makes it easier at query time to query only the
// parents or only the cover or both depending on whether it is a within, contains or intersects
// query.
func indexCells(g types.Geo) (parents, cover s2.CellUnion, err error) {
if g.Stride() != 2 {
return nil, nil, x.Errorf("Covering only available for 2D co-ordinates.")
}
switch v := g.T.(type) {
case *geom.Point:
p, c := indexCellsForPoint(v, MinCellLevel, MaxCellLevel)
return p, c, nil
case *geom.Polygon:
l, err := loopFromPolygon(v)
if err != nil {
return nil, nil, err
}
cover := coverLoop(l, MinCellLevel, MaxCellLevel, MaxCells)
parents := getParentCells(cover, MinCellLevel)
return parents, cover, nil
default:
return nil, nil, x.Errorf("Cannot index geometry of type %T", v)
}
}
func typeValueFromNQuad(nq *graph.NQuad) (types.Value, error) {
if nq.Value == nil || nq.Value.Val == nil {
return nil, nil
}
switch v := nq.Value.Val.(type) {
case *graph.Value_BytesVal:
b := types.Bytes(v.BytesVal)
return &b, nil
case *graph.Value_IntVal:
i := types.Int32(v.IntVal)
return &i, nil
case *graph.Value_StrVal:
s := types.String(v.StrVal)
return &s, nil
case *graph.Value_BoolVal:
b := types.Bool(v.BoolVal)
return &b, nil
case *graph.Value_DoubleVal:
f := types.Float(v.DoubleVal)
return &f, nil
case *graph.Value_GeoVal:
var geom types.Geo
err := geom.UnmarshalBinary(v.GeoVal)
if err != nil {
return nil, err
}
return &geom, nil
case nil:
log.Fatalf("Val being nil is already handled")
return nil, nil
default:
// Unknown type
return nil, x.Errorf("Unknown value type %T", v)
}
}
func populateGraph(t *testing.T) (string, string, *store.Store) {
// logrus.SetLevel(logrus.DebugLevel)
dir, err := ioutil.TempDir("", "storetest_")
require.NoError(t, err)
ps, err := store.NewStore(dir)
require.NoError(t, err)
schema.ParseBytes([]byte(schemaStr))
posting.Init(ps)
worker.Init(ps)
group.ParseGroupConfig("")
dir2, err := ioutil.TempDir("", "wal_")
require.NoError(t, err)
worker.StartRaftNodes(dir2)
// So, user we're interested in has uid: 1.
// She has 5 friends: 23, 24, 25, 31, and 101
addEdgeToUID(t, ps, "friend", 1, 23)
addEdgeToUID(t, ps, "friend", 1, 24)
addEdgeToUID(t, ps, "friend", 1, 25)
addEdgeToUID(t, ps, "friend", 1, 31)
addEdgeToUID(t, ps, "friend", 1, 101)
// Now let's add a few properties for the main user.
addEdgeToValue(t, ps, "name", 1, "Michonne")
addEdgeToValue(t, ps, "gender", 1, "female")
var coord types.Geo
err = coord.UnmarshalText([]byte("{\"Type\":\"Point\", \"Coordinates\":[1.1,2.0]}"))
require.NoError(t, err)
gData, err := coord.MarshalBinary()
require.NoError(t, err)
addEdgeToTypedValue(t, ps, "loc", 1, types.GeoID, gData)
data, err := types.Int32(15).MarshalBinary()
require.NoError(t, err)
addEdgeToTypedValue(t, ps, "age", 1, types.Int32ID, data)
addEdgeToValue(t, ps, "address", 1, "31, 32 street, Jupiter")
data, err = types.Bool(true).MarshalBinary()
require.NoError(t, err)
addEdgeToTypedValue(t, ps, "alive", 1, types.BoolID, data)
addEdgeToValue(t, ps, "age", 1, "38")
addEdgeToValue(t, ps, "survival_rate", 1, "98.99")
addEdgeToValue(t, ps, "sword_present", 1, "true")
addEdgeToValue(t, ps, "_xid_", 1, "mich")
// Now let's add a name for each of the friends, except 101.
addEdgeToTypedValue(t, ps, "name", 23, types.StringID, []byte("Rick Grimes"))
addEdgeToValue(t, ps, "age", 23, "15")
err = coord.UnmarshalText([]byte(`{"Type":"Polygon", "Coordinates":[[[0.0,0.0], [2.0,0.0], [2.0, 2.0], [0.0, 2.0]]]}`))
require.NoError(t, err)
gData, err = coord.MarshalBinary()
require.NoError(t, err)
addEdgeToTypedValue(t, ps, "loc", 23, types.GeoID, gData)
addEdgeToValue(t, ps, "address", 23, "21, mark street, Mars")
addEdgeToValue(t, ps, "name", 24, "Glenn Rhee")
addEdgeToValue(t, ps, "name", 25, "Daryl Dixon")
addEdgeToValue(t, ps, "name", 31, "Andrea")
addEdgeToValue(t, ps, "dob", 23, "1910-01-02")
addEdgeToValue(t, ps, "dob", 24, "1909-05-05")
addEdgeToValue(t, ps, "dob", 25, "1909-01-10")
addEdgeToValue(t, ps, "dob", 31, "1901-01-15")
return dir, dir2, ps
}
func TestSchema(t *testing.T) {
dir, dir2, _ := populateGraph(t)
defer os.RemoveAll(dir)
defer os.RemoveAll(dir2)
query := `
{
debug(_uid_:0x1) {
_xid_
name
gender
alive
loc
friend {
name
}
friend {
}
}
}
`
gq, _, err := gql.Parse(query)
require.NoError(t, err)
ctx := context.Background()
sg, err := ToSubGraph(ctx, gq)
require.NoError(t, err)
ch := make(chan error)
go ProcessGraph(ctx, sg, nil, ch)
err = <-ch
require.NoError(t, err)
var l Latency
gr, err := sg.ToProtocolBuffer(&l)
require.NoError(t, err)
require.EqualValues(t, "debug", gr.Children[0].Attribute)
require.EqualValues(t, 1, gr.Children[0].Uid)
require.EqualValues(t, "mich", gr.Children[0].Xid)
require.Len(t, gr.Children[0].Properties, 4)
require.EqualValues(t, "Michonne",
getProperty(gr.Children[0].Properties, "name").GetStrVal())
var g types.Geo
x.Check(g.UnmarshalBinary(getProperty(gr.Children[0].Properties, "loc").GetGeoVal()))
received, err := g.MarshalText()
require.EqualValues(t, "{'type':'Point','coordinates':[1.1,2]}", string(received))
require.Len(t, gr.Children[0].Children, 5)
child := gr.Children[0].Children[0]
require.EqualValues(t, 23, child.Uid)
require.EqualValues(t, "friend", child.Attribute)
require.Len(t, child.Properties, 1)
require.EqualValues(t, "Rick Grimes",
getProperty(child.Properties, "name").GetStrVal())
require.Empty(t, child.Children)
child = gr.Children[0].Children[1]
require.EqualValues(t, 24, child.Uid)
require.EqualValues(t, "friend", child.Attribute)
require.Len(t, child.Properties, 1)
require.EqualValues(t, "Glenn Rhee",
getProperty(child.Properties, "name").GetStrVal())
require.Empty(t, child.Children)
child = gr.Children[0].Children[4]
require.EqualValues(t, 101, child.Uid)
require.EqualValues(t, "friend", child.Attribute)
require.Len(t, child.Properties, 0)
}
// queryTokens returns the tokens to be used to look up the geo index for a given filter.
func queryTokens(qt QueryType, data string, maxDistance float64) ([]string, *QueryData, error) {
// Try to parse the data as geo type.
var g types.Geo
geoData := strings.Replace(data, "'", "\"", -1)
err := g.UnmarshalText([]byte(geoData))
if err != nil {
return nil, nil, x.Wrapf(err, "Cannot decode given geoJson input")
}
var l *s2.Loop
var pt *s2.Point
switch v := g.T.(type) {
case *geom.Point:
p := pointFromPoint(v)
pt = &p
case *geom.Polygon:
l, err = loopFromPolygon(v)
if err != nil {
return nil, nil, err
}
default:
return nil, nil, x.Errorf("Cannot query using a geometry of type %T", v)
}
x.AssertTruef(l != nil || pt != nil, "We should have a point or a loop.")
parents, cover, err := indexCells(g)
if err != nil {
return nil, nil, err
}
switch qt {
case QueryTypeWithin:
// For a within query we only need to look at the objects whose parents match our cover.
// So we take our cover and prefix with the parentPrefix to look in the index.
toks := toTokens(cover, parentPrefix)
return toks, &QueryData{pt: pt, loop: l, qtype: qt}, nil
case QueryTypeContains:
if l != nil {
return nil, nil, x.Errorf("Cannot use a polygon in a contains query")
}
// For a contains query, we only need to look at the objects whose cover matches our
// parents. So we take our parents and prefix with the coverPrefix to look in the index.
return toTokens(parents, coverPrefix), &QueryData{pt: pt, qtype: qt}, nil
case QueryTypeNear:
if l != nil {
return nil, nil, x.Errorf("Cannot use a polygon in a near query")
}
return nearQueryKeys(*pt, maxDistance)
case QueryTypeIntersects:
// An intersects query is essentially the union of contains and within. So we look at all
// the objects whose parents match our cover as well as all the objects whose cover matches
// our parents.
toks := parentCoverTokens(parents, cover)
return toks, &QueryData{pt: pt, loop: l, qtype: qt}, nil
default:
return nil, nil, x.Errorf("Unknown query type")
}
}