diff --git a/controller/collection/collection.go b/controller/collection/collection.go index eb51f74a..d2c72c3e 100644 --- a/controller/collection/collection.go +++ b/controller/collection/collection.go @@ -101,26 +101,63 @@ func (c *Collection) Bounds() (minX, minY, minZ, maxX, maxY, maxZ float64) { // The fields argument is optional. // The return values are the old object, the old fields, and the new fields func (c *Collection) ReplaceOrInsert(id string, obj geojson.Object, fields []string, values []float64) (oldObject geojson.Object, oldFields []float64, newFields []float64) { - nitem, oldItem := c.insertOrReplace(id, obj) - if oldItem != nil { + var oldItem *itemT + var newItem *itemT = &itemT{id: id, object: obj} + // add the new item to main btree and remove the old one if needed + oldItemPtr := c.items.ReplaceOrInsert(newItem) + if oldItemPtr != nil { + // the old item was removed, now let's remove from the rtree + // or strings tree. + oldItem = oldItemPtr.(*itemT) + if obj.IsGeometry() { + // geometry + c.index.Remove(oldItem) + c.objects-- + } else { + // string + c.values.Delete(oldItem) + c.nobjects-- + } + // decrement the point count + c.points -= oldItem.object.PositionCount() + + // decrement the weights + c.weight -= len(oldItem.fields) * 8 + c.weight -= oldItem.object.Weight() + len(oldItem.id) + + // references oldObject = oldItem.object oldFields = oldItem.fields - nitem.fields = oldFields - c.weight += len(nitem.fields) * 8 + newItem.fields = oldFields } - if fields == nil && len(values) > 0 { - // directly set the field values, update weight - c.weight -= len(nitem.fields) * 8 - nitem.fields = values - c.weight += len(nitem.fields) * 8 + // insert the new item into the rtree or strings tree. + if obj.IsGeometry() { + c.index.Insert(newItem) + c.objects++ + } else { + c.values.ReplaceOrInsert(newItem) + c.nobjects++ + } + // increment the point count + c.points += obj.PositionCount() + // add the new weights + c.weight += len(newItem.fields) * 8 + c.weight += obj.Weight() + len(id) + if fields == nil { + if len(values) > 0 { + // directly set the field values, update weight + c.weight -= len(newItem.fields) * 8 + newItem.fields = values + c.weight += len(newItem.fields) * 8 + } } else { // map field name to value for i, field := range fields { - c.setField(nitem, field, values[i]) + c.setField(newItem, field, values[i]) } } - return oldObject, oldFields, nitem.fields + return oldObject, oldFields, newItem.fields } func (c *Collection) remove(id string) (item *itemT, ok bool) { @@ -142,34 +179,19 @@ func (c *Collection) remove(id string) (item *itemT, ok bool) { return item, true } -func (c *Collection) insertOrReplace(id string, obj geojson.Object) (item, prev *itemT) { +func (c *Collection) insert(id string, obj geojson.Object) (item *itemT) { item = &itemT{id: id, object: obj} - old := c.items.ReplaceOrInsert(item) - if old != nil { - prev = old.(*itemT) - if item.object.IsGeometry() { - c.index.Remove(prev) - c.index.Insert(item) - } else { - c.values.ReplaceOrInsert(item) - } - c.weight -= len(prev.fields) * 8 - c.weight -= prev.object.Weight() + len(prev.id) - c.points -= prev.object.PositionCount() - c.weight += obj.Weight() + len(id) - c.points += obj.PositionCount() + if obj.IsGeometry() { + c.index.Insert(item) + c.objects++ } else { - if obj.IsGeometry() { - c.index.Insert(item) - c.objects++ - } else { - c.values.ReplaceOrInsert(item) - c.nobjects++ - } - c.weight += obj.Weight() + len(id) - c.points += obj.PositionCount() + c.values.ReplaceOrInsert(item) + c.nobjects++ } - return item, prev + c.items.ReplaceOrInsert(item) + c.weight += obj.Weight() + len(id) + c.points += obj.PositionCount() + return item } // Remove removes an object and returns it. diff --git a/tests/keys_test.go b/tests/keys_test.go index 09259fb4..8e2085be 100644 --- a/tests/keys_test.go +++ b/tests/keys_test.go @@ -29,6 +29,7 @@ func subTestKeys(t *testing.T, mc *mockServer) { runStep(t, mc, "TTL", keys_TTL_test) runStep(t, mc, "SET EX", keys_SET_EX_test) runStep(t, mc, "PDEL", keys_PDEL_test) + runStep(t, mc, "FIELDS", keys_FIELDS_test) } func keys_BOUNDS_test(mc *mockServer) error { @@ -325,6 +326,21 @@ func keys_SET_EX_test(mc *mockServer) (err error) { return nil } +func keys_FIELDS_test(mc *mockServer) error { + return mc.DoBatch([][]interface{}{ + {"SET", "mykey", "myid1a", "FIELD", "a", 1, "POINT", 33, -115}, {"OK"}, + {"GET", "mykey", "myid1a", "WITHFIELDS"}, {`[{"type":"Point","coordinates":[-115,33]} [a 1]]`}, + {"SET", "mykey", "myid1a", "FIELD", "a", "a", "POINT", 33, -115}, {"ERR invalid argument 'a'"}, + {"GET", "mykey", "myid1a", "WITHFIELDS"}, {`[{"type":"Point","coordinates":[-115,33]} [a 1]]`}, + {"SET", "mykey", "myid1a", "FIELD", "a", 1, "FIELD", "b", 2, "POINT", 33, -115}, {"OK"}, + {"GET", "mykey", "myid1a", "WITHFIELDS"}, {`[{"type":"Point","coordinates":[-115,33]} [a 1 b 2]]`}, + {"SET", "mykey", "myid1a", "FIELD", "b", 2, "POINT", 33, -115}, {"OK"}, + {"GET", "mykey", "myid1a", "WITHFIELDS"}, {`[{"type":"Point","coordinates":[-115,33]} [a 1 b 2]]`}, + {"SET", "mykey", "myid1a", "FIELD", "b", 2, "FIELD", "a", "1", "FIELD", "c", 3, "POINT", 33, -115}, {"OK"}, + {"GET", "mykey", "myid1a", "WITHFIELDS"}, {`[{"type":"Point","coordinates":[-115,33]} [a 1 b 2 c 3]]`}, + }) +} + func keys_PDEL_test(mc *mockServer) error { return mc.DoBatch([][]interface{}{ {"SET", "mykey", "myid1a", "POINT", 33, -115}, {"OK"}, diff --git a/tests/tests_test.go b/tests/tests_test.go index 266dfdb9..e38fa37e 100644 --- a/tests/tests_test.go +++ b/tests/tests_test.go @@ -57,7 +57,10 @@ func runStep(t *testing.T, mc *mockServer, name string, step func(mc *mockServer mc.ResetConn() defer mc.ResetConn() // clear the database so the test is consistent - if err := mc.DoBatch([][]interface{}{{"FLUSHDB"}, {"OK"}}); err != nil { + if err := mc.DoBatch([][]interface{}{ + {"OUTPUT", "resp"}, {"OK"}, + {"FLUSHDB"}, {"OK"}, + }); err != nil { return err } if err := step(mc); err != nil {