003_module_query.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "50a19ce1-bb8a-40c1-8ef3-ff78f60b71d9",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| default_exp query"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "885c918b-5721-4a0d-947e-62451a6d4c33",
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9ec96dcb-56ee-412f-bc63-253c5b1a19b3",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|hide\n",
    "from nbdev import nbdev_export\n",
    "from fastcore.test import test_eq, test_fail"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "14f00408-3ef3-475c-8ab4-eea43088ecae",
   "metadata": {},
   "source": [
    "# query\n",
    "> Reference API related to QueryModel related code"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7e79fb1c-ff24-4887-842d-6230f7c599f7",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export\n",
    "import os\n",
    "from typing import Optional, Dict, Callable, List, Tuple, Union\n",
    "from pandas import DataFrame\n",
    "from fastcore.utils import patch\n",
    "from vespa.io import VespaQueryResponse\n",
    "from vespa.application import Vespa"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "efbbdf6e-6c96-4b20-b6b5-4a8fb3ffd25e",
   "metadata": {},
   "source": [
    "## Match Filters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d3fca8bd-3b79-4d3d-a9d3-3a54a58d0bd2",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export\n",
    "class MatchFilter(object):\n",
    "    def __init__(self) -> None:    \n",
    "        \"Abstract class for match filters.\"\n",
    "        pass"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "84d88d28-4593-43e3-bcb6-f35cf49f5046",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export\n",
    "#|hide\n",
    "@patch\n",
    "def create_match_filter(\n",
    "    self: MatchFilter, \n",
    "    query: str  # Query input.\n",
    ") -> str:  # Part of the YQL expression related to the filter.\n",
    "    \"Abstract method to be implemented that creates part of the YQL expression related to the filter.\"\n",
    "    raise NotImplementedError"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2bfbe86e-d60e-4b12-96af-0c08ade2fbb6",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export\n",
    "#|hide\n",
    "@patch\n",
    "def get_query_properties(\n",
    "    self: MatchFilter, \n",
    "    query: Optional[str] = None  # Query input.\n",
    ") -> Dict:  # Contains the relevant request properties associated with the filter.\n",
    "    \"Abstract method to be implemented that get the relevant request properties associated with the filter.\"\n",
    "    raise NotImplementedError"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ffb3c718-0941-4ab8-bb30-82a61bbdd818",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export\n",
    "class AND(MatchFilter):\n",
    "    def __init__(self) -> None:\n",
    "        \"Filter that match document containing all the query terms.\"\n",
    "        super().__init__()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b96ffa07-0096-4993-9a59-c49938e74360",
   "metadata": {},
   "source": [
    "Usage: The `AND` filter is usually used when specifying query models."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "facca29b-b1f6-4996-a671-029e36956d79",
   "metadata": {},
   "outputs": [],
   "source": [
    "and_filter = AND()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "93f98a3c-01af-467d-91c0-b6a4ddd4538d",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export\n",
    "#|hide\n",
    "@patch\n",
    "def create_match_filter(\n",
    "    self: AND, \n",
    "    query: str  # Query input.  \n",
    ") -> str:  # Part of the YQL expression related to the AND filter.\n",
    "    \"Creates part of the YQL expression related to the AND filter\"\n",
    "    return '(userInput(\"{}\"))'.format(query)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0391132e-340c-4dde-a777-fcd7919d6cb3",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export\n",
    "#|hide\n",
    "@patch\n",
    "def get_query_properties(\n",
    "    self: AND, \n",
    "    query: Optional[str] = None  # Query input.\n",
    ") -> Dict:  # Get the relevant request properties associated with the AND filter.\n",
    "    \"Get the relevant request properties associated with the AND filter.\"\n",
    "    return {}\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d26c32f3-78de-4b39-80e3-bdf0683077de",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|hide\n",
    "query = \"this is  a test\"\n",
    "test_eq(\n",
    "    and_filter.create_match_filter(query),\n",
    "    '(userInput(\"this is  a test\"))',\n",
    ")\n",
    "test_eq(and_filter.get_query_properties(query), {})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c78a3e4f-a874-49cc-a4ad-9e8c1b545bb4",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export\n",
    "class OR(MatchFilter):\n",
    "    def __init__(self) -> None:\n",
    "        \"Filter that match any document containing at least one query term.\"\n",
    "        super().__init__()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bbf112de-fbdf-4cd8-b7ca-2b0f8e0a45c1",
   "metadata": {},
   "source": [
    "Usage: The `OR` filter is usually used when specifying query models."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fc9481c6-0314-4905-83ee-f76b27c7b833",
   "metadata": {},
   "outputs": [],
   "source": [
    "or_filter = OR()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "122c8920-60ae-4e20-9362-2a85f26a1f34",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export \n",
    "#|hide\n",
    "@patch\n",
    "def create_match_filter(\n",
    "    self: OR, \n",
    "    query: str  # Query input.\n",
    ") -> str:  # Part of the YQL expression related to the OR filter.\n",
    "    \"Creates part of the YQL expression related to the OR filter\"    \n",
    "    return '({{grammar: \"any\"}}userInput(\"{}\"))'.format(query)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fc2fb933-0110-4d5b-849c-ea860883d5f4",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export\n",
    "#|hide\n",
    "@patch\n",
    "def get_query_properties(\n",
    "    self: OR, \n",
    "    query: Optional[str] = None  # Query input.\n",
    ") -> Dict:  # Get the relevant request properties associated with the OR filter.\n",
    "    \"Get the relevant request properties associated with the OR filter.\"    \n",
    "    return {}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "32a60e35-7a05-45a1-ba67-0d7c1fc9ca76",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|hide\n",
    "query = \"this is  a test\"\n",
    "test_eq(\n",
    "    or_filter.create_match_filter(query=query),\n",
    "    '({grammar: \"any\"}userInput(\"this is  a test\"))',\n",
    ")\n",
    "test_eq(or_filter.get_query_properties(query=query), {})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "495f7ea7-15d4-4d25-900e-dd48f4120b41",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export\n",
    "class WeakAnd(MatchFilter):\n",
    "    def __init__(\n",
    "        self, \n",
    "        hits: int,  # Lower bound on the number of hits to be retrieved. \n",
    "        field: str = \"default\"  # Which Vespa field to search.\n",
    "    ) -> None:\n",
    "        \"\"\"\n",
    "        Match documents according to the weakAND algorithm.\n",
    "\n",
    "        Reference: [https://docs.vespa.ai/en/using-wand-with-vespa.html](https://docs.vespa.ai/en/using-wand-with-vespa.html)\n",
    "        \"\"\"\n",
    "        super().__init__()\n",
    "        self.hits = hits\n",
    "        self.field = field"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2f6cf082-f8ae-4d2e-83e0-8015e6ea8efc",
   "metadata": {},
   "source": [
    "Usage: The `WeakAnd` filter is usually used when specifying query models."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "89d17057-3d06-47b8-8f0a-cd96bdeaebb0",
   "metadata": {},
   "outputs": [],
   "source": [
    "weakand_filter = WeakAnd(hits=10, field=\"default\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "316b092c-bca6-4984-85b2-3107ced59a90",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export\n",
    "#|hide\n",
    "@patch\n",
    "def create_match_filter(\n",
    "    self: WeakAnd, \n",
    "    query: str  # Query input.\n",
    ") -> str:  # Part of the YQL expression related to the WeakAnd filter.\n",
    "    \"Creates part of the YQL expression related to the WeakAnd filter\"\n",
    "    return '({{grammar: \"weakAnd\", targetHits: {}, defaultIndex: \"{}\"}}userInput(\"{}\"))'.format(self.hits, self.field, query)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "79c380d1-e939-403e-a41c-c333f8f5d5ca",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export\n",
    "#|hide\n",
    "@patch\n",
    "def get_query_properties(\n",
    "    self: WeakAnd, \n",
    "    query: Optional[str] = None  # Query input.\n",
    ") -> Dict:  # Get the relevant request properties associated with the WeakAnd filter.\n",
    "    \"Get the relevant request properties associated with the WeakAnd filter.\"        \n",
    "    return {}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ec13431a-97ea-405d-86be-5ccb75e616b3",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|hide\n",
    "weakand_filter = WeakAnd(hits=10, field=\"field_name\")\n",
    "query = \"this is  a test\"\n",
    "test_eq(\n",
    "    weakand_filter.create_match_filter(query=query),\n",
    "    '({grammar: \"weakAnd\", targetHits: 10, defaultIndex: \"field_name\"}userInput(\"this is  a test\"))',\n",
    ")\n",
    "test_eq(weakand_filter.get_query_properties(query=query), {})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9caa69a6-3374-4bcf-ba66-297c76e6d705",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export\n",
    "class Tokenize(MatchFilter):\n",
    "    def __init__(\n",
    "        self, \n",
    "        hits: int,  # Lower bound on the number of hits to be retrieved. \n",
    "        field: str = \"default\"  # Which Vespa field to search.\n",
    "    ) -> None:\n",
    "        \"\"\"\n",
    "        Match documents according to the weakAND algorithm without parsing specials characters.\n",
    "\n",
    "        Reference: [https://docs.vespa.ai/en/reference/simple-query-language-reference.html](https://docs.vespa.ai/en/reference/simple-query-language-reference.html)\n",
    "        \"\"\"\n",
    "        super().__init__()\n",
    "        self.hits = hits\n",
    "        self.field = field"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "42d4a362-bb24-4188-a878-ac1bc8097d5d",
   "metadata": {},
   "source": [
    "Usage: The `Tokenize` filter is usually used when specifying query models."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a06729ff-1896-4842-bced-900c66d8716e",
   "metadata": {},
   "outputs": [],
   "source": [
    "tokenize_filter = Tokenize(hits=10, field=\"default\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5ee24942-77a7-4e2c-95f8-b6edb53ec917",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export\n",
    "#|hide\n",
    "@patch\n",
    "def create_match_filter(\n",
    "    self: Tokenize, \n",
    "    query: str  # Query input.\n",
    ") -> str:  # Part of the YQL expression related to the Tokenizer filter.\n",
    "    \"Creates part of the YQL expression related to the Tokenizer filter\"\n",
    "    return '({{grammar: \"tokenize\", targetHits: {}, defaultIndex: \"{}\"}}userInput(\"{}\"))'.format(self.hits, self.field, query)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d7a6ff6a-4ff1-4271-bb46-80a7f0590301",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export\n",
    "#|hide\n",
    "@patch\n",
    "def get_query_properties(\n",
    "    self: Tokenize, \n",
    "    query: Optional[str] = None  # Query input.\n",
    ") -> Dict:  # Get the relevant request properties associated with the Tokenize filter.\n",
    "    \"Get the relevant request properties associated with the Tokenize filter.\"        \n",
    "    return {}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ad06a82c-c8fa-47fd-b841-d59d312a491a",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|hide\n",
    "tokenize_filter = Tokenize(hits=10, field=\"field_name\")\n",
    "query = \"this is  a test\"\n",
    "test_eq(\n",
    "    tokenize_filter.create_match_filter(query=query),\n",
    "    '({grammar: \"tokenize\", targetHits: 10, defaultIndex: \"field_name\"}userInput(\"this is  a test\"))',\n",
    ")\n",
    "test_eq(tokenize_filter.get_query_properties(query=query), {})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eae83e4d-5e73-430b-89fa-6b8dd5b4b605",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export\n",
    "class ANN(MatchFilter):\n",
    "    def __init__(\n",
    "        self,\n",
    "        doc_vector: str,  # Name of the document field to be used in the distance calculation.\n",
    "        query_vector: str,  # Name of the query field to be used in the distance calculation.\n",
    "        hits: int,  # Lower bound on the number of hits to return.\n",
    "        label: str,  # A label to identify this specific operator instance.\n",
    "        approximate: bool = True,  # True to use approximate nearest neighbor and False to use brute force. Default to True.\n",
    "    ) -> None:\n",
    "        \"\"\"\n",
    "        Match documents according to the nearest neighbor operator.\n",
    "\n",
    "        Reference: [https://docs.vespa.ai/en/reference/query-language-reference.html](https://docs.vespa.ai/en/reference/query-language-reference.html)\n",
    "        \"\"\"\n",
    "        super().__init__()\n",
    "        self.doc_vector = doc_vector\n",
    "        self.query_vector = query_vector\n",
    "        self.hits = hits\n",
    "        self.label = label\n",
    "        self.approximate = approximate\n",
    "        self._approximate = \"true\" if self.approximate is True else \"false\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "aad13e04-afe4-43bf-8788-f32e46e5df7d",
   "metadata": {},
   "source": [
    "Usage: The `ANN` filter is usually used when specifying query models."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1402c564-3beb-42c1-8e30-24ab5289dadd",
   "metadata": {},
   "source": [
    "By default, the `ANN` operator uses approximate nearest neighbor:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e2194510-e419-4506-a2c8-7491e0693b67",
   "metadata": {},
   "outputs": [],
   "source": [
    "match_filter = ANN(\n",
    "    doc_vector=\"doc_vector\",\n",
    "    query_vector=\"query_vector\",\n",
    "    hits=10,\n",
    "    label=\"label\",\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "539c5802-d82f-4c07-a056-df6260962331",
   "metadata": {},
   "source": [
    "Brute-force can be used by specifying `approximate=False`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cf93d47c-7ee9-49b7-a736-cc8475b6bc5a",
   "metadata": {},
   "outputs": [],
   "source": [
    "ann_filter = ANN(\n",
    "    doc_vector=\"doc_vector\",\n",
    "    query_vector=\"query_vector\",\n",
    "    hits=10,\n",
    "    label=\"label\",\n",
    "    approximate=False,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d7a50b6b-3ca9-48c0-8e4e-66f6c9303a73",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export\n",
    "#|hide\n",
    "@patch\n",
    "def create_match_filter(\n",
    "    self: ANN, \n",
    "    query: str  # Query input is ignored in the ANN case.\n",
    ") -> str:  # Part of the YQL expression related to the ANN filter.\n",
    "    \"Creates part of the YQL expression related to the ANN filter\"    \n",
    "    return '({{targetHits: {}, label: \"{}\", approximate: {}}}nearestNeighbor({}, {}))'.format(\n",
    "        self.hits, self.label, self._approximate, self.doc_vector, self.query_vector\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dbc49f58-0371-4edb-a584-bc810cb005d5",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export\n",
    "#|hide\n",
    "@patch\n",
    "def get_query_properties(\n",
    "    self: ANN, \n",
    "    query: Optional[str] = None  # Query input is ignored in the ANN case.\n",
    ") -> Dict[str, str]:  # Get the relevant request properties associated with the ANN filter.\n",
    "    \"Get the relevant request properties associated with the ANN filter.\"            \n",
    "    return {}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cade5791-3ab6-40e5-b1b4-c11bb15c11d1",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|hide\n",
    "query = \"query string has no effect here\"\n",
    "match_filter = ANN(\n",
    "    doc_vector=\"doc_vector\",\n",
    "    query_vector=\"query_vector\",\n",
    "    hits=10,\n",
    "    label=\"label\",\n",
    ")\n",
    "test_eq(\n",
    "    match_filter.create_match_filter(query=query),\n",
    "    '({targetHits: 10, label: \"label\", approximate: true}nearestNeighbor(doc_vector, query_vector))',\n",
    ")\n",
    "test_eq(\n",
    "    match_filter.get_query_properties(query=query),\n",
    "    {},\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7c739915-ec55-4935-98fc-c0a4fa222e78",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|hide\n",
    "query = \"query string has no effect here\"\n",
    "match_filter = ANN(\n",
    "    doc_vector=\"doc_vector\",\n",
    "    query_vector=\"query_vector\",\n",
    "    hits=10,\n",
    "    label=\"label\",\n",
    "    approximate=False,\n",
    ")\n",
    "test_eq(\n",
    "    match_filter.create_match_filter(query=query),\n",
    "    '({targetHits: 10, label: \"label\", approximate: false}nearestNeighbor(doc_vector, query_vector))',\n",
    ")\n",
    "test_eq(\n",
    "    match_filter.get_query_properties(query=query),\n",
    "    {},\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7987b2ec-2a4e-4118-99d9-73857fae0495",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export\n",
    "class Union(MatchFilter):\n",
    "    def __init__(\n",
    "        self, \n",
    "        *args: MatchFilter  # Match filters to be taken the union of.\n",
    "    ) -> None:\n",
    "        \"Match documents that belongs to the union of many match filters.\"\n",
    "        super().__init__()\n",
    "        self.operators = args"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "95ec5c5f-31e2-40a1-b6cb-0261cf52f0e6",
   "metadata": {},
   "source": [
    "Usage: The `Union` filter is usually used when specifying query models."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "94b0c764-cc97-4153-b140-f2a693613a23",
   "metadata": {},
   "outputs": [],
   "source": [
    "union_filter = Union(\n",
    "    WeakAnd(hits=10, field=\"field_name\"),\n",
    "    ANN(\n",
    "        doc_vector=\"doc_vector\",\n",
    "        query_vector=\"query_vector\",\n",
    "        hits=10,\n",
    "        label=\"label\",\n",
    "    ),\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2b6273a2-9cf3-4db1-9f1c-ca7b42b75aa5",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export\n",
    "#|hide\n",
    "@patch\n",
    "def create_match_filter(\n",
    "    self: Union, \n",
    "    query: str  # Query input.\n",
    ") -> str:  # Part of the YQL expression related to the Union filter.\n",
    "    \"Creates part of the YQL expression related to the Union filter\"    \n",
    "    match_filters = []\n",
    "    for operator in self.operators:\n",
    "        match_filter = operator.create_match_filter(query=query)\n",
    "        if match_filter is not None:\n",
    "            match_filters.append(match_filter)\n",
    "    return \" or \".join(match_filters)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4af2916a-cced-4940-96a9-e673877b2b41",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export\n",
    "#|hide\n",
    "@patch\n",
    "def get_query_properties(\n",
    "    self: Union,  # Query input. \n",
    "    query: Optional[str] = None  # Get the relevant request properties associated with the Union filter.\n",
    ") -> Dict[str, str]:  # Get the relevant request properties associated with the Union filter.\n",
    "    query_properties = {}\n",
    "    for operator in self.operators:\n",
    "        query_properties.update(operator.get_query_properties(query=query))\n",
    "    return query_properties"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "226d79a4-9cef-430d-b2a3-83563e36632c",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|hide\n",
    "query = \"this is  a test\"\n",
    "match_filter = Union(\n",
    "    WeakAnd(hits=10, field=\"field_name\"),\n",
    "    ANN(\n",
    "        doc_vector=\"doc_vector\",\n",
    "        query_vector=\"query_vector\",\n",
    "        hits=10,\n",
    "        label=\"label\",\n",
    "    ),\n",
    ")\n",
    "test_eq(\n",
    "    match_filter.create_match_filter(query=query),\n",
    "    '({grammar: \"weakAnd\", targetHits: 10, defaultIndex: \"field_name\"}userInput(\"this is  a test\")) or '\n",
    "    '({targetHits: 10, label: \"label\", approximate: true}nearestNeighbor(doc_vector, query_vector))',\n",
    ")\n",
    "test_eq(\n",
    "    match_filter.get_query_properties(query=query),\n",
    "    {},\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fa7646f1-a275-485e-9e80-0e83652017be",
   "metadata": {},
   "source": [
    "## Ranking"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "48649f3c-4e72-439d-b783-e58acddc7fc7",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export\n",
    "class Ranking(object):\n",
    "    def __init__(\n",
    "        self, \n",
    "        name: str = \"default\",  # Name of the rank profile as defined in a Vespa search definition.\n",
    "        list_features: bool = False  # Should the ranking features be returned. Either 'true' or 'false'.\n",
    "    ) -> None:\n",
    "        \"Define the rank profile to be used during ranking.\"\n",
    "        self.name = name\n",
    "        self.list_features = \"false\"\n",
    "        if list_features:\n",
    "            self.list_features = \"true\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c7b328b6-fdae-4ca5-bfcd-a728660b9cad",
   "metadata": {},
   "source": [
    "Usage: `Ranking` is usually used when specifying query models."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f3ebc369-f585-4d1e-b1ae-b244e46b97e4",
   "metadata": {},
   "outputs": [],
   "source": [
    "ranking = Ranking(name=\"bm25\", list_features=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cfd02daf-68b3-4538-8c85-96dfb19c9144",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|hide\n",
    "ranking = Ranking(name=\"rank_profile\", list_features=True)\n",
    "test_eq(ranking.name, \"rank_profile\")\n",
    "test_eq(ranking.list_features, \"true\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b95a90b9-5349-47cd-a132-cf204798f58f",
   "metadata": {},
   "source": [
    "## Query properties"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a256f850-57f3-4550-9e21-a58ff39a968a",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export\n",
    "class QueryProperty(object):\n",
    "    def __init__(self) -> None:    \n",
    "        \"Abstract class for query property.\"\n",
    "        pass    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "93a459f4-f8e6-41bf-a9a2-648086ce405f",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export\n",
    "#|hide\n",
    "@patch\n",
    "def get_query_properties(\n",
    "    self: QueryProperty, \n",
    "    query: Optional[str] = None  # Query input.\n",
    ") -> Dict:  # Contains the relevant request properties to be included in the query.\n",
    "    \"Extract query property syntax.\"\n",
    "    raise NotImplementedError\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fd4fb173-39f3-40ad-83de-06598d29af9e",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export\n",
    "class QueryRankingFeature(QueryProperty):\n",
    "    def __init__(\n",
    "        self,\n",
    "        name: str,  # Name of the feature.\n",
    "        mapping: Callable[[str], List[float]],  # Function mapping a string to a list of floats.\n",
    "    ) -> None:\n",
    "        \"Include ranking.feature.query into a Vespa query.\"\n",
    "        super().__init__()\n",
    "        self.name = name\n",
    "        self.mapping = mapping"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b2b3f809-54bd-443c-adda-9e15ae66c555",
   "metadata": {},
   "source": [
    "Usage: `QueryRankingFeature` is usually used when specifying query models."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a58dc22b-f238-4d3b-801e-eb2f11c48e74",
   "metadata": {},
   "outputs": [],
   "source": [
    "query_property = QueryRankingFeature(\n",
    "    name=\"query_vector\", mapping=lambda x: [1, 2, 3]\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "375441bd-cdc6-4bb9-971b-3e22442cd720",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export\n",
    "#|hide\n",
    "@patch\n",
    "def get_query_properties(\n",
    "    self: QueryRankingFeature, \n",
    "    query: Optional[str] = None  # Query input.\n",
    ") -> Dict[str, str]:  # Contains the relevant request properties to be included in the query.\n",
    "    value = self.mapping(query)\n",
    "    return {\"ranking.features.query({})\".format(self.name): str(value)}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ebd9ffd0-344a-4f84-a3b6-0f84aef3079c",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|hide\n",
    "query = \"this is  a test\"\n",
    "query_property = QueryRankingFeature(\n",
    "    name=\"query_vector\", mapping=lambda x: [1, 2, 3]\n",
    ")\n",
    "test_eq(\n",
    "    query_property.get_query_properties(query=query),\n",
    "    {\"ranking.features.query(query_vector)\": \"[1, 2, 3]\"},\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f5ef348d-55f8-4ea4-a65f-df57db794a4a",
   "metadata": {},
   "source": [
    "## Query model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "86fa85c2-9708-49a0-b847-1454f3a55af4",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export\n",
    "class QueryModel(object):\n",
    "    def __init__(\n",
    "        self,\n",
    "        name: str = \"default_name\",  # Name of the query model. Used to tag model-related quantities, like evaluation metrics.\n",
    "        query_properties: Optional[List[QueryProperty]] = None,  # Query properties to be included in the queries.\n",
    "        match_phase: MatchFilter = AND(),  # Define the match criteria.\n",
    "        ranking: Ranking = Ranking(),  # Define the rank criteria.\n",
    "        body_function: Optional[Callable[[str], Dict]] = None,  # Function that take query as parameter and returns the body of a Vespa query.\n",
    "    ) -> None:\n",
    "        \"\"\"\n",
    "        Define a query model.\n",
    "\n",
    "        A `QueryModel` is an abstraction that encapsulates all the relevant information\n",
    "        controlling how a Vespa app matches and ranks documents.\n",
    "        \"\"\"\n",
    "        self.name = name\n",
    "        self.query_properties = query_properties if query_properties is not None else []\n",
    "        self.match_phase = match_phase\n",
    "        self.ranking = ranking\n",
    "        self.body_function = body_function\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bc616185-3d29-4c1c-aa8d-7680075f023b",
   "metadata": {},
   "source": [
    "Usage:"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1301c52c-7072-4647-89b5-52196e78f536",
   "metadata": {},
   "source": [
    "Specify a query model with default configurations:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a724cef8-4ba5-45e0-824e-e0ed56cbcab5",
   "metadata": {},
   "outputs": [],
   "source": [
    "query_model = QueryModel()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "61bc3c86-eb1a-4a99-a487-707bf1c19039",
   "metadata": {},
   "source": [
    "Specify match phase, ranking phase and properties used by them."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a72fdfce-9cc9-436f-b4f1-193b3a675d63",
   "metadata": {},
   "outputs": [],
   "source": [
    "query_model = QueryModel(\n",
    "    query_properties=[\n",
    "        QueryRankingFeature(name=\"query_embedding\", mapping=lambda x: [1, 2, 3])\n",
    "    ],\n",
    "    match_phase=ANN(\n",
    "        doc_vector=\"document_embedding\",\n",
    "        query_vector=\"query_embedding\",\n",
    "        hits=10,\n",
    "        label=\"label\",\n",
    "    ),\n",
    "    ranking=Ranking(name=\"bm25_plus_embeddings\", list_features=True),\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fcdaa003-b709-44ee-b771-4a54310a7d8f",
   "metadata": {},
   "source": [
    "Specify a query model based on a function that output Vespa YQL."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "30409603-be14-46f5-b330-c1c491072436",
   "metadata": {},
   "outputs": [],
   "source": [
    "def body_function(query):\n",
    "    body = {\n",
    "        \"yql\": \"select * from sources * where userQuery();\",\n",
    "        \"query\": query,\n",
    "        \"type\": \"any\",\n",
    "        \"ranking\": {\"profile\": \"bm25\", \"listFeatures\": \"true\"},\n",
    "    }\n",
    "    return body\n",
    "\n",
    "query_model = QueryModel(body_function=body_function)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6bd913d5-3fae-4de5-867e-12e092a1b029",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export\n",
    "#|hide\n",
    "@patch\n",
    "def create_body(\n",
    "    self: QueryModel, \n",
    "    query: str  # Query string.\n",
    ") -> Dict[str, str]:  # Request body\n",
    "    \"Create the appropriate request body to be sent to Vespa.\"\n",
    "\n",
    "    if self.body_function:\n",
    "        body = self.body_function(query)\n",
    "        return body\n",
    "\n",
    "    query_properties = {}\n",
    "    for query_property in self.query_properties:\n",
    "        query_properties.update(query_property.get_query_properties(query=query))\n",
    "    query_properties.update(self.match_phase.get_query_properties(query=query))\n",
    "\n",
    "    match_filter = self.match_phase.create_match_filter(query=query)\n",
    "\n",
    "    body = {\n",
    "        \"yql\": \"select * from sources * where {};\".format(match_filter),\n",
    "        \"ranking\": {\n",
    "            \"profile\": self.ranking.name,\n",
    "            \"listFeatures\": self.ranking.list_features,\n",
    "        },\n",
    "    }\n",
    "    body.update(query_properties)\n",
    "    return body"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bc9aba22-b4a9-4078-bebe-4f1158cdc587",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|hide\n",
    "query = \"this is  a test\"\n",
    "query_model = QueryModel()\n",
    "test_eq(\n",
    "    query_model.create_body(query=query),\n",
    "    {\n",
    "        \"yql\": 'select * from sources * where (userInput(\"this is  a test\"));',\n",
    "        \"ranking\": {\"profile\": \"default\", \"listFeatures\": \"false\"},\n",
    "    },\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "690ff95c-8ac0-4b2c-806e-bcf9fe8099dc",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|hide\n",
    "query = \"this is  a test\"\n",
    "def body_function(query):\n",
    "    body = {\n",
    "        \"yql\": \"select * from sources * where userQuery();\",\n",
    "        \"query\": query,\n",
    "        \"type\": \"any\",\n",
    "        \"ranking\": {\"profile\": \"bm25\", \"listFeatures\": \"true\"},\n",
    "    }\n",
    "    return body\n",
    "\n",
    "query_model = QueryModel(body_function=body_function)\n",
    "test_eq(\n",
    "    query_model.create_body(query=query),\n",
    "    {\n",
    "        \"yql\": \"select * from sources * where userQuery();\",\n",
    "        \"query\": \"this is  a test\",\n",
    "        \"type\": \"any\",\n",
    "        \"ranking\": {\"profile\": \"bm25\", \"listFeatures\": \"true\"},\n",
    "    },\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e1af818a-c1b3-47af-bb07-218b907d67da",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|hide\n",
    "query = \"this is  a test\"\n",
    "query_model = QueryModel(\n",
    "    query_properties=[\n",
    "        QueryRankingFeature(name=\"query_vector\", mapping=lambda x: [1, 2, 3])\n",
    "    ],\n",
    "    match_phase=OR(),\n",
    "    ranking=Ranking(name=\"bm25\", list_features=True),\n",
    ")\n",
    "test_eq(\n",
    "    query_model.create_body(query=query),\n",
    "    {\n",
    "        \"yql\": 'select * from sources * where ({grammar: \"any\"}userInput(\"this is  a test\"));',\n",
    "        \"ranking\": {\"profile\": \"bm25\", \"listFeatures\": \"true\"},\n",
    "        \"ranking.features.query(query_vector)\": \"[1, 2, 3]\",\n",
    "    },\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d9cf938e-ec21-4a05-bf8c-e5d746b9315f",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|hide\n",
    "query = \"this is  a test\"\n",
    "query_model = QueryModel(\n",
    "    query_properties=[\n",
    "        QueryRankingFeature(name=\"query_vector\", mapping=lambda x: [1, 2, 3])\n",
    "    ],\n",
    "    match_phase=ANN(\n",
    "        doc_vector=\"doc_vector\",\n",
    "        query_vector=\"query_vector\",\n",
    "        hits=10,\n",
    "        label=\"label\",\n",
    "    ),\n",
    "    ranking=Ranking(name=\"bm25\", list_features=True),\n",
    ")\n",
    "test_eq(\n",
    "    query_model.create_body(query=query),\n",
    "    {\n",
    "        \"yql\": 'select * from sources * where ({targetHits: 10, label: \"label\", approximate: true}nearestNeighbor(doc_vector, query_vector));',\n",
    "        \"ranking\": {\"profile\": \"bm25\", \"listFeatures\": \"true\"},\n",
    "        \"ranking.features.query(query_vector)\": \"[1, 2, 3]\",\n",
    "    },\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fb640d9a-78d9-4b57-ab53-74ab4cf02814",
   "metadata": {},
   "source": [
    "## Send query with QueryModel"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "74f45218-818c-4842-bf33-3cbc22666488",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|hide"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "917e3c3a-a5d1-4056-a417-1be074f8fb54",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export\n",
    "def _build_query_body(\n",
    "    query: str,\n",
    "    query_model: QueryModel,\n",
    "    recall: Optional[Tuple] = None,\n",
    "    **kwargs,\n",
    ") -> Dict:\n",
    "    assert query_model is not None, \"No 'query_model' specified.\"\n",
    "    body = query_model.create_body(query=query)\n",
    "    if recall is not None:\n",
    "        body.update(\n",
    "            {\n",
    "                \"recall\": \"+(\"\n",
    "                + \" \".join(\n",
    "                    [\"{}:{}\".format(recall[0], str(doc)) for doc in recall[1]]\n",
    "                )\n",
    "                + \")\"\n",
    "            }\n",
    "        )\n",
    "    body.update(kwargs)\n",
    "    return body"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d5b9888d-36e7-4ecf-97f9-9914f7729ac1",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export\n",
    "def send_query(\n",
    "    app: Vespa,  # Connection to a Vespa application\n",
    "    body: Optional[Dict] = None,  # Contains all the request parameters. None when using `query_model`.\n",
    "    query: Optional[str] = None,  # Query string. None when using `body`.\n",
    "    query_model: Optional[QueryModel] = None,  # Query model. None when using `body`.\n",
    "    debug_request: bool = False,  # Return request body for debugging instead of sending the request.\n",
    "    recall: Optional[Tuple] = None,  # Tuple of size 2 where the first element is the name of the field to use to recall and the second element is a list of the values to be recalled.\n",
    "    **kwargs,  # Additional parameters to be sent along the request.\n",
    ") -> VespaQueryResponse:  # Either the request body if debug_request is True or the result from the Vespa application.\n",
    "    \"\"\"\n",
    "    Send a query request to a Vespa application.\n",
    "\n",
    "    Either send 'body' containing all the request parameters or specify 'query' and 'query_model'.\n",
    "    \"\"\"\n",
    "    body = (\n",
    "        _build_query_body(query, query_model, recall, **kwargs)\n",
    "        if body is None\n",
    "        else body\n",
    "    )\n",
    "    if debug_request:\n",
    "        return VespaQueryResponse(\n",
    "            json={}, status_code=None, url=None, request_body=body\n",
    "        )\n",
    "    else:\n",
    "        return app.query(body=body)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "68cc8a08-bd1a-496a-9cb3-42fe3ad4fdda",
   "metadata": {},
   "source": [
    "Usage: Assume `app` is a Vespa connection."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5a211d10-bc34-4a3c-8509-d210b8cc2637",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|hide\n",
    "# setup a vespa app\n",
    "from vespa.package import ApplicationPackage, Field, FieldSet, RankProfile\n",
    "from vespa.deployment import VespaDocker"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8365236c-3b1d-40b8-aadf-42403e157d03",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Waiting for configuration server, 0/300 seconds...\n",
      "Waiting for configuration server, 5/300 seconds...\n",
      "Waiting for configuration server, 10/300 seconds...\n",
      "Waiting for application status, 0/300 seconds...\n",
      "Waiting for application status, 5/300 seconds...\n",
      "Waiting for application status, 10/300 seconds...\n",
      "Waiting for application status, 15/300 seconds...\n",
      "Waiting for application status, 20/300 seconds...\n",
      "Waiting for application status, 25/300 seconds...\n",
      "Waiting for application status, 30/300 seconds...\n",
      "Finished deployment.\n",
      "Successful documents fed: 10/10.\n",
      "Batch progress: 1/1.\n"
     ]
    }
   ],
   "source": [
    "#|hide\n",
    "app_package = ApplicationPackage(name=\"QueryApp\")\n",
    "app_package.schema.add_fields(\n",
    "    Field(name=\"doc_id\", type=\"string\", indexing=[\"attribute\", \"summary\"]),\n",
    "    Field(name=\"text\", type=\"string\", indexing=[\"index\", \"summary\"], index=\"enable-bm25\"),\n",
    ")\n",
    "app_package.schema.add_field_set(\n",
    "    FieldSet(name=\"default\", fields=[\"text\"])\n",
    ")\n",
    "app_package.schema.add_rank_profile(\n",
    "    RankProfile(\n",
    "        name=\"bm25\",\n",
    "        first_phase=\"bm25(text)\",\n",
    "        summary_features=[\"bm25(text)\"]\n",
    "    )\n",
    ")\n",
    "vespa_docker = VespaDocker(port=8083, cfgsrv_port=19073)\n",
    "app = vespa_docker.deploy(application_package=app_package)\n",
    "df = DataFrame(\n",
    "    data={\n",
    "        \"doc_id\": [f\"{i}\" for i in range(10)],\n",
    "        \"text\": [f\"this is title {i}\" for i in range(10)],\n",
    "    }\n",
    ")\n",
    "responses = app.feed_df(\n",
    "    df=df, \n",
    "    include_id=True, \n",
    "    id_field=\"doc_id\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "db28f211-3393-4a99-bdec-e42fd8088ac0",
   "metadata": {},
   "source": [
    "Send request body."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cf300632-d860-4786-a02d-324ff1249767",
   "metadata": {},
   "outputs": [],
   "source": [
    "body = {\"yql\": \"select * from sources * where test\"}\n",
    "result = send_query(app=app, body=body)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d750881e-2488-4f11-8879-ea057df565bf",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|hide\n",
    "app_test = Vespa(url=\"http://localhost\", port=8080)\n",
    "\n",
    "body = {\"yql\": \"select * from sources * where test\"}\n",
    "test_eq(\n",
    "    send_query(app=app_test, body=body, debug_request=True).request_body, body\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "abee0be3-a4d5-44db-b016-3058c37ce6bc",
   "metadata": {},
   "source": [
    "Use `query` and `query_model`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "89517c3b-1667-4fb6-bfeb-62f6a71ee099",
   "metadata": {},
   "outputs": [],
   "source": [
    "result = send_query(\n",
    "    app=app,\n",
    "    query=\"this is a test\",\n",
    "    query_model=QueryModel(\n",
    "        match_phase=OR(), \n",
    "        ranking=Ranking()\n",
    "    ),\n",
    "    hits=10,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "10999e9e-d89c-46f3-8492-2d81b2474c79",
   "metadata": {},
   "source": [
    "Debug the output of the `QueryModel` by setting `debug_request=True`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "306e05b3-f474-4c23-b881-63a95e1a43a0",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'yql': 'select * from sources * where ({grammar: \"any\"}userInput(\"this is a test\"));',\n",
       " 'ranking': {'profile': 'default', 'listFeatures': 'false'},\n",
       " 'hits': 10}"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "send_query(\n",
    "    app=app,\n",
    "    query=\"this is a test\",\n",
    "    query_model=QueryModel(match_phase=OR(), ranking=Ranking()),\n",
    "    debug_request=True,\n",
    "    hits=10,\n",
    ").request_body"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0ac49a61-3f4b-413b-aa34-da2653e9e3f4",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|hide\n",
    "test_eq(\n",
    "    send_query(\n",
    "        app=app_test,\n",
    "        query=\"this is a test\",\n",
    "        query_model=QueryModel(match_phase=OR(), ranking=Ranking()),\n",
    "        debug_request=True,\n",
    "        hits=10,\n",
    "    ).request_body,\n",
    "    {\n",
    "        \"yql\": 'select * from sources * where ({grammar: \"any\"}userInput(\"this is a test\"));',\n",
    "        \"ranking\": {\"profile\": \"default\", \"listFeatures\": \"false\"},\n",
    "        \"hits\": 10,\n",
    "    },\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2e310a9c-9a4e-46bb-9086-5a6726879719",
   "metadata": {},
   "source": [
    "Recall documents using the `id` field:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "34247c77-e7c1-49e8-9c1b-88185ad6fa4b",
   "metadata": {},
   "outputs": [],
   "source": [
    "result = send_query(\n",
    "    app=app,\n",
    "    query=\"this is a test\",\n",
    "    query_model=QueryModel(match_phase=OR(), ranking=Ranking()),\n",
    "    hits=10,\n",
    "    recall=(\"id\", [1, 5]),\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "341a3fb8-f5e6-485f-a6a6-4689a3f70070",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|hide\n",
    "test_eq(\n",
    "    send_query(\n",
    "        app=app_test,\n",
    "        query=\"this is a test\",\n",
    "        query_model=QueryModel(match_phase=OR(), ranking=Ranking()),\n",
    "        debug_request=True,\n",
    "        hits=10,\n",
    "        recall=(\"id\", [1, 5]),\n",
    "    ).request_body,\n",
    "    {\n",
    "        \"yql\": 'select * from sources * where ({grammar: \"any\"}userInput(\"this is a test\"));',\n",
    "        \"ranking\": {\"profile\": \"default\", \"listFeatures\": \"false\"},\n",
    "        \"hits\": 10,\n",
    "        \"recall\": \"+(id:1 id:5)\",\n",
    "    },\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "007aa33e-8d66-4c55-a0ad-360426223edb",
   "metadata": {},
   "source": [
    "Use a `body_function` to specify a `QueryModel`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "84320169-561f-4013-83d1-a0879eaf8bac",
   "metadata": {},
   "outputs": [],
   "source": [
    "def body_function(query):\n",
    "    body = {\n",
    "        \"yql\": \"select * from sources * where userQuery();\",\n",
    "        \"query\": query,\n",
    "        \"type\": \"any\",\n",
    "        \"ranking\": {\"profile\": \"bm25\", \"listFeatures\": \"true\"},\n",
    "    }\n",
    "    return body\n",
    "\n",
    "query_model = QueryModel(body_function=body_function)\n",
    "\n",
    "result = send_query(\n",
    "        app=app,\n",
    "        query=\"this is a test\",\n",
    "        query_model=query_model,\n",
    "        hits=10\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6f8c98a8-f959-4c66-851a-e1e8670c388a",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|hide\n",
    "def body_function(query):\n",
    "    body = {\n",
    "        \"yql\": \"select * from sources * where userQuery();\",\n",
    "        \"query\": query,\n",
    "        \"type\": \"any\",\n",
    "        \"ranking\": {\"profile\": \"bm25\", \"listFeatures\": \"true\"},\n",
    "    }\n",
    "    return body\n",
    "\n",
    "query_model = QueryModel(body_function=body_function)\n",
    "\n",
    "test_eq(\n",
    "    send_query(\n",
    "        app=app,\n",
    "        query=\"this is a test\",\n",
    "        query_model=query_model,\n",
    "        debug_request=True,\n",
    "        hits=10,\n",
    "        recall=(\"id\", [1, 5]),\n",
    "    ).request_body,\n",
    "    {\n",
    "        \"yql\": \"select * from sources * where userQuery();\",\n",
    "        \"query\": \"this is a test\",\n",
    "        \"type\": \"any\",\n",
    "        \"ranking\": {\"profile\": \"bm25\", \"listFeatures\": \"true\"},\n",
    "        \"hits\": 10,\n",
    "        \"recall\": \"+(id:1 id:5)\",\n",
    "    },\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6d927a8e-c50f-4c67-a5d3-3a33996fb402",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|hide\n",
    "test_fail(\n",
    "    send_query, \n",
    "    kwargs={\"app\": app, \"query\": \"this should not work\"}, \n",
    "    contains=\"No 'query_model' specified.\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "30e6f213-e2ba-4a2c-9424-da601fdb3efd",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export\n",
    "def send_query_batch(\n",
    "    app,  # Connection to a Vespa application\n",
    "    body_batch: Optional[List[Dict]] = None,  # Contains all the request parameters. Set to None if using 'query_batch'.\n",
    "    query_batch: Optional[List[str]] = None,  # Query strings. Set to None if using 'body_batch'.\n",
    "    query_model: Optional[QueryModel] = None,  # Query model to use when sending query strings. Set to None if using 'body_batch'.\n",
    "    recall_batch: Optional[List[Tuple]] = None,  # One tuple for each query. Tuple of size 2 where the first element is the name of the field to use to recall and the second element is a list of the values to be recalled.\n",
    "    asynchronous=True,  # Set True to send data in async mode. Default to True.\n",
    "    connections: Optional[int] = 100,  # Number of allowed concurrent connections, valid only if `asynchronous=True`.\n",
    "    total_timeout: int = 100,  # Total timeout in secs for each of the concurrent requests when using `asynchronous=True`.\n",
    "    **kwargs,  # Additional parameters to be sent along the request.\n",
    ") -> List[VespaQueryResponse]:  # HTTP POST responses.\n",
    "    \"Send queries in batch to a Vespa app.\"\n",
    "\n",
    "    if body_batch:\n",
    "        assert (\n",
    "            query_batch is None\n",
    "        ), \"'query_batch' has no effect if 'body_batch' is not None.\"\n",
    "    elif query_batch:\n",
    "        assert (\n",
    "            body_batch is None\n",
    "        ), \"'body_batch' has no effect if 'query_batch' is not None.\"\n",
    "        assert (\n",
    "            query_model is not None\n",
    "        ), \"Specify a 'query_model' when using 'query_batch' argument.\"\n",
    "        number_of_queries = len(query_batch)\n",
    "\n",
    "        if recall_batch:\n",
    "            assert (\n",
    "                len(recall_batch) == number_of_queries\n",
    "            ), \"Specify one recall tuple for each query in the batch.\"\n",
    "            body_batch = [\n",
    "                _build_query_body(\n",
    "                    query=query, \n",
    "                    query_model=query_model, \n",
    "                    recall=recall,\n",
    "                    **kwargs\n",
    "                ) for query, recall in zip(query_batch, recall_batch)\n",
    "            ]\n",
    "        else:\n",
    "            body_batch = [\n",
    "                _build_query_body(\n",
    "                    query=query, \n",
    "                    query_model=query_model, \n",
    "                    **kwargs\n",
    "                ) for query in query_batch\n",
    "            ]\n",
    "    else:\n",
    "        ValueError(\"Specify either 'query_batch' or 'body_batch'.\")\n",
    "\n",
    "    return app.query_batch(\n",
    "        body_batch=body_batch,\n",
    "        asynchronous=asynchronous,\n",
    "        connections=connections,\n",
    "        total_timeout=total_timeout,\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7574c485-bbc4-45ff-8e50-676f5e9915ec",
   "metadata": {},
   "source": [
    "Use `body_batch` to send a batch of body requests."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dc6fbaa0-926e-4218-8123-aba011cdc251",
   "metadata": {},
   "outputs": [],
   "source": [
    "body_batch = [\n",
    "    {\"yql\": \"select * from sources * where test\"},\n",
    "    {\"yql\": \"select * from sources * where test2\"}\n",
    "]\n",
    "result = send_query_batch(app=app, body_batch=body_batch)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6ea0de4c-82af-4c19-bdba-2cb99b13b8da",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|hide\n",
    "body_batch = [\n",
    "    {\"yql\": \"select * from sources * where test\"},\n",
    "    {\"yql\": \"select * from sources * where test2\"}\n",
    "]\n",
    "result = send_query_batch(app=app, body_batch=body_batch)\n",
    "test_eq(len(result), 2)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fc743a82-7198-46f7-9eb5-ed90b7f42294",
   "metadata": {},
   "source": [
    "Use `query_batch` to send a batch of query strings to be ranked according a `QueryModel`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8e5024cf-6185-4cfb-b954-14fb985bc253",
   "metadata": {},
   "outputs": [],
   "source": [
    "result = send_query_batch(\n",
    "    app=app,\n",
    "    query_batch=[\"this is a test\", \"this is a test 2\"],\n",
    "    query_model=QueryModel(\n",
    "        match_phase=OR(), \n",
    "        ranking=Ranking()\n",
    "    ),\n",
    "    hits=10,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "22c9f8f6-ecfa-48d4-9d77-06aeb778f9fc",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|hide\n",
    "result = send_query_batch(\n",
    "    app=app,\n",
    "    query_batch=[\"this is a test\", \"this is a test 2\"],\n",
    "    query_model=QueryModel(\n",
    "        match_phase=OR(), \n",
    "        ranking=Ranking()\n",
    "    ),\n",
    "    hits=10,\n",
    ")\n",
    "test_eq(len(result), 2)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ece540bf-fad6-4db8-b399-79d0466e1f81",
   "metadata": {},
   "source": [
    "Use `recall_batch` to send one tuple for each query in `query_batch`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fa9a9aec-cb62-4ed9-9d18-4adcd070a094",
   "metadata": {},
   "outputs": [],
   "source": [
    "result = send_query_batch(\n",
    "    app=app,\n",
    "    query_batch=[\"this is a test\", \"this is a test 2\"],\n",
    "    query_model=QueryModel(match_phase=OR(), ranking=Ranking()),\n",
    "    hits=10,\n",
    "    recall_batch=[(\"doc_id\", [2, 7]), (\"doc_id\", [0, 5])],\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b4159891-2b1e-443b-b439-35ef56dc7b5b",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|hide\n",
    "result = send_query_batch(\n",
    "    app=app,\n",
    "    query_batch=[\"this is a test 2\", \"this is a test 0\"],\n",
    "    query_model=QueryModel(match_phase=OR(), ranking=Ranking()),\n",
    "    hits=10,\n",
    "    recall_batch=[(\"doc_id\", [2, 7]), (\"doc_id\", [0, 5])],\n",
    ")\n",
    "test_eq(['2', '7', '0', '5'], [result[i].hits[j][\"fields\"][\"doc_id\"] for i in range(2) for j in range(2)])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "43a1abcf-6b96-42a1-bceb-71b58d3c47ac",
   "metadata": {},
   "source": [
    "## Collect Vespa features"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "157f6e10-3ba0-413d-8d95-ad84c1a82bd3",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export\n",
    "def _annotate_data(\n",
    "    hits, query_id, id_field, relevant_id, fields, relevant_score, default_score\n",
    "):\n",
    "    data = []\n",
    "    for h in hits:\n",
    "        record = {}\n",
    "        record.update({\"document_id\": h[\"fields\"][id_field]})\n",
    "        record.update({\"query_id\": query_id})\n",
    "        record.update(\n",
    "            {\n",
    "                \"label\": relevant_score\n",
    "                if h[\"fields\"][id_field] == relevant_id\n",
    "                else default_score\n",
    "            }\n",
    "        )\n",
    "        for field in fields:\n",
    "            field_value = h[\"fields\"].get(field, None)\n",
    "            if field_value:\n",
    "                if isinstance(field_value, dict):\n",
    "                    record.update(field_value)\n",
    "                else:\n",
    "                    record.update({field: field_value})\n",
    "        data.append(record)\n",
    "    return data\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0c16533b-fb1c-4453-9cf1-8bc66f24fc01",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export\n",
    "def _parse_labeled_data(\n",
    "    df: DataFrame  # DataFrame with the following required columns [\"qid\", \"query\", \"doc_id\", \"relevance\"].\n",
    ") -> List[Dict]:  # Concise representation of the labeled data, grouped by query_id and query.\n",
    "    \"Convert a DataFrame with labeled data to format used internally\"\n",
    "    required_columns = [\"qid\", \"query\", \"doc_id\", \"relevance\"]\n",
    "    assert all(\n",
    "        [x in list(df.columns) for x in required_columns]\n",
    "    ), \"DataFrame needs at least the following columns: {}\".format(required_columns)\n",
    "    qid_query = (\n",
    "        df[[\"qid\", \"query\"]].drop_duplicates([\"qid\", \"query\"]).to_dict(orient=\"records\")\n",
    "    )\n",
    "    labeled_data = []\n",
    "    for q in qid_query:\n",
    "        docid_relevance = df[(df[\"qid\"] == q[\"qid\"]) & (df[\"query\"] == q[\"query\"])][\n",
    "            [\"doc_id\", \"relevance\"]\n",
    "        ]\n",
    "        relevant_docs = []\n",
    "        for idx, row in docid_relevance.iterrows():\n",
    "            relevant_docs.append({\"id\": row[\"doc_id\"], \"score\": row[\"relevance\"]})\n",
    "        data_point = {\n",
    "            \"query_id\": q[\"qid\"],\n",
    "            \"query\": q[\"query\"],\n",
    "            \"relevant_docs\": relevant_docs,\n",
    "        }\n",
    "        labeled_data.append(data_point)\n",
    "    return labeled_data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5e6f74fa-470f-4697-baf4-410482d230a1",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|hide\n",
    "labeled_data_df = DataFrame(\n",
    "    data={\n",
    "        \"qid\": [0, 0, 1, 1],\n",
    "        \"query\": [\"Intrauterine virus infections and congenital heart disease\"]\n",
    "        * 2\n",
    "        + [\n",
    "            \"Clinical and immunologic studies in identical twins discordant for systemic lupus erythematosus\"\n",
    "        ]\n",
    "        * 2,\n",
    "        \"doc_id\": [0, 3, 1, 5],\n",
    "        \"relevance\": [1, 1, 1, 1],\n",
    "    }\n",
    ")\n",
    "labeled_data = _parse_labeled_data(df=labeled_data_df)\n",
    "expected_labeled_data = [\n",
    "    {\n",
    "        \"query_id\": 0,\n",
    "        \"query\": \"Intrauterine virus infections and congenital heart disease\",\n",
    "        \"relevant_docs\": [{\"id\": 0, \"score\": 1}, {\"id\": 3, \"score\": 1}],\n",
    "    },\n",
    "    {\n",
    "        \"query_id\": 1,\n",
    "        \"query\": \"Clinical and immunologic studies in identical twins discordant for systemic lupus erythematosus\",\n",
    "        \"relevant_docs\": [{\"id\": 1, \"score\": 1}, {\"id\": 5, \"score\": 1}],\n",
    "    },\n",
    "]\n",
    "test_eq(labeled_data, expected_labeled_data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "312efb89-60af-4c5b-b0d4-7e7ce1611c7d",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|hide\n",
    "# parse_labeled_data_with_wrong_columns\n",
    "labeled_data_df = DataFrame(\n",
    "    data={\n",
    "        \"qid\": [0, 0, 1, 1],\n",
    "        \"doc_id\": [0, 3, 1, 5],\n",
    "        \"relevance\": [1, 1, 1, 1],\n",
    "    }\n",
    ")\n",
    "test_fail(\n",
    "    _parse_labeled_data, \n",
    "    kwargs={\"df\":labeled_data_df}, \n",
    "    contains=\"DataFrame needs at least the following columns: ['qid', 'query', 'doc_id', 'relevance']\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0fa18f78-2bcf-4520-8977-ddb74f2fae0e",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|hide\n",
    "# labeled_data: Union[List[Dict], DataFrame]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e6841661-95b6-433c-a3ae-2b25292ff19b",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export\n",
    "def collect_vespa_features(\n",
    "    app: Vespa,  # Connection to a Vespa application.\n",
    "    labeled_data,  # Labelled data containing query, query_id and relevant ids. See examples about data format.\n",
    "    id_field: str,  # The Vespa field representing the document id.\n",
    "    query_model: QueryModel,  # Query model.\n",
    "    number_additional_docs: int,  # Number of additional documents to retrieve for each relevant document. Duplicate documents will be dropped.\n",
    "    fields: List[str],  # Vespa fields to collect, e.g. [\"rankfeatures\", \"summaryfeatures\"]\n",
    "    keep_features: Optional[List[str]] = None,  # List containing the names of the features that should be returned. Default to None, which return all the features contained in the 'fields' argument.\n",
    "    relevant_score: int = 1,  # Score to assign to relevant documents. Default to 1.\n",
    "    default_score: int = 0,  # Score to assign to the additional documents that are not relevant. Default to 0.\n",
    "    **kwargs,  # Extra keyword arguments to be included in the Vespa Query.\n",
    ") -> DataFrame:  # DataFrame containing document id (document_id), query id (query_id), scores (relevant) and vespa rank features returned by the Query model RankProfile used.\n",
    "    \"\"\"\n",
    "    Collect Vespa features based on a set of labelled data.\n",
    "    \"\"\"\n",
    "\n",
    "    if isinstance(labeled_data, DataFrame):\n",
    "        labeled_data = _parse_labeled_data(df=labeled_data)\n",
    "\n",
    "    flat_data = [\n",
    "        (\n",
    "            data[\"query_id\"],\n",
    "            data[\"query\"],\n",
    "            relevant_doc[\"id\"],\n",
    "            relevant_doc.get(\"score\", relevant_score),\n",
    "        )\n",
    "        for data in labeled_data\n",
    "        for relevant_doc in data[\"relevant_docs\"]\n",
    "    ]\n",
    "\n",
    "    queries = [x[1] for x in flat_data]\n",
    "    relevant_search = send_query_batch(\n",
    "        app=app,\n",
    "        query_batch=queries,\n",
    "        query_model=query_model,\n",
    "        recall_batch=[(id_field, [x[2]]) for x in flat_data],\n",
    "        **kwargs,\n",
    "    )\n",
    "    result = []\n",
    "    for ((query_id, query, relevant_id, relevant_score), query_result) in zip(\n",
    "        flat_data, relevant_search\n",
    "    ):\n",
    "        result.extend(\n",
    "            _annotate_data(\n",
    "                hits=query_result.hits,\n",
    "                query_id=query_id,\n",
    "                id_field=id_field,\n",
    "                relevant_id=relevant_id,\n",
    "                fields=fields,\n",
    "                relevant_score=relevant_score,\n",
    "                default_score=default_score,\n",
    "            )\n",
    "        )\n",
    "    if number_additional_docs > 0:\n",
    "        additional_hits_result = send_query_batch(\n",
    "            app=app,\n",
    "            query_batch=queries,\n",
    "            query_model=query_model,\n",
    "            hits=number_additional_docs,\n",
    "            **kwargs,\n",
    "        )\n",
    "        for ((query_id, query, relevant_id, relevant_score), query_result) in zip(\n",
    "            flat_data, additional_hits_result\n",
    "        ):\n",
    "            result.extend(\n",
    "                _annotate_data(\n",
    "                    hits=query_result.hits,\n",
    "                    query_id=query_id,\n",
    "                    id_field=id_field,\n",
    "                    relevant_id=relevant_id,\n",
    "                    fields=fields,\n",
    "                    relevant_score=relevant_score,\n",
    "                    default_score=default_score,\n",
    "                )\n",
    "            )\n",
    "    df = DataFrame.from_records(result)\n",
    "    df = df.drop_duplicates([\"document_id\", \"query_id\", \"label\"])\n",
    "    df = df.sort_values(\"query_id\")\n",
    "    if keep_features:\n",
    "        df = df[[\"document_id\", \"query_id\", \"label\"] + keep_features]\n",
    "    return df"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "02a94a6d-939d-4251-b4b0-4521d6c06e6f",
   "metadata": {},
   "source": [
    "Usage:"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1cd5b01b-bdc7-41de-8d2e-876ccdcef36d",
   "metadata": {},
   "source": [
    "Define `labeled_data` as a list of dict containing relevant documents:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0b5d3a8b-e939-4b0a-acda-6fe5352b02a8",
   "metadata": {},
   "outputs": [],
   "source": [
    "labeled_data = [\n",
    "    {\n",
    "        \"query_id\": 0,\n",
    "        \"query\": \"give me title 1\",\n",
    "        \"relevant_docs\": [{\"id\": \"1\", \"score\": 1}],\n",
    "    },\n",
    "    {\n",
    "        \"query_id\": 1,\n",
    "        \"query\": \"give me title 3\",\n",
    "        \"relevant_docs\": [{\"id\": \"3\", \"score\": 1}],\n",
    "    },\n",
    "]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "885ff001-29c3-4882-b900-f6f1fa822ec3",
   "metadata": {},
   "source": [
    "Collect vespa features:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d794b55b-8484-414b-8a8f-7b1d15f63a08",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>document_id</th>\n",
       "      <th>query_id</th>\n",
       "      <th>label</th>\n",
       "      <th>attributeMatch(doc_id)</th>\n",
       "      <th>attributeMatch(doc_id).averageWeight</th>\n",
       "      <th>attributeMatch(doc_id).completeness</th>\n",
       "      <th>attributeMatch(doc_id).fieldCompleteness</th>\n",
       "      <th>attributeMatch(doc_id).importance</th>\n",
       "      <th>attributeMatch(doc_id).matches</th>\n",
       "      <th>attributeMatch(doc_id).maxWeight</th>\n",
       "      <th>...</th>\n",
       "      <th>term(3).significance</th>\n",
       "      <th>term(3).weight</th>\n",
       "      <th>term(4).connectedness</th>\n",
       "      <th>term(4).significance</th>\n",
       "      <th>term(4).weight</th>\n",
       "      <th>textSimilarity(text).fieldCoverage</th>\n",
       "      <th>textSimilarity(text).order</th>\n",
       "      <th>textSimilarity(text).proximity</th>\n",
       "      <th>textSimilarity(text).queryCoverage</th>\n",
       "      <th>textSimilarity(text).score</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>...</td>\n",
       "      <td>0.583333</td>\n",
       "      <td>100.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.50</td>\n",
       "      <td>1.0</td>\n",
       "      <td>1.000000</td>\n",
       "      <td>0.50</td>\n",
       "      <td>0.750000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>7</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>...</td>\n",
       "      <td>0.583333</td>\n",
       "      <td>100.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.25</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.859375</td>\n",
       "      <td>0.25</td>\n",
       "      <td>0.425781</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>3</td>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>...</td>\n",
       "      <td>0.583333</td>\n",
       "      <td>100.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.50</td>\n",
       "      <td>1.0</td>\n",
       "      <td>1.000000</td>\n",
       "      <td>0.50</td>\n",
       "      <td>0.750000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5</th>\n",
       "      <td>7</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>...</td>\n",
       "      <td>0.583333</td>\n",
       "      <td>100.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.25</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.859375</td>\n",
       "      <td>0.25</td>\n",
       "      <td>0.425781</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>4 rows × 94 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "  document_id  query_id  label  attributeMatch(doc_id)  \\\n",
       "0           1         0      1                     0.0   \n",
       "3           7         0      0                     0.0   \n",
       "1           3         1      1                     0.0   \n",
       "5           7         1      0                     0.0   \n",
       "\n",
       "   attributeMatch(doc_id).averageWeight  attributeMatch(doc_id).completeness  \\\n",
       "0                                   0.0                                  0.0   \n",
       "3                                   0.0                                  0.0   \n",
       "1                                   0.0                                  0.0   \n",
       "5                                   0.0                                  0.0   \n",
       "\n",
       "   attributeMatch(doc_id).fieldCompleteness  \\\n",
       "0                                       0.0   \n",
       "3                                       0.0   \n",
       "1                                       0.0   \n",
       "5                                       0.0   \n",
       "\n",
       "   attributeMatch(doc_id).importance  attributeMatch(doc_id).matches  \\\n",
       "0                                0.0                             0.0   \n",
       "3                                0.0                             0.0   \n",
       "1                                0.0                             0.0   \n",
       "5                                0.0                             0.0   \n",
       "\n",
       "   attributeMatch(doc_id).maxWeight  ...  term(3).significance  \\\n",
       "0                               0.0  ...              0.583333   \n",
       "3                               0.0  ...              0.583333   \n",
       "1                               0.0  ...              0.583333   \n",
       "5                               0.0  ...              0.583333   \n",
       "\n",
       "   term(3).weight  term(4).connectedness  term(4).significance  \\\n",
       "0           100.0                    0.0                   0.0   \n",
       "3           100.0                    0.0                   0.0   \n",
       "1           100.0                    0.0                   0.0   \n",
       "5           100.0                    0.0                   0.0   \n",
       "\n",
       "   term(4).weight  textSimilarity(text).fieldCoverage  \\\n",
       "0             0.0                                0.50   \n",
       "3             0.0                                0.25   \n",
       "1             0.0                                0.50   \n",
       "5             0.0                                0.25   \n",
       "\n",
       "   textSimilarity(text).order  textSimilarity(text).proximity  \\\n",
       "0                         1.0                        1.000000   \n",
       "3                         0.0                        0.859375   \n",
       "1                         1.0                        1.000000   \n",
       "5                         0.0                        0.859375   \n",
       "\n",
       "   textSimilarity(text).queryCoverage  textSimilarity(text).score  \n",
       "0                                0.50                    0.750000  \n",
       "3                                0.25                    0.425781  \n",
       "1                                0.50                    0.750000  \n",
       "5                                0.25                    0.425781  \n",
       "\n",
       "[4 rows x 94 columns]"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "rank_features = collect_vespa_features(\n",
    "    app=app,\n",
    "    labeled_data=labeled_data,\n",
    "    id_field=\"doc_id\",\n",
    "    query_model=QueryModel(\n",
    "        match_phase=OR(), \n",
    "        ranking=Ranking(name=\"bm25\", list_features=True)\n",
    "    ),\n",
    "    number_additional_docs=2,\n",
    "    fields=[\"rankfeatures\"],\n",
    ")\n",
    "rank_features"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "14881801-c545-4d8d-ad04-ba31009a3a96",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|hide\n",
    "labeled_data = [\n",
    "    {\n",
    "        \"query_id\": 0,\n",
    "        \"query\": \"give me title 1\",\n",
    "        \"relevant_docs\": [{\"id\": \"1\", \"score\": 1}],\n",
    "    },\n",
    "    {\n",
    "        \"query_id\": 1,\n",
    "        \"query\": \"give me title 3\",\n",
    "        \"relevant_docs\": [{\"id\": \"3\", \"score\": 1}],\n",
    "    },\n",
    "]\n",
    "rank_features = collect_vespa_features(\n",
    "    app=app,\n",
    "    labeled_data=labeled_data,\n",
    "    id_field=\"doc_id\",\n",
    "    query_model=QueryModel(\n",
    "        match_phase=OR(), \n",
    "        ranking=Ranking(name=\"bm25\", list_features=True)\n",
    "    ),\n",
    "    number_additional_docs=2,\n",
    "    fields=[\"rankfeatures\"],\n",
    ")\n",
    "test_eq(rank_features.shape[0] > 2, True)\n",
    "test_eq(rank_features.shape[1] > 3, True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2cd9a9d1-b3e2-41f5-bf2e-41550010e9fc",
   "metadata": {},
   "source": [
    "Use a `DataFrame` for `labeled_data` instead of a list of dict:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "982b2ab2-a161-4453-9103-c7f8c4ee05d7",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>qid</th>\n",
       "      <th>query</th>\n",
       "      <th>doc_id</th>\n",
       "      <th>relevance</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>0</td>\n",
       "      <td>give me title 1</td>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>1</td>\n",
       "      <td>give me title 3</td>\n",
       "      <td>3</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   qid            query  doc_id  relevance\n",
       "0    0  give me title 1       1          1\n",
       "1    1  give me title 3       3          1"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "labeled_data = [\n",
    "    {\n",
    "        \"qid\": 0,\n",
    "        \"query\": \"give me title 1\",\n",
    "        \"doc_id\": 1, \n",
    "        \"relevance\": 1\n",
    "    },\n",
    "    {\n",
    "        \"qid\": 1,\n",
    "        \"query\": \"give me title 3\",\n",
    "        \"doc_id\": 3, \n",
    "        \"relevance\": 1\n",
    "    },\n",
    "]\n",
    "labeled_data_df = DataFrame.from_records(labeled_data)\n",
    "labeled_data_df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "063d9304-8d44-42b7-b7e1-1e4fc59ee505",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>document_id</th>\n",
       "      <th>query_id</th>\n",
       "      <th>label</th>\n",
       "      <th>attributeMatch(doc_id)</th>\n",
       "      <th>attributeMatch(doc_id).averageWeight</th>\n",
       "      <th>attributeMatch(doc_id).completeness</th>\n",
       "      <th>attributeMatch(doc_id).fieldCompleteness</th>\n",
       "      <th>attributeMatch(doc_id).importance</th>\n",
       "      <th>attributeMatch(doc_id).matches</th>\n",
       "      <th>attributeMatch(doc_id).maxWeight</th>\n",
       "      <th>...</th>\n",
       "      <th>term(3).significance</th>\n",
       "      <th>term(3).weight</th>\n",
       "      <th>term(4).connectedness</th>\n",
       "      <th>term(4).significance</th>\n",
       "      <th>term(4).weight</th>\n",
       "      <th>textSimilarity(text).fieldCoverage</th>\n",
       "      <th>textSimilarity(text).order</th>\n",
       "      <th>textSimilarity(text).proximity</th>\n",
       "      <th>textSimilarity(text).queryCoverage</th>\n",
       "      <th>textSimilarity(text).score</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>...</td>\n",
       "      <td>0.583333</td>\n",
       "      <td>100.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.50</td>\n",
       "      <td>1.0</td>\n",
       "      <td>1.000000</td>\n",
       "      <td>0.50</td>\n",
       "      <td>0.750000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>7</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>...</td>\n",
       "      <td>0.583333</td>\n",
       "      <td>100.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.25</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.859375</td>\n",
       "      <td>0.25</td>\n",
       "      <td>0.425781</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>3</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>...</td>\n",
       "      <td>0.583333</td>\n",
       "      <td>100.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.50</td>\n",
       "      <td>1.0</td>\n",
       "      <td>1.000000</td>\n",
       "      <td>0.50</td>\n",
       "      <td>0.750000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5</th>\n",
       "      <td>7</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>...</td>\n",
       "      <td>0.583333</td>\n",
       "      <td>100.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.25</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.859375</td>\n",
       "      <td>0.25</td>\n",
       "      <td>0.425781</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>4 rows × 94 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "  document_id  query_id  label  attributeMatch(doc_id)  \\\n",
       "0           1         0      0                     0.0   \n",
       "3           7         0      0                     0.0   \n",
       "1           3         1      0                     0.0   \n",
       "5           7         1      0                     0.0   \n",
       "\n",
       "   attributeMatch(doc_id).averageWeight  attributeMatch(doc_id).completeness  \\\n",
       "0                                   0.0                                  0.0   \n",
       "3                                   0.0                                  0.0   \n",
       "1                                   0.0                                  0.0   \n",
       "5                                   0.0                                  0.0   \n",
       "\n",
       "   attributeMatch(doc_id).fieldCompleteness  \\\n",
       "0                                       0.0   \n",
       "3                                       0.0   \n",
       "1                                       0.0   \n",
       "5                                       0.0   \n",
       "\n",
       "   attributeMatch(doc_id).importance  attributeMatch(doc_id).matches  \\\n",
       "0                                0.0                             0.0   \n",
       "3                                0.0                             0.0   \n",
       "1                                0.0                             0.0   \n",
       "5                                0.0                             0.0   \n",
       "\n",
       "   attributeMatch(doc_id).maxWeight  ...  term(3).significance  \\\n",
       "0                               0.0  ...              0.583333   \n",
       "3                               0.0  ...              0.583333   \n",
       "1                               0.0  ...              0.583333   \n",
       "5                               0.0  ...              0.583333   \n",
       "\n",
       "   term(3).weight  term(4).connectedness  term(4).significance  \\\n",
       "0           100.0                    0.0                   0.0   \n",
       "3           100.0                    0.0                   0.0   \n",
       "1           100.0                    0.0                   0.0   \n",
       "5           100.0                    0.0                   0.0   \n",
       "\n",
       "   term(4).weight  textSimilarity(text).fieldCoverage  \\\n",
       "0             0.0                                0.50   \n",
       "3             0.0                                0.25   \n",
       "1             0.0                                0.50   \n",
       "5             0.0                                0.25   \n",
       "\n",
       "   textSimilarity(text).order  textSimilarity(text).proximity  \\\n",
       "0                         1.0                        1.000000   \n",
       "3                         0.0                        0.859375   \n",
       "1                         1.0                        1.000000   \n",
       "5                         0.0                        0.859375   \n",
       "\n",
       "   textSimilarity(text).queryCoverage  textSimilarity(text).score  \n",
       "0                                0.50                    0.750000  \n",
       "3                                0.25                    0.425781  \n",
       "1                                0.50                    0.750000  \n",
       "5                                0.25                    0.425781  \n",
       "\n",
       "[4 rows x 94 columns]"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "rank_features = collect_vespa_features(\n",
    "    app=app,\n",
    "    labeled_data=labeled_data_df,\n",
    "    id_field=\"doc_id\",\n",
    "    query_model=QueryModel(\n",
    "        match_phase=OR(), ranking=Ranking(name=\"bm25\", list_features=True)\n",
    "    ),\n",
    "    number_additional_docs=2,\n",
    "    fields=[\"rankfeatures\"],\n",
    ")\n",
    "rank_features"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f2971fb1-2844-462c-b21b-3ba3484537d8",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|hide\n",
    "labeled_data = [\n",
    "    {\n",
    "        \"qid\": 0,\n",
    "        \"query\": \"give me title 1\",\n",
    "        \"doc_id\": 1, \n",
    "        \"relevance\": 1\n",
    "    },\n",
    "    {\n",
    "        \"qid\": 1,\n",
    "        \"query\": \"give me title 3\",\n",
    "        \"doc_id\": 3, \n",
    "        \"relevance\": 1\n",
    "    },\n",
    "]\n",
    "labeled_data_df = DataFrame.from_records(labeled_data)\n",
    "rank_features = collect_vespa_features(\n",
    "    app=app,\n",
    "    labeled_data=labeled_data_df,\n",
    "    id_field=\"doc_id\",\n",
    "    query_model=QueryModel(\n",
    "        match_phase=OR(), ranking=Ranking(name=\"bm25\", list_features=True)\n",
    "    ),\n",
    "    number_additional_docs=2,\n",
    "    fields=[\"rankfeatures\"],\n",
    ")\n",
    "test_eq(rank_features.shape[0] > 2, True)\n",
    "test_eq(rank_features.shape[1] > 3, True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a82696ab-de99-4010-9ef3-a4e8aca4d6cc",
   "metadata": {},
   "source": [
    "Keep only selected features by specifying their names in the `keep_features` argument:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "56b2be3b-8d79-4f7f-8713-2beac7bc7e88",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>document_id</th>\n",
       "      <th>query_id</th>\n",
       "      <th>label</th>\n",
       "      <th>textSimilarity(text).score</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0.750000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>7</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0.425781</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>3</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>0.750000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5</th>\n",
       "      <td>7</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>0.425781</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "  document_id  query_id  label  textSimilarity(text).score\n",
       "0           1         0      0                    0.750000\n",
       "3           7         0      0                    0.425781\n",
       "1           3         1      0                    0.750000\n",
       "5           7         1      0                    0.425781"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "rank_features = collect_vespa_features(\n",
    "    app=app,\n",
    "    labeled_data=labeled_data_df,\n",
    "    id_field=\"doc_id\",\n",
    "    query_model=QueryModel(\n",
    "        match_phase=OR(), ranking=Ranking(name=\"bm25\", list_features=True)\n",
    "    ),\n",
    "    number_additional_docs=2,\n",
    "    fields=[\"rankfeatures\"],\n",
    "    keep_features=[\"textSimilarity(text).score\"],\n",
    ")\n",
    "rank_features"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c255b6e7-9792-46e7-b15b-df90e5bd71dd",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|hide\n",
    "labeled_data = [\n",
    "    {\n",
    "        \"query_id\": 0,\n",
    "        \"query\": \"give me title 1\",\n",
    "        \"relevant_docs\": [{\"id\": \"1\", \"score\": 1}],\n",
    "    },\n",
    "    {\n",
    "        \"query_id\": 1,\n",
    "        \"query\": \"give me title 3\",\n",
    "        \"relevant_docs\": [{\"id\": \"3\", \"score\": 1}],\n",
    "    },\n",
    "]\n",
    "rank_features = collect_vespa_features(\n",
    "    app=app,\n",
    "    labeled_data=labeled_data,\n",
    "    id_field=\"doc_id\",\n",
    "    query_model=QueryModel(\n",
    "        match_phase=OR(), ranking=Ranking(name=\"bm25\", list_features=True)\n",
    "    ),\n",
    "    number_additional_docs=2,\n",
    "    fields=[\"rankfeatures\"],\n",
    "    keep_features=[\"textSimilarity(text).score\"],\n",
    ")\n",
    "test_eq(rank_features.shape[0] > 2, True)\n",
    "# It should have at least one rank feature in addition to document_id, query_id and label\n",
    "test_eq(rank_features.shape[1] == 4, True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "94c38660-4437-4928-8da1-aa6d6ca57028",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|hide\n",
    "labeled_data = [\n",
    "    {\n",
    "        \"query_id\": 0,\n",
    "        \"query\": \"give me title 1\",\n",
    "        \"relevant_docs\": [{\"id\": \"1\", \"score\": 1}],\n",
    "    }\n",
    "]\n",
    "rank_features = collect_vespa_features(\n",
    "    app=app,\n",
    "    labeled_data=labeled_data,\n",
    "    id_field=\"doc_id\",\n",
    "    query_model=QueryModel(\n",
    "        match_phase=OR(), ranking=Ranking(name=\"bm25\", list_features=True)\n",
    "    ),\n",
    "    number_additional_docs=12,\n",
    "    fields=[\"rankfeatures\"],\n",
    "    keep_features=[\"textSimilarity(text).score\"],\n",
    ")\n",
    "document_ids = rank_features.document_id.tolist()\n",
    "test_eq(len(document_ids), len(set(document_ids)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eb49420d-b182-481f-b2f5-f82fc288fecf",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export\n",
    "def store_vespa_features(\n",
    "    app: Vespa,  # Connection to a Vespa application.\n",
    "    output_file_path: str,  # Path of the .csv output file. It will create the file of it does not exist and append the vespa features to an pre-existing file.\n",
    "    labeled_data,  # Labelled data containing query, query_id and relevant ids. See details about data format.\n",
    "    id_field: str,  # The Vespa field representing the document id.\n",
    "    query_model: QueryModel,  # Query model.\n",
    "    number_additional_docs: int,  # Number of additional documents to retrieve for each relevant document.\n",
    "    fields: List[str],  # List of Vespa fields to collect, e.g. [\"rankfeatures\", \"summaryfeatures\"]\n",
    "    keep_features: Optional[List[str]] = None,  # List containing the names of the features that should be returned. Default to None, which return all the features contained in the 'fields' argument.\n",
    "    relevant_score: int = 1,  # Score to assign to relevant documents.\n",
    "    default_score: int = 0,  # Score to assign to the additional documents that are not relevant.\n",
    "    batch_size=1000,  # The size of the batch of labeled data points to be processed.\n",
    "    **kwargs,  # Extra keyword arguments to be included in the Vespa Query.\n",
    ") -> int:  # returns 0 upon success.\n",
    "    \"Retrieve Vespa rank features and store them in a .csv file.\"\n",
    "\n",
    "    if isinstance(labeled_data, DataFrame):\n",
    "        labeled_data = _parse_labeled_data(df=labeled_data)\n",
    "\n",
    "    mini_batches = [\n",
    "        labeled_data[i : i + batch_size]\n",
    "        for i in range(0, len(labeled_data), batch_size)\n",
    "    ]\n",
    "    for idx, mini_batch in enumerate(mini_batches):\n",
    "        vespa_features = collect_vespa_features(\n",
    "            app=app,\n",
    "            labeled_data=mini_batch,\n",
    "            id_field=id_field,\n",
    "            query_model=query_model,\n",
    "            number_additional_docs=number_additional_docs,\n",
    "            fields=fields,\n",
    "            keep_features=keep_features,\n",
    "            relevant_score=relevant_score,\n",
    "            default_score=default_score,\n",
    "            **kwargs,\n",
    "        )\n",
    "        if os.path.isfile(output_file_path):\n",
    "            vespa_features.to_csv(\n",
    "                path_or_buf=output_file_path, header=False, index=False, mode=\"a\"\n",
    "            )\n",
    "        else:\n",
    "            vespa_features.to_csv(\n",
    "                path_or_buf=output_file_path, header=True, index=False, mode=\"w\"\n",
    "            )\n",
    "        print(\n",
    "            \"Rows collected: {}.\\nBatch progress: {}/{}.\".format(\n",
    "                vespa_features.shape[0],\n",
    "                idx + 1,\n",
    "                len(mini_batches),\n",
    "            )\n",
    "        )\n",
    "    return 0\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dfa61f3c-71ea-4229-b8ca-cf9feabeec35",
   "metadata": {},
   "source": [
    "Usage:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1d2325e3-4400-4985-a197-b54b6f6778bd",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|hide\n",
    "from pandas import read_csv"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0441fec5-b626-4a84-80d2-05195f0df42f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Rows collected: 4.\n",
      "Batch progress: 1/1.\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>document_id</th>\n",
       "      <th>query_id</th>\n",
       "      <th>label</th>\n",
       "      <th>attributeMatch(doc_id)</th>\n",
       "      <th>attributeMatch(doc_id).averageWeight</th>\n",
       "      <th>attributeMatch(doc_id).completeness</th>\n",
       "      <th>attributeMatch(doc_id).fieldCompleteness</th>\n",
       "      <th>attributeMatch(doc_id).importance</th>\n",
       "      <th>attributeMatch(doc_id).matches</th>\n",
       "      <th>attributeMatch(doc_id).maxWeight</th>\n",
       "      <th>...</th>\n",
       "      <th>term(3).weight</th>\n",
       "      <th>term(4).connectedness</th>\n",
       "      <th>term(4).significance</th>\n",
       "      <th>term(4).weight</th>\n",
       "      <th>textSimilarity(text).fieldCoverage</th>\n",
       "      <th>textSimilarity(text).order</th>\n",
       "      <th>textSimilarity(text).proximity</th>\n",
       "      <th>textSimilarity(text).queryCoverage</th>\n",
       "      <th>textSimilarity(text).score</th>\n",
       "      <th>vespa.summaryFeatures.cached</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>...</td>\n",
       "      <td>100.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.50</td>\n",
       "      <td>1.0</td>\n",
       "      <td>1.000000</td>\n",
       "      <td>0.50</td>\n",
       "      <td>0.750000</td>\n",
       "      <td>0.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>7</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>...</td>\n",
       "      <td>100.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.25</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.859375</td>\n",
       "      <td>0.25</td>\n",
       "      <td>0.425781</td>\n",
       "      <td>0.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>3</td>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>...</td>\n",
       "      <td>100.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.50</td>\n",
       "      <td>1.0</td>\n",
       "      <td>1.000000</td>\n",
       "      <td>0.50</td>\n",
       "      <td>0.750000</td>\n",
       "      <td>0.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>7</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>...</td>\n",
       "      <td>100.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.25</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.859375</td>\n",
       "      <td>0.25</td>\n",
       "      <td>0.425781</td>\n",
       "      <td>0.0</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>4 rows × 95 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "   document_id  query_id  label  attributeMatch(doc_id)  \\\n",
       "0            1         0      1                     0.0   \n",
       "1            7         0      0                     0.0   \n",
       "2            3         1      1                     0.0   \n",
       "3            7         1      0                     0.0   \n",
       "\n",
       "   attributeMatch(doc_id).averageWeight  attributeMatch(doc_id).completeness  \\\n",
       "0                                   0.0                                  0.0   \n",
       "1                                   0.0                                  0.0   \n",
       "2                                   0.0                                  0.0   \n",
       "3                                   0.0                                  0.0   \n",
       "\n",
       "   attributeMatch(doc_id).fieldCompleteness  \\\n",
       "0                                       0.0   \n",
       "1                                       0.0   \n",
       "2                                       0.0   \n",
       "3                                       0.0   \n",
       "\n",
       "   attributeMatch(doc_id).importance  attributeMatch(doc_id).matches  \\\n",
       "0                                0.0                             0.0   \n",
       "1                                0.0                             0.0   \n",
       "2                                0.0                             0.0   \n",
       "3                                0.0                             0.0   \n",
       "\n",
       "   attributeMatch(doc_id).maxWeight  ...  term(3).weight  \\\n",
       "0                               0.0  ...           100.0   \n",
       "1                               0.0  ...           100.0   \n",
       "2                               0.0  ...           100.0   \n",
       "3                               0.0  ...           100.0   \n",
       "\n",
       "   term(4).connectedness  term(4).significance  term(4).weight  \\\n",
       "0                    0.0                   0.0             0.0   \n",
       "1                    0.0                   0.0             0.0   \n",
       "2                    0.0                   0.0             0.0   \n",
       "3                    0.0                   0.0             0.0   \n",
       "\n",
       "   textSimilarity(text).fieldCoverage  textSimilarity(text).order  \\\n",
       "0                                0.50                         1.0   \n",
       "1                                0.25                         0.0   \n",
       "2                                0.50                         1.0   \n",
       "3                                0.25                         0.0   \n",
       "\n",
       "   textSimilarity(text).proximity  textSimilarity(text).queryCoverage  \\\n",
       "0                        1.000000                                0.50   \n",
       "1                        0.859375                                0.25   \n",
       "2                        1.000000                                0.50   \n",
       "3                        0.859375                                0.25   \n",
       "\n",
       "   textSimilarity(text).score  vespa.summaryFeatures.cached  \n",
       "0                    0.750000                           0.0  \n",
       "1                    0.425781                           0.0  \n",
       "2                    0.750000                           0.0  \n",
       "3                    0.425781                           0.0  \n",
       "\n",
       "[4 rows x 95 columns]"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "labeled_data = [\n",
    "    {\n",
    "        \"query_id\": 0,\n",
    "        \"query\": \"give me title 1\",\n",
    "        \"relevant_docs\": [{\"id\": \"1\", \"score\": 1}],\n",
    "    },\n",
    "    {\n",
    "        \"query_id\": 1,\n",
    "        \"query\": \"give me title 3\",\n",
    "        \"relevant_docs\": [{\"id\": \"3\", \"score\": 1}],\n",
    "    },\n",
    "]\n",
    "\n",
    "store_vespa_features(\n",
    "    app=app,\n",
    "    output_file_path=\"vespa_features.csv\",\n",
    "    labeled_data=labeled_data,\n",
    "    id_field=\"doc_id\",\n",
    "    query_model=QueryModel(\n",
    "        match_phase=OR(), ranking=Ranking(name=\"bm25\", list_features=True)\n",
    "    ),\n",
    "    number_additional_docs=2,\n",
    "    fields=[\"rankfeatures\", \"summaryfeatures\"],\n",
    ")\n",
    "rank_features = read_csv(\"vespa_features.csv\")\n",
    "rank_features"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2a14a6c9-e939-4049-a49a-a426dbacd0ef",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Rows collected: 4.\n",
      "Batch progress: 1/1.\n"
     ]
    }
   ],
   "source": [
    "#|hide\n",
    "labeled_data = [\n",
    "    {\n",
    "        \"query_id\": 0,\n",
    "        \"query\": \"give me title 1\",\n",
    "        \"relevant_docs\": [{\"id\": \"1\", \"score\": 1}],\n",
    "    },\n",
    "    {\n",
    "        \"query_id\": 1,\n",
    "        \"query\": \"give me title 3\",\n",
    "        \"relevant_docs\": [{\"id\": \"3\", \"score\": 1}],\n",
    "    },\n",
    "]\n",
    "\n",
    "store_vespa_features(\n",
    "    app=app,\n",
    "    output_file_path=\"vespa_features.csv\",\n",
    "    labeled_data=labeled_data,\n",
    "    id_field=\"doc_id\",\n",
    "    query_model=QueryModel(\n",
    "        match_phase=OR(), ranking=Ranking(name=\"bm25\", list_features=True)\n",
    "    ),\n",
    "    number_additional_docs=2,\n",
    "    fields=[\"rankfeatures\", \"summaryfeatures\"],\n",
    ")\n",
    "rank_features = read_csv(\"vespa_features.csv\")\n",
    "# at least two relevant docs\n",
    "test_eq(rank_features.shape[0] > 2, True)\n",
    "# at least one feature besides document_id, query_id and label\n",
    "test_eq(rank_features.shape[1] > 3, True)\n",
    "os.remove(\"vespa_features.csv\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c75024d1-ab20-4662-8f83-52e457062816",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Rows collected: 4.\n",
      "Batch progress: 1/1.\n"
     ]
    }
   ],
   "source": [
    "#|hide\n",
    "labeled_data = [\n",
    "    {\n",
    "        \"qid\": 0,\n",
    "        \"query\": \"give me title 1\",\n",
    "        \"doc_id\": 1, \n",
    "        \"relevance\": 1\n",
    "    },\n",
    "    {\n",
    "        \"qid\": 1,\n",
    "        \"query\": \"give me title 3\",\n",
    "        \"doc_id\": 3, \n",
    "        \"relevance\": 1\n",
    "    },\n",
    "]\n",
    "labeled_data_df = DataFrame.from_records(labeled_data)\n",
    "store_vespa_features( \n",
    "    app=app,\n",
    "    output_file_path=\"vespa_features.csv\",\n",
    "    labeled_data=labeled_data_df,\n",
    "    id_field=\"doc_id\",\n",
    "    query_model=QueryModel(\n",
    "        match_phase=OR(), ranking=Ranking(name=\"bm25\", list_features=True)\n",
    "    ),\n",
    "    number_additional_docs=2,\n",
    "    fields=[\"rankfeatures\", \"summaryfeatures\"],\n",
    ")\n",
    "rank_features = read_csv(\"vespa_features.csv\")\n",
    "# at least two relevant docs\n",
    "test_eq(rank_features.shape[0] > 2, True)\n",
    "# at least one feature besides document_id, query_id and label\n",
    "test_eq(rank_features.shape[1] > 3, True)\n",
    "os.remove(\"vespa_features.csv\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8aa0e95c-f02a-429f-9389-7f02882d252a",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|hide\n",
    "vespa_docker.container.stop(timeout=600)\n",
    "vespa_docker.container.remove()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f0d77781-8f0d-46d0-92f6-22ef8b25719d",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|hide\n",
    "nbdev_export()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "learntorank",
   "language": "python",
   "name": "learntorank"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}