packages feed

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'