Fix bug in rule_query where text_expansion errored because it was not rewritten

[kderusso]

Query rules has a bug where a rule_query with an underlying organic text_expansion query will fail with the reason "failed to create query: text_expansion should have been rewritten to another query type".

The reason for this is that the text_expansion query will throw this as it expects to be rewritten as a boolean query of weighted tokens. However, when we change it to a pinned query during the rewrite phase, it never gets re-written again when the pinned query tries to generate a query.

Here's an example script to test the functionality. It requires an index pipeline set up using ELSER on a field called text.

// Index some sample data
POST /search-example/_doc/1?pipeline=search-example@ml-inference
{
  "text": "What is the weather in Jamaica?"
}
POST /search-example/_doc/2?pipeline=search-example@ml-inference
{
  "text": "Why is the sky blue?"
}
// Here we create a hidden document, with a different field. 
// This will not be returned by search results but will be returned by query rules 
POST /search-example/_doc/hidden
{
  "name": "Rule query pinned doc"
}

// Ruleset - we'll specify two rules: one that specifies IDs and one that specifies docs.
PUT /_query_rules/rules
{
  "ruleset_id": "rules",
  "rules": [
    {
      "rule_id": "pin_hidden_id",
      "type": "pinned",
      "criteria": [
        {
          "type": "exact",
          "metadata": "foo",
          "values": [
            "bar"
          ]
        }
      ],
      "actions": {
        "ids": [
          "hidden"
        ]
      }
    },
    {
      "rule_id": "pin_hidden_doc",
      "type": "pinned",
      "criteria": [
        {
          "type": "exact",
          "metadata": "foo",
          "values": [
            "baz"
          ]
        }
      ],
      "actions": {
        "docs": [
          {
            "_index": "search-example",
            "_id": "hidden"
          }
        ]
      }
    }
  ]
}

// Simple search examples verifying the rule query works with ID matches, docs matches, and no matches
POST search-example/_search
{
  "query": {
    "rule_query": {
      "organic": {
        "query_string": {
          "query": "blue",
          "default_field": "text"
        }
      },
      "ruleset_id": "rules",
      "match_criteria": {
        "foo": "bar"
      }
    }
  }, "_source": [ "text", "name" ]
}

POST search-example/_search
{
  "query": {
    "rule_query": {
      "organic": {
        "query_string": {
          "query": "blue",
          "default_field": "text"
        }
      },
      "ruleset_id": "rules",
      "match_criteria": {
        "foo": "baz"
      }
    }
  }, "_source": [ "text", "name" ]
}

POST search-example/_search
{
  "query": {
    "rule_query": {
      "organic": {
        "query_string": {
          "query": "blue",
          "default_field": "text"
        }
      },
      "ruleset_id": "rules",
      "match_criteria": {
        "foo": "nomatch"
      }
    }
  }, "_source": [ "text", "name" ]
}

// Example showing a rule query with a text_expansion query, this previously threw an error and now works - update the match criteria to test additional cases
POST search-example/_search
{
  "query": {
    "rule_query": {
      "organic": {
        "text_expansion": {
          "ml.inference.text_expanded.predicted_value": {
            "model_id": ".elser_model_2",
            "model_text": "Why is the sky blue?",
            "boost": 1
          }
        }
      },
      "ruleset_id": "rules",
      "match_criteria": {
        "foo": "bar"
      }
    }
  }, "_source": [ "text", "name" ]
}

[elasticsearchmachine]

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

[elasticsearchmachine]

Pinging @elastic/ent-search-eng (Team:Enterprise Search)

[saikatsarkar056]

LGTM. Just a non-blocking suggestion.

[jimczi]

I left some comments regarding the rewrite logic. I think we also need an integration test with an organic query that requires an async rewrite. If it's too difficult to test with an existing query with async rewrite, we can add a custom one for the test.

[jimczi]

I left more comments. We still needs an integration test to ensure that the bug is fixed. The unit tests are not enough here.

[jimczi]

Thanks for adding the yml tests! We're almost there ;). I left some comments to remove unneeded conditions and class members.
I also wonder if we can avoid the breaking change by being more lenient when dealing with queries coming from older nodes?

[kderusso]

@benwtrent had a great observation that these checks don't do anything because pinnedIds and docs are interstitial - moved these checks down into rewrite.

[benwtrent]

They are so interstitial, they are never set in this constructor, as once we have them (or not) we return a different kind of query.

Could we remove them from the ctor here? Its private, so not that big of a deal, but we always pass null, so seems like we could remove them.

[kderusso]

I was going to remove them, but then ran into the transport client changes, so figured I'd hold off for this PR...

[benwtrent]

Only one more concern around validation when both lists are non-empty. The rest is minor cleanup stuff. Looking much cleaner!

[benwtrent]

I don't see how these lists are ever null, down in the async action:

pinnedIdsSetOnce.set(appliedRules.pinnedIds().stream().distinct().toList());
pinnedDocsSetOnce.set(appliedRules.pinnedDocs().stream().distinct().toList());

If either were null, we would have thrown an NPE on List.stream()

So, once they are set, they aren't null. The only time the suppliers will return null from what I can tell is when the fetch hasn't finished yet.

[kderusso]

You're right, this is annoyingly defensive, I've cleaned it up a bit.

[benwtrent]

We should check

 if ((pinnedIds != null && pinnedIds.isEmpty() == false) && (pinnedDocs != null && pinnedDocs.isEmpty() == false)) {
            throw new IllegalArgumentException("applied rules contain both pinned ids and pinned docs, only one of ids or docs is allowed");
}

As that will never be checked and we arbitrarily (seemingly) identifiedPinnedIds first if both are supplied.

[benwtrent]

They are so interstitial, they are never set in this constructor, as once we have them (or not) we return a different kind of query.

Could we remove them from the ctor here? Its private, so not that big of a deal, but we always pass null, so seems like we could remove them.

[benwtrent]

There is still some minor cleanup to do, but I think this is fine to go.

Don't forget the test skips, I think you may need that for your new yaml test.

[benwtrent]

My guess is that this needs to have a skip version for the version that this is fixed in. Generally, how to handle this, is to make it restrictive on main (e.g. 8.13 only) (this may be even trickier if 8.13 branch is already cut, if so, this skip should initially be 8.14, and after backports, adjust). and in the back port adjust it to 8.12.2/3. Once the backport is merged, adjust main to 8.12.2/3.

[kderusso]

Thanks for this - I'll make sure to do this

[elasticsearchmachine]

💚 Backport successful

Status	Branch	Result
✅	8.13
✅	8.12