






from __future__ import annotations

import math
import random
from typing import Any






class 		lowercase_					:






      '''simple docstring'''




      def __init__(      self      :						List[str]  )    ->None:




       """simple docstring"""


       a		  =      []
       a		  =      0
       a		  =      0







      def       __lowerCAmelCase					(      self      :						str  )    ->bool:




       """simple docstring"""


       return self.head == self.tail







      def       __lowerCAmelCase					(      self      :						Tuple ,					__UpperCAmelCase      :						Any  )    ->None:




       """simple docstring"""


       self.data.append(__UpperCAmelCase  )
       a		  =      self.tail + 1







      def       __lowerCAmelCase					(      self      :						Optional[Any]  )    ->Any:




       """simple docstring"""


       a		  =      self.data[self.head]
       a		  =      self.head + 1
       return ret







      def       __lowerCAmelCase					(      self      :						List[Any]  )    ->int:




       """simple docstring"""


       return self.tail - self.head







      def       __lowerCAmelCase					(      self      :						Optional[int]  )    ->None:




       """simple docstring"""


       print(self.data  )
       print('''**************'''  )
       print(self.data[self.head : self.tail]  )






class 		lowercase_					:






      '''simple docstring'''




      def __init__(      self      :						int ,					__UpperCAmelCase      :						Any  )    ->None:




       """simple docstring"""


       a		  =      data
       a		  =      None
       a		  =      None
       a		  =      1







      def       __lowerCAmelCase					(      self      :						List[Any]  )    ->Any:




       """simple docstring"""


       return self.data







      def       __lowerCAmelCase					(      self      :						List[Any]  )    ->MyNode | None:




       """simple docstring"""


       return self.left







      def       __lowerCAmelCase					(      self      :						List[str]  )    ->MyNode | None:




       """simple docstring"""


       return self.right







      def       __lowerCAmelCase					(      self      :						Tuple  )    ->int:




       """simple docstring"""


       return self.height







      def       __lowerCAmelCase					(      self      :						List[str] ,					__UpperCAmelCase      :						Any  )    ->None:




       """simple docstring"""


       a		  =      data







      def       __lowerCAmelCase					(      self      :						List[Any] ,					__UpperCAmelCase      :						MyNode | None  )    ->None:




       """simple docstring"""


       a		  =      node







      def       __lowerCAmelCase					(      self      :						Union[str, Any] ,					__UpperCAmelCase      :						MyNode | None  )    ->None:




       """simple docstring"""


       a		  =      node







      def       __lowerCAmelCase					(      self      :						Any ,					__UpperCAmelCase      :						int  )    ->None:




       """simple docstring"""


       a		  =      height
def 						_a    (      a :MyNode | None )							->			int:
 if node is None:
  return 0
 return node.get_height()
def 						_a    (      a :int	,							a :int )							->			int:
 if a > b:
  return a
 return b
def 						_a    (      a :MyNode )							->			MyNode:
 print('''left rotation node:'''	,							node.get_data() )
 a		  =      node.get_left()
 assert ret is not None
 node.set_left(ret.get_right() )
 ret.set_right(a )
 a		  =      my_max(get_height(node.get_right() )	,							get_height(node.get_left() ) ) + 1
 node.set_height(a )
 a		  =      my_max(get_height(ret.get_right() )	,							get_height(ret.get_left() ) ) + 1
 ret.set_height(a )
 return ret
def 						_a    (      a :MyNode )							->			MyNode:
 print('''right rotation node:'''	,							node.get_data() )
 a		  =      node.get_right()
 assert ret is not None
 node.set_right(ret.get_left() )
 ret.set_left(a )
 a		  =      my_max(get_height(node.get_right() )	,							get_height(node.get_left() ) ) + 1
 node.set_height(a )
 a		  =      my_max(get_height(ret.get_right() )	,							get_height(ret.get_left() ) ) + 1
 ret.set_height(a )
 return ret
def 						_a    (      a :MyNode )							->			MyNode:
 a		  =      node.get_left()
 assert left_child is not None
 node.set_left(left_rotation(a ) )
 return right_rotation(a )
def 						_a    (      a :MyNode )							->			MyNode:
 a		  =      node.get_right()
 assert right_child is not None
 node.set_right(right_rotation(a ) )
 return left_rotation(a )
def 						_a    (      a :MyNode | None	,							a :Any )							->			MyNode | None:
 if node is None:
  return MyNode(a )
 if data < node.get_data():
  node.set_left(insert_node(node.get_left()	,							a ) )
  if (
      get_height(node.get_left() ) - get_height(node.get_right() ) == 2
  ):  # an unbalance detected
   a		  =      node.get_left()
   assert left_child is not None
   if (
       data < left_child.get_data()
   ):  # new node is the left child of the left child
    a		  =      right_rotation(a )
   else:
    a		  =      lr_rotation(a )
 else:
  node.set_right(insert_node(node.get_right()	,							a ) )
  if get_height(node.get_right() ) - get_height(node.get_left() ) == 2:
   a		  =      node.get_right()
   assert right_child is not None
   if data < right_child.get_data():
    a		  =      rl_rotation(a )
   else:
    a		  =      left_rotation(a )
 a		  =      my_max(get_height(node.get_right() )	,							get_height(node.get_left() ) ) + 1
 node.set_height(a )
 return node
def 						_a    (      a :MyNode )							->			Any:
 while True:
  a		  =      root.get_right()
  if right_child is None:
   break
  a		  =      right_child
 return root.get_data()
def 						_a    (      a :MyNode )							->			Any:
 while True:
  a		  =      root.get_left()
  if left_child is None:
   break
  a		  =      left_child
 return root.get_data()
def 						_a    (      a :MyNode	,							a :Any )							->			MyNode | None:
 a		  =      root.get_left()
 a		  =      root.get_right()
 if root.get_data() == data:
  if left_child is not None and right_child is not None:
   a		  =      get_left_most(a )
   root.set_data(a )
   root.set_right(del_node(a	,							a ) )
  elif left_child is not None:
   a		  =      left_child
  elif right_child is not None:
   a		  =      right_child
  else:
   return None
 elif root.get_data() > data:
  if left_child is None:
   print('''No such data''' )
   return root
  else:
   root.set_left(del_node(a	,							a ) )
 else:  # root.get_data() < data
  if right_child is None:
   return root
  else:
   root.set_right(del_node(a	,							a ) )

 if get_height(a ) - get_height(a ) == 2:
  assert right_child is not None
  if get_height(right_child.get_right() ) > get_height(right_child.get_left() ):
   a		  =      left_rotation(a )
  else:
   a		  =      rl_rotation(a )
 elif get_height(a ) - get_height(a ) == -2:
  assert left_child is not None
  if get_height(left_child.get_left() ) > get_height(left_child.get_right() ):
   a		  =      right_rotation(a )
  else:
   a		  =      lr_rotation(a )
 a		  =      my_max(get_height(root.get_right() )	,							get_height(root.get_left() ) ) + 1
 root.set_height(a )
 return root






class 		lowercase_					:






      '''simple docstring'''




      def __init__(      self      :						int  )    ->None:




       """simple docstring"""


       a		  =      None







      def       __lowerCAmelCase					(      self      :						Dict  )    ->int:




       """simple docstring"""


       return get_height(self.root  )







      def       __lowerCAmelCase					(      self      :						Dict ,					__UpperCAmelCase      :						Any  )    ->None:




       """simple docstring"""


       print('''insert:''' + str(__UpperCAmelCase  )  )
       a		  =      insert_node(self.root ,					__UpperCAmelCase  )







      def       __lowerCAmelCase					(      self      :						Any ,					__UpperCAmelCase      :						Any  )    ->None:




       """simple docstring"""


       print('''delete:''' + str(__UpperCAmelCase  )  )
       if self.root is None:
        print('''Tree is empty!'''  )
        return
       a		  =      del_node(self.root ,					__UpperCAmelCase  )







      def __str__(      self      :						List[Any] ,					)    ->str:  # a level traversale, gives a more intuitive look on the tree




       """simple docstring"""


       a		  =      ''''''
       a		  =      MyQueue()
       q.push(self.root  )
       a		  =      self.get_height()
       if layer == 0:
        return output
       a		  =      0
       while not q.is_empty():
        a		  =      q.pop()
        a		  =      ''' ''' * int(math.pow(2 ,					layer - 1  )  )
        output += space
        if node is None:
         output += "*"
         q.push(__UpperCAmelCase  )
         q.push(__UpperCAmelCase  )
        else:
         output += str(node.get_data()  )
         q.push(node.get_left()  )
         q.push(node.get_right()  )
        output += space
        a		  =      cnt + 1
        for i in range(100  ):
         if cnt == math.pow(2 ,					__UpperCAmelCase  ) - 1:
          a		  =      layer - 1
          if layer == 0:
           output += "\n*************************************"
           return output
          output += "\n"
          break
       output += "\n*************************************"
       return output
def 						_a    (      )							->			None:
 import doctest

 doctest.testmod()


if __name__ == "__main__":
     _test()
     UpperCAmelCase__				   =		AVLtree()
     UpperCAmelCase__				   =		list(range(10))
     random.shuffle(lst)
     for i in lst:
          t.insert(i)
          print(str(t))
     random.shuffle(lst)
     for i in lst:
          t.del_node(i)
          print(str(t))










import os
import tempfile
from functools import partial
from unittest import TestCase
from unittest.mock import patch

import numpy as np
import pytest

from datasets.arrow_dataset import Dataset
from datasets.search import ElasticSearchIndex, FaissIndex, MissingIndex

from .utils import require_elasticsearch, require_faiss


UpperCAmelCase__				   =		pytest.mark.integration






@require_faiss
class 		lowercase_					(	lowercase   ):






      '''simple docstring'''




      def       __lowerCAmelCase					(      self      :						Tuple  )    ->Any:




       """simple docstring"""


       a		  =      Dataset.from_dict({'''filename''': ['''my_name-train''' + '''_''' + str(__UpperCAmelCase  ) for x in np.arange(30  ).tolist()]}  )
       return dset







      def       __lowerCAmelCase					(      self      :						Tuple  )    ->Any:




       """simple docstring"""


       import faiss

       a		  =      self._create_dummy_dataset()
       a		  =      dset.map(
           lambda __UpperCAmelCase ,					__UpperCAmelCase  : {"vecs": i * np.ones(5 ,					dtype=np.floataa  )} ,					with_indices=__UpperCAmelCase ,					keep_in_memory=__UpperCAmelCase  )
       a		  =      dset.add_faiss_index('''vecs''' ,					batch_size=100 ,					metric_type=faiss.METRIC_INNER_PRODUCT  )
       a  ,     a		  =      dset.get_nearest_examples('''vecs''' ,					np.ones(5 ,					dtype=np.floataa  )  )
       self.assertEqual(examples['''filename'''][0] ,					'''my_name-train_29'''  )
       dset.drop_index('''vecs'''  )







      def       __lowerCAmelCase					(      self      :						Union[str, Any]  )    ->Optional[Any]:




       """simple docstring"""


       import faiss

       a		  =      self._create_dummy_dataset()
       dset.add_faiss_index_from_external_arrays(
           external_arrays=np.ones((30, 5)  ) * np.arange(30  ).reshape(-1 ,					1  ) ,					index_name='''vecs''' ,					batch_size=100 ,					metric_type=faiss.METRIC_INNER_PRODUCT ,					)
       a  ,     a		  =      dset.get_nearest_examples('''vecs''' ,					np.ones(5 ,					dtype=np.floataa  )  )
       self.assertEqual(examples['''filename'''][0] ,					'''my_name-train_29'''  )







      def       __lowerCAmelCase					(      self      :						Any  )    ->Union[str, Any]:




       """simple docstring"""


       import faiss

       a		  =      self._create_dummy_dataset()
       dset.add_faiss_index_from_external_arrays(
           external_arrays=np.ones((30, 5)  ) * np.arange(30  ).reshape(-1 ,					1  ) ,					index_name='''vecs''' ,					metric_type=faiss.METRIC_INNER_PRODUCT ,					)

       # Setting delete=False and unlinking manually is not pretty... but it is required on Windows to
       # ensure somewhat stable behaviour. If we don't, we get PermissionErrors. This is an age-old issue.
       # see https://bugs.python.org/issue14243 and
       # https://stackoverflow.com/questions/23212435/permission-denied-to-write-to-my-temporary-file/23212515
       with tempfile.NamedTemporaryFile(delete=__UpperCAmelCase  ) as tmp_file:
        dset.save_faiss_index('''vecs''' ,					tmp_file.name  )
        dset.load_faiss_index('''vecs2''' ,					tmp_file.name  )
       os.unlink(tmp_file.name  )

       a  ,     a		  =      dset.get_nearest_examples('''vecs2''' ,					np.ones(5 ,					dtype=np.floataa  )  )
       self.assertEqual(examples['''filename'''][0] ,					'''my_name-train_29'''  )







      def       __lowerCAmelCase					(      self      :						Optional[Any]  )    ->Any:




       """simple docstring"""


       a		  =      self._create_dummy_dataset()
       dset.add_faiss_index_from_external_arrays(
           external_arrays=np.ones((30, 5)  ) * np.arange(30  ).reshape(-1 ,					1  ) ,					index_name='''vecs'''  )
       dset.drop_index('''vecs'''  )
       self.assertRaises(__UpperCAmelCase ,					partial(dset.get_nearest_examples ,					'''vecs2''' ,					np.ones(5 ,					dtype=np.floataa  )  )  )







      def       __lowerCAmelCase					(      self      :						List[Any]  )    ->List[str]:




       """simple docstring"""


       from elasticsearch import Elasticsearch

       a		  =      self._create_dummy_dataset()
       with patch('''elasticsearch.Elasticsearch.search'''  ) as mocked_search, patch(
           '''elasticsearch.client.IndicesClient.create'''  ) as mocked_index_create, patch('''elasticsearch.helpers.streaming_bulk'''  ) as mocked_bulk:
        a		  =      {'''acknowledged''': True}
        mocked_bulk.return_value([(True, None)] * 30  )
        a		  =      {'''hits''': {'''hits''': [{'''_score''': 1, '''_id''': 29}]}}
        a		  =      Elasticsearch()

        dset.add_elasticsearch_index('''filename''' ,					es_client=__UpperCAmelCase  )
        a  ,     a		  =      dset.get_nearest_examples('''filename''' ,					'''my_name-train_29'''  )
        self.assertEqual(examples['''filename'''][0] ,					'''my_name-train_29'''  )






@require_faiss
class 		lowercase_					(	lowercase   ):






      '''simple docstring'''




      def       __lowerCAmelCase					(      self      :						Any  )    ->Any:




       """simple docstring"""


       import faiss

       a		  =      FaissIndex(metric_type=faiss.METRIC_INNER_PRODUCT  )

       # add vectors
       index.add_vectors(np.eye(5 ,					dtype=np.floataa  )  )
       self.assertIsNotNone(index.faiss_index  )
       self.assertEqual(index.faiss_index.ntotal ,					5  )
       index.add_vectors(np.zeros((5, 5) ,					dtype=np.floataa  )  )
       self.assertEqual(index.faiss_index.ntotal ,					10  )

       # single query
       a		  =      np.zeros(5 ,					dtype=np.floataa  )
       a		  =      1
       a  ,     a		  =      index.search(__UpperCAmelCase  )
       self.assertRaises(__UpperCAmelCase ,					index.search ,					query.reshape(-1 ,					1  )  )
       self.assertGreater(scores[0] ,					0  )
       self.assertEqual(indices[0] ,					1  )

       # batched queries
       a		  =      np.eye(5 ,					dtype=np.floataa  )[::-1]
       a  ,     a		  =      index.search_batch(__UpperCAmelCase  )
       self.assertRaises(__UpperCAmelCase ,					index.search_batch ,					queries[0]  )
       a		  =      [scores[0] for scores in total_scores]
       a		  =      [indices[0] for indices in total_indices]
       self.assertGreater(np.min(__UpperCAmelCase  ) ,					0  )
       self.assertListEqual([4, 3, 2, 1, 0] ,					__UpperCAmelCase  )







      def       __lowerCAmelCase					(      self      :						Any  )    ->List[Any]:




       """simple docstring"""


       import faiss

       a		  =      FaissIndex(string_factory='''Flat'''  )
       index.add_vectors(np.eye(5 ,					dtype=np.floataa  )  )
       self.assertIsInstance(index.faiss_index ,					faiss.IndexFlat  )
       a		  =      FaissIndex(string_factory='''LSH'''  )
       index.add_vectors(np.eye(5 ,					dtype=np.floataa  )  )
       self.assertIsInstance(index.faiss_index ,					faiss.IndexLSH  )
       with self.assertRaises(__UpperCAmelCase  ):
        a		  =      FaissIndex(string_factory='''Flat''' ,					custom_index=faiss.IndexFlat(5  )  )







      def       __lowerCAmelCase					(      self      :						int  )    ->Optional[Any]:




       """simple docstring"""


       import faiss

       a		  =      faiss.IndexFlat(5  )
       a		  =      FaissIndex(custom_index=__UpperCAmelCase  )
       index.add_vectors(np.eye(5 ,					dtype=np.floataa  )  )
       self.assertIsInstance(index.faiss_index ,					faiss.IndexFlat  )







      def       __lowerCAmelCase					(      self      :						int  )    ->Dict:




       """simple docstring"""


       import faiss

       a		  =      FaissIndex(metric_type=faiss.METRIC_INNER_PRODUCT  )
       index.add_vectors(np.eye(5 ,					dtype=np.floataa  )  )

       # Setting delete=False and unlinking manually is not pretty... but it is required on Windows to
       # ensure somewhat stable behaviour. If we don't, we get PermissionErrors. This is an age-old issue.
       # see https://bugs.python.org/issue14243 and
       # https://stackoverflow.com/questions/23212435/permission-denied-to-write-to-my-temporary-file/23212515
       with tempfile.NamedTemporaryFile(delete=__UpperCAmelCase  ) as tmp_file:
        index.save(tmp_file.name  )
        a		  =      FaissIndex.load(tmp_file.name  )
       os.unlink(tmp_file.name  )

       a		  =      np.zeros(5 ,					dtype=np.floataa  )
       a		  =      1
       a  ,     a		  =      index.search(__UpperCAmelCase  )
       self.assertGreater(scores[0] ,					0  )
       self.assertEqual(indices[0] ,					1  )
@require_faiss
def 						_a    (      a :Dict )							->			Any:
 import faiss

 a		  =      FaissIndex(metric_type=faiss.METRIC_INNER_PRODUCT )
 index.add_vectors(np.eye(5	,							dtype=np.floataa ) )

 a		  =      '''index.faiss'''
 a		  =      F"""mock://{index_name}"""
 index.save(a	,							storage_options=mockfs.storage_options )
 a		  =      FaissIndex.load(a	,							storage_options=mockfs.storage_options )

 a		  =      np.zeros(5	,							dtype=np.floataa )
 a		  =      1
 a  ,     a		  =      index.search(a )
 assert scores[0] > 0
 assert indices[0] == 1






@require_elasticsearch
class 		lowercase_					(	lowercase   ):






      '''simple docstring'''




      def       __lowerCAmelCase					(      self      :						int  )    ->List[Any]:




       """simple docstring"""


       from elasticsearch import Elasticsearch

       with patch('''elasticsearch.Elasticsearch.search'''  ) as mocked_search, patch(
           '''elasticsearch.client.IndicesClient.create'''  ) as mocked_index_create, patch('''elasticsearch.helpers.streaming_bulk'''  ) as mocked_bulk:
        a		  =      Elasticsearch()
        a		  =      {'''acknowledged''': True}
        a		  =      ElasticSearchIndex(es_client=__UpperCAmelCase  )
        mocked_bulk.return_value([(True, None)] * 3  )
        index.add_documents(['''foo''', '''bar''', '''foobar''']  )

        # single query
        a		  =      '''foo'''
        a		  =      {'''hits''': {'''hits''': [{'''_score''': 1, '''_id''': 0}]}}
        a  ,     a		  =      index.search(__UpperCAmelCase  )
        self.assertEqual(scores[0] ,					1  )
        self.assertEqual(indices[0] ,					0  )

        # single query with timeout
        a		  =      '''foo'''
        a		  =      {'''hits''': {'''hits''': [{'''_score''': 1, '''_id''': 0}]}}
        a  ,     a		  =      index.search(__UpperCAmelCase ,					request_timeout=30  )
        self.assertEqual(scores[0] ,					1  )
        self.assertEqual(indices[0] ,					0  )

        # batched queries
        a		  =      ['''foo''', '''bar''', '''foobar''']
        a		  =      {'''hits''': {'''hits''': [{'''_score''': 1, '''_id''': 1}]}}
        a  ,     a		  =      index.search_batch(__UpperCAmelCase  )
        a		  =      [scores[0] for scores in total_scores]
        a		  =      [indices[0] for indices in total_indices]
        self.assertGreater(np.min(__UpperCAmelCase  ) ,					0  )
        self.assertListEqual([1, 1, 1] ,					__UpperCAmelCase  )

        # batched queries with timeout
        a		  =      ['''foo''', '''bar''', '''foobar''']
        a		  =      {'''hits''': {'''hits''': [{'''_score''': 1, '''_id''': 1}]}}
        a  ,     a		  =      index.search_batch(__UpperCAmelCase ,					request_timeout=30  )
        a		  =      [scores[0] for scores in total_scores]
        a		  =      [indices[0] for indices in total_indices]
        self.assertGreater(np.min(__UpperCAmelCase  ) ,					0  )
        self.assertListEqual([1, 1, 1] ,					__UpperCAmelCase  )


