diff --git a/elastic4s-core/src/test/scala/com/sksamuel/elastic4s/requests/mappings/KnnVectorFieldTest.scala b/elastic4s-core/src/test/scala/com/sksamuel/elastic4s/requests/mappings/KnnVectorFieldTest.scala index 1c9d4810e..35c044171 100644 --- a/elastic4s-core/src/test/scala/com/sksamuel/elastic4s/requests/mappings/KnnVectorFieldTest.scala +++ b/elastic4s-core/src/test/scala/com/sksamuel/elastic4s/requests/mappings/KnnVectorFieldTest.scala @@ -128,7 +128,7 @@ class KnnVectorFieldTest extends AnyFlatSpec with Matchers with ElasticApi { |"dimension":512, |"method":{"name":"hnsw","engine":"faiss","space_type":"l2", |"parameters":{"ef_construction":100,"m":50,"ef_search":50, - |"encoder":{"name":"pq","encoder":{"m":50,"code_size":100}}}}}""".stripMargin.replace( + |"encoder":{"name":"pq","parameters":{"m":50,"code_size":100}}}}}""".stripMargin.replace( "\n", "" ) @@ -161,7 +161,7 @@ class KnnVectorFieldTest extends AnyFlatSpec with Matchers with ElasticApi { |"dimension":512, |"method":{"name":"hnsw","engine":"faiss","space_type":"innerproduct", |"parameters":{"ef_construction":100,"m":50,"ef_search":50, - |"encoder":{"name":"sq","encoder":{"clip":true,"type":"fp16"}}}}}""".stripMargin.replace( + |"encoder":{"name":"sq","parameters":{"clip":true,"type":"fp16"}}}}}""".stripMargin.replace( "\n", "" ) @@ -182,12 +182,11 @@ class KnnVectorFieldTest extends AnyFlatSpec with Matchers with ElasticApi { KnnVectorField( name = "myfield123", dimension = 512, - HnswParameters( + IvfParameters( engine = Some(KnnEngine.faiss), spaceType = Some(SpaceType.innerProduct), - efConstruction = Some(100), - m = Some(50), - efSearch = Some(50), + nlist = Some(4), + nprobes = Some(2), encoder = Some( FaissEncoder( Some(FaissEncoderName.sq), @@ -201,9 +200,9 @@ class KnnVectorFieldTest extends AnyFlatSpec with Matchers with ElasticApi { .string shouldBe """{"type":"knn_vector", |"dimension":512, - |"method":{"name":"hnsw","engine":"faiss","space_type":"innerproduct", - |"parameters":{"ef_construction":100,"m":50,"ef_search":50, - |"encoder":{"name":"sq","encoder":{"clip":true,"type":"fp16"}}}}}""".stripMargin.replace( + |"method":{"name":"ivf","engine":"faiss","space_type":"innerproduct", + |"parameters":{"nlist":4,"nprobes":2, + |"encoder":{"name":"sq","parameters":{"clip":true,"type":"fp16"}}}}}""".stripMargin.replace( "\n", "" ) diff --git a/elastic4s-domain/src/main/scala/com/sksamuel/elastic4s/fields/KnnVectorField.scala b/elastic4s-domain/src/main/scala/com/sksamuel/elastic4s/fields/KnnVectorField.scala index 039b8bbff..b4168d2f3 100644 --- a/elastic4s-domain/src/main/scala/com/sksamuel/elastic4s/fields/KnnVectorField.scala +++ b/elastic4s-domain/src/main/scala/com/sksamuel/elastic4s/fields/KnnVectorField.scala @@ -286,7 +286,7 @@ object IvfParameters { engine.get match { case KnnEngine.faiss => // do nothing as all parameters are accepted - if (!encoder.isInstanceOf[FaissEncoder]) { + if (encoder.nonEmpty && !encoder.get.isInstanceOf[FaissEncoder]) { throw new IllegalArgumentException( "encoder must be instance of FaissEncoder for faiss engine" ) diff --git a/elastic4s-handlers/src/main/scala/com/sksamuel/elastic4s/handlers/fields/KnnVectorFieldBuilderFn.scala b/elastic4s-handlers/src/main/scala/com/sksamuel/elastic4s/handlers/fields/KnnVectorFieldBuilderFn.scala index 1e0eb8c79..2cade8f77 100644 --- a/elastic4s-handlers/src/main/scala/com/sksamuel/elastic4s/handlers/fields/KnnVectorFieldBuilderFn.scala +++ b/elastic4s-handlers/src/main/scala/com/sksamuel/elastic4s/handlers/fields/KnnVectorFieldBuilderFn.scala @@ -60,13 +60,35 @@ object KnnVectorFieldBuilderFn { .get("name") .map(_.asInstanceOf[String]) .flatMap(v => FaissEncoderName.withName(v)) - val m = e.get("m").map(_.asInstanceOf[Int]) - val codeSize = e.get("code_size").map(_.asInstanceOf[Int]) - val sqClip = e.get("clip").map(_.asInstanceOf[Boolean]) + val m = e + .get("parameters") + .flatMap(p => + p.asInstanceOf[Map[String, Any]] + .get("m") + .map(_.asInstanceOf[Int]) + ) + val codeSize = e + .get("parameters") + .flatMap(p => + p.asInstanceOf[Map[String, Any]] + .get("code_size") + .map(_.asInstanceOf[Int]) + ) + val sqClip = e + .get("parameters") + .flatMap(p => + p.asInstanceOf[Map[String, Any]] + .get("clip") + .map(_.asInstanceOf[Boolean]) + ) val sqType = e - .get("type") - .map(_.asInstanceOf[String]) - .flatMap(v => FaissScalarQuantizationType.withName(v)) + .get("parameters") + .flatMap(p => + p.asInstanceOf[Map[String, Any]] + .get("type") + .map(_.asInstanceOf[String]) + flatMap (v => FaissScalarQuantizationType.withName(v)) + ) FaissEncoder( name = name, m = m, @@ -152,7 +174,7 @@ object KnnVectorFieldBuilderFn { encoder.name.foreach(v => builder.field("name", v.name)) // start of encoder `parameters` field - builder.startObject("encoder") + builder.startObject("parameters") encoder match { case e: FaissEncoder => e.m.foreach(v => builder.field("m", v)) @@ -174,7 +196,7 @@ object KnnVectorFieldBuilderFn { encoder.name.foreach(v => builder.field("name", v.name)) // start of encoder `parameters` field - builder.startObject("encoder") + builder.startObject("parameters") encoder match { case e: FaissEncoder => e.m.foreach(v => builder.field("m", v))