bloodhound-1.0.0.0: tests/Test/KnnOs3Spec.hs
{-# LANGUAGE OverloadedStrings #-}
module Test.KnnOs3Spec (spec) where
import Data.Aeson
import Data.Aeson.Key (Key)
import Data.Aeson.Key qualified as AK
import Data.Aeson.KeyMap qualified as KM
import Data.ByteString.Lazy.Char8 qualified as LBS
import Data.Maybe (fromJust, isJust)
import Database.Bloodhound.Client.Cluster (BackendType (..))
import Database.Bloodhound.OpenSearch3.Client
import Database.Bloodhound.OpenSearch3.Requests qualified as OS3Requests
import TestsUtils.Common
import TestsUtils.Import
vectorField :: FieldName
vectorField = FieldName "vector"
queryVector :: [Double]
queryVector = [0.1, 0.2, 0.3, 0.4, 0.5]
hasKey :: Key -> LBS.ByteString -> Bool
hasKey k bs = case decode bs of
Just (Object obj) -> k `KM.member` obj
_ -> False
-- | Walk into a nested dict path, returning the inner Value if reachable.
walk :: [Key] -> Value -> Maybe Value
walk [] v = Just v
walk (k : ks) (Object obj) = case KM.lookup k obj of
Just v -> walk ks v
Nothing -> Nothing
walk _ _ = Nothing
-- | True iff the encoded OpenSearchKnnQuery produces the OS-correct shape
-- @{"knn": {<field>: {...}}}@ with the field name as a dict key.
hasKnnNestedUnderField :: FieldName -> LBS.ByteString -> Bool
hasKnnNestedUnderField (FieldName f) bs =
case decode bs of
Just (Object obj) -> case KM.lookup "knn" obj of
Just (Object knnObj) -> KM.member (AK.fromText f) knnObj
_ -> False
_ -> False
spec :: Spec
spec = describe "k-NN Search API - OpenSearch 3" $ do
describe "OpenSearchKnnQuery body shape" $ do
it "serializes as query.knn.<field> with the field name as the dict key" $ do
let clause = mkOpenSearchKnnQuery vectorField queryVector 3
let json = encode clause
json `shouldSatisfy` hasKey "knn"
json `shouldSatisfy` hasKnnNestedUnderField vectorField
-- No ES-style wrapper.
json `shouldNotSatisfy` hasKey "field"
json `shouldNotSatisfy` hasKey "query_vector"
it "omits Nothing optional fields from the clause body" $ do
let clause = mkOpenSearchKnnQuery vectorField queryVector 3
let Just (Object knnDict) = walk ["knn", AK.fromText "vector"] (fromJust (decode (encode clause)))
knnDict `shouldSatisfy` not . KM.member "filter"
knnDict `shouldSatisfy` not . KM.member "boost"
knnDict `shouldSatisfy` not . KM.member "method_parameters"
it "includes filter, boost, method_parameters when set" $ do
let clause =
(mkOpenSearchKnnQuery vectorField queryVector 5)
{ osknnFilter = Just (Filter (TermQuery (Term "label" "alpha") Nothing)),
osknnBoost = Just 1.2,
osknnMethodParameters = Just (KnnHnswParams 100)
}
let Just (Object knnDict) = walk ["knn", AK.fromText "vector"] (fromJust (decode (encode clause)))
knnDict `shouldSatisfy` KM.member "filter"
knnDict `shouldSatisfy` KM.member "boost"
knnDict `shouldSatisfy` KM.member "method_parameters"
it "serializes HNSW method_parameters as {\"ef_search\":N} only" $ do
let clause =
(mkOpenSearchKnnQuery vectorField queryVector 5)
{ osknnMethodParameters = Just (KnnHnswParams 100)
}
let Just (Object methodParams) = walk ["knn", AK.fromText "vector", "method_parameters"] (fromJust (decode (encode clause)))
KM.lookup "ef_search" methodParams `shouldBe` Just (Number 100)
methodParams `shouldSatisfy` not . KM.member "nprobes"
it "serializes IVF method_parameters as {\"nprobes\":N} only" $ do
let clause =
(mkOpenSearchKnnQuery vectorField queryVector 5)
{ osknnMethodParameters = Just (KnnIvfParams 10)
}
let Just (Object methodParams) = walk ["knn", AK.fromText "vector", "method_parameters"] (fromJust (decode (encode clause)))
KM.lookup "nprobes" methodParams `shouldBe` Just (Number 10)
methodParams `shouldSatisfy` not . KM.member "ef_search"
it "round-trips KnnHnswParams and KnnIvfParams through ToJSON/FromJSON" $ do
decode (encode (KnnHnswParams 100)) `shouldBe` Just (KnnHnswParams 100 :: OpenSearchKnnMethodParameters)
decode (encode (KnnIvfParams 10)) `shouldBe` Just (KnnIvfParams 10 :: OpenSearchKnnMethodParameters)
it "rejects a method_parameters object with neither expected key" $ do
let v = encode (object ["other" .= Number 1])
decode v `shouldBe` (Nothing :: Maybe OpenSearchKnnMethodParameters)
it "round-trips a minimal OpenSearchKnnQuery through ToJSON/FromJSON" $ do
let clause = mkOpenSearchKnnQuery vectorField queryVector 3
decode (encode clause) `shouldBe` Just clause
it "round-trips a fully-populated OpenSearchKnnQuery through ToJSON/FromJSON" $ do
let clause =
(mkOpenSearchKnnQuery vectorField queryVector 5)
{ osknnFilter = Just (Filter (TermQuery (Term "label" "alpha") Nothing)),
osknnBoost = Just 1.2,
osknnMethodParameters = Just (KnnHnswParams 100)
}
decode (encode clause) `shouldBe` Just clause
it "serializes OpenSearchKnnEngine as lowercase faiss/nmslib/lucene" $ do
encode KnnEngineFaiss `shouldBe` "\"faiss\""
encode KnnEngineNmslib `shouldBe` "\"nmslib\""
encode KnnEngineLucene `shouldBe` "\"lucene\""
decode "\"faiss\"" `shouldBe` Just (KnnEngineFaiss :: OpenSearchKnnEngine)
describe "OpenSearchKnnQuery radial-search variants" $ do
it "max_distance variant emits max_distance and no k in the clause body" $ do
let clause = mkOpenSearchKnnQueryByMaxDistance vectorField queryVector 1.5
let json = encode clause
json `shouldSatisfy` hasKey "knn"
json `shouldSatisfy` hasKnnNestedUnderField vectorField
let Just (Object knnDict) = walk ["knn", AK.fromText "vector"] (fromJust (decode json))
KM.lookup "max_distance" knnDict `shouldBe` Just (Number 1.5)
knnDict `shouldSatisfy` not . KM.member "k"
knnDict `shouldSatisfy` not . KM.member "min_score"
it "min_score variant emits min_score and no k in the clause body" $ do
let clause = mkOpenSearchKnnQueryByMinScore vectorField queryVector 0.8
let json = encode clause
json `shouldSatisfy` hasKey "knn"
json `shouldSatisfy` hasKnnNestedUnderField vectorField
let Just (Object knnDict) = walk ["knn", AK.fromText "vector"] (fromJust (decode json))
KM.lookup "min_score" knnDict `shouldBe` Just (Number 0.8)
knnDict `shouldSatisfy` not . KM.member "k"
knnDict `shouldSatisfy` not . KM.member "max_distance"
it "k variant still emits k (regression)" $ do
let clause = mkOpenSearchKnnQuery vectorField queryVector 3
let Just (Object knnDict) = walk ["knn", AK.fromText "vector"] (fromJust (decode (encode clause)))
KM.lookup "k" knnDict `shouldBe` Just (Number 3)
knnDict `shouldSatisfy` not . KM.member "max_distance"
knnDict `shouldSatisfy` not . KM.member "min_score"
it "preserves optional fields on the max_distance variant" $ do
let clause =
(mkOpenSearchKnnQueryByMaxDistance vectorField queryVector 1.5)
{ osknnFilter = Just (Filter (TermQuery (Term "label" "alpha") Nothing)),
osknnBoost = Just 1.2,
osknnMethodParameters = Just (KnnHnswParams 100)
}
let Just (Object knnDict) = walk ["knn", AK.fromText "vector"] (fromJust (decode (encode clause)))
KM.lookup "max_distance" knnDict `shouldBe` Just (Number 1.5)
knnDict `shouldSatisfy` KM.member "filter"
knnDict `shouldSatisfy` KM.member "boost"
knnDict `shouldSatisfy` KM.member "method_parameters"
it "round-trips a minimal max_distance query" $ do
let clause = mkOpenSearchKnnQueryByMaxDistance vectorField queryVector 1.5
decode (encode clause) `shouldBe` Just clause
it "round-trips a fully-populated min_score query" $ do
let clause =
(mkOpenSearchKnnQueryByMinScore vectorField queryVector 0.8)
{ osknnFilter = Just (Filter (TermQuery (Term "label" "alpha") Nothing)),
osknnBoost = Just 1.2,
osknnMethodParameters = Just (KnnIvfParams 10)
}
decode (encode clause) `shouldBe` Just clause
it "FromJSON picks KnnTerminationByK / ByMaxDistance / ByMinScore by present key" $ do
decode (encode (object ["k" .= Number 3]))
`shouldBe` Just (KnnTerminationByK 3 :: OpenSearchKnnQueryTermination)
decode (encode (object ["max_distance" .= Number 1.5]))
`shouldBe` Just (KnnTerminationByMaxDistance 1.5 :: OpenSearchKnnQueryTermination)
decode (encode (object ["min_score" .= Number 0.8]))
`shouldBe` Just (KnnTerminationByMinScore 0.8 :: OpenSearchKnnQueryTermination)
it "ToJSON OpenSearchKnnQueryTermination emits the single expected key" $ do
encode (KnnTerminationByK 3) `shouldBe` "{\"k\":3}"
encode (KnnTerminationByMaxDistance 1.5) `shouldBe` "{\"max_distance\":1.5}"
encode (KnnTerminationByMinScore 0.8) `shouldBe` "{\"min_score\":0.8}"
it "rejects a termination object with none of the expected keys" $ do
let v = encode (object ["other" .= Number 1])
decode v `shouldBe` (Nothing :: Maybe OpenSearchKnnQueryTermination)
describe "Search.osKnnBody integration" $ do
it "when osKnnBody is set, emits it under the top-level query slot" $ do
let clause = mkOpenSearchKnnQuery vectorField queryVector 3
let search = (mkSearch Nothing Nothing) {osKnnBody = Just clause}
let json = encode search
json `shouldSatisfy` hasKey "query"
let Just (Object queryObj) = walk ["query"] (fromJust (decode json))
queryObj `shouldSatisfy` KM.member "knn"
it "when osKnnBody is Nothing, falls back to queryBody/filterBody logic" $ do
let search = mkSearch Nothing Nothing
let json = encode search
json `shouldNotSatisfy` hasKey "query"
describe "endpoint shape" $ do
it "POSTs to /{index}/_search (not /_knn_search)" $ do
let bareReq = OS3Requests.knnSearch @KnnDoc osKnnTestIndex (mkOpenSearchKnnQuery vectorField queryVector 3) (mkSearch Nothing Nothing)
getRawEndpoint (bhRequestEndpoint bareReq)
`shouldBe` [unIndexName osKnnTestIndex, "_search"]
it "carries no ?index= query string (the old bug)" $ do
let bareReq = OS3Requests.knnSearch @KnnDoc osKnnTestIndex (mkOpenSearchKnnQuery vectorField queryVector 3) (mkSearch Nothing Nothing)
getRawEndpointQueries (bhRequestEndpoint bareReq) `shouldBe` []
it "embeds the knn clause in the request body as query.knn.<field>" $ do
let bareReq = OS3Requests.knnSearch @KnnDoc osKnnTestIndex (mkOpenSearchKnnQuery vectorField queryVector 3) (mkSearch Nothing Nothing)
let body = bhRequestBody bareReq
body `shouldSatisfy` isJust
let Just bodyBytes = body
case decode bodyBytes of
Just (Object obj) -> case KM.lookup "query" obj of
Just (Object q) -> case KM.lookup "knn" q of
Just (Object knn) -> KM.member (AK.fromText "vector") knn `shouldBe` True
_ -> expectationFailure "missing knn key inside query"
_ -> expectationFailure "missing query key"
_ -> expectationFailure "body did not decode to Object"
describe "live integration (requires OS3 backend with index.knn=true)" $
backendSpecific [OpenSearch3] $ do
it "returns the nearest neighbour document" $
withTestEnv $ do
insertOsKnnData
let search = mkSearch Nothing Nothing
clause = mkOpenSearchKnnQuery vectorField queryVector 3
searchResult <- knnSearch osKnnTestIndex clause search :: BH IO (SearchResult KnnDoc)
let topHit = headMay $ hits $ searchHits searchResult
liftIO $ (hitDocId <$> topHit) `shouldBe` Just (Just (DocId "a"))
it "supports optional filter and boost fields" $
withTestEnv $ do
insertOsKnnData
let search = mkSearch Nothing Nothing
clause =
(mkOpenSearchKnnQuery vectorField queryVector 5)
{ osknnBoost = Just 1.0,
osknnFilter = Just (Filter (TermQuery (Term "label" "alpha") Nothing))
}
searchResult <- knnSearch osKnnTestIndex clause search :: BH IO (SearchResult KnnDoc)
let topHit = headMay $ hits $ searchHits searchResult
liftIO $ (hitDocId <$> topHit) `shouldBe` Just (Just (DocId "a"))
it "radial search by max_distance returns the nearest neighbour" $
withTestEnv $ do
insertOsKnnData
let search = mkSearch Nothing Nothing
clause = mkOpenSearchKnnQueryByMaxDistance vectorField queryVector 2.0
searchResult <- knnSearch osKnnTestIndex clause search :: BH IO (SearchResult KnnDoc)
let topHit = headMay $ hits $ searchHits searchResult
liftIO $ (hitDocId <$> topHit) `shouldBe` Just (Just (DocId "a"))
it "radial search by min_score returns the nearest neighbour" $
withTestEnv $ do
insertOsKnnData
let search = mkSearch Nothing Nothing
clause = mkOpenSearchKnnQueryByMinScore vectorField queryVector 0.5
searchResult <- knnSearch osKnnTestIndex clause search :: BH IO (SearchResult KnnDoc)
let topHit = headMay $ hits $ searchHits searchResult
liftIO $ (hitDocId <$> topHit) `shouldBe` Just (Just (DocId "a"))