Merge branch 'meetrajesh-support_for_aggregations'

Andrew Kane
2 parents 54e5e5a6 aa6812c4
Showing 4 changed files with 136 additions and 1 deletions Show diff stats
CHANGELOG.md
lib/searchkick/query.rb
lib/searchkick/results.rb
test/aggs_test.rb
 ## 0.9.2 [unreleased]
+- Added support for aggregations
 - Added ability to use misspellings for partial matches
 - Added `fragment_size` option for highlight
 - Added `took` method to results
@@ -266,7 +266,7 @@ module Searchkick
         # filters
         filters = where_filters(options[:where])
         if filters.any?
-          if options[:facets]
+          if options[:facets] || options[:aggs]
             payload[:filter] = {
               and: filters
             }
@@ -336,6 +336,43 @@ module Searchkick
           end
         end
+        # aggregations
+        if options[:aggs]
+          aggs = options[:aggs]
+          payload[:aggs] = {}
+
+          if aggs.is_a?(Array) # convert to more advanced syntax
+            aggs = aggs.map { |f| [f, {}] }.to_h
+          end
+
+          aggs.each do |field, agg_options|
+            size = agg_options[:limit] ? agg_options[:limit] : 100_000
+
+            payload[:aggs][field] = {
+              terms: {
+                field: agg_options[:field] || field,
+                size: size
+              }
+            }
+
+            agg_options.deep_merge!(where: options.fetch(:where, {}).reject { |k| k == field } ) if options[:smart_aggs] == true
+            agg_filters = where_filters(agg_options[:where])
+
+            if agg_filters.any?
+              payload[:aggs][field] = {
+                filter: {
+                  bool: {
+                    must: agg_filters
+                  }
+                },
+                aggs: {
+                  field => payload[:aggs][field]
+                }
+              }
+            end
+          end
+        end
+
         # suggestions
         if options[:suggest]
           suggest_fields = (searchkick_options[:suggest] || []).map(&:to_s)
@@ -75,6 +75,19 @@ module Searchkick
       response["facets"]
     end
+    def aggs
+      @aggs ||= begin
+        response["aggregations"].dup.each do |field, filtered_agg|
+          buckets = filtered_agg[field]
+          # move the buckets one level above into the field hash
+          if buckets
+            filtered_agg.delete(field)
+            filtered_agg.merge!(buckets)
+          end
+        end
+      end
+    end
+
     def took
       response["took"]
     end
@@ -0,0 +1,84 @@
+require_relative "test_helper"
+require "active_support/core_ext"
+
+class TestAggs < Minitest::Test
+
+  def setup
+    super
+    store [
+      {name: "Product Show", latitude: 37.7833, longitude: 12.4167, store_id: 1, in_stock: true, color: "blue", price: 21, created_at: 2.days.ago},
+      {name: "Product Hide", latitude: 29.4167, longitude: -98.5000, store_id: 2, in_stock: false, color: "green", price: 25, created_at: 2.days.from_now},
+      {name: "Product B", latitude: 43.9333, longitude: -122.4667, store_id: 2, in_stock: false, color: "red", price: 5},
+      {name: "Foo", latitude: 43.9333, longitude: 12.4667, store_id: 3, in_stock: false, color: "yellow", price: 15}
+    ]
+  end
+
+  def test_basic
+    assert_equal ({1 => 1, 2 => 2}), store_agg(aggs: [:store_id])
+  end
+
+  def test_where
+    assert_equal ({1 => 1}), store_agg(aggs: {store_id: {where: {in_stock: true}}})
+  end
+
+  def test_field
+    assert_equal ({1 => 1, 2 => 2}), store_agg(aggs: {store_id: {}})
+    assert_equal ({1 => 1, 2 => 2}), store_agg(aggs: {store_id: {field: "store_id"}})
+    assert_equal ({1 => 1, 2 => 2}), store_agg({aggs: {store_id_new: {field: "store_id"}}}, "store_id_new")
+  end
+
+  def test_limit
+    agg = Product.search("Product", aggs: {store_id: {limit: 1}}).aggs["store_id"]
+    assert_equal 1, agg["buckets"].size
+    # assert_equal 3, agg["doc_count"]
+    assert_equal 1, agg["sum_other_doc_count"]
+  end
+
+  def test_where_no_smart_aggs
+    assert_equal ({2 => 2}), store_agg(where: {color: "red"}, aggs: {store_id: {where: {in_stock: false}}})
+    assert_equal ({2 => 2}), store_agg(where: {color: "blue"}, aggs: {store_id: {where: {in_stock: false}}})
+  end
+
+  def test_smart_aggs
+    assert_equal ({1 => 1}), store_agg(where: {in_stock: true}, aggs: [:store_id], smart_aggs: true)
+  end
+
+  def test_smart_aggs_where
+    assert_equal ({2 => 1}), store_agg(where: {color: "red"}, aggs: {store_id: {where: {in_stock: false}}}, smart_aggs: true)
+  end
+
+  def test_smart_aggs_where_override
+    assert_equal ({2 => 1}), store_agg(where: {color: "red"}, aggs: {store_id: {where: {in_stock: false, color: "blue"}}}, smart_aggs: true)
+    assert_equal ({}), store_agg(where: {color: "blue"}, aggs: {store_id: {where: {in_stock: false, color: "red"}}}, smart_aggs: true)
+  end
+
+  def test_smart_aggs_skip_agg
+    assert_equal ({1 => 1, 2 => 2}), store_agg(where: {store_id: 2}, aggs: [:store_id], smart_aggs: true)
+  end
+
+  def test_smart_aggs_skip_agg_complex
+    assert_equal ({1 => 1, 2 => 1}), store_agg(where: {store_id: 2, price: {gt: 5}}, aggs: [:store_id], smart_aggs: true)
+  end
+
+  def test_multiple_aggs
+    assert_equal ({"store_id" => {1 => 1, 2 => 2}, "color" => {"blue" => 1, "green" => 1, "red" => 1}}), store_multiple_aggs(aggs: [:store_id, :color])
+  end
+
+  protected
+
+  def buckets_as_hash(agg)
+    agg["buckets"].map { |v| [v["key"], v["doc_count"]] }.to_h
+  end
+
+  def store_agg(options, agg_key = "store_id")
+    buckets = Product.search("Product", options).aggs[agg_key]
+    buckets_as_hash(buckets)
+  end
+
+  def store_multiple_aggs(options)
+    Product.search("Product", options).aggs.map do |field, filtered_agg|
+      [field, buckets_as_hash(filtered_agg)]
+    end.to_h
+  end
+
+end
1	## 0.9.2 [unreleased]	1	## 0.9.2 [unreleased]
2		2
		3	+- Added support for aggregations
3	- Added ability to use misspellings for partial matches	4	- Added ability to use misspellings for partial matches
4	- Added `fragment_size` option for highlight	5	- Added `fragment_size` option for highlight
5	- Added `took` method to results	6	- Added `took` method to results
	@@ -266,7 +266,7 @@ module Searchkick		@@ -266,7 +266,7 @@ module Searchkick
266	# filters	266	# filters
267	filters = where_filters(options[:where])	267	filters = where_filters(options[:where])
268	if filters.any?	268	if filters.any?
269	- if options[:facets]	269	+ if options[:facets] \|\| options[:aggs]
270	payload[:filter] = {	270	payload[:filter] = {
271	and: filters	271	and: filters
272	}	272	}
	@@ -336,6 +336,43 @@ module Searchkick		@@ -336,6 +336,43 @@ module Searchkick
336	end	336	end
337	end	337	end
338		338
		339	+ # aggregations
		340	+ if options[:aggs]
		341	+ aggs = options[:aggs]
		342	+ payload[:aggs] = {}
		343	+
		344	+ if aggs.is_a?(Array) # convert to more advanced syntax
		345	+ aggs = aggs.map { \|f\| [f, {}] }.to_h
		346	+ end
		347	+
		348	+ aggs.each do \|field, agg_options\|
		349	+ size = agg_options[:limit] ? agg_options[:limit] : 100_000
		350	+
		351	+ payload[:aggs][field] = {
		352	+ terms: {
		353	+ field: agg_options[:field] \|\| field,
		354	+ size: size
		355	+ }
		356	+ }
		357	+
		358	+ agg_options.deep_merge!(where: options.fetch(:where, {}).reject { \|k\| k == field } ) if options[:smart_aggs] == true
		359	+ agg_filters = where_filters(agg_options[:where])
		360	+
		361	+ if agg_filters.any?
		362	+ payload[:aggs][field] = {
		363	+ filter: {
		364	+ bool: {
		365	+ must: agg_filters
		366	+ }
		367	+ },
		368	+ aggs: {
		369	+ field => payload[:aggs][field]
		370	+ }
		371	+ }
		372	+ end
		373	+ end
		374	+ end
		375	+
339	# suggestions	376	# suggestions
340	if options[:suggest]	377	if options[:suggest]
341	suggest_fields = (searchkick_options[:suggest] \|\| []).map(&:to_s)	378	suggest_fields = (searchkick_options[:suggest] \|\| []).map(&:to_s)
	@@ -75,6 +75,19 @@ module Searchkick		@@ -75,6 +75,19 @@ module Searchkick
75	response["facets"]	75	response["facets"]
76	end	76	end
77		77
		78	+ def aggs
		79	+ @aggs \|\|= begin
		80	+ response["aggregations"].dup.each do \|field, filtered_agg\|
		81	+ buckets = filtered_agg[field]
		82	+ # move the buckets one level above into the field hash
		83	+ if buckets
		84	+ filtered_agg.delete(field)
		85	+ filtered_agg.merge!(buckets)
		86	+ end
		87	+ end
		88	+ end
		89	+ end
		90	+
78	def took	91	def took
79	response["took"]	92	response["took"]
80	end	93	end
@@ -0,0 +1,84 @@		@@ -0,0 +1,84 @@
	1	+require_relative "test_helper"
	2	+require "active_support/core_ext"
	3	+
	4	+class TestAggs < Minitest::Test
	5	+
	6	+ def setup
	7	+ super
	8	+ store [
	9	+ {name: "Product Show", latitude: 37.7833, longitude: 12.4167, store_id: 1, in_stock: true, color: "blue", price: 21, created_at: 2.days.ago},
	10	+ {name: "Product Hide", latitude: 29.4167, longitude: -98.5000, store_id: 2, in_stock: false, color: "green", price: 25, created_at: 2.days.from_now},
	11	+ {name: "Product B", latitude: 43.9333, longitude: -122.4667, store_id: 2, in_stock: false, color: "red", price: 5},
	12	+ {name: "Foo", latitude: 43.9333, longitude: 12.4667, store_id: 3, in_stock: false, color: "yellow", price: 15}
	13	+ ]
	14	+ end
	15	+
	16	+ def test_basic
	17	+ assert_equal ({1 => 1, 2 => 2}), store_agg(aggs: [:store_id])
	18	+ end
	19	+
	20	+ def test_where
	21	+ assert_equal ({1 => 1}), store_agg(aggs: {store_id: {where: {in_stock: true}}})
	22	+ end
	23	+
	24	+ def test_field
	25	+ assert_equal ({1 => 1, 2 => 2}), store_agg(aggs: {store_id: {}})
	26	+ assert_equal ({1 => 1, 2 => 2}), store_agg(aggs: {store_id: {field: "store_id"}})
	27	+ assert_equal ({1 => 1, 2 => 2}), store_agg({aggs: {store_id_new: {field: "store_id"}}}, "store_id_new")
	28	+ end
	29	+
	30	+ def test_limit
	31	+ agg = Product.search("Product", aggs: {store_id: {limit: 1}}).aggs["store_id"]
	32	+ assert_equal 1, agg["buckets"].size
	33	+ # assert_equal 3, agg["doc_count"]
	34	+ assert_equal 1, agg["sum_other_doc_count"]
	35	+ end
	36	+
	37	+ def test_where_no_smart_aggs
	38	+ assert_equal ({2 => 2}), store_agg(where: {color: "red"}, aggs: {store_id: {where: {in_stock: false}}})
	39	+ assert_equal ({2 => 2}), store_agg(where: {color: "blue"}, aggs: {store_id: {where: {in_stock: false}}})
	40	+ end
	41	+
	42	+ def test_smart_aggs
	43	+ assert_equal ({1 => 1}), store_agg(where: {in_stock: true}, aggs: [:store_id], smart_aggs: true)
	44	+ end
	45	+
	46	+ def test_smart_aggs_where
	47	+ assert_equal ({2 => 1}), store_agg(where: {color: "red"}, aggs: {store_id: {where: {in_stock: false}}}, smart_aggs: true)
	48	+ end
	49	+
	50	+ def test_smart_aggs_where_override
	51	+ assert_equal ({2 => 1}), store_agg(where: {color: "red"}, aggs: {store_id: {where: {in_stock: false, color: "blue"}}}, smart_aggs: true)
	52	+ assert_equal ({}), store_agg(where: {color: "blue"}, aggs: {store_id: {where: {in_stock: false, color: "red"}}}, smart_aggs: true)
	53	+ end
	54	+
	55	+ def test_smart_aggs_skip_agg
	56	+ assert_equal ({1 => 1, 2 => 2}), store_agg(where: {store_id: 2}, aggs: [:store_id], smart_aggs: true)
	57	+ end
	58	+
	59	+ def test_smart_aggs_skip_agg_complex
	60	+ assert_equal ({1 => 1, 2 => 1}), store_agg(where: {store_id: 2, price: {gt: 5}}, aggs: [:store_id], smart_aggs: true)
	61	+ end
	62	+
	63	+ def test_multiple_aggs
	64	+ assert_equal ({"store_id" => {1 => 1, 2 => 2}, "color" => {"blue" => 1, "green" => 1, "red" => 1}}), store_multiple_aggs(aggs: [:store_id, :color])
	65	+ end
	66	+
	67	+ protected
	68	+
	69	+ def buckets_as_hash(agg)
	70	+ agg["buckets"].map { \|v\| [v["key"], v["doc_count"]] }.to_h
	71	+ end
	72	+
	73	+ def store_agg(options, agg_key = "store_id")
	74	+ buckets = Product.search("Product", options).aggs[agg_key]
	75	+ buckets_as_hash(buckets)
	76	+ end
	77	+
	78	+ def store_multiple_aggs(options)
	79	+ Product.search("Product", options).aggs.map do \|field, filtered_agg\|
	80	+ [field, buckets_as_hash(filtered_agg)]
	81	+ end.to_h
	82	+ end
	83	+
	84	+end