code: def faiss_search(
      faiss_index_pkl,
      item_dict_pkl,
      cluster_mode,
      batch_size=65536,
      k=200,
  ):
      from bigdl.friesian.utils import SafePickle
      import faiss
  
      if cluster_mode == "yarn":
          load_SPARK = True
      else:
          load_SPARK = False
  
      def do_search(partition):
          with open_pickle(faiss_index_pkl, load_SPARK) as index_pkl:
              faiss_idx = faiss.deserialize_index(
                  SafePickle.load(index_pkl)
              )
  
          with open_pickle(item_dict_pkl, load_SPARK) as f:
              item_dict = SafePickle.load(f)
  
          buffer = []
          for record in partition:
              if len(buffer) == batch_size:
                  s1 = time.time()
                  seed_ids = [row[0] for row in buffer]
                  embeddings = [row[1] for row in buffer]
                  buffer = [record]
                  q_vec = np.stack(embeddings).astype(np.float32)
                  similarity_array, idx_array = faiss_idx.search(
                      q_vec, k=k
                  )
                  e1 = time.time()
                  print("Search time: ", e1 - s1)
  
                  for i in range(batch_size):
                      seed_idx = int(seed_ids[i])
                      seed_item = str(item_dict[seed_idx])
                      for n, (score, rec_id) in enumerate(
                          zip(similarity_array[i], idx_array[i])
                      ):
                          rec_id = int(rec_id)
                          yield (
                              seed_item,
                              str(item_dict[rec_id]),
                              int(n),
                              float(score),
                          )
  
              else:
                  buffer.append(record)
  
          remain_size = len(buffer)
          if remain_size > 0:
              seed_ids = [row[0] for row in buffer]
              embeddings = [row[1] for row in buffer]
              q_vec = np.stack(embeddings).astype(np.float32)
              similarity_array, idx_array = faiss_idx.search(q_vec, k=k)
  
              for i in range(remain_size):
                  seed_idx = int(seed_ids[i])
                  seed_item = str(item_dict[seed_idx])
                  for n, (score, rec_id) in enumerate(
                      zip(similarity_array[i], idx_array[i])
                  ):
                      rec_id = int(rec_id)
                      yield (
                          seed_item,
                          str(item_dict[rec_id]),
                          int(n),
                          float(score),
                      )
  
      return do_search
  
file_path: BigDL/python/friesian/example/recall/search.py
imports: import sys
  import argparse
  import numpy
  from bigdl.dllib.utils.log4Error import *
  from bigdl.orca import init_orca_context, OrcaContext, stop_orca_context
  from pyspark.sql.types import StructType, StructField, IntegerType, FloatType, StringType
  from utils import *
label_desc: Use image node_id 1 for calling a global function with example usage: faiss_search(faiss_index_pkl, item_dict_pkl, cluster_mode, batch_size, k) and returns: do_search