Merge pull request #60 from Soostone/value-count-aggregation

Add value_count aggregation support

src/Database/Bloodhound/Types.hs

@@ -103,6 +103,7 @@ module Database.Bloodhound.Types
        , RegexpFlags(..)
        , RegexpFlag(..)
        , FieldName(..)
+       , Script(..)
        , IndexName(..)
        , TemplateName(..)
        , TemplatePattern(..)
@@ -201,6 +202,7 @@ module Database.Bloodhound.Types
        , Bucket(..)
        , BucketAggregation(..)
        , TermsAggregation(..)
+       , ValueCountAggregation(..)
        , DateHistogramAggregation(..)
 
        , Highlights(..)
@@ -607,6 +609,12 @@ newtype QueryString = QueryString Text deriving (Eq, Generic, Show)
 -}
 newtype FieldName = FieldName Text deriving (Eq, Show)
 
+
+{-| 'Script' is often used in place of 'FieldName' to specify more
+complex ways of extracting a value from a document.
+-}
+newtype Script = Script { scriptText :: Text } deriving (Eq, Show)
+
 {-| 'CacheName' is used in 'RegexpFilter' for describing the
     'CacheKey' keyed caching behavior.
 -}
@@ -1310,7 +1318,8 @@ data Interval = Year
               | FractionalInterval Float TimeInterval deriving (Eq, Show)
 
 data Aggregation = TermsAgg TermsAggregation
-                 | DateHistogramAgg DateHistogramAggregation deriving (Eq, Show)
+                 | DateHistogramAgg DateHistogramAggregation
+                 | ValueCountAgg ValueCountAggregation deriving (Eq, Show)
 
 
 data TermsAggregation = TermsAggregation { term              :: Either Text Text
@@ -1336,6 +1345,8 @@ data DateHistogramAggregation = DateHistogramAggregation { dateField      :: Fie
                                                          , dateAggs       :: Maybe Aggregations
                                                          } deriving (Eq, Show)
 
+data ValueCountAggregation = FieldValueCount FieldName
+                           | ScriptValueCount Script deriving (Eq, Show)
 
 mkTermsAggregation :: Text -> TermsAggregation
 mkTermsAggregation t = TermsAggregation (Left t) Nothing Nothing Nothing Nothing Nothing Nothing Nothing Nothing Nothing
@@ -1409,6 +1420,10 @@ instance ToJSON Aggregation where
                                                "post_offset" .= postOffset
                                              ],
                "aggs"           .= dateHistoAggs ]
+  toJSON (ValueCountAgg a) = object ["value_count" .= v]
+    where v = case a of
+                (FieldValueCount (FieldName n)) -> object ["field" .= n]
+                (ScriptValueCount (Script s))   -> object ["script" .= s]
 
 type AggregationResults = M.Map Text Value
 

tests/tests.hs

@@ -14,6 +14,7 @@ import           Data.List                       (nub)
 import           Data.List.NonEmpty              (NonEmpty (..))
 import qualified Data.List.NonEmpty              as NE
 import qualified Data.Map.Strict                 as M
+import           Data.Monoid
 import           Data.Text                       (Text)
 import qualified Data.Text                       as T
 import           Data.Time.Calendar              (Day (..))
@@ -174,15 +175,16 @@ insertOther = do
 
 searchTweet :: Search -> BH IO (Either String Tweet)
 searchTweet search = do
-  reply <- searchByIndex testIndex search
-  let result = eitherDecode (responseBody reply) :: Either String (SearchResult Tweet)
+  result <- searchTweets search
   let myTweet = fmap (hitSource . head . hits . searchHits) result
   return myTweet
 
+searchTweets :: Search -> BH IO (Either String (SearchResult Tweet))
+searchTweets search = eitherDecode . responseBody <$> searchByIndex testIndex search
+
 searchExpectNoResults :: Search -> BH IO ()
 searchExpectNoResults search = do
-  reply <- searchByIndex testIndex search
-  let result = eitherDecode (responseBody reply) :: Either String (SearchResult Tweet)
+  result <- searchTweets search
   let emptyHits = fmap (hits . searchHits) result
   liftIO $
     emptyHits `shouldBe` Right []
@@ -214,8 +216,7 @@ searchTermsAggHint hints = do
 
 searchTweetHighlight :: Search -> BH IO (Either String (Maybe HitHighlight))
 searchTweetHighlight search = do
-  reply <- searchByIndex testIndex search
-  let result = eitherDecode (responseBody reply) :: Either String (SearchResult Tweet)
+  result <- searchTweets search
   let myHighlight = fmap (hitHighlight . head . hits . searchHits) result
   return myHighlight
 
@@ -435,8 +436,7 @@ main = hspec $ do
       let search = Search Nothing
                    (Just IdentityFilter) (Just [sortSpec]) Nothing Nothing
                    False (From 0) (Size 10) Nothing
-      reply <- searchByIndex testIndex search
-      let result = eitherDecode (responseBody reply) :: Either String (SearchResult Tweet)
+      result <- searchTweets search
       let myTweet = fmap (hitSource . head . hits . searchHits) result
       liftIO $
         myTweet `shouldBe` Right otherTweet
@@ -628,6 +628,19 @@ main = hspec $ do
       _ <- insertData
       searchTermsAggHint [GlobalOrdinals, GlobalOrdinalsHash, GlobalOrdinalsLowCardinality, Map]
 
+
+    it "can execute value_count aggregations" $ withTestEnv $ do
+      _ <- insertData
+      _ <- insertOther
+      let ags = mkAggregations "user_count" (ValueCountAgg (FieldValueCount (FieldName "user"))) <>
+                mkAggregations "bogus_count" (ValueCountAgg (FieldValueCount (FieldName "bogus")))
+      let search = mkAggregateSearch Nothing ags
+      let countPair k n = (k, object ["value" .= Number n])
+      res <- searchTweets search
+      liftIO $
+        fmap aggregations res `shouldBe` Right (Just (M.fromList [ countPair "user_count" 2
+                                                                 , countPair "bogus_count" 0
+                                                                 ]))
     -- Interaction of date serialization and date histogram aggregation is broken.
     -- it "returns date histogram aggregation results" $ withTestEnv $ do
     --   _ <- insertData