Skip to content

Commit 33518a2

Browse files
mkaruzamkaruza
committed
Address PR review / Search inside box
1 parent c59479f commit 33518a2

File tree

6 files changed

+73
-67
lines changed

6 files changed

+73
-67
lines changed

‎src/core/search/ast_expr.cc‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -20,7 +20,7 @@ AstRangeNode::AstRangeNode(double lo, bool lo_excl, double hi, bool hi_excl)
2020
: lo{lo_excl ? nextafter(lo, hi) : lo}, hi{hi_excl ? nextafter(hi, lo) : hi} {
2121
}
2222

23-
AstGeoNode::AstGeoNode(double lat, double lon, size_t radius, std::string unit)
23+
AstGeoNode::AstGeoNode(double lat, double lon, double radius, std::string unit)
2424
: lat(lat), lon(lon), radius(radius), unit(std::move(unit)) {
2525
}
2626

‎src/core/search/ast_expr.h‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -46,9 +46,9 @@ struct AstRangeNode {
4646
};
4747

4848
struct AstGeoNode {
49-
AstGeoNode(double lat, double lon, size_t radius, std::string unit);
49+
AstGeoNode(double lat, double lon, double radius, std::string unit);
5050
double lat, lon;
51-
size_t radius;
51+
double radius;
5252
std::string unit;
5353
};
5454

‎src/core/search/indices.cc‎

Lines changed: 61 additions & 52 deletions
Original file line numberDiff line numberDiff line change
@@ -94,6 +94,39 @@ void IterateAllSuffixes(const absl::flat_hash_set<string>& words,
9494
}
9595
}
9696

97+
double ConvertToRadiusInMeters(size_t radius, std::string_view arg) {
98+
const std::string unit = absl::AsciiStrToUpper(arg);
99+
if (unit == "M") {
100+
return radius * 1;
101+
} else if (unit == "KM") {
102+
return radius * 1000;
103+
} else if (unit == "FT") {
104+
return radius * 0.3048;
105+
} else if (unit == "MI") {
106+
return radius * 1609.34;
107+
} else {
108+
return -1;
109+
}
110+
}
111+
112+
std::optional<GeoIndex::point> GetGeoPoint(const DocumentAccessor& doc, string_view field) {
113+
auto element = doc.GetStrings(field);
114+
115+
if (!element)
116+
return std::nullopt;
117+
118+
absl::InlinedVector<string_view, 2> coordinates = absl::StrSplit(element.value()[0], ",");
119+
120+
if (coordinates.size() != 2)
121+
return std::nullopt;
122+
123+
double lat, lon;
124+
if (!absl::SimpleAtod(coordinates[0], &lat) || !absl::SimpleAtod(coordinates[1], &lon))
125+
return nullopt;
126+
127+
return GeoIndex::point{lat, lon};
128+
}
129+
97130
}; // namespace
98131

99132
class RangeTreeAdapter : public NumericIndex::RangeTreeBase {
@@ -638,15 +671,18 @@ void GeoIndex::Remove(DocId id, const DocumentAccessor& doc, string_view field)
638671
rtree_->remove({doc_point.value(), id});
639672
}
640673

641-
std::vector<DocId> GeoIndex::RadiusSearch(double lat, double lon, double radius) const {
674+
std::vector<DocId> GeoIndex::RadiusSearch(double lat, double lon, double radius,
675+
std::string_view unit) {
642676
std::vector<DocId> results;
643677

644-
// Declare the geographic_point_circle strategy (with 36 points)
645-
// Default template arguments (taking Andoyer strategy)
646-
boost::geometry::strategy::buffer::geographic_point_circle<> point_strategy(36);
678+
// Get radius in meters
679+
double converted_radius = ConvertToRadiusInMeters(radius, unit);
680+
681+
// Declare the geographic_point_circle strategy with 4 points
682+
boost::geometry::strategy::buffer::geographic_point_circle<> point_strategy(4);
647683

648684
// Declare the distance strategy in meters around the point
649-
boost::geometry::strategy::buffer::distance_symmetric<double> distance_strategy(radius);
685+
boost::geometry::strategy::buffer::distance_symmetric<double> distance_strategy(converted_radius);
650686

651687
// Declare other necessary strategies, unused for point
652688
boost::geometry::strategy::buffer::join_round join_strategy;
@@ -655,59 +691,32 @@ std::vector<DocId> GeoIndex::RadiusSearch(double lat, double lon, double radius)
655691

656692
point p{lat, lon};
657693

658-
// Create the buffer of a point on the Earth
659-
boost::geometry::model::multi_polygon<boost::geometry::model::polygon<point>> result;
694+
// Create polygon with 4 point around point
695+
boost::geometry::model::multi_polygon<boost::geometry::model::polygon<point>> buffer_polygon;
660696

661-
boost::geometry::buffer(p, result, distance_strategy, side_strategy, join_strategy, end_strategy,
662-
point_strategy);
697+
boost::geometry::buffer(p, buffer_polygon, distance_strategy, side_strategy, join_strategy,
698+
end_strategy, point_strategy);
663699

664-
rtree_->query(boost::geometry::index::within(result),
665-
boost::make_function_output_iterator(
666-
[&results](auto const& val) { results.push_back(val.second); }));
700+
// Create bouding box around polygon to include all possible points
701+
boost::geometry::model::box<point> box;
702+
boost::geometry::envelope(buffer_polygon, box);
667703

668-
return results;
669-
}
704+
rtree_->query(
705+
boost::geometry::index::within(box),
706+
boost::make_function_output_iterator([&results, &p, &converted_radius](auto const& val) {
707+
if (boost::geometry::distance(val.first, p) <= converted_radius) {
708+
results.push_back(val.second);
709+
}
710+
}));
670711

671-
double GeoIndex::ConvertToRadiusInMeters(size_t radius, std::string_view arg) {
672-
const std::string unit = absl::AsciiStrToUpper(arg);
673-
if (unit == "M") {
674-
return radius * 1;
675-
} else if (unit == "KM") {
676-
return radius * 1000;
677-
} else if (unit == "FT") {
678-
return radius * 0.3048;
679-
} else if (unit == "MI") {
680-
return radius * 1609.34;
681-
} else {
682-
return -1;
683-
}
712+
return results;
684713
}
685714

686-
std::optional<GeoIndex::point> GeoIndex::GetGeoPoint(const DocumentAccessor& doc,
687-
string_view field) {
688-
auto element = doc.GetStrings(field);
689-
690-
if (!element) {
691-
return std::nullopt;
692-
}
693-
694-
absl::InlinedVector<string_view, 2> coordinates = absl::StrSplit(element.value()[0], ",");
695-
696-
if (coordinates.size() != 2) {
697-
return std::nullopt;
698-
}
699-
700-
double lat;
701-
if (!absl::SimpleAtod(coordinates[0], &lat)) {
702-
return std::nullopt;
703-
}
704-
705-
double lon;
706-
if (!absl::SimpleAtod(coordinates[1], &lon)) {
707-
return std::nullopt;
708-
}
709-
710-
return point{lat, lon};
715+
std::vector<DocId> GeoIndex::GetAllDocsWithNonNullValues() const {
716+
std::vector<DocId> results;
717+
std::for_each(boost::geometry::index::begin(*rtree_), boost::geometry::index::end(*rtree_),
718+
[&results](auto const& val) { results.push_back(val.second); });
719+
return results;
711720
}
712721

713722
} // namespace dfly::search

‎src/core/search/indices.h‎

Lines changed: 5 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,8 @@
99
#include <absl/container/flat_hash_set.h>
1010

1111
#include <boost/geometry.hpp>
12+
#include <boost/geometry/core/cs.hpp>
13+
#include <boost/geometry/index/parameters.hpp>
1214
#include <map>
1315
#include <memory>
1416
#include <optional>
@@ -217,17 +219,11 @@ struct GeoIndex : public BaseIndex {
217219

218220
bool Add(DocId id, const DocumentAccessor& doc, std::string_view field) override;
219221
void Remove(DocId id, const DocumentAccessor& doc, std::string_view field) override;
220-
std::vector<DocId> RadiusSearch(double lat, double lon, double radius) const;
221-
222-
std::vector<DocId> GetAllDocsWithNonNullValues() const override {
223-
return std::vector<DocId>{};
224-
}
225-
226-
static double ConvertToRadiusInMeters(size_t radius, std::string_view arg);
222+
std::vector<DocId> RadiusSearch(double lat, double lon, double radius, std::string_view arg);
223+
std::vector<DocId> GetAllDocsWithNonNullValues() const override;
227224

228225
private:
229-
std::optional<point> GetGeoPoint(const DocumentAccessor& doc, std::string_view field);
230-
using rtree = boost::geometry::index::rtree<index_entry, boost::geometry::index::quadratic<16>>;
226+
using rtree = boost::geometry::index::rtree<index_entry, boost::geometry::index::linear<16>>;
231227
std::unique_ptr<rtree> rtree_;
232228
};
233229

‎src/core/search/parser.y‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -184,8 +184,9 @@ bracket_filter_expr:
184184
| LPAREN UINT32 COMMA UINT32 { $$ = AstRangeNode(toUint32($2), true, toUint32($4), false); }
185185
| UINT32 COMMA LPAREN UINT32 { $$ = AstRangeNode(toUint32($1), false, toUint32($4), true); }
186186
| LPAREN UINT32 COMMA LPAREN UINT32 { $$ = AstRangeNode(toUint32($2), true, toUint32($5), true); }
187-
/* This is rule for GEO filter */
187+
/* GEO filter */
188188
| DOUBLE DOUBLE UINT32 TERM { $$ = AstGeoNode(toDouble($1), toDouble($2), toUint32($3), std::move($4)); }
189+
| DOUBLE DOUBLE DOUBLE TERM { $$ = AstGeoNode(toDouble($1), toDouble($2), toDouble($3), std::move($4)); }
189190

190191
field_cond_expr:
191192
field_unary_expr { $$ = std::move($1); }

‎src/core/search/search.cc‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -281,9 +281,8 @@ struct BasicSearch {
281281

282282
IndexResult Search(const AstGeoNode& node, string_view active_field) {
283283
DCHECK(!active_field.empty());
284-
double converted_radius = GeoIndex::ConvertToRadiusInMeters(node.radius, node.unit);
285284
if (auto* index = GetIndex<GeoIndex>(active_field); index) {
286-
return IndexResult{index->RadiusSearch(node.lat, node.lon, converted_radius)};
285+
return IndexResult{index->RadiusSearch(node.lat, node.lon, node.radius, node.unit)};
287286
}
288287
return IndexResult{};
289288
}
@@ -408,6 +407,7 @@ struct BasicSearch {
408407
// used by knn
409408

410409
DCHECK(top_level || holds_alternative<AstKnnNode>(node.Variant()) ||
410+
holds_alternative<AstGeoNode>(node.Variant()) ||
411411
visit([](auto* set) { return is_sorted(set->begin(), set->end()); }, result.Borrowed()));
412412

413413
if (profile_builder_)

0 commit comments

Comments
 (0)