git.iko.soy/typeable/bloodhound
on master help sign in
Merge pull request #14 from chrisguiney/aggregations-merge

Add terms and date histogram aggregations

commit1dde598ee7a96c63cc2d9de5d454ae823544010c authorChris Allen <cma@bitemyapp.com> dateSat 11 Oct 2014, 16:10 (12 years ago) parentf4d0e1065e6be2
4 files changed+359−15
compare with master
unifiedsplit
Database/Bloodhound/Client.hs+5−1 Database/Bloodhound/Types.hs+232−12 bloodhound.cabal+2 tests/tests.hs+120−2
Database/Bloodhound/Client.hs+5−1
@@ -15,6 +15,7 @@ module Database.Bloodhound.Client
1515 , searchByType
1616 , refreshIndex
1717 , mkSearch
18+ , mkAggregateSearch
1819 , bulk
1920 , pageSearch
2021 , mkShardCount
@@ -244,7 +245,10 @@ searchByType (Server server) (IndexName indexName)
244245 url = joinPath [server, indexName, mappingName, "_search"]
245246
246247mkSearch :: Maybe Query -> Maybe Filter -> Search
247−mkSearch query filter = Search query filter Nothing False 0 10
248+mkSearch query filter = Search query filter Nothing Nothing False 0 10
249+
250+mkAggregateSearch :: Maybe Query -> Aggregations -> Search
251+mkAggregateSearch query aggregations = Search query Nothing Nothing (Just aggregations) False 0 0
248252
249253pageSearch :: Int -> Int -> Search -> Search
250254pageSearch pageFrom pageSize search = search { from = pageFrom, size = pageSize }
Database/Bloodhound/Types.hs+232−12
@@ -22,6 +22,12 @@ module Database.Bloodhound.Types
2222 , mkBoolQuery
2323 , mkRangeQuery
2424 , mkQueryStringQuery
25+ , mkAggregations
26+ , mkTermsAggregation
27+ , mkTermsScriptAggregation
28+ , mkDateHistogram
29+ , toTerms
30+ , toDateHistogram
2531 , omitNulls
2632 , Version(..)
2733 , Status(..)
@@ -143,16 +149,33 @@ module Database.Bloodhound.Types
143149 , FieldOrFields(..)
144150 , Monoid(..)
145151 , ToJSON(..)
152+ , Interval(..)
153+ , TimeInterval(..)
154+ , ExecutionHint(..)
155+ , CollectionMode(..)
156+ , TermOrder(..)
157+ , TermInclusion(..)
158+ , Aggregation(..)
159+ , Aggregations
160+ , AggregationResults
161+ , Bucket(..)
162+ , BucketAggregation(..)
163+ , TermsAggregation(..)
164+ , DateHistogramAggregation(..)
165+ , TermsResult
166+ , DateHistogramResult
146167 ) where
147168
148169import Control.Applicative
149170import Data.Aeson
171+import Data.Aeson.Types (parseMaybe)
150172import qualified Data.ByteString.Lazy.Char8 as L
151173import Data.List (nub)
152174import Data.List.NonEmpty (NonEmpty (..))
153175import Data.Monoid
154176import Data.Text (Text)
155177import qualified Data.Text as T
178+import qualified Data.Map.Strict as M
156179import Data.Time.Clock (UTCTime)
157180import qualified Data.Vector as V
158181import GHC.Generics (Generic)
@@ -473,9 +496,10 @@ type TrackSortScores = Bool
473496type From = Int
474497type Size = Int
475498
476−data Search = Search { queryBody :: Maybe Query
477− , filterBody :: Maybe Filter
478− , sortBody :: Maybe Sort
499+data Search = Search { queryBody :: Maybe Query
500+ , filterBody :: Maybe Filter
501+ , sortBody :: Maybe Sort
502+ , aggBody :: Maybe Aggregations
479503 -- default False
480504 , trackSortScores :: TrackSortScores
481505 , from :: From
@@ -915,10 +939,11 @@ data DistanceRange =
915939 , distanceTo :: Distance } deriving (Eq, Show)
916940
917941data FromJSON a => SearchResult a =
918− SearchResult { took :: Int
919− , timedOut :: Bool
920− , shards :: ShardResult
921− , searchHits :: SearchHits a } deriving (Eq, Show)
942+ SearchResult { took :: Int
943+ , timedOut :: Bool
944+ , shards :: ShardResult
945+ , searchHits :: SearchHits a
946+ , aggregations :: Maybe AggregationResults } deriving (Eq, Show)
922947
923948type Score = Double
924949
@@ -942,6 +967,199 @@ data ShardResult =
942967showText :: Show a => a -> Text
943968showText = T.pack . show
944969
970+type Aggregations = M.Map Text Aggregation
971+
972+emptyAggregations :: Aggregations
973+emptyAggregations = M.empty
974+
975+mkAggregations :: Text -> Aggregation -> Aggregations
976+mkAggregations name aggregation = M.insert name aggregation emptyAggregations
977+
978+data TermOrder = TermOrder{ termSortField :: Text
979+ , termSortOrder :: SortOrder } deriving (Eq, Show)
980+
981+data TermInclusion = TermInclusion Text
982+ | TermPattern Text Text deriving (Eq, Show)
983+
984+data CollectionMode = BreadthFirst
985+ | DepthFirst deriving (Eq, Show)
986+
987+data ExecutionHint = Ordinals
988+ | GlobalOrdinals
989+ | GlobalOrdinalsHash
990+ | GlobalOrdinalsLowCardinality
991+ | Map deriving(Eq, Show)
992+
993+data TimeInterval = Weeks
994+ | Days
995+ | Hours
996+ | Minutes
997+ | Seconds deriving (Eq)
998+
999+data Interval = Year
1000+ | Quarter
1001+ | Month
1002+ | Week
1003+ | Day
1004+ | Hour
1005+ | Minute
1006+ | Second
1007+ | FractionalInterval Float TimeInterval deriving (Eq, Show)
1008+
1009+data Aggregation = TermsAgg TermsAggregation
1010+ | DateHistogramAgg DateHistogramAggregation deriving (Eq, Show)
1011+
1012+
1013+data TermsAggregation = TermsAggregation { term :: Either Text Text
1014+ , termInclude :: Maybe TermInclusion
1015+ , termExclude :: Maybe TermInclusion
1016+ , termOrder :: Maybe TermOrder
1017+ , termMinDocCount :: Maybe Int
1018+ , termSize :: Maybe Int
1019+ , termShardSize :: Maybe Int
1020+ , termCollectMode :: Maybe CollectionMode
1021+ , termExecutionHint :: Maybe ExecutionHint
1022+ , termAggs :: Maybe Aggregations
1023+ } deriving (Eq, Show)
1024+
1025+data DateHistogramAggregation = DateHistogramAggregation { dateField :: FieldName
1026+ , dateInterval :: Interval
1027+ , dateFormat :: Maybe Text
1028+ , datePreZone :: Maybe Text
1029+ , datePostZone :: Maybe Text
1030+ , datePreOffset :: Maybe Text
1031+ , datePostOffset :: Maybe Text
1032+ , dateAggs :: Maybe Aggregations
1033+ } deriving (Eq, Show)
1034+
1035+
1036+mkTermsAggregation :: Text -> TermsAggregation
1037+mkTermsAggregation t = TermsAggregation (Left t) Nothing Nothing Nothing Nothing Nothing Nothing Nothing Nothing Nothing
1038+
1039+mkTermsScriptAggregation :: Text -> TermsAggregation
1040+mkTermsScriptAggregation t = TermsAggregation (Right t) Nothing Nothing Nothing Nothing Nothing Nothing Nothing Nothing Nothing
1041+
1042+mkDateHistogram :: FieldName -> Interval -> DateHistogramAggregation
1043+mkDateHistogram t i = DateHistogramAggregation t i Nothing Nothing Nothing Nothing Nothing Nothing
1044+
1045+instance ToJSON TermOrder where
1046+ toJSON (TermOrder termSortField termSortOrder) = object [termSortField .= termSortOrder]
1047+
1048+instance ToJSON TermInclusion where
1049+ toJSON (TermInclusion x) = toJSON x
1050+ toJSON (TermPattern pattern flags) = omitNulls [ "pattern" .= pattern,
1051+ "flags" .= flags]
1052+
1053+instance ToJSON CollectionMode where
1054+ toJSON BreadthFirst = "breadth_first"
1055+ toJSON DepthFirst = "depth_first"
1056+
1057+instance ToJSON ExecutionHint where
1058+ toJSON Ordinals = "ordinals"
1059+ toJSON GlobalOrdinals = "global_ordinals"
1060+ toJSON GlobalOrdinalsHash = "global_ordinals_hash"
1061+ toJSON GlobalOrdinalsLowCardinality = "global_ordinals_low_cardinality"
1062+ toJSON Map = "map"
1063+
1064+instance ToJSON Interval where
1065+ toJSON Year = "year"
1066+ toJSON Quarter = "quarter"
1067+ toJSON Month = "month"
1068+ toJSON Week = "week"
1069+ toJSON Day = "day"
1070+ toJSON Hour = "hour"
1071+ toJSON Minute = "minute"
1072+ toJSON Second = "second"
1073+ toJSON (FractionalInterval fraction interval) = toJSON $ show fraction ++ show interval
1074+
1075+instance Show TimeInterval where
1076+ show Weeks = "w"
1077+ show Days = "d"
1078+ show Hours = "h"
1079+ show Minutes = "m"
1080+ show Seconds = "s"
1081+
1082+instance ToJSON Aggregation where
1083+ toJSON (TermsAgg (TermsAggregation term include exclude order minDocCount size shardSize collectMode executionHint aggs)) =
1084+ omitNulls ["terms" .= omitNulls [ toJSON' term,
1085+ "include" .= toJSON include,
1086+ "exclude" .= toJSON exclude,
1087+ "order" .= toJSON order,
1088+ "min_doc_count" .= toJSON minDocCount,
1089+ "size" .= toJSON size,
1090+ "shard_size" .= toJSON shardSize,
1091+ "collect_mode" .= toJSON collectMode,
1092+ "execution_hint" .= toJSON executionHint
1093+ ],
1094+ "aggs" .= toJSON aggs ]
1095+ where
1096+ toJSON' x = case x of { Left y -> "field" .= toJSON y; Right y -> "script" .= toJSON y }
1097+
1098+ toJSON (DateHistogramAgg (DateHistogramAggregation field interval format preZone postZone preOffset postOffset aggs)) =
1099+ omitNulls ["date_histogram" .= omitNulls [ "field" .= toJSON field,
1100+ "interval" .= toJSON interval,
1101+ "format" .= toJSON format,
1102+ "pre_zone" .= toJSON preZone,
1103+ "post_zone" .= toJSON postZone,
1104+ "pre_offset" .= toJSON preOffset,
1105+ "post_offset" .= toJSON postOffset
1106+ ],
1107+ "aggs" .= toJSON aggs ]
1108+
1109+type AggregationResults = M.Map Text Value
1110+
1111+class BucketAggregation a where
1112+ key :: a -> Text
1113+ docCount :: a -> Int
1114+ aggs :: a -> Maybe AggregationResults
1115+
1116+
1117+data (FromJSON a, BucketAggregation a) => Bucket a = Bucket { buckets :: [a]} deriving (Show)
1118+
1119+data TermsResult = TermsResult { termKey :: Text
1120+ , termsDocCount :: Int
1121+ , termsAggs :: Maybe AggregationResults } deriving (Show)
1122+
1123+data DateHistogramResult = DateHistogramResult { dateKey :: Int
1124+ , dateKeyStr :: Maybe Text
1125+ , dateDocCount :: Int
1126+ , dateHistogramAggs :: Maybe AggregationResults } deriving (Show)
1127+
1128+toTerms :: Text -> AggregationResults -> Maybe (Bucket TermsResult)
1129+toTerms t a = M.lookup t a >>= deserialize
1130+ where deserialize = parseMaybe parseJSON
1131+
1132+toDateHistogram :: Text -> AggregationResults -> Maybe (Bucket DateHistogramResult)
1133+toDateHistogram t a = M.lookup t a >>= deserialize
1134+ where deserialize = parseMaybe parseJSON
1135+
1136+instance BucketAggregation TermsResult where
1137+ key = termKey
1138+ docCount = termsDocCount
1139+ aggs = termsAggs
1140+
1141+instance BucketAggregation DateHistogramResult where
1142+ key = showText . dateKey
1143+ docCount = dateDocCount
1144+ aggs = dateHistogramAggs
1145+
1146+instance (FromJSON a, BucketAggregation a) => FromJSON (Bucket a) where
1147+ parseJSON (Object v) = Bucket <$>
1148+ v .: "buckets"
1149+ parseJSON _ = mempty
1150+
1151+instance FromJSON TermsResult where
1152+ parseJSON (Object v) = TermsResult <$>
1153+ v .: "key" <*>
1154+ v .: "doc_count" <*>
1155+ v .:? "aggregations"
1156+
1157+instance FromJSON DateHistogramResult where
1158+ parseJSON (Object v) = DateHistogramResult <$>
1159+ v .: "key" <*>
1160+ v .:? "key_as_string" <*>
1161+ v .: "doc_count" <*>
1162+ v .:? "aggregations"
9451163
9461164instance Monoid Filter where
9471165 mempty = IdentityFilter
@@ -1536,10 +1754,11 @@ instance (FromJSON a) => FromJSON (EsResult a) where
15361754
15371755
15381756instance ToJSON Search where
1539− toJSON (Search query sFilter sort sTrackSortScores sFrom sSize) =
1757+ toJSON (Search query sFilter sort aggs sTrackSortScores sFrom sSize) =
15401758 omitNulls [ "query" .= query
15411759 , "filter" .= sFilter
15421760 , "sort" .= sort
1761+ , "aggregations" .= aggs
15431762 , "from" .= sFrom
15441763 , "size" .= sSize
15451764 , "track_scores" .= sTrackSortScores ]
@@ -1677,10 +1896,11 @@ instance ToJSON BoolMatch where
16771896
16781897instance (FromJSON a) => FromJSON (SearchResult a) where
16791898 parseJSON (Object v) = SearchResult <$>
1680− v .: "took" <*>
1681− v .: "timed_out" <*>
1682− v .: "_shards" <*>
1683− v .: "hits"
1899+ v .: "took" <*>
1900+ v .: "timed_out" <*>
1901+ v .: "_shards" <*>
1902+ v .: "hits" <*>
1903+ v .:? "aggregations"
16841904 parseJSON _ = empty
16851905
16861906instance (FromJSON a) => FromJSON (SearchHits a) where
bloodhound.cabal+2
@@ -25,6 +25,7 @@ library
2525 Database.Bloodhound.Types.Class
2626 build-depends: base >= 4.3 && <5,
2727 bytestring >= 0.10.2,
28+ containers >= 0.5.0.0,
2829 aeson >= 0.7,
2930 conduit >= 1.0,
3031 http-client >= 0.3,
@@ -45,6 +46,7 @@ test-suite tests
4546 bloodhound,
4647 http-client,
4748 http-types,
49+ containers >= 0.5.0.0,
4850 hspec >= 1.8,
4951 text >= 0.11,
5052 time >= 1.4,
tests/tests.hs+120−2
@@ -4,10 +4,13 @@
44module Main where
55
66import Control.Applicative
7+import Control.Monad
78import Data.Aeson
9+import Data.Aeson.Types (parseMaybe)
810import Data.HashMap.Strict (fromList)
911import Data.List (nub)
1012import Data.List.NonEmpty (NonEmpty (..))
13+import qualified Data.Map.Strict as M
1114import Data.Text (Text)
1215import qualified Data.Text as T
1316import Data.Time.Calendar (Day (..))
@@ -39,6 +42,55 @@ createExampleIndex = createIndex testServer defaultIndexSettings testIndex
3942deleteExampleIndex :: IO Reply
4043deleteExampleIndex = deleteIndex testServer testIndex
4144
45+data ServerVersion = ServerVersion Int Int Int deriving (Show, Eq, Ord)
46+
47+es14 :: ServerVersion
48+es14 = ServerVersion 1 4 0
49+
50+es13 :: ServerVersion
51+es13 = ServerVersion 1 3 0
52+
53+es12 :: ServerVersion
54+es12 = ServerVersion 1 2 0
55+
56+es11 :: ServerVersion
57+es11 = ServerVersion 1 1 0
58+
59+es10 :: ServerVersion
60+es10 = ServerVersion 1 0 0
61+
62+serverBranch :: ServerVersion -> ServerVersion
63+serverBranch (ServerVersion maj min patch) = ServerVersion maj min 0
64+
65+mkServerVersion :: [Int] -> Maybe ServerVersion
66+mkServerVersion [maj, min, patch] = Just (ServerVersion maj min patch)
67+mkServerVersion _ = Nothing
68+
69+getServerVersion :: Server -> IO (Maybe ServerVersion)
70+getServerVersion s = liftM extractVersion (getStatus s)
71+ where
72+ version' = T.splitOn "." . number . version
73+ toInt = read . T.unpack
74+ parseVersion v = map toInt (version' v)
75+ extractVersion = join . liftM (mkServerVersion . parseVersion)
76+
77+
78+
79+testServerBranch :: IO (Maybe ServerVersion)
80+testServerBranch = getServerVersion testServer >>= \v -> return $ liftM serverBranch v
81+
82+atleast :: ServerVersion -> IO Bool
83+atleast v = testServerBranch >>= \x -> return $ x >= Just (serverBranch v)
84+
85+atmost :: ServerVersion -> IO Bool
86+atmost v = testServerBranch >>= \x -> return $ x <= Just (serverBranch v)
87+
88+is :: ServerVersion -> IO Bool
89+is v = testServerBranch >>= \x -> return $ x == Just (serverBranch v)
90+
91+when' :: Monad m => m Bool -> m () -> m ()
92+when' b f = b >>= \x -> when x f
93+
4294data Location = Location { lat :: Double
4395 , lon :: Double } deriving (Eq, Generic, Show)
4496
@@ -109,6 +161,28 @@ searchExpectNoResults search = do
109161 let emptyHits = fmap (hits . searchHits) result
110162 emptyHits `shouldBe` Right []
111163
164+searchExpectAggs :: Search -> IO ()
165+searchExpectAggs search = do
166+ reply <- searchAll testServer search
167+ let isEmpty x = return (M.null x)
168+ let result = decode (responseBody reply) :: Maybe (SearchResult Tweet)
169+ (result >>= aggregations >>= isEmpty) `shouldBe` Just False
170+
171+searchValidBucketAgg :: (BucketAggregation a, FromJSON a, Show a) => Search -> Text -> (Text -> AggregationResults -> Maybe (Bucket a)) -> IO ()
172+searchValidBucketAgg search aggKey extractor = do
173+ reply <- searchAll testServer search
174+ let bucketDocs = docCount . head . buckets
175+ let result = decode (responseBody reply) :: Maybe (SearchResult Tweet)
176+ let count = result >>= aggregations >>= extractor aggKey >>= \x -> return (bucketDocs x)
177+ count `shouldBe` Just 1
178+
179+searchTermsAggHint :: [ExecutionHint] -> IO ()
180+searchTermsAggHint hints = do
181+ let terms hint = TermsAgg $ (mkTermsAggregation "user") { termExecutionHint = Just hint }
182+ let search hint = mkAggregateSearch Nothing $ mkAggregations "users" $ terms hint
183+ forM_ hints $ searchExpectAggs . search
184+ forM_ hints (\x -> searchValidBucketAgg (search x) "users" toTerms)
185+
112186data BulkTest = BulkTest { name :: Text } deriving (Eq, Generic, Show)
113187instance FromJSON BulkTest
114188instance ToJSON BulkTest
@@ -237,7 +311,7 @@ main = hspec $ do
237311 _ <- insertOther
238312 let sortSpec = DefaultSortSpec $ mkSort (FieldName "age") Ascending
239313 let search = Search Nothing
240− (Just IdentityFilter) (Just [sortSpec])
314+ (Just IdentityFilter) (Just [sortSpec]) Nothing
241315 False 0 10
242316 reply <- searchByIndex testServer testIndex search
243317 let result = eitherDecode (responseBody reply) :: Either String (SearchResult Tweet)
@@ -364,6 +438,51 @@ main = hspec $ do
364438 let search = mkSearch Nothing (Just filter)
365439 searchExpectNoResults search
366440
441+ describe "Aggregation API" $ do
442+ it "returns term aggregation results" $ do
443+ _ <- insertData
444+ let terms = TermsAgg $ mkTermsAggregation "user"
445+ let search = mkAggregateSearch Nothing $ mkAggregations "users" terms
446+ searchExpectAggs search
447+ searchValidBucketAgg search "users" toTerms
448+
449+ it "can give collection hint parameters to term aggregations" $ when' (atleast es13) $ do
450+ _ <- insertData
451+ let terms = TermsAgg $ (mkTermsAggregation "user") { termCollectMode = Just BreadthFirst }
452+ let search = mkAggregateSearch Nothing $ mkAggregations "users" terms
453+ searchExpectAggs search
454+ searchValidBucketAgg search "users" toTerms
455+
456+ it "can give execution hint paramters to term aggregations" $ when' (atmost es11) $ do
457+ _ <- insertData
458+ searchTermsAggHint [Map, Ordinals]
459+
460+ it "can give execution hint paramters to term aggregations" $ when' (is es12) $ do
461+ _ <- insertData
462+ searchTermsAggHint [GlobalOrdinals, GlobalOrdinalsHash, GlobalOrdinalsLowCardinality, Map, Ordinals]
463+
464+ it "can give execution hint paramters to term aggregations" $ when' (atleast es12) $ do
465+ _ <- insertData
466+ searchTermsAggHint [GlobalOrdinals, GlobalOrdinalsHash, GlobalOrdinalsLowCardinality, Map]
467+
468+ it "returns date histogram aggregation results" $ do
469+ _ <- insertData
470+ let histogram = DateHistogramAgg $ mkDateHistogram (FieldName "postDate") Minute
471+ let search = mkAggregateSearch Nothing (mkAggregations "byDate" histogram)
472+ searchExpectAggs search
473+ searchValidBucketAgg search "byDate" toDateHistogram
474+
475+ it "returns date histogram using fractional date" $ do
476+ _ <- insertData
477+ let periods = [Year, Quarter, Month, Week, Day, Hour, Minute, Second]
478+ let fractionals = map (FractionalInterval 1.5) [Weeks, Days, Hours, Minutes, Seconds]
479+ let intervals = periods ++ fractionals
480+ let histogram = mkDateHistogram (FieldName "postDate")
481+ let search interval = mkAggregateSearch Nothing $ mkAggregations "byDate" $ DateHistogramAgg (histogram interval)
482+ let expect interval = searchExpectAggs (search interval)
483+ let valid interval = searchValidBucketAgg (search interval) "byDate" toDateHistogram
484+ forM_ intervals expect
485+ forM_ intervals valid
367486
368487 describe "ToJSON RegexpFlags" $ do
369488 it "generates the correct JSON for AllRegexpFlags" $
@@ -385,7 +504,6 @@ main = hspec $ do
385504 flagStrs = T.splitOn "|" str
386505 in noDuplicates flagStrs
387506
388−
389507 describe "omitNulls" $ do
390508 it "checks that omitNulls drops list elements when it should" $
391509 let dropped = omitNulls $ [ "test1" .= (toJSON ([] :: [Int]))