Filtering by geo arbitrary rectangle (#64)

machinio · Jan 17, 2024 · 829daa5 · 829daa5
1 parent 28f66cb
commit 829daa5
Show file tree

Hide file tree

Showing 14 changed files with 131 additions and 35 deletions.
diff --git a/.rubocop.yml b/.rubocop.yml
@@ -6,7 +6,7 @@ AllCops:
   TargetRubyVersion: '2.5'
   DisabledByDefault: true
 Layout/LineLength:
-  Max: 120
+  Max: 160
 Style/Documentation:
   Enabled: false
 Style/FrozenStringLiteralComment:

diff --git a/README.md b/README.md
@@ -308,6 +308,35 @@ end
   request.run(page: 1, page_size: 10)
 ```
 
+### Filtering by a Geofilt
+
+```ruby
+  spatial_point = Solr::SpatialPoint.new(lat: 40.0, lon: -120.0)
+
+  filters = [
+    Solr::Query::Request::Geofilt.new(field: :location, spatial_point: spatial_point, spatial_radius: 100)
+  ]
+
+  request = Solr::Query::Request.new(search_term: 'term', filters: filters)
+  request.run(page: 1, page_size: 10)
+```
+
+### Filtering by an Arbitrary Rectangle
+
+```ruby
+  spatial_rectangle = Solr::SpatialRectangle.new(
+    top_left: Solr::SpatialPoint.new(lat: 40.0, lon: -120.0),
+    bottom_right: Solr::SpatialPoint.new(lat: 30.0, lon: -110.0)
+  )
+
+  filters = [
+    Solr::Query::Request::Filter.new(type: :equal, field: :location, value: spatial_rectangle)
+  ]
+
+  request = Solr::Query::Request.new(search_term: 'term', filters: filters)
+  request.run(page: 1, page_size: 10)
+```
+
 ## Query with sorting
 
 ```ruby
@@ -374,7 +403,9 @@ Default sorting logic is following: nulls last, not-nulls first.
   )
   request.run(page: 1, page_size: 10)
 ```
+
 ### Dictionary boosting function
+
 Sometimes you want to do a dictionary-style boosting
 example: given a hash (dictionary)
 

diff --git a/lib/solr.rb b/lib/solr.rb
@@ -9,6 +9,8 @@
 require 'solr/document_collection'
 require 'solr/grouped_document_collection'
 require 'solr/response'
+require 'solr/spatial_point'
+require 'solr/spatial_rectangle'
 require 'solr/request/http_request'
 require 'solr/request/runner'
 require 'solr/query/request'

diff --git a/lib/solr/query/request.rb b/lib/solr/query/request.rb
@@ -1,7 +1,7 @@
 require 'solr/query/request/facet'
 require 'solr/query/request/filter'
 require 'solr/query/request/boost_magnitude'
-require 'solr/query/request/geo_filter'
+require 'solr/query/request/geofilt'
 require 'solr/query/request/sorting'
 require 'solr/query/request/edismax_adapter'
 require 'solr/query/request/grouping'

diff --git a/lib/solr/query/request/filter.rb b/lib/solr/query/request/filter.rb
@@ -30,6 +30,8 @@ def solr_value
             value.map { |v| to_primitive_solr_value(v) }.join(' OR ')
           elsif value.is_a?(::Range)
             to_interval_solr_value(value)
+          elsif value.is_a?(Solr::SpatialRectangle)
+            value.to_solr_s
           else
             to_primitive_solr_value(value)
           end

diff --git a/lib/solr/query/request/geo_filter.rb b/lib/solr/query/request/geo_filter.rb
diff --git a/lib/solr/query/request/geofilt.rb b/lib/solr/query/request/geofilt.rb
@@ -0,0 +1,28 @@
+module Solr
+  module Query
+    class Request
+      class Geofilt
+        include Solr::Support::SchemaHelper
+
+        attr_reader :field, :spatial_point, :spatial_radius, :score
+
+        def initialize(field:, spatial_point:, spatial_radius:, score: nil)
+          raise ArgumentError, 'spatial_point must be a Solr::SpatialPoint' unless spatial_point.is_a?(Solr::SpatialPoint)
+
+          @field = field
+          @spatial_point = spatial_point
+          @spatial_radius = spatial_radius
+          @score = score
+          freeze
+        end
+
+        def to_solr_s
+          solr_field = solarize_field(@field)
+          filter = "{!geofilt sfield=#{solr_field} pt=#{spatial_point.to_solr_s} d=#{spatial_radius}"
+          filter << " score=#{score}" if score
+          filter + "}"
+        end
+      end
+    end
+  end
+end
diff --git a/lib/solr/spatial_point.rb b/lib/solr/spatial_point.rb
@@ -0,0 +1,14 @@
+module Solr
+  class SpatialPoint
+    attr_reader :lat, :lon
+
+    def initialize(lat:, lon:)
+      @lat = lat
+      @lon = lon
+    end
+
+    def to_solr_s
+      "#{lat},#{lon}"
+    end
+  end
+end
diff --git a/lib/solr/spatial_rectangle.rb b/lib/solr/spatial_rectangle.rb
@@ -0,0 +1,17 @@
+module Solr
+  class SpatialRectangle
+    attr_reader :top_right, :bottom_left
+
+    def initialize(top_right:, bottom_left:)
+      raise ArgumentError, 'top_right must be a Solr::SpatialPoint' unless top_right.is_a?(Solr::SpatialPoint)
+      raise ArgumentError, 'bottom_left must be a Solr::SpatialPoint' unless bottom_left.is_a?(Solr::SpatialPoint)
+
+      @top_right = top_right
+      @bottom_left = bottom_left
+    end
+
+    def to_solr_s
+      "[#{bottom_left.to_solr_s} TO #{top_right.to_solr_s}]"
+    end
+  end
+end
diff --git a/spec/query/request/filter_spec.rb b/spec/query/request/filter_spec.rb
@@ -25,6 +25,15 @@
       end
     end
 
+    context 'when value is spatial rectangle' do
+      let(:top_right) { Solr::SpatialPoint.new(lat: 1.0, lon: 2.0) }
+      let(:bottom_left) { Solr::SpatialPoint.new(lat: 3.0, lon: 4.0) }
+      let(:spatial_rectangle) { Solr::SpatialRectangle.new(top_right: top_right, bottom_left: bottom_left) }
+
+      subject { described_class.new(type: :equal, field: :field, value: spatial_rectangle).to_solr_s }
+      it { is_expected.to eq('field:([3.0,4.0 TO 1.0,2.0])') }
+    end
+
     context 'when value is date' do
       subject { described_class.new(type: :equal, field: :field, value: Date.parse('2018-08-01')).to_solr_s }
       it { is_expected.to eq('field:(2018-08-01T00:00:00Z)') }

diff --git a/spec/query/request/geo_filter_spec.rb b/spec/query/request/geo_filter_spec.rb
diff --git a/spec/query/request/geofilt_spec.rb b/spec/query/request/geofilt_spec.rb
@@ -0,0 +1,9 @@
+RSpec.describe Solr::Query::Request::Geofilt do
+  describe '.to_solr_s' do
+    let(:spatial_point) { Solr::SpatialPoint.new(lat: -25.429692, lon: -49.271265) }
+
+    subject { described_class.new(field: :machine_type, spatial_point: spatial_point, spatial_radius: 161) }
+
+    it { expect(subject.to_solr_s).to eq('{!geofilt sfield=machine_type pt=-25.429692,-49.271265 d=161}') }
+  end
+end
diff --git a/spec/spatial_point_spec.rb b/spec/spatial_point_spec.rb
@@ -0,0 +1,7 @@
+RSpec.describe Solr::SpatialPoint do
+  describe '#to_solr_s' do
+    it 'returns a solr string' do
+      expect(described_class.new(lat: 1.0, lon: 2.0).to_solr_s).to eq('1.0,2.0')
+    end
+  end
+end
diff --git a/spec/spatial_rectangle_spec.rb b/spec/spatial_rectangle_spec.rb
@@ -0,0 +1,10 @@
+RSpec.describe Solr::SpatialRectangle do
+  describe '#to_solr_s' do
+    it 'returns a solr string' do
+      expect(described_class.new(
+        top_right: Solr::SpatialPoint.new(lat: 1.0, lon: 2.0),
+        bottom_left: Solr::SpatialPoint.new(lat: 3.0, lon: 4.0)
+      ).to_solr_s).to eq('[3.0,4.0 TO 1.0,2.0]')
+    end
+  end
+end