add modern location iterator that can set rawscore

2 files changed, 111 insertions, 1 deletions

diff --git a/searchlib/src/vespa/searchlib/common/locationiterators.cpp b/searchlib/src/vespa/searchlib/common/locationiterators.cpp
index d90ed3b41f3..ea7f5315122 100644
--- a/searchlib/src/vespa/searchlib/common/locationiterators.cpp
+++ b/searchlib/src/vespa/searchlib/common/locationiterators.cpp
@@ -7,6 +7,106 @@
 #include <vespa/log/log.h>
 LOG_SETUP(".searchlib.common.locationiterators");
 
+namespace search::common {
+
+class LocationIterator : public search::queryeval::SearchIterator
+{
+private:
+    static constexpr double pi = 3.14159265358979323846;
+    // microdegrees -> degrees -> radians -> km (using Earth mean radius)
+    static constexpr double udeg_to_km = 1.0e-6 * (pi / 180.0) * 6371.0088;
+    search::fef::TermFieldMatchData & _tfmd;
+    const unsigned int _numDocs;
+    const bool         _strict;
+    const Location &   _location;
+    std::vector<search::AttributeVector::largeint_t> _pos;
+
+    void doSeek(uint32_t docId) override;
+    void doUnpack(uint32_t docId) override;
+public:
+    LocationIterator(search::fef::TermFieldMatchData &tfmd,
+                     unsigned int numDocs,
+                     bool strict,
+                     const Location & location);
+    ~LocationIterator() override;
+};
+
+LocationIterator::LocationIterator(search::fef::TermFieldMatchData &tfmd,
+                                   unsigned int numDocs,
+                                   bool strict,
+                                   const Location & location)
+  : SearchIterator(),
+    _tfmd(tfmd),
+    _numDocs(numDocs),
+    _strict(strict),
+    _location(location),
+    _pos()
+{
+    _pos.resize(1);  //Need at least 1 entry as the singlevalue attributes does not honour given size.
+    LOG(debug, "created LocationIterator(numDocs=%u)\n", numDocs);
+};
+
+
+LocationIterator::~LocationIterator() = default;
+
+void
+LocationIterator::doSeek(uint32_t docId)
+{
+    while (__builtin_expect(docId < getEndId(), true)) {
+        if (__builtin_expect(docId >= _numDocs, false)) {
+            break;
+        }
+        uint32_t numValues = _location.getVec()->get(docId, &_pos[0], _pos.size());
+        if (numValues > _pos.size()) {
+            _pos.resize(numValues);
+            numValues = _location.getVec()->get(docId, &_pos[0], _pos.size());
+        }
+        for (uint32_t i = 0; i < numValues; i++) {
+            int64_t docxy(_pos[i]);
+            if (_location.inside_limit(docxy)) {
+                setDocId(docId);
+                return;
+            }
+        }
+        if (!_strict) {
+            return;
+        }
+        ++docId;
+    }
+    setAtEnd();
+}
+
+void
+LocationIterator::doUnpack(uint32_t docId)
+{
+    uint64_t sqabsdist = std::numeric_limits<uint64_t>::max();
+    int32_t docx = 0;
+    int32_t docy = 0;
+    uint32_t numValues = _location.getVec()->get(docId, &_pos[0], _pos.size());
+    for (uint32_t i = 0; i < numValues; i++) {
+        int64_t docxy(_pos[i]);
+        vespalib::geo::ZCurve::decode(docxy, &docx, &docy);
+        uint64_t sqdist = _location.sq_distance_to({docx, docy});
+        if (sqdist < sqabsdist) {
+            sqabsdist = sqdist;
+        }
+    }
+    double dist = std::sqrt(double(sqabsdist));
+    double score = 1.0 / (1.0 + (udeg_to_km * dist));
+    LOG(debug, "unpack LI(%u) score %f\n", docId, score);
+    LOG(debug, "distance: %f micro-degrees ~= %f km", dist, udeg_to_km * dist);
+    _tfmd.setRawScore(docId, score);
+}
+
+std::unique_ptr<search::queryeval::SearchIterator>
+create_location_iterator(search::fef::TermFieldMatchData &tfmd, unsigned int numDocs,
+                         bool strict, const Location & location)
+{
+    return std::make_unique<LocationIterator>(tfmd, numDocs, strict, location);
+}
+
+} // namespace
+
 using namespace search::common;
 
 class FastS_2DZLocationIterator : public search::queryeval::SearchIterator
diff --git a/searchlib/src/vespa/searchlib/common/locationiterators.h b/searchlib/src/vespa/searchlib/common/locationiterators.h
index e345bcae4fe..d963ac2e479 100644
--- a/searchlib/src/vespa/searchlib/common/locationiterators.h
+++ b/searchlib/src/vespa/searchlib/common/locationiterators.h
@@ -4,9 +4,19 @@
 
 #include <vespa/searchlib/queryeval/searchiterator.h>
 #include <vespa/searchlib/common/location.h>
+#include <vespa/searchlib/fef/termfieldmatchdata.h>
+
+namespace search::common {
+
+std::unique_ptr<search::queryeval::SearchIterator>
+create_location_iterator(search::fef::TermFieldMatchData &tfmd,
+                         unsigned int numDocs,
+                         bool strict,
+                         const Location & location);
+
+} // namespace
 
 std::unique_ptr<search::queryeval::SearchIterator>
 FastS_AllocLocationIterator(unsigned int numDocs,
                             bool strict,
                             const search::common::Location & location);
-