Add support for aggregations

Rajesh Kumar
1 parent 1a0ff891
Showing 3 changed files with 127 additions and 1 deletions Show diff stats
lib/searchkick/query.rb
lib/searchkick/results.rb
test/aggs_test.rb
@@ -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,44 @@ 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|
+            payload[:aggs][field] = {
+              terms: {
+                field: agg_options[:field] || field
+              }
+            }
+
+            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 => {
+                    terms: {
+                      field: field
+                    }
+                  }
+                }
+              }
+            end
+          end
+        end
+
         # suggestions
         if options[:suggest]
           suggest_fields = (searchkick_options[:suggest] || []).map(&:to_s)
@@ -75,6 +75,17 @@ module Searchkick
       response["facets"]
     end
+    def aggs
+      response["aggregations"].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
+
     def took
       response["took"]
     end
@@ -0,0 +1,77 @@
+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_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
	@@ -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,44 @@ module Searchkick		@@ -336,6 +336,44 @@ 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	+ payload[:aggs][field] = {
		350	+ terms: {
		351	+ field: agg_options[:field] \|\| field
		352	+ }
		353	+ }
		354	+
		355	+ agg_options.deep_merge!(where: options.fetch(:where, {}).reject { \|k\| k == field } ) if options[:smart_aggs] == true
		356	+ agg_filters = where_filters(agg_options[:where])
		357	+
		358	+ if agg_filters.any?
		359	+ payload[:aggs][field] = {
		360	+ filter: {
		361	+ bool: {
		362	+ must: agg_filters
		363	+ }
		364	+ },
		365	+ aggs: {
		366	+ field => {
		367	+ terms: {
		368	+ field: field
		369	+ }
		370	+ }
		371	+ }
		372	+ }
		373	+ end
		374	+ end
		375	+ end
		376	+
339	# suggestions	377	# suggestions
340	if options[:suggest]	378	if options[:suggest]
341	suggest_fields = (searchkick_options[:suggest] \|\| []).map(&:to_s)	379	suggest_fields = (searchkick_options[:suggest] \|\| []).map(&:to_s)
	@@ -75,6 +75,17 @@ module Searchkick		@@ -75,6 +75,17 @@ module Searchkick
75	response["facets"]	75	response["facets"]
76	end	76	end
77		77
		78	+ def aggs
		79	+ response["aggregations"].each do \|field, filtered_agg\|
		80	+ buckets = filtered_agg[field]
		81	+ # move the buckets one level above into the field hash
		82	+ if buckets
		83	+ filtered_agg.delete(field)
		84	+ filtered_agg.merge!(buckets)
		85	+ end
		86	+ end
		87	+ end
		88	+
78	def took	89	def took
79	response["took"]	90	response["took"]
80	end	91	end
@@ -0,0 +1,77 @@		@@ -0,0 +1,77 @@
	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_where_no_smart_aggs
	31	+ assert_equal ({2 => 2}), store_agg(where: {color: "red"}, aggs: {store_id: {where: {in_stock: false}}})
	32	+ assert_equal ({2 => 2}), store_agg(where: {color: "blue"}, aggs: {store_id: {where: {in_stock: false}}})
	33	+ end
	34	+
	35	+ def test_smart_aggs
	36	+ assert_equal ({1 => 1}), store_agg(where: {in_stock: true}, aggs: [:store_id], smart_aggs: true)
	37	+ end
	38	+
	39	+ def test_smart_aggs_where
	40	+ assert_equal ({2 => 1}), store_agg(where: {color: "red"}, aggs: {store_id: {where: {in_stock: false}}}, smart_aggs: true)
	41	+ end
	42	+
	43	+ def test_smart_aggs_where_override
	44	+ assert_equal ({2 => 1}), store_agg(where: {color: "red"}, aggs: {store_id: {where: {in_stock: false, color: "blue"}}}, smart_aggs: true)
	45	+ assert_equal ({}), store_agg(where: {color: "blue"}, aggs: {store_id: {where: {in_stock: false, color: "red"}}}, smart_aggs: true)
	46	+ end
	47	+
	48	+ def test_smart_aggs_skip_agg
	49	+ assert_equal ({1 => 1, 2 => 2}), store_agg(where: {store_id: 2}, aggs: [:store_id], smart_aggs: true)
	50	+ end
	51	+
	52	+ def test_smart_aggs_skip_agg_complex
	53	+ assert_equal ({1 => 1, 2 => 1}), store_agg(where: {store_id: 2, price: {gt: 5}}, aggs: [:store_id], smart_aggs: true)
	54	+ end
	55	+
	56	+ def test_multiple_aggs
	57	+ assert_equal ({"store_id" => {1 => 1, 2 => 2}, "color" => {"blue" => 1, "green" => 1, "red" => 1}}), store_multiple_aggs(aggs: [:store_id, :color])
	58	+ end
	59	+
	60	+ protected
	61	+
	62	+ def buckets_as_hash(agg)
	63	+ agg["buckets"].map { \|v\| [v["key"], v["doc_count"]] }.to_h
	64	+ end
	65	+
	66	+ def store_agg(options, agg_key = "store_id")
	67	+ buckets = Product.search("Product", options).aggs[agg_key]
	68	+ buckets_as_hash(buckets)
	69	+ end
	70	+
	71	+ def store_multiple_aggs(options)
	72	+ Product.search("Product", options).aggs.map do \|field, filtered_agg\|
	73	+ [field, buckets_as_hash(filtered_agg)]
	74	+ end.to_h
	75	+ end
	76	+
	77	+end