diff --git a/include/rocksdb/utilities/geo_db.h b/include/rocksdb/utilities/geo_db.h index 41c0f1408..d603c5770 100644 --- a/include/rocksdb/utilities/geo_db.h +++ b/include/rocksdb/utilities/geo_db.h @@ -59,6 +59,16 @@ class GeoObject { } }; +class GeoIterator { + public: + GeoIterator() = default; + virtual ~GeoIterator() {} + virtual void Next() = 0; + virtual bool Valid() const = 0; + virtual const GeoObject& geo_object() = 0; + virtual Status status() const = 0; +}; + // // Stack your DB with GeoDB to be able to get geo-spatial support // @@ -91,14 +101,13 @@ class GeoDB : public StackableDB { // Delete the specified object virtual Status Remove(const Slice& id) = 0; - // Returns a list of all items within a circular radius from the + // Returns an iterator for the items within a circular radius from the // specified gps location. If 'number_of_values' is specified, - // then this call returns at most that many number of objects. + // then the iterator is capped to that number of objects. // The radius is specified in 'meters'. - virtual Status SearchRadial(const GeoPosition& pos, - double radius, - std::vector* values, - int number_of_values = INT_MAX) = 0; + virtual GeoIterator* SearchRadial(const GeoPosition& pos, + double radius, + int number_of_values = INT_MAX) = 0; }; } // namespace rocksdb diff --git a/utilities/geodb/geodb_impl.cc b/utilities/geodb/geodb_impl.cc index afa2f4c91..f509e277e 100644 --- a/utilities/geodb/geodb_impl.cc +++ b/utilities/geodb/geodb_impl.cc @@ -159,15 +159,63 @@ Status GeoDBImpl::Remove(const Slice& id) { return db_->Write(woptions_, &batch); } -Status GeoDBImpl::SearchRadial(const GeoPosition& pos, +class GeoIteratorImpl : public GeoIterator { + private: + std::vector values_; + std::vector::iterator iter_; + public: + explicit GeoIteratorImpl(std::vector values) + : values_(std::move(values)) { + iter_ = values_.begin(); + } + virtual void Next() override; + virtual bool Valid() const override; + virtual const GeoObject& geo_object() override; + virtual Status status() const override; +}; + +class GeoErrorIterator : public GeoIterator { + private: + Status status_; + public: + explicit GeoErrorIterator(Status s) : status_(s) {} + virtual void Next() override {}; + virtual bool Valid() const override { return false; } + virtual const GeoObject& geo_object() override { + GeoObject* g = new GeoObject(); + return *g; + } + virtual Status status() const override { return status_; } +}; + +void GeoIteratorImpl::Next() { + assert(Valid()); + iter_++; +} + +bool GeoIteratorImpl::Valid() const { + return iter_ != values_.end(); +} + +const GeoObject& GeoIteratorImpl::geo_object() { + assert(Valid()); + return *iter_; +} + +Status GeoIteratorImpl::status() const { + return Status::OK(); +} + +GeoIterator* GeoDBImpl::SearchRadial(const GeoPosition& pos, double radius, - std::vector* values, int number_of_values) { + std::vector values; + // Gather all bounding quadkeys std::vector qids; Status s = searchQuadIds(pos, radius, &qids); if (!s.ok()) { - return s; + return new GeoErrorIterator(s); } // create an iterator @@ -200,7 +248,7 @@ Status GeoDBImpl::SearchRadial(const GeoPosition& pos, if (res.first == qid.end()) { GeoPosition obj_pos(atof(parts[3].c_str()), atof(parts[4].c_str())); GeoObject obj(obj_pos, parts[4], iter->value().ToString()); - values->push_back(obj); + values.push_back(obj); number_of_values--; } else { break; @@ -208,7 +256,7 @@ Status GeoDBImpl::SearchRadial(const GeoPosition& pos, } } delete iter; - return Status::OK(); + return new GeoIteratorImpl(std::move(values)); } std::string GeoDBImpl::MakeKey1(const GeoPosition& pos, Slice id, diff --git a/utilities/geodb/geodb_impl.h b/utilities/geodb/geodb_impl.h index aaf3a25ef..d63102856 100644 --- a/utilities/geodb/geodb_impl.h +++ b/utilities/geodb/geodb_impl.h @@ -46,9 +46,8 @@ class GeoDBImpl : public GeoDB { // Returns a list of all items within a circular radius from the // specified gps location - virtual Status SearchRadial(const GeoPosition& pos, double radius, - std::vector* values, - int number_of_values) override; + virtual GeoIterator* SearchRadial(const GeoPosition& pos, double radius, + int number_of_values) override; private: DB* db_; diff --git a/utilities/geodb/geodb_test.cc b/utilities/geodb/geodb_test.cc index 503e5331b..98cda370b 100644 --- a/utilities/geodb/geodb_test.cc +++ b/utilities/geodb/geodb_test.cc @@ -103,17 +103,21 @@ TEST_F(GeoDBTest, Search) { // search all objects centered at 46 degree latitude with // a radius of 200 kilometers. We should find the one object that // we inserted earlier. - std::vector values; - status = getdb()->SearchRadial(GeoPosition(46, 46), 200000, &values); + GeoIterator* iter = getdb()->SearchRadial(GeoPosition(46, 46), 200000); ASSERT_TRUE(status.ok()); - ASSERT_EQ(values.size(), 1U); + ASSERT_EQ(iter->geo_object().value, "midvalue1"); + uint size = 0; + while (iter->Valid()) { + size++; + iter->Next(); + } + ASSERT_EQ(size, 1U); // search all objects centered at 46 degree latitude with // a radius of 2 kilometers. There should be none. - values.clear(); - status = getdb()->SearchRadial(GeoPosition(46, 46), 2, &values); + iter = getdb()->SearchRadial(GeoPosition(46, 46), 2); ASSERT_TRUE(status.ok()); - ASSERT_EQ(values.size(), 0U); + ASSERT_FALSE(iter->Valid()); } } // namespace rocksdb