#!/usr/bin/env python
# coding: utf-8

# # OpenSearch の基本的な検索機能

# ## 概要
# 本ラボでは、OpenSearch が提供する多様な検索機能について、実際に検索クエリを実行しながら確認していきます。
# 
# ### ラボの構成
# 本ラボでは、ノートブック環境(JupyterLab) および Amazon OpenSearch Service を使用します。
# <img src="./img/architecture.png" width="1024"> 
# 
# ### 使用するデータセット
# 本ラボでは、Amazon OpenSearch Service [デベロッパーガイド](https://docs.aws.amazon.com/ja_jp/opensearch-service/latest/developerguide/what-is.html)の[チュートリアル](https://docs.aws.amazon.com/ja_jp/opensearch-service/latest/developerguide/search-example.html)内でも使用している[サンプルムービーデータセット](https://docs.aws.amazon.com/opensearch-service/latest/developerguide/samples/sample-movies.zip)を使用します。
# 

# ## 事前作業

# ### パッケージインストール

# In[1]:


get_ipython().system('pip install opensearch-py requests-aws4auth "awswrangler[opensearch]" --quiet')


# ### インポート

# In[2]:


import boto3
import json
import logging
import awswrangler as wr
import pandas as pd
from opensearchpy import OpenSearch, RequestsHttpConnection, AWSV4SignerAuth


# ### ヘルパー関数の定義
# 以降の処理を実行する際に必要なヘルパー関数を定義しておきます。

# In[3]:


def search_cloudformation_output(stackname, key):
    cloudformation_client = boto3.client("cloudformation", region_name=default_region)
    for output in cloudformation_client.describe_stacks(StackName=stackname)["Stacks"][0]["Outputs"]:
        if output["OutputKey"] == key:
            return output["OutputValue"]
    raise ValueError(f"{key} is not found in outputs of {stackname}.")


# ### 共通変数のセット

# In[4]:


default_region = boto3.Session().region_name
logging.getLogger().setLevel(logging.ERROR)


# ### OpenSearch クラスターへの接続確認
# 
# OpenSearch クラスターへのネットワーク接続性が確保されており、OpenSearch の Security 機能により API リクエストが許可されているかを確認します。
# 
# レスポンスに cluster_name や cluster_uuid が含まれていれば、接続確認が無事完了したと判断できます

# In[5]:


cloudformation_stack_name = "search-lab-jp"
opensearch_cluster_endpoint = search_cloudformation_output(cloudformation_stack_name, "OpenSearchDomainEndpoint")

credentials = boto3.Session().get_credentials()
service_code = "es"
auth = AWSV4SignerAuth(credentials=credentials, region=default_region, service=service_code)
opensearch_client = OpenSearch(
    hosts=[{"host": opensearch_cluster_endpoint, "port": 443}],
    http_compress=True, 
    http_auth=auth,
    use_ssl=True,
    verify_certs=True,
    connection_class = RequestsHttpConnection
)
opensearch_client.info()


# ## インデックスの作成・データロード
# インデックスを作成し、検索用のサンプルデータを格納します。

# ### サンプルデータの準備

# In[6]:


dataset_dir = "./dataset/sample-movies"

get_ipython().run_line_magic('mkdir', '-p $dataset_dir')

# ファイルダウンロード
get_ipython().system('curl -s -o $dataset_dir/sample-movies.zip https://docs.aws.amazon.com/opensearch-service/latest/developerguide/samples/sample-movies.zip')

# zip ファイルから sample-movies.bulk のみを展開
get_ipython().system('unzip -oq $dataset_dir/sample-movies.zip sample-movies.bulk -d $dataset_dir')

# .bulk ファイルから実データ行だけを抜き出して jsonl ファイルとして保存
get_ipython().system('grep -v "_index" $dataset_dir/sample-movies.bulk > $dataset_dir/sample-movies.jsonl')


# ### インデックスの作成
# 本ラボでは、
# [OpenSearch の基本操作と基本概念](../introduction-to-opensearch/introduction-to-opensearch.ipynb) をベースに、全文検索を意識した [Analyzer][analyzer] や [Normalizer][normalizer] の設定を追加しています。詳細は以降のセクションで順次解説していきます。
# 
# [analyzer]: https://opensearch.org/docs/latest/analyzers/supported-analyzers/index/
# [normalizer]: https://opensearch.org/docs/latest/analyzers/normalizers/

# In[7]:


index_name = "movies"

payload = {
  "settings": {
    "index": {
      "number_of_shards": 1,
      "number_of_replicas": 0
    },
    "analysis": {
      "normalizer": {
        "lowercase_normalizer": {
          "type": "custom",
          "filter": ["lowercase"]
        }
      }
    }
  },
  "mappings": {
    "properties": {
      "id":  { "type": "keyword" },
      "directors":  { "type": "text" },
      "release_date": {"type": "date"},
      "rating": {"type": "scaled_float", "scaling_factor": 10},
      "genres":  {
        "type": "keyword",
        "normalizer": "lowercase_normalizer"
      },
      "image_url":  { "type": "keyword" },
      "plot":  {
        "type": "text",
        "analyzer": "english"
      },
      "title":  { "type": "text" },
      "rank": { "type": "integer" },
      "running_time_secs": { "type": "integer" },
      "actors":  { "type": "text" },
      "year": {"type": "short"},
      "type": { "type": "keyword" } 
    }
  }
}

try:
    # 既に同名のインデックスが存在する場合、いったん削除を行う
    print("# delete index")
    response = opensearch_client.indices.delete(index=index_name)
    print(json.dumps(response, indent=2))
except Exception as e:
    print(e)

# インデックスの作成を行う
print("# create index")
response = opensearch_client.indices.create(index=index_name, body=payload)
print(json.dumps(response, indent=2))


# ### ドキュメントのロード
# ドキュメントのロードを行います。ドキュメントのロードは "OpenSearch の基本概念・基本操作の理解" でも解説した通り bulk API を使用することで効率よく進められますが、データ処理フレームワークを利用することでより簡単にデータを取り込むことも可能です。本ワークショップでは、[AWS SDK for Pandas][aws-sdk-pandas] を使用したデータ取り込みを行います。
# 
# [aws-sdk-pandas]: https://github.com/aws/aws-sdk-pandas
# 

# In[8]:


get_ipython().run_cell_magic('time', '', 'index_name = "movies"\nfile_path = f"{dataset_dir}/sample-movies.jsonl"\n\nresponse = wr.opensearch.index_json(\n    client=opensearch_client,\n    path=file_path,\n    use_threads=True, #クライアントの CPU 数に応じて自動で書き込みを並列化\n    id_keys=["id"],\n    index=index_name,\n    bulk_size=200, # 200 件ずつ書き込み,\n    refresh=False # 書き込み処理実行中の refresh (セグメントの書き出し) を無効化\n)\n')


# response["success"] の値が jsonl ファイルの行数と一致しているかを確認します。True が表示される場合は全件登録に成功していると判断できます。

# In[9]:


response["success"] == len(open(file_path).readlines())


# 本ラボではデータ登録時に意図的に Refresh オプションを無効化しているため、念のため Refresh API を実行し、登録されたドキュメントが確実に検索可能となるようにします

# In[10]:


index_name = "movies"
response = opensearch_client.indices.refresh(index_name)
response = opensearch_client.indices.forcemerge(index_name, max_num_segments=1)


# ## ドキュメント検索
# ドキュメント検索は [Search][search] API を使用します。Search API は様々なパラメーターを持ちます。
# 基本的なパラメーターはリクエストボディに含めることが可能です。以下は代表的なパラメーターです。
# 
# - query: 検索条件を記載します
# - aggs: 集計条件を記載します
# - sort: ソート条件を記載します
# - size: 取得する検索結果の件数を指定します
# 
# 以降のセクションでは、検索要件に応じた query/aggs/sort の書き方を解説します。
# 
# [search]: https://opensearch.org/docs/latest/api-reference/search/

# ### [全文検索][full-text]
# 全文検索は OpenSearch の根幹をなす機能です。OpenSearch では文字列をトークンと呼ばれる単位に分割してインデックスに格納することで、高速な全文検索(部分一致検索)を実現しています。
# 
# [full-text]: https://opensearch.org/docs/latest/query-dsl/full-text/index/

# #### 逐次検索と転置インデックス
# 逐次検索とは、検索対象の複数のデータを順次照合することで対象となるデータを抽出する方法です。事前にインデックスを作成せずに検索を行うことができますが、順次データの走査を行うためデータサイズに応じて検索時間が伸びていきます。UNIX の grep コマンドは逐次検索にあたります。
# 
# <img src="./img/grep.png" width=1024>
# 
# 
# 一方、OpenSearch は全文検索向けに、Apache Lucene が提供する転置インデックスを使用しています。転置インデックスは特定の見出し語に対応する文書 ID と出現位置の情報を持っています。これにより、特定語句にマッチする検索結果を素早く返すことを可能としています。
# 
# <img src="./img/inverted_index.png" width=1024>
# 

# #### text 型とアナライザー
# 全文検索の対象となるデータは、text 型のフィールドに格納される必要があります。text 型のフィールドに格納されたデータは、以下の流れで処理され、転置インデックスに登録されます
# 
# 1. Character Filter によるテキスト内の文字の置換、削除 (オプション)
# 2. Tokenizer によるテキストのトークン分割 (必須)
# 3. Token Filter によるトークンの正規化処理 (オプション)
# 
# これらの処理を通じて作成されたトークンがインデックスに格納されることで、キーワードによる文字列の検索が可能となります。この一連の処理をアナライズといい、これらのコンポーネントの集合体をアナライザーと呼びます。
# 
# <img src="./img/analyzer.png" width=1024>
# 
# アナライザーは、インデックスにドキュメントを格納する際の文字列の処理、また検索時の検索キーワードの処理に用いられます。一般的には格納と検索に同じアナライザーを用いますが、オートコンプリートなど一部のユースケースでは格納時と検索時で異なるアナライザーを使用します。
# 
# アナライザーは text 型フィールドの新規定義時にのみ指定が可能です。text 型フィールドに対してアナライザーを明示的に指定しない場合、OpenSearch はデフォルトで standard アナライザーを設定します。フィールドに設定されたアナライザーを後から変更することはできないため、事前に検索要件を確認した上でカスタムアナライザーを定義することがより良い検索体験の提供に繋がります。
# 
# OpenSearch では、いくつかの言語向けにアナライザーが標準で提供されています。これらのアナライザーは、予め決められた Character Filter, Tokenizer, Token Filter で構成されています。
# 
# デフォルトのアナライザーである Standard Analyzer は以下のコンポーネントで構成されており、主に英語などの半角スペースでターム(語句)が区切られている言語の文章解析に使用できます。トークンは Lower Case Token Filter によりすべて小文字に変換され格納されるため、大文字・小文字の違いに影響されない検索が可能です。
# 
# - Character Filter: なし
# - Tokenizer: Standard Tokenizer
# - Token Filter: Lower Case Token Filter
# 
# OpenSearch では英語やフランス語、日本語など各言語に応じたアナライザーのプリセットを提供しています。本ワークショップでは plot フィールドに english アナライザーをセットすることで、解説文による部分検索の精度を高めています。
# 

# <div class="alert alert-block alert-info"> 
# <b>Tips: カスタムアナライザ</b>
# 
# 標準で提供されているアナライザーが要件を満たさない場合は、任意の Character Filter, Tokenizer, Token Filter を組み合わせたカスタムアナライザーを使用できます。
# 
# 例えば、standard analyzer に加えて以下の処理を追加したい場合は、カスタムアナライザーを定義します。
# 
# トークンのステミング(sampling -> sample や cars -> car といった語幹の統一)を行う
# a や the などのストップワードを除去する。ストップワード判定の際、大文字小文字の違いは無視する
# カスタムアナライザーは、インデックス作成時の anaysis オプション配下で定義可能です。カスタムアナライザーをテストする場合は、カスタムアナライザーが定義されているインデックスに対して　_analyze API を発行します。
# 
# standard アナライザーに対する _analyze API 結果と比較してみると、on などの不要なトークンが削除され、意味のある単語のみが抽出されていることが分かります。また一部の単語についてはステミングされていることも確認できます。
# 
# 
# </div>

# #### アナライザーの動作確認
# _analyze API を使用することで、アナライザーの動作確認を行えます。standard と english など、アナライザー毎の処理結果の違いを見ることにも役立ちます
# 

# In[11]:


payload = {
  "text": "OpenSearch is a distributed search and analytics engine based on Apache Lucene.",
  "analyzer": "standard"
}

response = opensearch_client.indices.analyze(
  body=payload
)
df_standard = pd.json_normalize(response["tokens"])
print(json.dumps(response, indent=2))


# In[12]:


payload = {
  "text": "OpenSearch is a distributed search and analytics engine based on Apache Lucene.",
  "analyzer": "english"
}

response = opensearch_client.indices.analyze(
  body=payload
)

df_english = pd.json_normalize(response["tokens"])
print(json.dumps(response, indent=2))


# 各アナライザーによって出力されるトークンを表形式で比較してみます。
# 
# standard analyzer はシンプルに文章の単語分割が行われています。一方、 english analyzer では is や a といった検索上意味を持たないワード(ストップワード)の除去やステミング(distrubuted -> distribut など、単語の変化しない前方部分のみの抽出)なども行われています。
# 
# english analyzer の方が一見すると検索精度が良いようにみえますが、"To be, or not to be, that is the question" を english analyzer で処理すると "question" しか残らないなど、対象のテキストによっては思わぬ副作用を生みます。
# 
# 例えば映画のタイトル検索では、"avengers" と "avenger" は区別したいという要件がある場合、english analyzer ではなく standard analyzer を使用することが望ましいと考えられます。一方で映画のプロット検索ではそこまで厳密な区別が要求されないと考えられること、ストップワードを除去した方がノイズが少ないと考えられることから、english analyzer を使用することが望ましいと考えられます。
# 
# 以上の点を踏まえて、本ラボでは、title フィールドについては standard analyzer を、plot フィールドについては english analyzer をセットしていきます。

# In[13]:


pd.merge(df_standard, df_english, on=["start_offset","end_offset", "position", "type"], how="left").rename(columns={"token_x": "token_standard", "token_y": "token_english"}).reindex(["start_offset","end_offset","position","type","token_standard","token_english"],axis=1).fillna("")


# #### match query による全文検索
# match query は、単一フィールドに対する全文検索を実行するものです。以下は "avengers" を含むタイトルの映画を検索するクエリです。

# In[14]:


index_name = "movies"

payload = {
  "query": {
    "match": {
      "title": "avengers"
    }
  }
}

response = opensearch_client.search(
    index=index_name,
    body=payload
)

print(json.dumps(response, indent=2))


# OpenSearch の検索 API の実行結果は JSON 形式で返却されます。これを表形式に変換、出力した結果は以下の通りです。
# 本ワークショップでは、以降見やすさを重視して結果を表形式で出力します。`pd.json_normalize` から始まる行をコメントアウトし、コメントアウトされている `print` から始まる行のコメントを解除することで JSON による出力結果を確認することもできます
# 
# ```
# # print(json.dumps(response, indent=2))
# pd.json_normalize(response["hits"]["hits"])
# ```

# In[15]:


pd.json_normalize(response["hits"]["hits"])


# title フィールドについては avenger と avengers は分けて扱われるため、avengers ではなく avenger で検索すると異なる結果が得られます

# In[16]:


index_name = "movies"

payload = {
  "query": {
    "match": {
      "title": "avenger"
    }
  }
}

response = opensearch_client.search(
    index=index_name,
    body=payload
)

# print(json.dumps(response, indent=2))
pd.json_normalize(response["hits"]["hits"])


# 一方、plot フィールドは english analyzer によりステミングが行われるため、superhero で検索した際に superheros を含むドキュメントもヒットします
# 
# [Highlight][highlight] と呼ばれる機能を使うことで、フィールド内のどの文字列(トークン)でヒットしたかを確認することができます。
# 
# デフォルトでは、ヒットしたトークンは em タグで囲まれて出力されます。タグは変更が可能です。
# 
# [highlight]: https://opensearch.org/docs/latest/search-plugins/searching-data/highlight/

# In[17]:


index_name = "movies"

payload = {
  "size": 3,
  "query": {
    "match": {
      "plot": "superhero"
    }
  },
  "highlight": {
    "fields": {
      "plot": {}
    }
  } 
}

response = opensearch_client.search(
    index=index_name,
    body=payload
)

#print(json.dumps(response, indent=2))
pd.json_normalize(response["hits"]["hits"])


# #### match query における複数キーワードによる検索 (OR / AND)
# テキスト検索を行う際、複数のキーワードによる OR or AND 検索を行う要件があります。match query も複数キーワード検索に対応しています。
# 複数キーワードを入力した場合、デフォルトでは OR 検索となります。

# In[18]:


index_name = "movies"

payload = {
  "query": {
    "match": {
      "title": {
        "query": "avenger avengers"
      }
    }
  },
  "highlight": {
    "fields": {
      "plot": {}
    }
  } 
}

response = opensearch_client.search(
    index=index_name,
    body=payload
)

#print(json.dumps(response, indent=2))
pd.json_normalize(response["hits"]["hits"])


# AND 検索は、operator オプションに明示的に and と指定することで実装可能です。同オプションのデフォルトは or となっています。

# In[19]:


index_name = "movies"

payload = {
  "query": {
    "match": {
      "plot": {
        "query": "superhero transform academy",
        "operator": "and"
      }
    }
  },
  "highlight": {
    "fields": {
      "plot": {}
    }
  } 
}

response = opensearch_client.search(
    index=index_name,
    body=payload
)

#print(json.dumps(response, indent=2))
pd.json_normalize(response["hits"]["hits"])


# [minimum_should_match][minimum_should_match] オプションを使うことで、クエリに含まれるキーワードにおいて、特定数、もしくは特定割合のワード数マッチすればヒットしたとみなすことも可能です。
# 
# [minimum_should_match]: https://opensearch.org/docs/latest/query-dsl/minimum-should-match/

# In[20]:


index_name = "movies"

payload = {
  "query": {
    "match": {
      "plot": {
        "query": "superhero transform academy",
        "minimum_should_match": "2"
      }
    }
  },
  "highlight": {
    "fields": {
      "plot": {}
    }
  } 
}

response = opensearch_client.search(
    index=index_name,
    body=payload
)

pd.json_normalize(response["hits"]["hits"])
#print(json.dumps(response, indent=2))


# #### multi match query によるフィールド横断の全文検索
# 同一の検索条件で複数フィールドにまたがった検索を実行したい場合は、[Multi-match][multi-match] クエリを使用します。
# 
# title もしくは plot フィールドに wind の文字列を含むドキュメントを検索する場合、Multi-match を利用して以下のように記述することができます。
# 
# [multi-match]: https://opensearch.org/docs/latest/query-dsl/full-text/multi-match/

# In[21]:


index_name = "movies"

payload = {
  "size": 3,
  "query": {
    "multi_match": {
      "query": "wind",
      "fields": ["title", "plot"]
    }
  },
  "highlight": {
    "fields": {
      "title": {},
      "plot": {}
    }
  } 
}

response = opensearch_client.search(
    index=index_name,
    body=payload,
    filter_path="hits.hits"
)

# print(json.dumps(response, indent=2))
pd.json_normalize(response["hits"]["hits"])


# #### フレーズ検索
# 語句が特定の順序で並んでいるドキュメントのみを検索したい場合、フレーズ検索の機能が有用です。
# 以下のクエリは "iron man" を match query で検索した例ですが、"Iron man" シリーズだけではなく "The Man with the Iron Fists" などもヒットしています。

# In[22]:


index_name = "movies"

payload = {
  "query": {
    "match": {
      "title": {
        "query": "iron man",
        "operator": "and"
      }
    }
  },
  "size": 5
}

response = opensearch_client.search(
    index=index_name,
    body=payload
)

#print(json.dumps(response, indent=2))
pd.json_normalize(response["hits"]["hits"])


# [match_phrase][match_phrase] クエリを使用することで、特定の並び順でトークンが配置されているドキュメントのみを検索することが可能です。
# 
# [match_phrase]: https://opensearch.org/docs/latest/query-dsl/full-text/match-phrase/

# In[23]:


index_name = "movies"

payload = {
  "query": {
    "match_phrase": {
      "title": {
        "query": "iron man",
      }
    }
  },
  "size": 5
}

response = opensearch_client.search(
    index=index_name,
    body=payload
)

#print(json.dumps(response, indent=2))
pd.json_normalize(response["hits"]["hits"])


# slop パラメーターを変更することで、ワードの順序が入れ替わっていてもフレーズ検索にマッチさせることが可能です。デフォルトは 0 であるため、厳密なマッチングが要求されます。
# 
# slop を 2 に変更すると、2 つのキーワードが入れ替えの対象になります。

# In[24]:


index_name = "movies"

payload = {
  "query": {
    "match_phrase": {
      "title": {
        "query": "man iron 2",
        "slop": 2
      }
    }
  },
  "size": 5
}

response = opensearch_client.search(
    index=index_name,
    body=payload
)

#print(json.dumps(response, indent=2))
pd.json_normalize(response["hits"]["hits"])


# 3 つのキーワードを入れ替えたい場合は、slop を 3 にセットします。

# In[25]:


index_name = "movies"

payload = {
  "query": {
    "match_phrase": {
      "title": {
        "query": "man 2 iron",
        "slop": 3
      }
    }
  },
  "size": 5
}

response = opensearch_client.search(
    index=index_name,
    body=payload
)

#print(json.dumps(response, indent=2))
pd.json_normalize(response["hits"]["hits"])


# ### 完全一致検索 (Term query)
# 検索キーワードに完全に一致するドキュメントを取得する場合は、Term query を使用します。主に keyword タイプのフィールドに対する検索で使用します。
# 以下のクエリは、genres フィールドに "Comedy" という値を持つドキュメントを取得しています。検索結果の件数は size と呼ばれるパラメーターで制御可能です。

# In[26]:


index_name = "movies"

payload = {
  "size": 3,
  "query": {
    "term": {
      "genres": {
        "value": "Comedy"
      }
    }
  }
}

response = opensearch_client.search(
    index=index_name,
    body=payload
)

#print(json.dumps(response, indent=2))
pd.json_normalize(response["hits"]["hits"])


# #### normalizer による正規化
# 完全一致検索の中でも、case-sensitive (大文字小文字の区別) の要否が分かれる場合があります。今回のようにジャンル名での検索であれば、Comedy ではなく comedy でも同様の結果を得られた方が好ましい場合があります。こうした要件に対応するために、OpenSearch では [Normalizer][normalizer] と呼ばれる機能を提供しています。Normalizer はデータ格納時、およびクエリ実行時に文字列の正規化を行う機能です。
# 
# 本ワークショップでは、genres フィールドに lowercase_normalizer と呼ばれる Normalizer をセットしています。この Normalizer は lowercase filter により入力された文字を自動的に小文字に変換するよう設定されています。
# 
# Normalizer はインデックスの設定(settings/analysis)内で定義します。定義した Normalizer は、フィールドの normalizer オプションで指定することができます。以下はインデックス設定の抜粋です。
# 
# ```json
# {
#   "settings": {
#     "analysis": {
#       "normalizer": {
#         "lowercase_normalizer": {
#           "type": "custom",
#           "filter": ["lowercase"]
#         }
#       }
#     }
#   },
#   "mappings": {
#     "properties": {
#       "genres":  {
#         "type": "keyword",
#         "normalizer": "lowercase_normalizer"
#       }
#     }
#   }
# }
# ```
# 
# Normalizer により入力された文字が小文字に統一されるため、以下のように "cOmEdY" という文字列で一致検索を行っても、"Comedy" で検索した時と同様の結果が得られます。
# 
# [normalizer]: https://opensearch.org/docs/latest/analyzers/normalizers/

# In[27]:


index_name = "movies"

payload = {
  "size": 3,
  "query": {
    "term": {
      "genres": {
        "value": "cOmEdY"
      }
    }
  }
}

response = opensearch_client.search(
    index=index_name,
    body=payload
)

print(json.dumps(response, indent=2))


# #### Terms
# 指定した複数の検索条件のいずれかに完全一致するドキュメントを取得したい場合は、term ではなく [terms][terms] クエリを使用します
# 
# 以下のクエリでは comedy もしくは drama のジャンルの映画を検索しています
# 
# [terms]: https://opensearch.org/docs/latest/query-dsl/term/terms/

# In[28]:


index_name = "movies"

payload = {
  "size": 3,
  "query": {
    "terms": {
      "genres": ["comedy", "drama"]
    }
  }
}

response = opensearch_client.search(
    index=index_name,
    body=payload
)

#print(json.dumps(response, indent=2))
pd.json_normalize(response["hits"]["hits"])


# #### Terms set
# Terms は配列内のいずれかの文字列と完全一致していればマッチした文書とみなして検索結果を返していました。一方で、配列内のすべて、もしくは一部の文字列と完全一致している場合にマッチしたとみなしたいケースもあります。
# こうしたケースでは Terms set クエリが有用です。以下は comedy, drama, sci-fi, family の要素をすべて含むジャンルの映画を検索するものです

# In[29]:


index_name = "movies"

payload = {
  "size": 3,
  "query": {
    "terms_set": {
      "genres": {
        "terms": ["comedy", "drama", "sci-fi", "family"],
        "minimum_should_match_script": {
          "source": "params.num_terms"
        }
      }
    }
  }
}

response = opensearch_client.search(
    index=index_name,
    body=payload
)

#print(json.dumps(response, indent=2))
pd.json_normalize(response["hits"]["hits"])


# #### Terms lookup
# [Terms lookup][terms-lookup] を利用することで、他インデックスのデータを使用した Terms query が実行可能です。 
# 
# 以下のサンプルコードでは、ユーザーごとのお気に入り情報を格納した movie-users インデックスを作成・データを格納し、ユーザーに対応したお気に入りジャンルの映画を検索しています
# 
# [terms-lookup]: https://opensearch.org/docs/latest/query-dsl/term/terms/#terms-lookup

# In[30]:


lookup_index_name = "movie-users"

payload = {
  "settings": {
    "index": {
      "number_of_shards": 1,
      "number_of_replicas": 0
    }
  },
  "mappings": {
    "properties": {
      "userid":  { "type": "keyword" },
      "favorite-genres":  { "type": "keyword" }
    }
  }
}

try:
    # 既に同名のインデックスが存在する場合、いったん削除を行う
    print("# delete index")
    response = opensearch_client.indices.delete(index=lookup_index_name)
    print(json.dumps(response, indent=2))
except Exception as e:
    print(e)

# インデックスの作成を行う
print("# create index")
response = opensearch_client.indices.create(index=lookup_index_name, body=payload)
print(json.dumps(response, indent=2))


# In[31]:


lookup_index_name = "movie-users"

payload = {
  "userid": "00000001",
  "favorite-genres": ["comedy", "drama"]
}

response = opensearch_client.index(
    index = lookup_index_name,
    body = payload,
    id = "00000001",
    refresh = True
)

print(json.dumps(response, indent=2))


# In[32]:


index_name = "movies"
lookup_index_name = "movie-users"

payload = {
  "size": 3,
  "query": {
    "terms": {
      "genres": {
        "index": lookup_index_name,
        "id": "00000001",
        "path": "favorite-genres"
      }
    }
  }
}

response = opensearch_client.search(
    index=index_name,
    body=payload
)

#print(json.dumps(response, indent=2))
pd.json_normalize(response["hits"]["hits"])


# <div class="alert alert-block alert-info"> 
# <b>Tips: 参照させる要素数が多い場合の対応</b>
#     
# Terms lookup で与える要素数が 10,000 を超える場合は、パフォーマンス向上のために [Bitmap Filtering][bitmap-filtering] の利用を検討してください。
# 
# [bitmap-filtering]: https://opensearch.org/docs/latest/query-dsl/term/terms/#bitmap-filtering
# </div>
# 
# 

# ### 部分一致検索・あいまい検索
# Keyword フィールドに対して部分的に一致する検索条件を記載したい場合や、多少の表記ゆれや誤字脱字をフォローした検索を行いたい場合の検索手法について解説します。
# 
# 全ドキュメントに対して横断的に検索が実行されるため、通常の term クエリと比較してパフォーマンスは大きく劣ります。
# 
# <div class="alert alert-block alert-warning"> 
# <b>あいまい検索利用時の注意事項</b>
# 
# これらの検索は、インデックス内の全ドキュメントが検索対象となる可能性があり、パフォーマンスへの影響があります。
# 
# 対象フィールドに対する検索クエリの大部分を部分一致検索やあいまい検索が占める場合は、text 型フィールドを使用した全文検索に切り替えるか、対象のフィールドを [wildcard][wildcard] 型に変更することを検討してください。
# 
# </div>
# 
# 
# [wildcard]: https://opensearch.org/docs/latest/field-types/supported-field-types/wildcard/

# #### Prefix query
# [Prefix query][prefix] は、前方一致に基づく検索を行うクエリです。特定の文字列から始まるテキストを検索する際に役立ちます。
# 
# 以下のサンプルクエリでは、com から始まる genres の映画を検索します。
# 
# [prefix]: https://opensearch.org/docs/latest/query-dsl/term/prefix/

# In[33]:


index_name = "movies"

payload = {
  "size": 3,
  "query": {
    "prefix": {
      "genres": "com"
    }
  }
}

response = opensearch_client.search(
    index=index_name,
    body=payload
)

# print(json.dumps(response, indent=2))
pd.json_normalize(response["hits"]["hits"])


# #### Regexp query
# Regexp query は、正規表現に基づく検索を行うクエリです。特定のパターンに一致する文字列を抽出する際に使用できます。
# 
# 以下のサンプルクエリでは、アルファベット 5 文字の genres の映画を検索します。
# 
# [regexp]: https://opensearch.org/docs/latest/query-dsl/term/regexp/

# In[34]:


index_name = "movies"

payload = {
  "size": 1,
  "query": {
    "regexp": {
      "genres": "[a-z]{5}"
    }
  }
}

response = opensearch_client.search(
    index=index_name,
    body=payload
)

# print(json.dumps(response, indent=2))
pd.json_normalize(response["hits"]["hits"])


# #### Wildcard query
# [Wildcard query][wildcard] は、ワイルドカードを使用した部分一致検索を提供します。以下のクエリでは、 c から始まって y で終わる genres の映画を検索します。
# 
# [wildcard]: https://opensearch.org/docs/latest/query-dsl/term/wildcard/

# In[35]:


index_name = "movies"

payload = {
  "size": 1,
  "query": {
    "wildcard": {
      "genres": "c*y"
    }
  }
}

response = opensearch_client.search(
    index=index_name,
    body=payload
)

# print(json.dumps(response, indent=2))
pd.json_normalize(response["hits"]["hits"])


# #### Fuzzy query
# [Fuzzy query][fuzzy] は、以下のような表記ゆれをクエリ側で吸収する機能です。
# 
# - 誤字: **c**at to **b**at
# - 不要文字の混入: cat to cat**s**
# - 脱字: **c**at to at
# - タイプミス(順序の入れ替わり): **ca**t to **ac**t
# 
# [fuzzy]: https://opensearch.org/docs/latest/query-dsl/term/fuzzy/

# 以下のサンプルクエリでは、comedy から二文字かけた cmdy でも検索にヒットするように fuziness を 2 にセットして Fuzzy query を実行しています。
# 
# fuziness はデフォルトで `AUTO` という値がセットされており、OpenSearch 側で語句の長さから自動的に補完文字数をセットする仕様になっています。fuziness に数値を明示的にセットすることで、何文字まで補完されるかをコントロールすることが可能です。

# In[36]:


index_name = "movies"

payload = {
  "size": 1,
  "query": {
    "fuzzy": {
      "genres": {
        "value": "cmdy",
        "fuzziness": 2
      }
    }
  }
}

response = opensearch_client.search(
    index=index_name,
    body=payload
)

# print(json.dumps(response, indent=2))
pd.json_normalize(response["hits"]["hits"])


# 以下のサンプルクエリでは、comedy のタイプミスである ocmedy で検索しています。
# 
# タイプミスにより隣り合った文字が入れ替わったケースに対応するかどうかは、**transpositions** オプションで制御します。同オプションのデフォルト値は **true** となっているため、Fuzzy query のデフォルトの動作としては、タイプミスをフォローする形になっているといえます。

# In[37]:


index_name = "movies"

payload = {
  "size": 1,
  "query": {
    "fuzzy": {
      "genres": {
        "value": "ocmedy",
      }
    }
  }
}

response = opensearch_client.search(
    index=index_name,
    body=payload
)

# print(json.dumps(response, indent=2))
pd.json_normalize(response["hits"]["hits"])


# 文字の入れ替わりを何か所まで許容するかは、**fuzziness** パラメーターで制御可能です。
# 
# ocemdy で comedy が含まれるドキュメントをヒットさせたい場合、**fuzziness** が `1` ではヒットしません。`2` 以上ではヒットします

# In[38]:


index_name = "movies"

payload = {
  "size": 1,
  "query": {
    "fuzzy": {
      "genres": {
        "value": "ocemdy",
        "fuzziness": 1
      }
    }
  }
}

response = opensearch_client.search(
    index=index_name,
    body=payload
)

# print(json.dumps(response, indent=2))
pd.json_normalize(response["hits"]["hits"])


# In[39]:


index_name = "movies"

payload = {
  "size": 1,
  "query": {
    "fuzzy": {
      "genres": {
        "value": "ocemdy",
        "fuzziness": 2
      }
    }
  }
}

response = opensearch_client.search(
    index=index_name,
    body=payload
)

# print(json.dumps(response, indent=2))
pd.json_normalize(response["hits"]["hits"])


# タイプミスを Fuzzy query でフォローしない場合、**transpositions** に `false` をセットします。

# In[40]:


index_name = "movies"

payload = {
  "size": 1,
  "query": {
    "fuzzy": {
      "genres": {
        "value": "ocemdy",
        "fuzziness": 2,
        "transpositions": False
      }
    }
  }
}

response = opensearch_client.search(
    index=index_name,
    body=payload
)

#print(json.dumps(response, indent=2))
pd.json_normalize(response["hits"]["hits"])


# transpositions による補正は隣り合った文字同士でのみ機能します。mocedy のように飛び石で入れ替わってしまったケースは対応できません。

# In[41]:


index_name = "movies"

payload = {
  "size": 1,
  "query": {
    "fuzzy": {
      "genres": {
        "value": "mocedy",
        "fuzziness": 1,
        "transpositions": True
      }
    }
  }
}

response = opensearch_client.search(
    index=index_name,
    body=payload
)

# print(json.dumps(response, indent=2))
pd.json_normalize(response["hits"]["hits"])


# このようなケースでも、Fuzzy query は誤字の補正対応で対応可能です。 fuzziness の値に対応した文字数までであれば、入れ替わりや誤字に対して対処可能です。

# In[42]:


index_name = "movies"

payload = {
  "size": 1,
  "query": {
    "fuzzy": {
      "genres": {
        "value": "mocedy",
        "fuzziness": 2,
        "transpositions": False
      }
    }
  }
}

response = opensearch_client.search(
    index=index_name,
    body=payload
)

#print(json.dumps(response, indent=2))
pd.json_normalize(response["hits"]["hits"])


# 以下のケースでは、入れ替わりに加えて m ではなく n のタイプミスも加わっていますが、**n**o**c**edy のうち n と c の 2 文字が fuzziness = 2 の設定により補完され、**c**o**m**edy とマッチすると判定されています。

# In[43]:


index_name = "movies"

payload = {
  "size": 1,
  "query": {
    "fuzzy": {
      "genres": {
        "value": "nocedy",
        "fuzziness": 2,
        "transpositions": False
      }
    }
  }
}

response = opensearch_client.search(
    index=index_name,
    body=payload
)

#print(json.dumps(response, indent=2))
pd.json_normalize(response["hits"]["hits"])


# ### 検索結果のソート
# 本セクションでは、検索結果の並べ替え手法について解説します。

# #### スコアによるソート
# OpenSearch では、[キーワード検索][keyword-search]における文書の関連度計算に [Okapi BM25][okapi-bm25] を使用しています。BM25 は、クエリ内に出現する単語の語彙検索を実行するキーワードベースのアルゴリズムです。
# 
# 文書の関連性を判断する際、BM25 は用語頻度(TF - Term Frequency) および 逆文書頻度 (IDF - Inverse Document Frequency) を考慮します。
# 
# TF(用語頻度)は、特定文書における、特定の単語の出現頻度を示します。検索対象の用語がより頻繁に出現する文書を、関連性が高い文書として扱います。
# 一方、IDF(逆文書頻度)は、コーパス内のすべての文書に共通して頻出する単語の重みを低くするものです。the や a のような冠詞が該当します。
# 
# <img src="./img/bm25.png" width=1024>
# 
# [keyword-search]: https://opensearch.org/docs/latest/search-plugins/keyword-search/
# [okapi-bm25]:https://en.wikipedia.org/wiki/Okapi_BM25
# 
# 計算後の score は "_score" フィールドに格納されています。Search API 実行時に [explain][explain] オプションを付与することで、詳細なスコア計算の過程を確認することが可能です。
# 
# [explain]: https://opensearch.org/docs/latest/api-reference/explain/

# In[44]:


index_name = "movies"

payload = {
  "size": 3,
  "query": {
    "match": {
      "plot": "superhero"
    }
  },
  "highlight": {
    "fields": {
      "plot": {}
    }
  },
  "explain": True
}

response = opensearch_client.search(
    index=index_name,
    body=payload
)

# print(json.dumps(response, indent=2))
pd.json_normalize(response["hits"]["hits"])


# #### 重みづけによるスコアの調整
# Multi match query など、複数フィールドに対して横断的に検索を行う場合、フィールドごとに重みづけを設定することが可能です。以下の例では、title と plot に対する横断検索を行う際、title フィールドの重みを 3 倍にしています。

# In[45]:


index_name = "movies"

payload = {
  "size": 3,
  "query": {
    "multi_match": {
      "query": "wind",
      "fields": ["title^3", "plot"]
    }
  },
  "highlight": {
    "fields": {
      "title": {},
      "plot": {}
    }
  },
  "explain": True
}

response = opensearch_client.search(
    index=index_name,
    body=payload
)

#print(json.dumps(response, indent=2))
pd.json_normalize(response["hits"]["hits"])


# #### スクリプトによるスコアの調整
# [Script score][script-score] クエリを使用することで、スコア計算を script ベースで行うことも可能です。以下の例では、関連度に基づいて算出された score を rank の値で割った値を最終的な score としています。
# 
# [script-score]: https://opensearch.org/docs/latest/query-dsl/specialized/script-score/

# In[46]:


index_name = "movies"

payload = {
  "size": 3,
  "query": {
    "script_score": {
      "query": {
        "match": { 
            "plot": "superhero" 
        }
      },
      "script": {
        "source": "_score / doc['rank'].value"
      }
    }
  }
}

response = opensearch_client.search(
    index=index_name,
    body=payload
)

#print(json.dumps(response, indent=2))
pd.json_normalize(response["hits"]["hits"])


# Script score では sigmoid といったいくつかの関数も提供しています。

# In[47]:


index_name = "movies"

payload = {
  "size": 3,
  "query": {
    "script_score": {
      "query": {
        "match": { 
            "plot": "superhero" 
        }
      },
      "script": {
        "source": "sigmoid(_score, 2, 1)"
      }
    }
  }
}

response = opensearch_client.search(
    index=index_name,
    body=payload
)

#print(json.dumps(response, indent=2))
pd.json_normalize(response["hits"]["hits"])


# #### キーワードによる Boosting
# スコアの補正は検索キーワードに対しても適用することが可能です。[Boosting][boosting] クエリを使用することで、negative に指定したキーワードにマッチするドキュメントのスコアを下げて検索順位を変動させることができます。以下の例では、private という文字列を plot フィールドに含むドキュメントのスコアを 0.9 倍にセットしています。
# 
# [boosting]: https://opensearch.org/docs/latest/query-dsl/compound/boosting/

# In[48]:


index_name = "movies"

payload = {
  "size": 3,
  "query": {
    "boosting": {
      "positive": {
        "match": {
          "plot": "superhero"
        }
      },
      "negative": {
        "match": {
          "plot": "private"
        }
      },
      "negative_boost": 0.9
    }
  }
}

response = opensearch_client.search(
    index=index_name,
    body=payload
)

#print(json.dumps(response, indent=2))
pd.json_normalize(response["hits"]["hits"])


# #### 特定フィールドによるソート
# スコアではなく特定のフィールドに基づいてソートを行いたい場合は、[sort][sort] オプションを検索 API に追加します。
# 以下のサンプルでは、公開年度が新しい映画の上位 3 位を取得しています。query オプションで指定している match_all は、全てのドキュメントが取得対象であることを示します。また上位 5 位だけを取得するために、size オプションに 5 を設定しています。
# 
# [sort]: https://opensearch.org/docs/latest/search-plugins/searching-data/sort/

# In[49]:


index_name = "movies"

payload = {
  "size": 5,
  "query": {
    "match_all": {}
  },
  "sort": [
    {
      "year": {
        "order": "desc"
      }
    }
  ]
}

response = opensearch_client.search(
    index=index_name,
    body=payload
)

#print(json.dumps(response, indent=2))
pd.json_normalize(response["hits"]["hits"])


# sort には複数条件を記載することができます。上記のクエリでは 2016 年の映画が下位 3 件を占めていました。2016 年に公開された映画を ID 順に並び替えたいと思います。
# 以下のように year によるソート条件の下に id によるソート条件を追加することで、year によるソートの後に id によるソートを実行することが可能です。複数条件で sort を行う場合は、先に実行したいソート条件を上に記載します。

# In[50]:


index_name = "movies"

payload = {
  "size": 5,
  "query": {
    "match_all": {}
  },
  "sort": [
    {
      "year": {
        "order": "desc"
      }
    },
    {
      "id": {
        "order": "asc"
      }
    }
  ]
}

response = opensearch_client.search(
    index=index_name,
    body=payload
)

#print(json.dumps(response, indent=2))
pd.json_normalize(response["hits"]["hits"])


# ### 範囲検索
# 数値や日付を元に連続した範囲で検索を行いたい場合は [Range query][range] を使用します。 以下の例では 1920 年から 1923 年に公開された映画を検索しています。
# 
# [range]: https://opensearch.org/docs/latest/query-dsl/term/range/

# In[51]:


index_name = "movies"

payload = {
  "query": {
    "range": {
      "year": {
        "gte": 1920,
        "lte": 1923
      }
    }
  }
}

response = opensearch_client.search(
    index=index_name,
    body=payload
)

#print(json.dumps(response, indent=2))
pd.json_normalize(response["hits"]["hits"])


# ### 複合検索
# [Boolean query][bool] を利用することで、複数の条件を組み合わせた検索を行うことが可能です。以下の要素を使用します。
# 
# - must: 条件に全て合致するドキュメントを返す
# - should: 1 つ以上の条件に合致するドキュメントを返す
# - filter: 条件に全て合致するドキュメントを返す。must と違って score 計算に関与しない。
# - must_not: 条件に合致するドキュメントを検索結果から除外する
# 
# [bool]: https://opensearch.org/docs/latest/query-dsl/compound/bool/
# 
# 

# #### Boolean query による複合検索
# 実際に複合条件による検索を行っていきます。まず、クエリに含まれる条件を列挙します
# 
# - must: plot に school が含まれること
# - must: genres に Music が必ず含まれること。
# - should: genres に Romance または Comedy のいずれか 1 つ (minimum_should_match にて指定)が含まれていること
# - must_not: genres に Thriller または Horror を含まないこと
# - filter: year が 2000 から 2009 の範囲であること
# - sort: 検索結果を rating フィールドで降順に並べ替えること
# - size: 検索結果のうち上位 5 件を出力すること
# 
# 

# In[52]:


index_name = "movies"

payload = {
  "size": 5,
  "query": {
    "bool": {
      "must": [
        {
          "match": {
            "plot": "school"
          }
        },
        {
          "term": {
            "genres": {
              "value": "Music"
            }
          }
        }
      ],
      "should": [
        {
          "term": {
            "genres": {
              "value": "Romance"
            }
          }
        },
        {
          "term": {
            "genres": {
              "value": "Comedy"
            }
          }
        }
      ],
      "minimum_should_match": 1,
      "filter": [
        {
          "range": {
            "year": {
              "gte": 2000,
              "lte": 2009
            }
          }
        }
      ],
      "must_not": [
        {
          "terms": {
            "genres": ["Thriller", "Horror"]
          }
        }
      ]
    }
  },
  "highlight": {
    "fields": {
      "plot": {}
    }
  },
  "sort": [
    {
       "rating": {
          "order": "desc"
      }
    }
  ]
}

response = opensearch_client.search(
    index=index_name,
    body=payload
)

#print(json.dumps(response, indent=2))
pd.json_normalize(response["hits"]["hits"])


# #### Zero terms query の取り扱い
# 
# 検索においては、match クエリに与えられるトークンが実質的に 0 件という状況が発生します。例えば以下のようなケースが考えられます。
# 
# - ユーザーが入力するクエリ文字列がストップワードのみで構成されているため、アナライザーによって全てのトークンが削除されてしまう ("To be or not to be" など)
# - そもそもクエリが入力されておらず、フィルタ条件のみが指定されている
# 
# この場合、OpenSearch のデフォルトの挙動は検索を行わない、つまり 0 件ヒットとなります。
# 
# しかしながら、キーワードが入力されない場合はフィルタ条件として入力したジャンルや公開年に基づいて絞り込まれた映画の、rating が高いものを返却するのが自然な挙動と考えられます。
# この挙動は、[zero_terms_query][zero_terms_query] のオプションに all と指定することで実現可能です。以下 2 つのクエリを実行することで、zero_terms_query の設定による結果の差異を確認することが可能です。1 つ目のクエリは zero_terms_query オプションにデフォルトの none を、2 つ目のクエリは all をセットしています。
# 
# </div>
# 
# [zero_terms_query]: https://opensearch.org/docs/latest/query-dsl/full-text/match/#empty-query

# In[53]:


index_name = "movies"

payload = {
  "size": 5,
  "query": {
    "bool": {
      "must": [
        {
          "match": {
            "plot": {
              "query": "",
              "zero_terms_query": "none"
            }
          }
        },
        {
          "term": {
            "genres": {
              "value": "Music"
            }
          }
        }
      ],
      "should": [
        {
          "term": {
            "genres": {
              "value": "Romance"
            }
          }
        },
        {
          "term": {
            "genres": {
              "value": "Comedy"
            }
          }
        }
      ],
      "minimum_should_match": 1,
      "filter": [
        {
          "range": {
            "year": {
              "gte": 2000,
              "lte": 2009
            }
          }
        }
      ],
      "must_not": [
        {
          "terms": {
            "genres": ["Thriller", "Horror"]
          }
        }
      ]
    }
  },
  "highlight": {
    "fields": {
      "plot": {}
    }
  },
  "sort": [
    {
       "rating": {
          "order": "desc"
      }
    }
  ]
}

response = opensearch_client.search(
    index=index_name,
    body=payload
)

#print(json.dumps(response, indent=2))
pd.json_normalize(response["hits"]["hits"])


# In[54]:


index_name = "movies"

payload = {
  "size": 5,
  "query": {
    "bool": {
      "must": [
        {
          "match": {
            "plot": {
              "query": "",
              "zero_terms_query": "all"
            }
          }
        },
        {
          "term": {
            "genres": {
              "value": "Music"
            }
          }
        }
      ],
      "should": [
        {
          "term": {
            "genres": {
              "value": "Romance"
            }
          }
        },
        {
          "term": {
            "genres": {
              "value": "Comedy"
            }
          }
        }
      ],
      "minimum_should_match": 1,
      "filter": [
        {
          "range": {
            "year": {
              "gte": 2000,
              "lte": 2009
            }
          }
        }
      ],
      "must_not": [
        {
          "terms": {
            "genres": ["Thriller", "Horror"]
          }
        }
      ]
    }
  },
  "highlight": {
    "fields": {
      "plot": {}
    }
  },
  "sort": [
    {
       "rating": {
          "order": "desc"
      }
    }
  ]
}

response = opensearch_client.search(
    index=index_name,
    body=payload
)

#print(json.dumps(response, indent=2))
pd.json_normalize(response["hits"]["hits"])


# ### 集計 (Aggregation)
# ユーザーが検索を行う際、キーワードを入力して検索する方法とは別に、ジャンルなどの分類を元に絞り込んでいく方法があります。このような検索方法をファセットナビゲーションやファセット検索と呼びます。
# 
# ファセットナビゲーションにより、ユーザーはどのような分類が存在するのか、各分類ごとに検索対象のドキュメント件数がどの程度存在するかを事前に確認することが可能です。またおおまかな検索結果に対してファセットナビゲーションを併用することで検索結果のフィルタリングを行うことも可能です。
# 
# [AWS のドキュメント検索](https://docs.aws.amazon.com/search/doc-search.html?searchPath=documentation&searchQuery=OpenSearch) でもファセットナビゲーションが採用されています。上部のプルダウンメニューには検索結果から得られたサービス一覧などがあり、このメニューからサービスごとのドキュメントに絞った検索を行うことができます。
# 
# <img src="./img/facet.png" width=1024>
# 
# OpenSearch では、Aggregations[aggregations] クエリで集計を実行します。Aggregations は、ファセット検索で有用なキーワードの集計から、平均値(avg) や最大値(max) など統計値を出力するなど様々な集計方法を提供しています。Aggregations クエリを使用する場合、Search API で aggs オプションを使用します
# 
# 以下のサンプルでは、plot に school を含みかつ genres に Music を持つ映画について、ジャンルごとのドキュメント件数を集計しています。
# 
# [aggregations]: https://opensearch.org/docs/latest/aggregations/
# 

# In[55]:


index_name = "movies"

payload = {
  "size": 5,
  "query": {
    "bool": {
      "must": [
        {
          "match": {
            "plot": {
              "query": "school"
            }
          }
        },
        {
          "term": {
            "genres": {
              "value": "Music"
            }
          }
        }
      ]
    }
  },
  "highlight": {
    "fields": {
      "plot": {}
    }
  },
  "sort": [
    {
       "rating": {
          "order": "desc"
      }
    }
  ],
  "aggs": {
    "genres": {
      "terms": {
        "field": "genres"
      }
    }
  }
}

response = opensearch_client.search(
    index=index_name,
    body=payload
)

#print(json.dumps(response, indent=2))
pd.json_normalize(response["aggregations"]["genres"]["buckets"])


# 上記の集計結果を見て、family もフィルタリング対象のジャンルに追加したとします。この場合、以下のようにクエリを書き換えて実行することで、ファセットによるドリルダウン相当の処理を実現できます。

# In[56]:


index_name = "movies"

payload = {
  "size": 5,
  "query": {
    "bool": {
      "must": [
        {
          "match": {
            "plot": {
              "query": "school"
            }
          }
        },
        {
          "terms_set": {
            "genres": {
              "terms": ["music", "family"],
              "minimum_should_match_script": {
                "source": "params.num_terms"
              }
            }
          }
        }
      ]
    }
  },
  "highlight": {
    "fields": {
      "plot": {}
    }
  },
  "sort": [
    {
       "rating": {
          "order": "desc"
      }
    }
  ],
  "aggs": {
    "genres": {
      "terms": {
        "field": "genres"
      }
    }
  }
}

response = opensearch_client.search(
    index=index_name,
    body=payload
)

#print(json.dumps(response, indent=2))
pd.json_normalize(response["aggregations"]["genres"]["buckets"])



# ### ページング
# アプリケーション要件によっては、数千件の検索結果を扱うケースがあります。多量のデータを一括で返却することはパフォーマンス上の悪影響があるため、Opensearch では多数の結果を少量ずつ取得する[ページング処理][paginate]を実装することを推奨しています。
# 
# ページングの実装方法は以下の 4 パターンが存在します。
# 
# - from + size: 10000 件以下の検索結果における簡易なページング
# - scroll : 多数のドキュメントをバッチ処理
# - search_after
# - search_after + PIT
# 
# [paginate]: https://opensearch.org/docs/latest/search-plugins/searching-data/paginate/

# #### from + size によるページング
# [from + size][the-from-and-size-parameters] は最もシンプルな実装方法です。 from に指定した開始位置から size 件数分のデータを返却します。from に 100、size に 20 と指定した場合は、先頭から数えて 100 件目から 20 件のデータが返却されます。SQL の OFFSET と LIMIT の組み合わせと似ているところがあります。
# 
# 以下のサンプルでは、8 件の検索結果を 2 つのクエリで分割取得しています。1 つめのクエリで from に 0(デフォルト)、size に 4 を指定し先頭 4 件の検索結果を、2 つめのクエリで from に 5、size に 4 を指定して後続の 4 件の検索結果を取得しています。
# 
# [the-from-and-size-parameters]: https://opensearch.org/docs/latest/search-plugins/searching-data/paginate/#the-from-and-size-parameters

# In[57]:


index_name = "movies"

payload = {
  "from": 0,
  "size": 4,
  "query": {
    "bool": {
      "must": [
        {
          "match": {
            "plot": "school"
          }
        },
        {
          "term": {
            "genres": {
              "value": "Music"
            }
          }
        }
      ],
      "should": [
        {
          "term": {
            "genres": {
              "value": "Romance"
            }
          }
        },
        {
          "term": {
            "genres": {
              "value": "Comedy"
            }
          }
        }
      ],
      "minimum_should_match": 1,
      "filter": [
        {
          "range": {
            "year": {
              "gte": 2000,
              "lte": 2009
            }
          }
        }
      ],
      "must_not": [
        {
          "terms": {
            "genres": ["Thriller", "Horror"]
          }
        }
      ]
    }
  },
  "highlight": {
    "fields": {
      "plot": {}
    }
  },
  "sort": [
    {
       "rating": {
          "order": "desc"
      }
    }
  ]
}

response = opensearch_client.search(
    index=index_name,
    body=payload
)

#print(json.dumps(response, indent=2))
pd.json_normalize(response["hits"]["hits"])


# In[58]:


payload = {
  "from": 4,
  "size": 4,
  "query": {
    "bool": {
      "must": [
        {
          "match": {
            "plot": "school"
          }
        },
        {
          "term": {
            "genres": {
              "value": "Music"
            }
          }
        }
      ],
      "should": [
        {
          "term": {
            "genres": {
              "value": "Romance"
            }
          }
        },
        {
          "term": {
            "genres": {
              "value": "Comedy"
            }
          }
        }
      ],
      "minimum_should_match": 1,
      "filter": [
        {
          "range": {
            "year": {
              "gte": 2000,
              "lte": 2009
            }
          }
        }
      ],
      "must_not": [
        {
          "terms": {
            "genres": ["Thriller", "Horror"]
          }
        }
      ]
    }
  },
  "highlight": {
    "fields": {
      "plot": {}
    }
  },
  "sort": [
    {
       "rating": {
          "order": "desc"
      }
    }
  ]
}

response = opensearch_client.search(
    index=index_name,
    body=payload
)

#print(json.dumps(response, indent=2))
pd.json_normalize(response["hits"]["hits"])


# このように最もシンプルにページングを実装できる from + to ですが、以下の懸念事項があります
# 
# - 検索は Search API の都度新規に実行されるため、処理に一貫性はありません。前回の Search API の実行からデータの変化がある場合、paging の結果にも影響します
# - 検索は from と size の範囲に関係なく実行されるため、結果を分割取得したとしても、Search API を発行する都度、毎回同じ検索負荷が生じます
# - デフォルトでは、from + size は 10000 を超えることができません。`index.max_result_window` の[インデックス設定][index-settings]を変更することでこの制限を緩和できますが、パフォーマンスへの影響があるため一般的には推奨されません。
# 
# 数万件の結果を取得し paging を行う場合は、異なる手法を検討する必要があります。
# 
# [index-settings]: https://opensearch.org/docs/latest/install-and-configure/configuring-opensearch/index-settings/

# In[59]:


index_name = "movies"

payload = {
  "from": 9999,
  "size": 2,
  "query": {
    "bool": {
      "must": [
        {
          "match": {
            "plot": "school"
          }
        },
        {
          "term": {
            "genres": {
              "value": "Music"
            }
          }
        }
      ],
      "should": [
        {
          "term": {
            "genres": {
              "value": "Romance"
            }
          }
        },
        {
          "term": {
            "genres": {
              "value": "Comedy"
            }
          }
        }
      ],
      "minimum_should_match": 1,
      "filter": [
        {
          "range": {
            "year": {
              "gte": 2000,
              "lte": 2009
            }
          }
        }
      ],
      "must_not": [
        {
          "terms": {
            "genres": ["Thriller", "Horror"]
          }
        }
      ]
    }
  },
  "highlight": {
    "fields": {
      "plot": {}
    }
  },
  "sort": [
    {
       "rating": {
          "order": "desc"
      }
    }
  ]
}

try:
    response = opensearch_client.search(
        index=index_name,
        body=payload
    )
    print(json.dumps(response, indent=2))
except Exception as e:
    print(e)


# #### scroll によるページング
# scroll は機械学習ジョブといった PB クラスのデータを取得するバッチ処理に向いている手法です。リクエストヘッダに scroll パラメーターを付与することで有効化されます。scroll はリクエスト時点のスナップショットを取得するため、メモリを大量に消費する場合があります。この特性上、頻繁に実行されるクエリには向いていません。
# 
# scroll オプションを使用して search クエリを実行すると、実行結果に "_scroll_id" が付与されます。_scroll_id を次回のリクエストに付与することで、後続の結果を得ることができます。

# In[60]:


index_name = "movies"

payload = {
  "size": 5,
  "query": {
    "bool": {
      "must": [
        {
          "match": {
            "plot": {
              "query": "school"
            }
          }
        }
      ]
    }
  },
  "sort": [
    {
       "rating": {
          "order": "desc"
      }
    }
  ]
}

response = opensearch_client.search(
    index=index_name,
    body=payload,
    scroll="10m"
)

scroll_id = response["_scroll_id"]
#print(json.dumps(response, indent=2))
print("ScrollId: " + scroll_id)
pd.json_normalize(response["hits"]["hits"])



# In[61]:


payload = {
  scroll_id: scroll_id,
  "size": 5,
  "query": {
    "bool": {
      "must": [
        {
          "match": {
            "plot": {
              "query": "school"
            }
          }
        }
      ]
    }
  },
  "sort": [
    {
       "rating": {
          "order": "desc"
      }
    }
  ]
}

response = opensearch_client.search(
    index=index_name,
    scroll="10m",
)

scroll_id = response["_scroll_id"]
#print(json.dumps(response, indent=2))
print("ScrollId: " + scroll_id)
pd.json_normalize(response["hits"]["hits"])





# #### search_after によるページング
# [Search after][search_after] は from + size に似たアプローチですが、前回取得した範囲は検索結果から除外できるメリットがあります。検索結果を全件取得してから部分的に切り出しを行う from + size と比較して、ページングによるオーバーヘッドを抑えられます。
# 
# 以下は Search after のサンプルです。全部で 8 件の検索結果を、4 件ずつに分けて取得しています。
# 
# 1 つ目のクエリでは rating フィールドで降順ソートされた結果から size = 4 で先頭 4 件を取得しています。2 つ目のクエリでは、search_after オプションに 1 つ目のクエリ実行結果の 4 件目の ratings の値を入れています。これにより、1 つ目のクエリで取得した結果より後の結果から 4 件のドキュメントを取得することに成功しています。
# 
# [search_after]: https://opensearch.org/docs/latest/search-plugins/searching-data/paginate/#the-search_after-parameter

# In[62]:


index_name = "movies"
payload = {
  "size": 4,
  "query": {
    "bool": {
      "must": [
        {
          "match": {
            "plot": "school"
          }
        },
        {
          "term": {
            "genres": {
              "value": "Music"
            }
          }
        }
      ],
      "should": [
        {
          "term": {
            "genres": {
              "value": "Romance"
            }
          }
        },
        {
          "term": {
            "genres": {
              "value": "Comedy"
            }
          }
        }
      ],
      "minimum_should_match": 1,
      "filter": [
        {
          "range": {
            "year": {
              "gte": 2000,
              "lte": 2009
            }
          }
        }
      ],
      "must_not": [
        {
          "terms": {
            "genres": ["Thriller", "Horror"]
          }
        }
      ]
    }
  },
  "highlight": {
    "fields": {
      "plot": {}
    }
  },
  "sort": [
    {
       "rating": {
          "order": "desc"
      }
    }
  ]
}

response = opensearch_client.search(
    index=index_name,
    body=payload
)

search_after_rating = response["hits"]["hits"][-1]["sort"]
search_after_rating

#print(json.dumps(response, indent=2))
print("search_after_rating: " + str(search_after_rating))
pd.json_normalize(response["hits"]["hits"])


# In[63]:


index_name = "movies"

payload = {
  "size": 4,
  "query": {
    "bool": {
      "must": [
        {
          "match": {
            "plot": "school"
          }
        },
        {
          "term": {
            "genres": {
              "value": "Music"
            }
          }
        }
      ],
      "should": [
        {
          "term": {
            "genres": {
              "value": "Romance"
            }
          }
        },
        {
          "term": {
            "genres": {
              "value": "Comedy"
            }
          }
        }
      ],
      "minimum_should_match": 1,
      "filter": [
        {
          "range": {
            "year": {
              "gte": 2000,
              "lte": 2009
            }
          }
        }
      ],
      "must_not": [
        {
          "terms": {
            "genres": ["Thriller", "Horror"]
          }
        }
      ]
    }
  },
  "highlight": {
    "fields": {
      "plot": {}
    }
  },
  "sort": [
    {
       "rating": {
          "order": "desc"
      }
    }
  ],
  "search_after": search_after_rating
}

response = opensearch_client.search(
    index=index_name,
    body=payload
)

search_after_rating = response["hits"]["hits"][-1]["sort"]
search_after_rating

#print(json.dumps(response, indent=2))
print("search_after_rating: " + str(search_after_rating))
pd.json_normalize(response["hits"]["hits"])


# <div class="alert alert-block alert-info"> 
# <b>Tips: sort 条件を明示的に指定しない場合に Search after を利用する方法</b>
# 
# デフォルトでは OpenSearch はスコア順にドキュメントのソートを行います。明示的にソート対象のフィールドに "_score" を指定した場合と挙動は同じです。この特性を利用して、sort に _score による降順ソートの条件を書くことでスコアによるソートと Search after を両立できます。ただし、ソート条件が _score だけだとエラーになってしまうため、id など何かしらのフィールドとセットで sort 条件を書く必要があります。
# 
# </div>

# In[64]:


index_name = "movies"

payload = {
  "size": 4,
  "query": {
    "bool": {
      "must": [
        {
          "match": {
            "plot": "school"
          }
        },
        {
          "term": {
            "genres": {
              "value": "Music"
            }
          }
        }
      ]
    }
  },
  "sort": [
    {
       "_score": {
          "order": "desc"
      }
    },
    {
        "id": {
            "order": "asc"
        }
    }
  ]
}

response = opensearch_client.search(
    index=index_name,
    body=payload
)

search_after_score = response["hits"]["hits"][-1]["sort"]

#print(json.dumps(response, indent=2))
print("search_after_score: " + str(search_after_score))
pd.json_normalize(response["hits"]["hits"])


# In[65]:


index_name = "movies"

payload = {
  "size": 4,
  "query": {
    "bool": {
      "must": [
        {
          "match": {
            "plot": "school"
          }
        },
        {
          "term": {
            "genres": {
              "value": "Music"
            }
          }
        }
      ]
    }
  },
  "sort": [
    {
      "_score": {
          "order": "desc"
      },
      "id": {
          "order": "asc"
      }
    }
  ],
  "search_after": search_after_score
}

response = opensearch_client.search(
    index=index_name,
    body=payload
)

search_after_score = response["hits"]["hits"][-1]["sort"]

#print(json.dumps(response, indent=2))
print("search_after_score: " + str(search_after_score))
pd.json_normalize(response["hits"]["hits"])


# #### Search after + PIT(Point in Time) によるページング
# [PIT(Point in Time)][point-in-time] は、インデックスのある時点の固定された状態を作り出す機能です。PIT と Search after を組み合わせることで一貫性のあるページングされた結果を取得することができます。
# 
# PIT を使用した検索の流れは以下の通りです。PIT 作成時に keep_alive パラメーターで PIT の保持時間を指定するため、PIT の削除は任意です。
# 
# 1. PIT の作成
# 2. PIT を使用した検索の実行
# 3. (option) PIT の削除
# 
# PIT を使用する場合、PIT 側に Index の情報が入っているため、Search API 実行時にインデックス名の指定は行いません。リクエストパラメーターにインデックス名を含めてはいけない点に注意が必要です。
# 
# [point-in-time]: https://opensearch.org/docs/latest/search-plugins/searching-data/point-in-time/

# In[66]:


response = opensearch_client.create_point_in_time(
    index=index_name,
    keep_alive="10m"
)

pit_id = response.get("pit_id")
print('\n Point in time ID: '+ pit_id)


# In[67]:


index_name = "movies"

payload = {
  "size": 4,
  "query": {
    "bool": {
      "must": [
        {
          "match": {
            "plot": "school"
          }
        },
        {
          "term": {
            "genres": {
              "value": "Music"
            }
          }
        }
      ]
    }
  },
  "sort": [
    {
      "_score": {
          "order": "desc"
      },
      "id": {
          "order": "asc"
      }
    }
  ],
  "pit": {
    "id":  pit_id, 
    "keep_alive": "10m"
  },
}

response = opensearch_client.search(
    body=payload
)

search_after_score = response["hits"]["hits"][-1]["sort"]

print("search_after_score: " + str(search_after_score))
#print(json.dumps(response, indent=2))
pd.json_normalize(response["hits"]["hits"])


# In[68]:


index_name = "movies"

payload = {
  "size": 4,
  "query": {
    "bool": {
      "must": [
        {
          "match": {
            "plot": "school"
          }
        },
        {
          "term": {
            "genres": {
              "value": "Music"
            }
          }
        }
      ]
    }
  },
  "sort": [
    {
      "_score": {
          "order": "desc"
      },
      "id": {
          "order": "asc"
      }
    }
  ],
  "pit": {
    "id":  pit_id, 
    "keep_alive": "10m"
  },
  "search_after": search_after_score
}

response = opensearch_client.search(
    body=payload
)

search_after_score = response["hits"]["hits"][-1]["sort"]

print("search_after_score: " + str(search_after_score))
#print(json.dumps(response, indent=2))
pd.json_normalize(response["hits"]["hits"])


# ### 非同期検索
# 大規模なデータ検索、特にウォームノードや複数のリモートクラスターにまたがって検索が実行される場合、完了までに時間がかかることがあります。
# 完了までクライアントが待機するにはタイムアウトを延長するなどの措置が必要ですが、何らかの問題で接続が切断された場合は改めて検索リクエストを発行する必要があります。
# 
# このような課題に対処するために、OpenSearch の[非同期検索][async]を使用することができます。非同期検索は、バックグラウンドで実行される検索リクエストを送信できます。検索の進行状況は監視可能であり、結果を段階的に取得することも可能です。検索結果は任意の期間保存することが可能であり、後から取得することもできます。
# 
# [async]: https://opensearch.org/docs/latest/search-plugins/async/index/

# In[69]:


index_name = "movies"

payload = {
  "index": index_name
}

response = opensearch_client.http.post(
    url = "/_plugins/_asynchronous_search?index=" + index_name
)

asynchronous_search_id = response["id"]

#print(json.dumps(response, indent=2))
print("asynchronous_search_id: " + asynchronous_search_id)
pd.json_normalize(response["response"]["hits"]["hits"])


# In[70]:


index_name = "movies"

payload = {
  "index": index_name
}

response = opensearch_client.http.get(
    url = "/_plugins/_asynchronous_search/stats"
)

print(json.dumps(response, indent=2))


# ### 検索と組み合わせたドキュメント処理
# OpenSearch では、query で抽出したドキュメントの一括更新や削除を行うことが可能です。以降のセクションで、具体的な実行方法を確認していきます。

# #### 複数ドキュメントの一括更新
# [Update by query][update-by-query] を使用することで、特定条件に一致する複数ドキュメントのデータ更新をまとめて行うことが可能です。条件は Search API と同様に query で指定します。
# 
# 以下のサンプルでは、year が 1920 から 1923 までの movies 内のドキュメントについて、rating の値を 0.1 増加させています。レスポンスの total フィールドに、処理対象となったドキュメントの件数が記載されています。
# 
# [update-by-query]: https://opensearch.org/docs/latest/api-reference/document-apis/update-by-query/

# In[71]:


index_name = "movies"

payload = {
  "query": {
    "range": {
      "year": {
        "gte": 1920,
        "lte": 1923
      }
    }
  },
  "script" : {
    "source": "ctx._source.rating += params.delta",
    "lang": "painless",
    "params" : {
      "delta" : 0.1
    }
  }
}

response = opensearch_client.update_by_query(
    index = index_name,
    body = payload,
    refresh = True
)

print(json.dumps(response, indent=2))


# #### 複数ドキュメントの一括削除
# [Delete by query][delete-by-query] は、特定条件に一致する複数ドキュメントを一括で削除します。条件は Search API と同様に query で指定します。
# 
# 以下のサンプルでは、year が 1989 までの movies 内のドキュメントを削除します。レスポンスの total フィールドに、処理対象となったドキュメントの件数が記載されています。
# 
# [delete-by-query]: https://opensearch.org/docs/latest/api-reference/document-apis/delete-by-query/

# In[72]:


index_name = "movies"

payload = {
  "query": {
    "range": {
      "year": {
        "lte": 1989
      }
    }
  }
}

response = opensearch_client.delete_by_query(
    index = index_name,
    body = payload,
    refresh = True
)

print(json.dumps(response, indent=2))


# ## まとめ
# 本ラボでは、OpenSearch の基本的な検索クエリについて、ユースケースと実際の使い方を解説しました。本ラボで学習した内容を元に、次のステップとして以下のラボを実行してみましょう。
# 
# - [日本語全文検索の実装](../full-text-search/full-text-search-jp.ipynb)
# - [ベクトル検索の実装 (Amazon SageMaker 編)](../vector-search/vector-search-with-sagemaker.ipynb)

# ## 後片付け

# ### インデックス削除
# 本ワークショップで使用したインデックスを削除します。インデックスの削除は Delete index API で行います。インデックスを削除するとインデックス内のドキュメントも削除されます。

# In[73]:


index_name = "movies"

try:
    response = opensearch_client.indices.delete(index=index_name)
    print(json.dumps(response, indent=2))
except Exception as e:
    print(e)


# ### データセット削除
# ダウンロードしたデータセットを削除します。./dataset ディレクトリ配下に何もない場合は、./dataset ディレクトリも合わせて削除します。

# In[74]:


get_ipython().run_line_magic('rm', '-rf {dataset_dir}')


# In[75]:


get_ipython().run_line_magic('rmdir', './dataset')