feat(search): Implement GEO index and support for RADIUS search (#5854)

mkaruza · web-flow · commit 1da4817e341e · 2025-10-01T10:55:21.000+02:00
Implementation of GEO index use boost::geometry::index::rtree as structure to store/search/remove points. Search is done in a way that we construct polygon of 4 points with 90 degrees angle difference around search point with fixed distance and bounding box around that polygon to cover whole area of possible points that fit in radius. For each point we calculate haversine distance between found point and search point. Fixes #4536 Signed-off-by: mkaruza <mario@dragonflydb.io>
diff --git a/src/core/search/ast_expr.cc b/src/core/search/ast_expr.cc
@@ -20,6 +20,10 @@ AstRangeNode::AstRangeNode(double lo, bool lo_excl, double hi, bool hi_excl)
     : lo{lo_excl ? nextafter(lo, hi) : lo}, hi{hi_excl ? nextafter(hi, lo) : hi} {
 }
 
+AstGeoNode::AstGeoNode(double lon, double lat, double radius, std::string unit)
+    : lon(lon), lat(lat), radius(radius), unit(std::move(unit)) {
+}
+
 AstNegateNode::AstNegateNode(AstNode&& node) : node{make_unique<AstNode>(std::move(node))} {
 }
 
diff --git a/src/core/search/ast_expr.h b/src/core/search/ast_expr.h
@@ -45,6 +45,13 @@ struct AstRangeNode {
   double lo, hi;
 };
 
+struct AstGeoNode {
+  AstGeoNode(double lon, double lat, double radius, std::string unit);
+  double lon, lat;
+  double radius;
+  std::string unit;
+};
+
 // Negates subtree
 struct AstNegateNode {
   AstNegateNode(AstNode&& node);
@@ -112,7 +119,7 @@ struct AstKnnNode {
 using NodeVariants =
     std::variant<std::monostate, AstStarNode, AstStarFieldNode, AstTermNode, AstPrefixNode,
                  AstSuffixNode, AstInfixNode, AstRangeNode, AstNegateNode, AstLogicalNode,
-                 AstFieldNode, AstTagsNode, AstKnnNode>;
+                 AstFieldNode, AstTagsNode, AstKnnNode, AstGeoNode>;
 
 struct AstNode : public NodeVariants {
   using variant::variant;
diff --git a/src/core/search/indices.cc b/src/core/search/indices.cc
@@ -10,6 +10,8 @@
 #include <absl/strings/str_join.h>
 #include <absl/strings/str_split.h>
 
+#include <boost/iterator/function_output_iterator.hpp>
+
 #define UNI_ALGO_DISABLE_NFKC_NFKD
 
 #include <hnswlib/hnswalg.h>
@@ -90,6 +92,42 @@ void IterateAllSuffixes(const absl::flat_hash_set<string>& words,
   }
 }
 
+// Haversine with earth radius in meters. Used to calculate distance.
+boost::geometry::strategy::distance::haversine haversine_(6372797.560856);
+
+double ConvertToRadiusInMeters(size_t radius, std::string_view arg) {
+  const std::string unit = absl::AsciiStrToUpper(arg);
+  if (unit == "M") {
+    return radius * 1;
+  } else if (unit == "KM") {
+    return radius * 1000;
+  } else if (unit == "FT") {
+    return radius * 0.3048;
+  } else if (unit == "MI") {
+    return radius * 1609.34;
+  } else {
+    return -1;
+  }
+}
+
+std::optional<GeoIndex::point> GetGeoPoint(const DocumentAccessor& doc, string_view field) {
+  auto element = doc.GetStrings(field);
+
+  if (!element)
+    return std::nullopt;
+
+  absl::InlinedVector<string_view, 2> coordinates = absl::StrSplit(element.value()[0], ",");
+
+  if (coordinates.size() != 2)
+    return std::nullopt;
+
+  double lon, lat;
+  if (!absl::SimpleAtod(coordinates[0], &lon) || !absl::SimpleAtod(coordinates[1], &lat))
+    return nullopt;
+
+  return GeoIndex::point{lon, lat};
+}
+
 };  // namespace
 
 class RangeTreeAdapter : public NumericIndex::RangeTreeBase {
@@ -614,4 +652,73 @@ void HnswVectorIndex::Remove(DocId id, const DocumentAccessor& doc, string_view
   adapter_->Remove(id);
 }
 
+GeoIndex::GeoIndex(PMR_NS::memory_resource* mr) : rtree_(make_unique<rtree>()) {
+}
+
+GeoIndex::~GeoIndex() {
+}
+
+bool GeoIndex::Add(DocId id, const DocumentAccessor& doc, std::string_view field) {
+  auto doc_point = GetGeoPoint(doc, field);
+  if (!doc_point) {
+    return false;
+  }
+  rtree_->insert({doc_point.value(), id});
+  return true;
+}
+
+void GeoIndex::Remove(DocId id, const DocumentAccessor& doc, string_view field) {
+  auto doc_point = GetGeoPoint(doc, field);
+  rtree_->remove({doc_point.value(), id});
+}
+
+std::vector<DocId> GeoIndex::RadiusSearch(double lon, double lat, double radius,
+                                          std::string_view unit) {
+  std::vector<DocId> results;
+
+  // Get radius in meters
+  double converted_radius = ConvertToRadiusInMeters(radius, unit);
+
+  // Declare the geographic_point_circle strategy with 4 points
+  boost::geometry::strategy::buffer::geographic_point_circle<> point_strategy(4);
+
+  // Declare the distance strategy in meters around the point
+  boost::geometry::strategy::buffer::distance_symmetric<double> distance_strategy(converted_radius);
+
+  // Declare other necessary strategies, unused for point
+  boost::geometry::strategy::buffer::join_round join_strategy;
+  boost::geometry::strategy::buffer::end_round end_strategy;
+  boost::geometry::strategy::buffer::side_straight side_strategy;
+
+  point p{lon, lat};
+
+  // Create polygon with 4 point around point
+  boost::geometry::model::multi_polygon<boost::geometry::model::polygon<point>> buffer_polygon;
+
+  boost::geometry::buffer(p, buffer_polygon, distance_strategy, side_strategy, join_strategy,
+                          end_strategy, point_strategy);
+
+  // Create bouding box around polygon to include all possible points
+  boost::geometry::model::box<point> box;
+  boost::geometry::envelope(buffer_polygon, box);
+
+  rtree_->query(
+      boost::geometry::index::within(box),
+      boost::make_function_output_iterator([&results, &p, &converted_radius](auto const& val) {
+        if (haversine_.apply(val.first, p) <= converted_radius) {
+          results.push_back(val.second);
+        }
+      }));
+
+  // TODO: we should return sorted results by radius distance
+  return results;
+}
+
+std::vector<DocId> GeoIndex::GetAllDocsWithNonNullValues() const {
+  std::vector<DocId> results;
+  std::for_each(boost::geometry::index::begin(*rtree_), boost::geometry::index::end(*rtree_),
+                [&results](auto const& val) { results.push_back(val.second); });
+  return results;
+}
+
 }  // namespace dfly::search
diff --git a/src/core/search/indices.h b/src/core/search/indices.h
@@ -8,6 +8,7 @@
 #include <absl/container/flat_hash_map.h>
 #include <absl/container/flat_hash_set.h>
 
+#include <boost/geometry.hpp>
 #include <map>
 #include <memory>
 #include <optional>
@@ -205,4 +206,23 @@ struct HnswVectorIndex : public BaseVectorIndex {
   std::unique_ptr<HnswlibAdapter> adapter_;
 };
 
+struct GeoIndex : public BaseIndex {
+  using point =
+      boost::geometry::model::point<double, 2,
+                                    boost::geometry::cs::geographic<boost::geometry::degree>>;
+  using index_entry = std::pair<point, DocId>;
+
+  explicit GeoIndex(PMR_NS::memory_resource* mr);
+  ~GeoIndex();
+
+  bool Add(DocId id, const DocumentAccessor& doc, std::string_view field) override;
+  void Remove(DocId id, const DocumentAccessor& doc, std::string_view field) override;
+  std::vector<DocId> RadiusSearch(double lon, double lat, double radius, std::string_view arg);
+  std::vector<DocId> GetAllDocsWithNonNullValues() const override;
+
+ private:
+  using rtree = boost::geometry::index::rtree<index_entry, boost::geometry::index::linear<16>>;
+  std::unique_ptr<rtree> rtree_;
+};
+
 }  // namespace dfly::search
diff --git a/src/core/search/lexer.lex b/src/core/search/lexer.lex
@@ -68,6 +68,10 @@ astrsk_ch  \*
 "KNN"          return Parser::make_KNN (loc());
 "AS"           return Parser::make_AS (loc());
 "EF_RUNTIME"   return Parser::make_EF_RUNTIME (loc());
+"M"            return Parser::make_GEOUNIT_M (loc());
+"KM"           return Parser::make_GEOUNIT_KM (loc());
+"MI"           return Parser::make_GEOUNIT_MI (loc());
+"FT"           return Parser::make_GEOUNIT_FT (loc());
 
 [0-9]{1,9}                          return Parser::make_UINT32(str(), loc());
 [+-]?(([0-9]*[.])?[0-9]+|inf)       return Parser::make_DOUBLE(str(), loc());
diff --git a/src/core/search/parser.y b/src/core/search/parser.y
@@ -61,6 +61,10 @@ double toDouble(string_view src);
   KNN         "KNN"
   AS          "AS"
   EF_RUNTIME  "EF_RUNTIME"
+  GEOUNIT_M   "GEOUNIT_M"
+  GEOUNIT_KM  "GEOUNIT_KM"
+  GEOUNIT_MI  "GEOUNIT_MI"
+  GEOUNIT_FT  "GEOUNIT_FT"
 ;
 
 %token AND_OP
@@ -77,14 +81,13 @@ double toDouble(string_view src);
 
 %token <std::string> DOUBLE "double"
 %token <std::string> UINT32 "uint32"
-%nterm <double> generic_number
-%nterm <bool> opt_lparen
-%nterm <AstExpr> final_query filter search_expr search_unary_expr search_or_expr search_and_expr numeric_filter_expr
+%nterm <AstExpr> final_query filter search_expr search_unary_expr search_or_expr search_and_expr bracket_filter_expr
 %nterm <AstExpr> field_cond field_cond_expr field_unary_expr field_or_expr field_and_expr tag_list
 %nterm <AstTagsNode::TagValueProxy> tag_list_element
 
 %nterm <AstKnnNode> knn_query
 %nterm <std::string> opt_knn_alias
+%nterm <std::string> geounit
 %nterm <std::optional<size_t>> opt_ef_runtime
 
 %printer { yyo << $$; } <*>;
@@ -149,23 +152,55 @@ field_cond:
   | STAR                                                { $$ = AstStarFieldNode();           }
   | NOT_OP field_cond                                   { $$ = AstNegateNode(std::move($2)); }
   | LPAREN field_cond_expr RPAREN                       { $$ = std::move($2); }
-  | LBRACKET numeric_filter_expr RBRACKET               { $$ = std::move($2); }
+  | LBRACKET bracket_filter_expr RBRACKET               { $$ = std::move($2); }
   | LCURLBR tag_list RCURLBR                            { $$ = std::move($2); }
   | PREFIX                                              { $$ = AstPrefixNode(std::move($1)); }
   | SUFFIX                                              { $$ = AstSuffixNode(std::move($1)); }
   | INFIX                                               { $$ = AstInfixNode(std::move($1));  }
 
-numeric_filter_expr:
-  opt_lparen generic_number opt_lparen generic_number         { $$ = AstRangeNode($2, $1, $4, $3); }
-  | opt_lparen generic_number COMMA opt_lparen generic_number { $$ = AstRangeNode($2, $1, $5, $4); }
-
-generic_number:
-  DOUBLE { $$ = toDouble($1); }
-  | UINT32 { $$ = toUint32($1); }
-
-opt_lparen:
-  /* empty */ { $$ = false; }
-  | LPAREN { $$ = true; }
+bracket_filter_expr:
+  /* Numeric filter has form [(] UINT32|DOUBLE [COMMA] [(] UINT32|DOUBLE */
+  DOUBLE DOUBLE                                { $$ = AstRangeNode(toDouble($1), false, toDouble($2), false); }
+  | LPAREN DOUBLE DOUBLE                       { $$ = AstRangeNode(toDouble($2), true, toDouble($3), false); }
+  | DOUBLE LPAREN DOUBLE                       { $$ = AstRangeNode(toDouble($1), false, toDouble($3), true); }
+  | LPAREN DOUBLE LPAREN DOUBLE                { $$ = AstRangeNode(toDouble($2), true, toDouble($4), true); }
+  | DOUBLE UINT32                              { $$ = AstRangeNode(toDouble($1), false, toUint32($2), false); }
+  | LPAREN DOUBLE UINT32                       { $$ = AstRangeNode(toDouble($2), true, toUint32($3), false); }
+  | DOUBLE LPAREN UINT32                       { $$ = AstRangeNode(toDouble($1), false, toUint32($3), true); }
+  | LPAREN DOUBLE LPAREN UINT32                { $$ = AstRangeNode(toDouble($2), true, toUint32($4), true); }
+  | UINT32 DOUBLE                              { $$ = AstRangeNode(toUint32($1), false, toDouble($2), false); }
+  | LPAREN UINT32 DOUBLE                       { $$ = AstRangeNode(toUint32($2), true, toDouble($3), false); }
+  | UINT32 LPAREN DOUBLE                       { $$ = AstRangeNode(toUint32($1), false, toDouble($3), true); }
+  | LPAREN UINT32 LPAREN DOUBLE                { $$ = AstRangeNode(toUint32($2), true, toDouble($4), true); }
+  | UINT32 UINT32                              { $$ = AstRangeNode(toUint32($1), false, toUint32($2), false); }
+  | LPAREN UINT32 UINT32                       { $$ = AstRangeNode(toUint32($2), true, toUint32($3), false); }
+  | UINT32 LPAREN UINT32                       { $$ = AstRangeNode(toUint32($1), false, toUint32($3), true); }
+  | LPAREN UINT32 LPAREN UINT32                { $$ = AstRangeNode(toUint32($2), true, toUint32($4), true); }
+  | DOUBLE COMMA DOUBLE                        { $$ = AstRangeNode(toDouble($1), false, toDouble($3), false); }
+  | DOUBLE COMMA UINT32                        { $$ = AstRangeNode(toDouble($1), false, toUint32($3), false); }
+  | UINT32 COMMA DOUBLE                        { $$ = AstRangeNode(toUint32($1), false, toDouble($3), false); }
+  | UINT32 COMMA UINT32                        { $$ = AstRangeNode(toUint32($1), false, toUint32($3), false); }
+  | LPAREN DOUBLE COMMA DOUBLE                 { $$ = AstRangeNode(toDouble($2), true, toDouble($4), false); }
+  | DOUBLE COMMA LPAREN DOUBLE                 { $$ = AstRangeNode(toDouble($1), false, toDouble($4), true); }
+  | LPAREN DOUBLE COMMA LPAREN DOUBLE          { $$ = AstRangeNode(toDouble($2), true, toDouble($5), true); }
+  | LPAREN DOUBLE COMMA UINT32                 { $$ = AstRangeNode(toDouble($2), true, toUint32($4), false); }
+  | DOUBLE COMMA LPAREN UINT32                 { $$ = AstRangeNode(toDouble($1), false, toUint32($4), true); }
+  | LPAREN DOUBLE COMMA LPAREN UINT32          { $$ = AstRangeNode(toDouble($2), true, toUint32($5), true); }
+  | LPAREN UINT32 COMMA DOUBLE                 { $$ = AstRangeNode(toUint32($2), true, toDouble($4), false); }
+  | UINT32 COMMA LPAREN DOUBLE                 { $$ = AstRangeNode(toUint32($1), false, toDouble($4), true); }
+  | LPAREN UINT32 COMMA LPAREN DOUBLE          { $$ = AstRangeNode(toUint32($2), true, toDouble($5), true); }
+  | LPAREN UINT32 COMMA UINT32                 { $$ = AstRangeNode(toUint32($2), true, toUint32($4), false); }
+  | UINT32 COMMA LPAREN UINT32                 { $$ = AstRangeNode(toUint32($1), false, toUint32($4), true); }
+  | LPAREN UINT32 COMMA LPAREN UINT32          { $$ = AstRangeNode(toUint32($2), true, toUint32($5), true); }
+  /* GEO filter */
+  | DOUBLE DOUBLE UINT32 geounit               { $$ = AstGeoNode(toDouble($1), toDouble($2), toUint32($3), std::move($4)); }
+  | DOUBLE DOUBLE DOUBLE geounit               { $$ = AstGeoNode(toDouble($1), toDouble($2), toDouble($3), std::move($4)); }
+
+geounit:
+  GEOUNIT_M     { $$ = "M"; }
+  | GEOUNIT_KM  { $$ = "KM"; }
+  | GEOUNIT_MI  { $$ = "MI"; }
+  | GEOUNIT_FT  { $$ = "FT"; }
 
 field_cond_expr:
   field_unary_expr { $$ = std::move($1); }
diff --git a/src/core/search/search.cc b/src/core/search/search.cc
@@ -79,6 +79,9 @@ struct ProfileBuilder {
         [](const AstNegateNode& n) { return absl::StrCat("Negate{}"); },
         [](const AstStarNode& n) { return absl::StrCat("Star{}"); },
         [](const AstStarFieldNode& n) { return absl::StrCat("StarField{}"); },
+        [](const AstGeoNode& n) {
+          return absl::StrCat("Geo{", n.lat, " ", n.lon, " ", n.radius, " ", n.unit, "}");
+        },
     };
     return visit(node_info, node.Variant());
   }
@@ -276,6 +279,14 @@ struct BasicSearch {
     return IndexResult{};
   }
 
+  IndexResult Search(const AstGeoNode& node, string_view active_field) {
+    DCHECK(!active_field.empty());
+    if (auto* index = GetIndex<GeoIndex>(active_field); index) {
+      return IndexResult{index->RadiusSearch(node.lon, node.lat, node.radius, node.unit)};
+    }
+    return IndexResult{};
+  }
+
   // negate -(*subquery*): explicitly compute result complement. Needs further optimizations
   IndexResult Search(const AstNegateNode& node, string_view active_field) {
     vector<DocId> matched = SearchGeneric(*node.node, active_field).Take();
@@ -396,6 +407,7 @@ struct BasicSearch {
     // used by knn
 
     DCHECK(top_level || holds_alternative<AstKnnNode>(node.Variant()) ||
+           holds_alternative<AstGeoNode>(node.Variant()) ||
            visit([](auto* set) { return is_sorted(set->begin(), set->end()); }, result.Borrowed()));
 
     if (profile_builder_)
@@ -500,6 +512,10 @@ void FieldIndices::CreateIndices(PMR_NS::memory_resource* mr) {
         indices_[field_ident] = std::move(vector_index);
         break;
       }
+      case SchemaField::GEO: {
+        indices_[field_ident] = make_unique<GeoIndex>(mr);
+        break;
+      }
     }
   }
 }
@@ -518,6 +534,7 @@ void FieldIndices::CreateSortIndices(PMR_NS::memory_resource* mr) {
         sort_indices_[field_ident] = make_unique<NumericSortIndex>(mr);
         break;
       case SchemaField::VECTOR:
+      case SchemaField::GEO:
         break;
     }
   }
diff --git a/src/core/search/search.h b/src/core/search/search.h
@@ -55,7 +55,7 @@ struct OptionalNumericFilter : public OptionalFilterBase {
 
 // Describes a specific index field
 struct SchemaField {
-  enum FieldType { TAG, TEXT, NUMERIC, VECTOR };
+  enum FieldType { TAG, TEXT, NUMERIC, VECTOR, GEO };
   enum FieldFlags : uint8_t { NOINDEX = 1 << 0, SORTABLE = 1 << 1 };
 
   struct VectorParams {
diff --git a/src/core/search/search_test.cc b/src/core/search/search_test.cc
diff --git a/src/server/search/doc_index.cc b/src/server/search/doc_index.cc
diff --git a/src/server/search/search_family.cc b/src/server/search/search_family.cc

Original file line number	Diff line number	Diff line change
`@@ -20,6 +20,10 @@ AstRangeNode::AstRangeNode(double lo, bool lo_excl, double hi, bool hi_excl)`
`20`	`20`	`: lo{lo_excl ? nextafter(lo, hi) : lo}, hi{hi_excl ? nextafter(hi, lo) : hi} {`
`21`	`21`	`}`
`22`	`22`
	`23`	`+AstGeoNode::AstGeoNode(double lon, double lat, double radius, std::string unit)`
	`24`	`+ : lon(lon), lat(lat), radius(radius), unit(std::move(unit)) {`
	`25`	`+}`
	`26`	`+`
`23`	`27`	`AstNegateNode::AstNegateNode(AstNode&& node) : node{make_unique<AstNode>(std::move(node))} {`
`24`	`28`	`}`
`25`	`29`