diff --git a/elastic4s-domain/src/main/scala/com/sksamuel/elastic4s/fields/SparseVectorField.scala b/elastic4s-domain/src/main/scala/com/sksamuel/elastic4s/fields/SparseVectorField.scala index f273957ef..86e8440eb 100644 --- a/elastic4s-domain/src/main/scala/com/sksamuel/elastic4s/fields/SparseVectorField.scala +++ b/elastic4s-domain/src/main/scala/com/sksamuel/elastic4s/fields/SparseVectorField.scala @@ -1,9 +1,17 @@ package com.sksamuel.elastic4s.fields +import com.sksamuel.elastic4s.requests.searches.queries.PruningConfig + object SparseVectorField { val `type`: String = "sparse_vector" } -case class SparseVectorField(name: String) extends ElasticField { +case class SparseVectorIndexOptions(prune: Boolean, pruningConfig: Option[PruningConfig]) + +case class SparseVectorField( + name: String, + store: Boolean = false, + indexOptions: Option[SparseVectorIndexOptions] = None +) extends ElasticField { override def `type`: String = SparseVectorField.`type` } diff --git a/elastic4s-handlers/src/main/scala/com/sksamuel/elastic4s/handlers/fields/SparseVectorFieldBuilderFn.scala b/elastic4s-handlers/src/main/scala/com/sksamuel/elastic4s/handlers/fields/SparseVectorFieldBuilderFn.scala index d8def1dc2..c7abbe528 100644 --- a/elastic4s-handlers/src/main/scala/com/sksamuel/elastic4s/handlers/fields/SparseVectorFieldBuilderFn.scala +++ b/elastic4s-handlers/src/main/scala/com/sksamuel/elastic4s/handlers/fields/SparseVectorFieldBuilderFn.scala @@ -2,13 +2,45 @@ package com.sksamuel.elastic4s.handlers.fields import com.sksamuel.elastic4s.fields._ import com.sksamuel.elastic4s.json.{XContentBuilder, XContentFactory} +import com.sksamuel.elastic4s.requests.searches.queries.PruningConfig object SparseVectorFieldBuilderFn { - def toField(name: String, values: Map[String, Any]): SparseVectorField = SparseVectorField(name) + private def getPruningConfig(values: Map[String, Any]): PruningConfig = + PruningConfig( + values.get("tokens_freq_ratio_threshold").map(_.asInstanceOf[Int]), + values.get("tokens_weight_threshold").map(_.asInstanceOf[Double].toFloat) + ) + + private def getIndexOptions(values: Map[String, Any]): SparseVectorIndexOptions = { + SparseVectorIndexOptions( + values("prune").asInstanceOf[Boolean], + pruningConfig = values.get("pruning_config").map(_.asInstanceOf[Map[String, Any]]).map(getPruningConfig) + ) + } + + def toField(name: String, values: Map[String, Any]): SparseVectorField = SparseVectorField( + name, + values("store").asInstanceOf[Boolean], + indexOptions = values.get("index_options").map(_.asInstanceOf[Map[String, Any]]).map(getIndexOptions) + ) def build(field: SparseVectorField): XContentBuilder = { val builder = XContentFactory.jsonBuilder() builder.field("type", field.`type`) + builder.field("store", field.store) + field.indexOptions.foreach { indexOptions => + builder.startObject("index_options") + builder.field("prune", indexOptions.prune) + indexOptions.pruningConfig.foreach { pruningConfig => + if (pruningConfig.tokensFreqRatioThreshold.nonEmpty || pruningConfig.tokensWeighThreshold.nonEmpty) { + builder.startObject("pruning_config") + pruningConfig.tokensFreqRatioThreshold.foreach(builder.field("tokens_freq_ratio_threshold", _)) + pruningConfig.tokensWeighThreshold.foreach(builder.field("tokens_weight_threshold", _)) + builder.endObject() + } + } + builder.endObject() + } builder.endObject() } } diff --git a/elastic4s-tests/src/test/scala/com/sksamuel/elastic4s/fields/ElasticFieldBuilderFnTest.scala b/elastic4s-tests/src/test/scala/com/sksamuel/elastic4s/fields/ElasticFieldBuilderFnTest.scala index 05c3a33a1..6c2cf02aa 100644 --- a/elastic4s-tests/src/test/scala/com/sksamuel/elastic4s/fields/ElasticFieldBuilderFnTest.scala +++ b/elastic4s-tests/src/test/scala/com/sksamuel/elastic4s/fields/ElasticFieldBuilderFnTest.scala @@ -246,16 +246,6 @@ class ElasticFieldBuilderFnTest extends AnyWordSpec with Matchers { JacksonSupport.mapper.readValue[Map[String, Any]](jsonString) ) shouldBe field } - - "support SparseVectorField" in { - val field = SparseVectorField("sparse_vector_field") - val jsonString = """{"type":"sparse_vector"}""" - ElasticFieldBuilderFn(field).string shouldBe jsonString - ElasticFieldBuilderFn.construct( - field.name, - JacksonSupport.mapper.readValue[Map[String, Any]](jsonString) - ) shouldBe field - } "support RankFeaturesField with positive_score_impact" in { val field = RankFeaturesField("rank_features_field", positiveScoreImpact = Some(false)) val jsonString = """{"type":"rank_features","positive_score_impact":false}""" diff --git a/elastic4s-tests/src/test/scala/com/sksamuel/elastic4s/fields/SparseVectorFieldTest.scala b/elastic4s-tests/src/test/scala/com/sksamuel/elastic4s/fields/SparseVectorFieldTest.scala new file mode 100644 index 000000000..490e06eb9 --- /dev/null +++ b/elastic4s-tests/src/test/scala/com/sksamuel/elastic4s/fields/SparseVectorFieldTest.scala @@ -0,0 +1,33 @@ +package com.sksamuel.elastic4s.fields + +import com.sksamuel.elastic4s.ElasticApi +import com.sksamuel.elastic4s.handlers.fields.{ElasticFieldBuilderFn, SparseVectorFieldBuilderFn} +import com.sksamuel.elastic4s.jackson.JacksonSupport +import com.sksamuel.elastic4s.requests.searches.queries.PruningConfig +import org.scalatest.flatspec.AnyFlatSpec +import org.scalatest.matchers.should.Matchers + +class SparseVectorFieldTest extends AnyFlatSpec with Matchers with ElasticApi { + "A SparseVectorField" should "be minimally supported" in { + val field = SparseVectorField("test.tokens") + val jsonString = SparseVectorFieldBuilderFn.build(field).string + jsonString shouldBe """{"type":"sparse_vector","store":false}""" + ElasticFieldBuilderFn.construct( + field.name, + JacksonSupport.mapper.readValue[Map[String, Any]](jsonString) + ) shouldBe field + } + + it should "support all fields" in { + val sparseVectorIndexOptions = + SparseVectorIndexOptions(prune = true, pruningConfig = Some(PruningConfig(Some(5), Some(1.0F)))) + + val field = SparseVectorField("test.tokens", store = true, indexOptions = Some(sparseVectorIndexOptions)) + val jsonString = SparseVectorFieldBuilderFn.build(field).string + jsonString shouldBe """{"type":"sparse_vector","store":true,"index_options":{"prune":true,"pruning_config":{"tokens_freq_ratio_threshold":5,"tokens_weight_threshold":1.0}}}""" + ElasticFieldBuilderFn.construct( + field.name, + JacksonSupport.mapper.readValue[Map[String, Any]](jsonString) + ) shouldBe field + } +}