Add support for knn vector queries on semantic_text fields by kderusso · Pull Request #119011 · elastic/elasticsearch

kderusso · 2024-12-18T19:29:09Z

Supports knn queries in the query DSL against semantic_text fields. Note that top level knn search using the knn section is not supported.

Example script:


PUT _inference/text_embedding/my-e5-model
{
  "service": "elasticsearch",
  "service_settings": {
    "num_allocations": 1,
    "num_threads": 1,
    "model_id": ".multilingual-e5-small"
  }
}

PUT my-index-1
{
    "mappings": {
        "properties": {
            "inference_field": {
                "type": "semantic_text",
                "inference_id": "my-e5-model"
            },
            "text_field": {
                "type": "text"
            }
        }
    }
}

POST my-index-1/_doc/1
{
  "inference_field": "a test value",
  "text_field": "a test value"
}

// KNN against a semantic text field using inference
GET my-index-1/_search
{
  "query": {
    "knn": {
      "field": "inference_field",
      "k": 10,
      "num_candidates": 100,
      "query_vector_builder": {
        "text_embedding": {
          "model_id": "my-e5-model",
          "model_text": "test"
        }
      }
    }
  },
  "_source": {
    "exclude": "inference_field.inference.chunks"
  }
}

// KNN against a semantic text field not specifying inference
GET my-index-1/_search?error_trace
{
  "query": {
    "knn": {
      "field": "inference_field",
      "k": 10,
      "num_candidates": 100,
      "query_vector_builder": {
        "text_embedding": {
          "model_text": "test"
        }
      }
    }
  },
  "_source": {
    "exclude": "inference_field.inference.chunks"
  }
}

// KNN against a semantic text field using query vectors
GET my-index-1/_search
{
  "query": {
    "knn": {
      "field": "inference_field",
      "k": 10,
      "num_candidates": 100,
      "query_vector": [ 0.007934765, -0.015913952, -0.043057524, -0.10636589, 0.011301906, -0.015008567, 0.012158697, 0.06435651, 0.05081897, 0.017386554, 0.042363416, 0.028541336, 0.08033845, -0.022166155, -0.051751476, 0.02091861, 0.09317102, -0.03462641, -0.06918183, -0.03136944, 0.010598254, -0.025934521, -0.031956308, 0.059287872, 0.05378277, 0.06513278, -0.04364164, -0.025723265, 0.046480402, -0.049092673, -0.021595225, -0.0032410314, 0.031168742, -0.0739943, 0.048144553, 0.08345922, -0.074975915, -0.060630694, 0.07822045, -0.07663624, -0.07743249, -0.035663426, 0.011769564, 0.06594226, 0.043934587, 0.033019762, -0.059856612, 0.017330708, -0.04865905, -0.028922942, -0.06462057, 0.029822953, -0.0023164605, 0.02530421, 0.069650486, -0.073910154, -0.004165947, -0.05491859, -0.085993074, 0.072383724, 0.07173003, -0.015900861, -0.046266597, 0.05065501, 0.0779681, 0.0891497, 0.03152424, 0.023294332, -0.04338317, -0.017575819, -0.059077594, 0.018678484, 0.020382047, -0.015786897, 0.043291267, 0.013902389, 0.00058899104, -0.08444341, 0.06319704, -0.052247103, -0.05578253, -0.06578668, -0.025945362, 0.017614782, -0.047085304, 0.057253387, 0.033382308, -0.08257727, 0.047596436, -0.037106443, 0.074864075, 0.052347176, -0.060556035, -0.046495277, -0.08871017, -0.046331394, -0.056618407, 0.054185994, 0.036283605, -0.052142534, 0.05231953, -0.046478417, 0.006274289, -0.09101135, -0.020674909, 0.050636552, 0.077042066, -0.0431401, 0.040718045, -0.041647647, 0.030319883, 0.026050929, 0.085587256, 0.028968774, -0.050006703, -0.00958342, -0.004807631, -0.0480301, 0.0633695, -0.09563107, 0.07481829, -0.025702624, -0.05917035, -0.03292549, -0.045559984, -0.021310199, 0.016081415, 0.06298191, 0.053787895, 0.010182924, 0.04226957, 0.037372373, 0.05062729, 0.0406053, 0.043430693, 0.054629993, -0.016552584, 0.060336597, -0.005067171, -0.0627392, -0.026641937, 0.047697227, -0.061994646, 0.053344723, 0.06145448, 0.035446793, 0.028673936, -0.00889342, 0.065559044, -0.07846254, 0.041033193, -0.027498009, 0.046198983, 0.004036812, 0.058940724, -0.03080122, -0.06507373, -0.0674627, 0.061292157, 0.04102403, -0.06849838, -0.08000652, -0.03143754, -0.035949968, -0.068278916, -0.07951413, 0.008518893, 0.07328911, -0.078702435, -0.056579445, -0.049364146, 0.032737307, -0.0018469784, 0.0382715, -0.0045158467, 0.06932468, -0.05570215, 0.056880653, 0.08039553, 0.032512806, -0.059217937, -0.039542492, -0.08205577, -0.07043393, -0.036309723, -0.05440946, -0.05358154, 0.01491165, 0.016233535, -0.0011302672, -0.0003497226, 0.044191226, -0.034170866, -0.06359364, -0.041531622, 0.0039237407, -0.06328019, 0.071654275, 0.056353293, -0.017949982, 0.009957283, -0.03611878, 0.015675843, 0.042649325, 0.05194668, 0.07188313, -0.086326554, 0.023199761, -0.05551281, 0.031496353, 0.04902332, -0.045349505, -0.08469811, 0.06258924, -0.0802043, -0.050104994, 0.008524371, 0.0867973, -0.009479119, -0.005370896, 0.06352957, -0.0040729647, 0.015049631, -0.04841483, -0.04341567, 0.056807704, 0.09788039, -0.11132506, -0.0172819, 0.05246852, -0.025508488, -0.053615686, -0.13030247, -0.07100205, -0.0426781, -0.02301062, -0.013061844, 0.02653803, 0.055927757, -0.06096397, -0.032301165, -0.04471813, 0.051232424, -0.08266851, 0.06761164, -0.053912845, -0.0014131556, 0.0330563, 0.015858745, 0.043734655, -0.013962677, -0.042700607, -0.033999793, -0.07273025, -0.025160242, 0.03139777, 0.06657497, 0.05504236, -0.05545997, 0.039574854, 0.038956527, -0.026925432, 0.079213314, 0.05498472, 0.048151035, 0.060903218, -0.024777696, -0.0387965, -0.105480194, -0.040409025, -0.064113736, 0.03950475, 0.07954551, -0.043298814, -0.020724343, -0.054619316, 0.09759665, 0.10228246, -0.032124214, -0.04095005, 0.05043362, -0.052726258, 0.05982946, 0.04796253, 0.06194914, 0.010950634, -0.030110938, 0.03283955, -0.0002696275, -0.0098341135, -0.022008123, -0.074884586, 0.033381447, -0.040135764, 0.072015, 0.033126105, -0.0070500704, 0.044344094, -0.033021137, 0.041409004, -0.019219067, -0.057634152, 0.014906292, 0.06305334, -0.00911885, 0.051950686, -0.0076152226, 0.0006077424, 0.007057235, 0.0767627, 0.0671153, 0.06547248, -0.04329299, -0.052492123, 0.009058822, 0.087170124, -0.033063732, 0.014588411, -0.02944014, -0.034347236, -0.030877389, -0.045103014, -0.025955923, -0.053345636, 0.054236747, 0.045914765, -0.034200758, -0.022325076, 0.056416932, -0.040778212, -0.012222789, -0.032742303, -0.058351234, 0.0128486985, -0.028538564, -0.04633306, -0.0015151277, 0.03103562, -0.049697988, -0.06285182, 0.015321202, 0.07325691, -0.028998345, 0.060792767, -0.023697915, -0.008528026, 0.03551946, -0.021368505, -0.030065581, 0.07350493, 0.05271221, -0.106231354, 0.017700348, 0.08116126, -0.07330424, 0.005489799, -0.065025605, 0.00013129896, 0.028777555, 0.05560677, -0.029058384, -0.06901476, 0.016285766, 0.034541924, 0.08150406, 0.0714121, -0.024971426, -0.04193428, 0.056124993, -0.06255153, 0.09743034, 0.049381103, -0.028551068, -0.015498741, 0.009420205, -0.08466289, -0.05514481, 0.05644234, -0.027993482, -0.01833721, 0.039042473, 0.050628796, 0.071234666, 0.046101592 ]
    }
  },
  "_source": {
    "exclude": "inference_field.inference.chunks"
  }
}

// Add another index with a dense_vector field and ingest a compatible document
PUT my-index-2
{
  "mappings": {
    "properties": {
      "inference_field": {
        "type": "dense_vector"
      },
      "text_field": { 
        "type": "text"
      }
    }
  }
}

POST my-index-2/_doc/1
{
  "inference_field": [ 0.007934765, -0.015913952, -0.043057524, -0.10636589, 0.011301906, -0.015008567, 0.012158697, 0.06435651, 0.05081897, 0.017386554, 0.042363416, 0.028541336, 0.08033845, -0.022166155, -0.051751476, 0.02091861, 0.09317102, -0.03462641, -0.06918183, -0.03136944, 0.010598254, -0.025934521, -0.031956308, 0.059287872, 0.05378277, 0.06513278, -0.04364164, -0.025723265, 0.046480402, -0.049092673, -0.021595225, -0.0032410314, 0.031168742, -0.0739943, 0.048144553, 0.08345922, -0.074975915, -0.060630694, 0.07822045, -0.07663624, -0.07743249, -0.035663426, 0.011769564, 0.06594226, 0.043934587, 0.033019762, -0.059856612, 0.017330708, -0.04865905, -0.028922942, -0.06462057, 0.029822953, -0.0023164605, 0.02530421, 0.069650486, -0.073910154, -0.004165947, -0.05491859, -0.085993074, 0.072383724, 0.07173003, -0.015900861, -0.046266597, 0.05065501, 0.0779681, 0.0891497, 0.03152424, 0.023294332, -0.04338317, -0.017575819, -0.059077594, 0.018678484, 0.020382047, -0.015786897, 0.043291267, 0.013902389, 0.00058899104, -0.08444341, 0.06319704, -0.052247103, -0.05578253, -0.06578668, -0.025945362, 0.017614782, -0.047085304, 0.057253387, 0.033382308, -0.08257727, 0.047596436, -0.037106443, 0.074864075, 0.052347176, -0.060556035, -0.046495277, -0.08871017, -0.046331394, -0.056618407, 0.054185994, 0.036283605, -0.052142534, 0.05231953, -0.046478417, 0.006274289, -0.09101135, -0.020674909, 0.050636552, 0.077042066, -0.0431401, 0.040718045, -0.041647647, 0.030319883, 0.026050929, 0.085587256, 0.028968774, -0.050006703, -0.00958342, -0.004807631, -0.0480301, 0.0633695, -0.09563107, 0.07481829, -0.025702624, -0.05917035, -0.03292549, -0.045559984, -0.021310199, 0.016081415, 0.06298191, 0.053787895, 0.010182924, 0.04226957, 0.037372373, 0.05062729, 0.0406053, 0.043430693, 0.054629993, -0.016552584, 0.060336597, -0.005067171, -0.0627392, -0.026641937, 0.047697227, -0.061994646, 0.053344723, 0.06145448, 0.035446793, 0.028673936, -0.00889342, 0.065559044, -0.07846254, 0.041033193, -0.027498009, 0.046198983, 0.004036812, 0.058940724, -0.03080122, -0.06507373, -0.0674627, 0.061292157, 0.04102403, -0.06849838, -0.08000652, -0.03143754, -0.035949968, -0.068278916, -0.07951413, 0.008518893, 0.07328911, -0.078702435, -0.056579445, -0.049364146, 0.032737307, -0.0018469784, 0.0382715, -0.0045158467, 0.06932468, -0.05570215, 0.056880653, 0.08039553, 0.032512806, -0.059217937, -0.039542492, -0.08205577, -0.07043393, -0.036309723, -0.05440946, -0.05358154, 0.01491165, 0.016233535, -0.0011302672, -0.0003497226, 0.044191226, -0.034170866, -0.06359364, -0.041531622, 0.0039237407, -0.06328019, 0.071654275, 0.056353293, -0.017949982, 0.009957283, -0.03611878, 0.015675843, 0.042649325, 0.05194668, 0.07188313, -0.086326554, 0.023199761, -0.05551281, 0.031496353, 0.04902332, -0.045349505, -0.08469811, 0.06258924, -0.0802043, -0.050104994, 0.008524371, 0.0867973, -0.009479119, -0.005370896, 0.06352957, -0.0040729647, 0.015049631, -0.04841483, -0.04341567, 0.056807704, 0.09788039, -0.11132506, -0.0172819, 0.05246852, -0.025508488, -0.053615686, -0.13030247, -0.07100205, -0.0426781, -0.02301062, -0.013061844, 0.02653803, 0.055927757, -0.06096397, -0.032301165, -0.04471813, 0.051232424, -0.08266851, 0.06761164, -0.053912845, -0.0014131556, 0.0330563, 0.015858745, 0.043734655, -0.013962677, -0.042700607, -0.033999793, -0.07273025, -0.025160242, 0.03139777, 0.06657497, 0.05504236, -0.05545997, 0.039574854, 0.038956527, -0.026925432, 0.079213314, 0.05498472, 0.048151035, 0.060903218, -0.024777696, -0.0387965, -0.105480194, -0.040409025, -0.064113736, 0.03950475, 0.07954551, -0.043298814, -0.020724343, -0.054619316, 0.09759665, 0.10228246, -0.032124214, -0.04095005, 0.05043362, -0.052726258, 0.05982946, 0.04796253, 0.06194914, 0.010950634, -0.030110938, 0.03283955, -0.0002696275, -0.0098341135, -0.022008123, -0.074884586, 0.033381447, -0.040135764, 0.072015, 0.033126105, -0.0070500704, 0.044344094, -0.033021137, 0.041409004, -0.019219067, -0.057634152, 0.014906292, 0.06305334, -0.00911885, 0.051950686, -0.0076152226, 0.0006077424, 0.007057235, 0.0767627, 0.0671153, 0.06547248, -0.04329299, -0.052492123, 0.009058822, 0.087170124, -0.033063732, 0.014588411, -0.02944014, -0.034347236, -0.030877389, -0.045103014, -0.025955923, -0.053345636, 0.054236747, 0.045914765, -0.034200758, -0.022325076, 0.056416932, -0.040778212, -0.012222789, -0.032742303, -0.058351234, 0.0128486985, -0.028538564, -0.04633306, -0.0015151277, 0.03103562, -0.049697988, -0.06285182, 0.015321202, 0.07325691, -0.028998345, 0.060792767, -0.023697915, -0.008528026, 0.03551946, -0.021368505, -0.030065581, 0.07350493, 0.05271221, -0.106231354, 0.017700348, 0.08116126, -0.07330424, 0.005489799, -0.065025605, 0.00013129896, 0.028777555, 0.05560677, -0.029058384, -0.06901476, 0.016285766, 0.034541924, 0.08150406, 0.0714121, -0.024971426, -0.04193428, 0.056124993, -0.06255153, 0.09743034, 0.049381103, -0.028551068, -0.015498741, 0.009420205, -0.08466289, -0.05514481, 0.05644234, -0.027993482, -0.01833721, 0.039042473, 0.050628796, 0.071234666, 0.046101592 ],
  "text_field": "another test value"
}


// Combined search using inference
GET my-index-*/_search
{
  "query": {
    "knn": {
      "field": "inference_field",
      "k": 10,
      "num_candidates": 100,
      "query_vector_builder": {
        "text_embedding": {
          "model_id": "my-e5-model",
          "model_text": "test"
        }
      }
    }
  },
  "_source": {
    "exclude": "inference_field.inference.chunks"
  }
}

// Combined search using default inference
// will error because dense_vector fields don't have inference built in 
GET my-index-*/_search
{
  "query": {
    "knn": {
      "field": "inference_field",
      "k": 10,
      "num_candidates": 100,
      "query_vector_builder": {
        "text_embedding": {
          "model_text": "test"
        }
      }
    }
  }
}

// Combined search using query vectors 
GET my-index-*/_search
{
  "query": {
    "knn": {
      "field": "inference_field",
      "k": 10,
      "num_candidates": 100,
      "query_vector": [ 0.007934765, -0.015913952, -0.043057524, -0.10636589, 0.011301906, -0.015008567, 0.012158697, 0.06435651, 0.05081897, 0.017386554, 0.042363416, 0.028541336, 0.08033845, -0.022166155, -0.051751476, 0.02091861, 0.09317102, -0.03462641, -0.06918183, -0.03136944, 0.010598254, -0.025934521, -0.031956308, 0.059287872, 0.05378277, 0.06513278, -0.04364164, -0.025723265, 0.046480402, -0.049092673, -0.021595225, -0.0032410314, 0.031168742, -0.0739943, 0.048144553, 0.08345922, -0.074975915, -0.060630694, 0.07822045, -0.07663624, -0.07743249, -0.035663426, 0.011769564, 0.06594226, 0.043934587, 0.033019762, -0.059856612, 0.017330708, -0.04865905, -0.028922942, -0.06462057, 0.029822953, -0.0023164605, 0.02530421, 0.069650486, -0.073910154, -0.004165947, -0.05491859, -0.085993074, 0.072383724, 0.07173003, -0.015900861, -0.046266597, 0.05065501, 0.0779681, 0.0891497, 0.03152424, 0.023294332, -0.04338317, -0.017575819, -0.059077594, 0.018678484, 0.020382047, -0.015786897, 0.043291267, 0.013902389, 0.00058899104, -0.08444341, 0.06319704, -0.052247103, -0.05578253, -0.06578668, -0.025945362, 0.017614782, -0.047085304, 0.057253387, 0.033382308, -0.08257727, 0.047596436, -0.037106443, 0.074864075, 0.052347176, -0.060556035, -0.046495277, -0.08871017, -0.046331394, -0.056618407, 0.054185994, 0.036283605, -0.052142534, 0.05231953, -0.046478417, 0.006274289, -0.09101135, -0.020674909, 0.050636552, 0.077042066, -0.0431401, 0.040718045, -0.041647647, 0.030319883, 0.026050929, 0.085587256, 0.028968774, -0.050006703, -0.00958342, -0.004807631, -0.0480301, 0.0633695, -0.09563107, 0.07481829, -0.025702624, -0.05917035, -0.03292549, -0.045559984, -0.021310199, 0.016081415, 0.06298191, 0.053787895, 0.010182924, 0.04226957, 0.037372373, 0.05062729, 0.0406053, 0.043430693, 0.054629993, -0.016552584, 0.060336597, -0.005067171, -0.0627392, -0.026641937, 0.047697227, -0.061994646, 0.053344723, 0.06145448, 0.035446793, 0.028673936, -0.00889342, 0.065559044, -0.07846254, 0.041033193, -0.027498009, 0.046198983, 0.004036812, 0.058940724, -0.03080122, -0.06507373, -0.0674627, 0.061292157, 0.04102403, -0.06849838, -0.08000652, -0.03143754, -0.035949968, -0.068278916, -0.07951413, 0.008518893, 0.07328911, -0.078702435, -0.056579445, -0.049364146, 0.032737307, -0.0018469784, 0.0382715, -0.0045158467, 0.06932468, -0.05570215, 0.056880653, 0.08039553, 0.032512806, -0.059217937, -0.039542492, -0.08205577, -0.07043393, -0.036309723, -0.05440946, -0.05358154, 0.01491165, 0.016233535, -0.0011302672, -0.0003497226, 0.044191226, -0.034170866, -0.06359364, -0.041531622, 0.0039237407, -0.06328019, 0.071654275, 0.056353293, -0.017949982, 0.009957283, -0.03611878, 0.015675843, 0.042649325, 0.05194668, 0.07188313, -0.086326554, 0.023199761, -0.05551281, 0.031496353, 0.04902332, -0.045349505, -0.08469811, 0.06258924, -0.0802043, -0.050104994, 0.008524371, 0.0867973, -0.009479119, -0.005370896, 0.06352957, -0.0040729647, 0.015049631, -0.04841483, -0.04341567, 0.056807704, 0.09788039, -0.11132506, -0.0172819, 0.05246852, -0.025508488, -0.053615686, -0.13030247, -0.07100205, -0.0426781, -0.02301062, -0.013061844, 0.02653803, 0.055927757, -0.06096397, -0.032301165, -0.04471813, 0.051232424, -0.08266851, 0.06761164, -0.053912845, -0.0014131556, 0.0330563, 0.015858745, 0.043734655, -0.013962677, -0.042700607, -0.033999793, -0.07273025, -0.025160242, 0.03139777, 0.06657497, 0.05504236, -0.05545997, 0.039574854, 0.038956527, -0.026925432, 0.079213314, 0.05498472, 0.048151035, 0.060903218, -0.024777696, -0.0387965, -0.105480194, -0.040409025, -0.064113736, 0.03950475, 0.07954551, -0.043298814, -0.020724343, -0.054619316, 0.09759665, 0.10228246, -0.032124214, -0.04095005, 0.05043362, -0.052726258, 0.05982946, 0.04796253, 0.06194914, 0.010950634, -0.030110938, 0.03283955, -0.0002696275, -0.0098341135, -0.022008123, -0.074884586, 0.033381447, -0.040135764, 0.072015, 0.033126105, -0.0070500704, 0.044344094, -0.033021137, 0.041409004, -0.019219067, -0.057634152, 0.014906292, 0.06305334, -0.00911885, 0.051950686, -0.0076152226, 0.0006077424, 0.007057235, 0.0767627, 0.0671153, 0.06547248, -0.04329299, -0.052492123, 0.009058822, 0.087170124, -0.033063732, 0.014588411, -0.02944014, -0.034347236, -0.030877389, -0.045103014, -0.025955923, -0.053345636, 0.054236747, 0.045914765, -0.034200758, -0.022325076, 0.056416932, -0.040778212, -0.012222789, -0.032742303, -0.058351234, 0.0128486985, -0.028538564, -0.04633306, -0.0015151277, 0.03103562, -0.049697988, -0.06285182, 0.015321202, 0.07325691, -0.028998345, 0.060792767, -0.023697915, -0.008528026, 0.03551946, -0.021368505, -0.030065581, 0.07350493, 0.05271221, -0.106231354, 0.017700348, 0.08116126, -0.07330424, 0.005489799, -0.065025605, 0.00013129896, 0.028777555, 0.05560677, -0.029058384, -0.06901476, 0.016285766, 0.034541924, 0.08150406, 0.0714121, -0.024971426, -0.04193428, 0.056124993, -0.06255153, 0.09743034, 0.049381103, -0.028551068, -0.015498741, 0.009420205, -0.08466289, -0.05514481, 0.05644234, -0.027993482, -0.01833721, 0.039042473, 0.050628796, 0.071234666, 0.046101592 ]
    }
  },
  "_source": {
    "exclude": "inference_field.inference.chunks"
  }
}

elasticsearchmachine · 2024-12-18T20:26:35Z

Hi @kderusso, I've created a changelog YAML for you.

kderusso · 2024-12-19T18:15:53Z

...e/src/main/java/org/elasticsearch/xpack/core/ml/vectors/TextEmbeddingQueryVectorBuilder.java

 */

-package org.elasticsearch.xpack.ml.vectors;
+package org.elasticsearch.xpack.core.ml.vectors;


This had to move from the ml plugin in order to access it in the inference plugin. While I could have moved it to inference, core seemed like a better fit for this standalone class. I did keep the tests in the ml plugin however, because they rely on the ML test runner.

Mikep86

Partial review. Great start to this! As far as edge cases, I think we need some test cases for the following:

Missing both query vector and query vector builder
Querying multiple indices that use different inference IDs, particularly when each index has more than k potential matches

...e/src/main/java/org/elasticsearch/xpack/core/ml/vectors/TextEmbeddingQueryVectorBuilder.java

...java/org/elasticsearch/xpack/inference/queries/SemanticKnnVectorQueryRewriteInterceptor.java

elasticsearchmachine · 2025-01-03T13:16:18Z

Pinging @elastic/es-search-relevance (Team:Search Relevance)

elasticsearchmachine · 2025-01-03T13:16:18Z

Pinging @elastic/search-eng (Team:SearchOrg)

elasticsearchmachine · 2025-01-03T13:16:19Z

Pinging @elastic/search-relevance (Team:Search - Relevance)

carlosdelest

LGTM 💯

Some questions about pulling up some common methods - can be done in follow up work

...e/src/main/java/org/elasticsearch/xpack/core/ml/vectors/TextEmbeddingQueryVectorBuilder.java

...java/org/elasticsearch/xpack/inference/queries/SemanticKnnVectorQueryRewriteInterceptor.java

carlosdelest · 2025-01-03T17:01:34Z

...c/test/java/org/elasticsearch/index/query/SemanticKnnVectorQueryRewriteInterceptorTests.java

+import java.io.IOException;
+import java.util.Map;
+
+public class SemanticKnnVectorQueryRewriteInterceptorTests extends ESTestCase {


We may find a common structure for these tests and create a base class for them in follow up work

Yes, I would like to have this organized better in a followup. Unfortunately, since the interceptor is a QueryBuilder we can't take advantage of built in abstract query builder tests. I was trying to focus more on test functionality first, so we don't back ourselves into a corner. Right now I'm thinking this would be a better followup optimization if the test cases are complete though. WDYT?

Yes, I'm totally fine with doing that as follow up work 👍

...n/inference/src/yamlRestTest/resources/rest-api-spec/test/inference/47_semantic_text_knn.yml

Mikep86

Partial review of everything except for the tests. Good work on this, I have some suggestions about how we can make the intercept logic simpler and more efficient.

...e/src/main/java/org/elasticsearch/xpack/core/ml/vectors/TextEmbeddingQueryVectorBuilder.java

Mikep86 · 2025-01-03T16:52:31Z

...java/org/elasticsearch/xpack/inference/queries/SemanticKnnVectorQueryRewriteInterceptor.java

+                createSubQueryForIndices(
+                    inferenceIdsIndices.get(inferenceId),
+                    buildNestedQueryFromKnnVectorQuery(queryBuilder, inferenceId)
+                )


I think we could update buildNestedQueryFromKnnVectorQuery to take the index list and build the nested knn query that filters on the indices directly in it. This would both simplify the code and be more efficient, as it looks like right now you're effectively using the query builder returned from buildNestedQueryFromKnnVectorQuery as a wrapper for data that you use to build the actual nested query in createSubQueryForIndices.

Another benefit of this approach is that we remove the index post-filtering (via the boolean query), which is no longer necessary since we are pre-filtering on index in the knn query.

Good call on adding these as a prefilter, however the post filtering can't be removed without breakinng functionality for dense vector combined fields. I think your suggestions helped to clean it up a bit though.

...java/org/elasticsearch/xpack/inference/queries/SemanticKnnVectorQueryRewriteInterceptor.java

...c/test/java/org/elasticsearch/index/query/SemanticKnnVectorQueryRewriteInterceptorTests.java

...n/inference/src/yamlRestTest/resources/rest-api-spec/test/inference/47_semantic_text_knn.yml

Mikep86

Thanks for making the changes! There's one more follow-up, I think we can avoid index post-filtering when not using nested queries. Please correct me if I'm wrong though.

...java/org/elasticsearch/xpack/inference/queries/SemanticKnnVectorQueryRewriteInterceptor.java

Mikep86

LGTM, thanks for the iterations!

…19011) * First cut at KNN interceptor * Move TextEmbeddingQueryVectorBuilder to xpack-core * Infer model ID * Fix test compilation errors * Update docs/changelog/119011.yaml * Update changelog * Update test * Cleanup * PR feedback * Add yaml test * Cleanup tests * Update test * Minor PR feedback * refactor pre filter indices for knn * PR feedback

…119648) * First cut at KNN interceptor * Move TextEmbeddingQueryVectorBuilder to xpack-core * Infer model ID * Fix test compilation errors * Update docs/changelog/119011.yaml * Update changelog * Update test * Cleanup * PR feedback * Add yaml test * Cleanup tests * Update test * Minor PR feedback * refactor pre filter indices for knn * PR feedback

kderusso added 3 commits December 18, 2024 13:26

First cut at KNN interceptor

678c43d

Move TextEmbeddingQueryVectorBuilder to xpack-core

e30bf6b

Infer model ID

a136ec1

elasticsearchmachine added the v9.0.0 label Dec 18, 2024

Fix test compilation errors

9bd3796

kderusso added >enhancement :SearchOrg/Relevance Label for the Search (solution/org) Relevance team labels Dec 18, 2024

kderusso added 2 commits December 18, 2024 15:26

Update docs/changelog/119011.yaml

5729e3b

Update changelog

ed2ab65

kderusso added the :Search Relevance/Search Catch all for Search Relevance label Dec 18, 2024

kderusso added 3 commits December 18, 2024 16:35

Update test

ea2c60f

Merge from main

a34be31

Merge branch 'main' into kderusso/semantic-text-knn-query

1f7b54f

kderusso commented Dec 19, 2024

View reviewed changes

kderusso added 2 commits December 19, 2024 13:46

Cleanup

ef06879

Merge branch 'main' into kderusso/semantic-text-knn-query

dea977d

kderusso changed the title ~~WIP: Add support for knn vector queries on semantic_text fields~~ Dec 19, 2024

Mikep86 reviewed Dec 19, 2024

View reviewed changes

kderusso added 5 commits January 2, 2025 15:22

PR feedback

055e8ed

Add yaml test

88ef614

Merge from main

792ef55

Cleanup tests

f779b41

Update test

1ab6517

kderusso marked this pull request as ready for review January 3, 2025 13:15

kderusso requested a review from Mikep86 January 3, 2025 13:16

kderusso requested a review from jimczi January 3, 2025 13:16

kderusso requested a review from a team January 3, 2025 13:16

kderusso added auto-backport Automatically create backport pull requests when merged v8.18.0 labels Jan 3, 2025

kderusso mentioned this pull request Jan 3, 2025

[Docs] Add docs for new semantic text query functionality #119520

Merged

carlosdelest approved these changes Jan 3, 2025

View reviewed changes

Mikep86 reviewed Jan 3, 2025

View reviewed changes

kderusso added 3 commits January 3, 2025 14:06

Minor PR feedback

5fa97c2

refactor pre filter indices for knn

1152df6

Merge from main

561cc7e

Mikep86 reviewed Jan 3, 2025

View reviewed changes

...java/org/elasticsearch/xpack/inference/queries/SemanticKnnVectorQueryRewriteInterceptor.java Outdated Show resolved Hide resolved

...java/org/elasticsearch/xpack/inference/queries/SemanticKnnVectorQueryRewriteInterceptor.java Outdated Show resolved Hide resolved

PR feedback

0d5765a

kderusso requested a review from Mikep86 January 6, 2025 19:41

Merge branch 'main' into kderusso/semantic-text-knn-query

c42f5aa

Mikep86 approved these changes Jan 6, 2025

View reviewed changes

Merge branch 'main' into kderusso/semantic-text-knn-query

ccca93e

kderusso merged commit 3f99211 into elastic:main Jan 7, 2025

kderusso mentioned this pull request Jan 7, 2025

[8.x] Add support for knn vector queries on semantic_text fields (#119011) #119648

Merged

kderusso deleted the kderusso/semantic-text-knn-query branch January 24, 2025 13:52

pquentin mentioned this pull request Sep 30, 2025

Specification fixes for 9.2 elastic/elasticsearch-specification#5372

Closed

9 tasks

kderusso mentioned this pull request Oct 1, 2025

Support inferred model ID for knn queries over semantic text fields elastic/elasticsearch-specification#5390

Merged

Conversation

kderusso commented Dec 18, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

elasticsearchmachine commented Dec 18, 2024

kderusso Dec 19, 2024

Choose a reason for hiding this comment

Mikep86 left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

elasticsearchmachine commented Jan 3, 2025

elasticsearchmachine commented Jan 3, 2025

elasticsearchmachine commented Jan 3, 2025

carlosdelest left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

carlosdelest Jan 3, 2025

Choose a reason for hiding this comment

kderusso Jan 3, 2025

Choose a reason for hiding this comment

carlosdelest Jan 7, 2025

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Mikep86 left a comment

Choose a reason for hiding this comment

Uh oh!

Mikep86 Jan 3, 2025

Choose a reason for hiding this comment

kderusso Jan 3, 2025

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Mikep86 left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Mikep86 left a comment

Choose a reason for hiding this comment

Labels

4 participants

kderusso commented Dec 18, 2024 •

edited

Loading