bloodhound-1.0.0.0: tests/Test/AggregationSpec.hs
{-# LANGUAGE OverloadedStrings #-}
module Test.AggregationSpec (spec) where
import Control.Error (fmapL, note)
import Data.Map qualified as M
import Data.Text qualified as T
import Database.Bloodhound qualified
import Database.Bloodhound.Types (BucketsPath)
import TestsUtils.Common
import TestsUtils.Import
-- | Extract a single aggregation value from a tweet-search result that may
-- itself have failed. Returns 'Nothing' if the request failed, no
-- aggregations were returned, or the named aggregation is absent. Used by
-- the pipeline-aggregation specs (bloodhound-5vp) to assert just the
-- pipeline value without depending on the (brittle) parent bucket shape.
lookupAgg :: Key -> Either EsError (SearchResult a) -> Maybe Value
lookupAgg name = either (const Nothing) (aggregations >=> M.lookup name)
spec :: Spec
spec =
describe "Aggregation API" $ do
it "returns term aggregation results" $
withTestEnv $ do
_ <- insertData
let terms = TermsAgg $ mkTermsAggregation $ FieldName "user"
let search = mkAggregateSearch Nothing $ mkAggregations "users" terms
searchExpectAggs search
searchValidBucketAgg search "users" toTerms
it "return sub-aggregation results" $
withTestEnv $ do
_ <- insertData
let subaggs = mkAggregations "age_agg" . TermsAgg $ mkTermsAggregation $ FieldName "age"
agg = TermsAgg $ (mkTermsAggregation $ FieldName "user") {termAggs = Just subaggs}
search = mkAggregateSearch Nothing $ mkAggregations "users" agg
result <- performBHRequest $ searchByIndex @Tweet testIndex search
let usersAggResults = aggregations result >>= toTerms "users"
subAggResults = usersAggResults >>= (listToMaybe . buckets) >>= termsAggs >>= toTerms "age_agg"
subAddResultsExists = isJust subAggResults
liftIO $ subAddResultsExists `shouldBe` True
it "returns cardinality aggregation results" $
withTestEnv $ do
_ <- insertData
let cardinality = CardinalityAgg $ mkCardinalityAggregation $ FieldName "user"
let search = mkAggregateSearch Nothing $ mkAggregations "users" cardinality
let search' = search {Database.Bloodhound.from = From 0, size = Size 0}
searchExpectAggs search'
let docCountPair k n = (k, object ["value" .= Number n])
res <- searchTweets search'
liftIO $
fmap aggregations res `shouldBe` Right (Just (M.fromList [docCountPair "users" 1]))
it "returns stats aggregation results" $
withTestEnv $ do
_ <- insertData
let stats = StatsAgg $ mkStatsAggregation $ FieldName "age"
let search = mkAggregateSearch Nothing $ mkAggregations "users" stats
let search' = search {Database.Bloodhound.from = From 0, size = Size 0}
searchExpectAggs search'
let statsAggRes k n =
( k,
object
[ "max" .= Number n,
"avg" .= Number n,
"count" .= Number 1,
"min" .= Number n,
"sum" .= Number n
]
)
res <- searchTweets search'
liftIO $
fmap aggregations res `shouldBe` Right (Just (M.fromList [statsAggRes "users" 10000]))
it "can give collection hint parameters to term aggregations" $
withTestEnv $ do
_ <- insertData
let terms = TermsAgg $ (mkTermsAggregation $ FieldName "user") {termCollectMode = Just BreadthFirst}
let search = mkAggregateSearch Nothing $ mkAggregations "users" terms
searchExpectAggs search
searchValidBucketAgg search "users" toTerms
it "can give execution hint parameters to term aggregations" $
withTestEnv $ do
_ <- insertData
searchTermsAggHint [GlobalOrdinals, Map]
-- One of the above.
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 docCountPair k n = (k, object ["value" .= Number n])
res <- searchTweets search
liftIO $
fmap aggregations res
`shouldBe` Right
( Just
( M.fromList
[ docCountPair "user_count" 2,
docCountPair "bogus_count" 0
]
)
)
it "can execute date_range aggregations" $
withTestEnv $ do
let now = fromGregorian 2015 3 14
let ltAMonthAgo = UTCTime (fromGregorian 2015 3 1) 0
let ltAWeekAgo = UTCTime (fromGregorian 2015 3 10) 0
let oldDoc = exampleTweet {postDate = ltAMonthAgo}
let newDoc = exampleTweet {postDate = ltAWeekAgo}
_ <- performBHRequest $ indexDocument testIndex defaultIndexDocumentSettings oldDoc (DocId "1")
_ <- performBHRequest $ indexDocument testIndex defaultIndexDocumentSettings newDoc (DocId "2")
_ <- performBHRequest $ refreshIndex testIndex
let thisMonth = DateRangeFrom (DateMathExpr (DMDate now) [SubtractTime 1 DMMonth])
let thisWeek = DateRangeFrom (DateMathExpr (DMDate now) [SubtractTime 1 DMWeek])
let agg = DateRangeAggregation (FieldName "postDate") Nothing (thisMonth :| [thisWeek])
let ags = mkAggregations "date_ranges" (DateRangeAgg agg)
let search = mkAggregateSearch Nothing ags
res <- searchTweets search
liftIO $ hitsTotal . searchHits <$> res `shouldBe` Right (Just (HitsTotal 2 HTR_EQ))
let bucks = do
magrs <- fmapL show (aggregations <$> res)
agrs <- note "no aggregations returned" magrs
rawBucks <- note "no date_ranges aggregation" $ M.lookup "date_ranges" agrs
parseEither parseJSON rawBucks
let fromMonthT = UTCTime (fromGregorian 2015 2 14) 0
let fromWeekT = UTCTime (fromGregorian 2015 3 7) 0
liftIO $
buckets
<$> bucks
`shouldBe` Right
[ DateRangeResult
"2015-02-14T00:00:00.000Z-*"
(Just fromMonthT)
(Just "2015-02-14T00:00:00.000Z")
Nothing
Nothing
2
Nothing,
DateRangeResult
"2015-03-07T00:00:00.000Z-*"
(Just fromWeekT)
(Just "2015-03-07T00:00:00.000Z")
Nothing
Nothing
1
Nothing
]
it "returns date histogram aggregation results" $
withTestEnv $ do
_ <- insertData
let histogram = DateHistogramAgg (mkDateHistogram (FieldName "postDate")) {dateFixedInterval = Just (FixedInterval 1 Minutes)}
let search = mkAggregateSearch Nothing (mkAggregations "byDate" histogram)
searchExpectAggs search
searchValidBucketAgg search "byDate" toDateHistogram
it "can execute missing aggregations" $
withTestEnv $ do
_ <- insertData
_ <- insertExtra
let ags = mkAggregations "missing_agg" (MissingAgg (MissingAggregation "extra"))
let search = mkAggregateSearch Nothing ags
let docCountPair k n = (k, object ["doc_count" .= Number n])
res <- searchTweets search
liftIO $
fmap aggregations res `shouldBe` Right (Just (M.fromList [docCountPair "missing_agg" 1]))
-- With a single document indexed (exampleTweet, age = 10000), min == max
-- == 10000. Mirrors the sibling "max aggregation" spec below.
it "can execute min aggregation" $
withTestEnv $ do
_ <- insertData
let minAgg = MinAgg (mkMinAggregation (FieldName "age"))
let search = mkAggregateSearch Nothing (mkAggregations "min_age" minAgg)
let search' = search {Database.Bloodhound.from = From 0, size = Size 0}
searchExpectAggs search'
res <- searchTweets search'
liftIO $
fmap aggregations res `shouldBe` Right (Just (M.fromList [("min_age", object ["value" .= Number 10000])]))
it "can execute max aggregation" $
withTestEnv $ do
_ <- insertData
let maxAgg = MaxAgg (mkMaxAggregation (FieldName "age"))
let search = mkAggregateSearch Nothing (mkAggregations "max_age" maxAgg)
let search' = search {Database.Bloodhound.from = From 0, size = Size 0}
searchExpectAggs search'
res <- searchTweets search'
liftIO $
fmap aggregations res `shouldBe` Right (Just (M.fromList [("max_age", object ["value" .= Number 10000])]))
-- Bucket aggregations return a @buckets@ array, not a top-level
-- @doc_count@. Single doc (age = 10000), interval 1000 -> one bucket
-- keyed at 10000.0.
it "can execute histogram aggregation" $
withTestEnv $ do
_ <- insertData
let histogram = HistogramAgg (mkHistogramAggregation (FieldName "age") 1000)
let search = mkAggregateSearch Nothing (mkAggregations "age_histogram" histogram)
let search' = search {Database.Bloodhound.from = From 0, size = Size 0}
searchExpectAggs search'
res <- searchTweets search'
liftIO $
lookupAgg "age_histogram" res `shouldBe` Just (object ["buckets" .= [object ["key" .= Number 10000, "doc_count" .= Number 1]]])
-- Bucket-shape drift as above: @buckets@ array with one entry per range.
-- Ranges: 0.0-*, 0.0-1000.0, 10000.0-*. Single doc age = 10000 falls in
-- both unbounded-from-above ranges. ES orders buckets by @from@ asc,
-- bounded-before-unbounded on ties, hence 0.0-1000.0 precedes 0.0-*.
it "can execute range aggregation" $
withTestEnv $ do
_ <- insertData
let rangeAgg = RangeAgg (mkRangeAggregation (FieldName "age") (RangeFrom 0 :| [RangeFromTo 0 1000, RangeFrom 10000]))
let search = mkAggregateSearch Nothing (mkAggregations "age_ranges" rangeAgg)
let search' = search {Database.Bloodhound.from = From 0, size = Size 0}
searchExpectAggs search'
res <- searchTweets search'
liftIO $
lookupAgg "age_ranges" res `shouldBe` Just (object ["buckets" .= [object ["key" .= String "0.0-1000.0", "from" .= Number 0, "to" .= Number 1000, "doc_count" .= Number 0], object ["key" .= String "0.0-*", "from" .= Number 0, "doc_count" .= Number 1], object ["key" .= String "10000.0-*", "from" .= Number 10000, "doc_count" .= Number 1]]])
-- The Tweet fixture has no @comments@ field, so the nested aggregation
-- correctly reports doc_count = 0.
it "can execute nested aggregation" $
withTestEnv $ do
_ <- insertData
let nestedAgg = NestedAgg (mkNestedAggregation (FieldName "comments"))
let search = mkAggregateSearch Nothing (mkAggregations "comments_nested" nestedAgg)
let search' = search {Database.Bloodhound.from = From 0, size = Size 0}
searchExpectAggs search'
res <- searchTweets search'
liftIO $
lookupAgg "comments_nested" res `shouldBe` Just (object ["doc_count" .= Number 0])
-- With a single document, every requested percentile equals age = 10000.
-- ES defaults to the 7 percentiles below.
it "can execute percentile aggregation" $
withTestEnv $ do
_ <- insertData
let percentileAgg = PercentileAgg (mkPercentileAggregation (FieldName "age"))
let search = mkAggregateSearch Nothing (mkAggregations "age_percentiles" percentileAgg)
let search' = search {Database.Bloodhound.from = From 0, size = Size 0}
searchExpectAggs search'
res <- searchTweets search'
liftIO $
lookupAgg "age_percentiles" res `shouldBe` Just (object ["values" .= object ["1.0" .= Number 10000, "5.0" .= Number 10000, "25.0" .= Number 10000, "50.0" .= Number 10000, "75.0" .= Number 10000, "95.0" .= Number 10000, "99.0" .= Number 10000]])
-- The pipeline-aggregation tests below were restructured in bloodhound-5vp.
-- Previous shape placed the metric agg (SumAgg/AvgAgg/...) directly under
-- the @ages@ name as a sibling of the pipeline agg, then pointed
-- @buckets_path@ at @ages>sum_age@. Modern Elasticsearch validates the
-- path and rejects it because @ages@ has no bucket children. The fix is to
-- make @ages@ a real bucket aggregation (TermsAgg on @user@) with the
-- metric as a named sub-aggregation, so the path "ages>metric"
-- resolves. With a single document indexed (exampleTweet, age = 10000),
-- the bucket contains one value, 10000, so every sibling pipeline
-- aggregator reduces to that same number.
it "can execute avg bucket aggregation" $
withTestEnv $ do
_ <- insertData
let sumAgg = SumAgg (mkSumAggregation (FieldName "age"))
let agesAgg = TermsAgg $ (mkTermsAggregation (FieldName "user")) {termAggs = Just (mkAggregations "sum_age" sumAgg)}
let avgBucketAgg = AvgBucketAgg (mkAvgBucketAggregation (BucketsPath "ages>sum_age"))
let search = mkAggregateSearch Nothing (mkAggregations "ages" agesAgg <> mkAggregations "avg_ages" avgBucketAgg)
let search' = search {Database.Bloodhound.from = From 0, size = Size 0}
searchExpectAggs search'
res <- searchTweets search'
liftIO $
lookupAgg "avg_ages" res `shouldBe` Just (object ["value" .= Number 10000])
it "can execute sum bucket aggregation" $
withTestEnv $ do
_ <- insertData
let avgAgg = AvgAgg (mkAvgAggregation (FieldName "age"))
let agesAgg = TermsAgg $ (mkTermsAggregation (FieldName "user")) {termAggs = Just (mkAggregations "avg_age" avgAgg)}
let sumBucketAgg = SumBucketAgg (mkSumBucketAggregation (BucketsPath "ages>avg_age"))
let search = mkAggregateSearch Nothing (mkAggregations "ages" agesAgg <> mkAggregations "sum_ages" sumBucketAgg)
let search' = search {Database.Bloodhound.from = From 0, size = Size 0}
searchExpectAggs search'
res <- searchTweets search'
liftIO $
lookupAgg "sum_ages" res `shouldBe` Just (object ["value" .= Number 10000])
it "can execute min bucket aggregation" $
withTestEnv $ do
_ <- insertData
let maxAgg = MaxAgg (mkMaxAggregation (FieldName "age"))
let agesAgg = TermsAgg $ (mkTermsAggregation (FieldName "user")) {termAggs = Just (mkAggregations "max_age" maxAgg)}
let minBucketAgg = MinBucketAgg (mkMinBucketAggregation (BucketsPath "ages>max_age"))
let search = mkAggregateSearch Nothing (mkAggregations "ages" agesAgg <> mkAggregations "min_ages" minBucketAgg)
let search' = search {Database.Bloodhound.from = From 0, size = Size 0}
searchExpectAggs search'
res <- searchTweets search'
liftIO $
lookupAgg "min_ages" res `shouldBe` Just (object ["keys" .= ["bitemyapp" :: Text], "value" .= Number 10000])
it "can execute max bucket aggregation" $
withTestEnv $ do
_ <- insertData
let minAgg = MinAgg (mkMinAggregation (FieldName "age"))
let agesAgg = TermsAgg $ (mkTermsAggregation (FieldName "user")) {termAggs = Just (mkAggregations "min_age" minAgg)}
let maxBucketAgg = MaxBucketAgg (mkMaxBucketAggregation (BucketsPath "ages>min_age"))
let search = mkAggregateSearch Nothing (mkAggregations "ages" agesAgg <> mkAggregations "max_ages" maxBucketAgg)
let search' = search {Database.Bloodhound.from = From 0, size = Size 0}
searchExpectAggs search'
res <- searchTweets search'
liftIO $
lookupAgg "max_ages" res `shouldBe` Just (object ["keys" .= ["bitemyapp" :: Text], "value" .= Number 10000])
it "can execute stats bucket aggregation" $
withTestEnv $ do
_ <- insertData
let sumAgg = SumAgg (mkSumAggregation (FieldName "age"))
let agesAgg = TermsAgg $ (mkTermsAggregation (FieldName "user")) {termAggs = Just (mkAggregations "sum_age" sumAgg)}
let statsBucketAgg = StatsBucketAgg (mkStatsBucketAggregation (BucketsPath "ages>sum_age"))
let search = mkAggregateSearch Nothing (mkAggregations "ages" agesAgg <> mkAggregations "stats_ages" statsBucketAgg)
let search' = search {Database.Bloodhound.from = From 0, size = Size 0}
searchExpectAggs search'
res <- searchTweets search'
liftIO $
lookupAgg "stats_ages" res `shouldBe` Just (object ["count" .= Number 1, "min" .= Number 10000, "max" .= Number 10000, "avg" .= Number 10000, "sum" .= Number 10000])
-- derivative and cumulative_sum are *parent* pipeline aggregations:
-- Elasticsearch requires them to live inside a histogram-family bucket
-- aggregation (histogram / date_histogram / auto_date_histogram), with a
-- path relative to each bucket. Two documents are indexed so derivative
-- has more than one bucket to diff over; we assert only that the request
-- succeeds (searchExpectAggs) since per-bucket values are ordering
-- dependent.
it "can execute derivative aggregation" $
withTestEnv $ do
_ <- insertData
_ <- insertOther
let sumAgg = SumAgg (mkSumAggregation (FieldName "age"))
let derivativeAgg = DerivativeAgg (mkDerivativeAggregation (BucketsPath "sum_age"))
let subAggs = mkAggregations "sum_age" sumAgg <> mkAggregations "derivatives" derivativeAgg
let histAgg = HistogramAgg $ (mkHistogramAggregation (FieldName "age") 1000) {histogramMinDocCount = Just 0, histogramAggs = Just subAggs}
let search = mkAggregateSearch Nothing (mkAggregations "ages" histAgg)
let search' = search {Database.Bloodhound.from = From 0, size = Size 0}
searchExpectAggs search'
it "can execute cumulative sum aggregation" $
withTestEnv $ do
_ <- insertData
_ <- insertOther
let sumAgg = SumAgg (mkSumAggregation (FieldName "age"))
let cumulativeSumAgg = CumulativeSumAgg (mkCumulativeSumAggregation (BucketsPath "sum_age"))
let subAggs = mkAggregations "sum_age" sumAgg <> mkAggregations "cumulative" cumulativeSumAgg
let histAgg = HistogramAgg $ (mkHistogramAggregation (FieldName "age") 1000) {histogramMinDocCount = Just 0, histogramAggs = Just subAggs}
let search = mkAggregateSearch Nothing (mkAggregations "ages" histAgg)
let search' = search {Database.Bloodhound.from = From 0, size = Size 0}
searchExpectAggs search'