code: def generate_data(args):
      print("create emb_vecs and item_dict data >>>>>>")
      from bigdl.orca import (
          init_orca_context,
          OrcaContext,
          stop_orca_context,
      )
      from pyspark.sql.types import Row
      from pyspark.sql.types import (
          StructType,
          StructField,
          IntegerType,
          ArrayType,
          FloatType,
      )
  
      print(">>>>>> using spark >>>>>>")
      sc = init_orca_context(
          cores="*",
          memory="16g",
          conf={"spark.driver.maxResultSize": "10g"},
      )
      spark = OrcaContext.get_spark_session()
      rdd = sc.parallelize(range(args.row_nums))
      rdd = rdd.map(lambda row: gen_vector(row, args.vec_dim))
      rdd.cache()  # If not cache, will run random generation again when saving to parquet.
      data = rdd.collect()
      emb_vecs = np.array([row[1] for row in data], dtype=np.float32)
      print("vector shape: ", emb_vecs.shape)
      # np.save("data.npy", vectors)
      schema = StructType(
          [
              StructField("id", IntegerType(), False),
              StructField("embedding", ArrayType(FloatType()), False),
          ]
      )
      df = spark.createDataFrame(
          rdd.map(lambda x: Row(x[0], x[1].tolist())), schema=schema
      )
      df.write.mode("overwrite").parquet(args.parquet_path)
      print("Finished")
      stop_orca_context()
      item_dict = create_dummy_data(
          args.row_nums,
          args.vec_dim,
          args.verbose,
          rnd_seed=42,
          HEADER_LEN=args.header_len,
          only_item=True,
      )
  
      print("create index for faiss >>>>>>")
      if args.index_type == "FlatL2":
          print("index_type: ", "# FlatL2")
          index_faiss = faiss.IndexFlatL2(args.vec_dim)
          # index_flat = faiss_flat.IndexFlatIP(VEC_DIM)
          print(
              "index_flat.is_trained: {}".format(index_faiss.is_trained)
          )
          index_faiss.add(emb_vecs)
          print("index_faiss.ntotal: ", index_faiss.ntotal)
      elif args.index_type == "IVFFlatL2":
          print("index_type: ", "# IVFFlatL2")
          # The number of cells (space partition). Typical value is sqrt(N)
          NLIST = int(math.sqrt(args.row_nums))
          print("NLIST: ", NLIST)
          quantizer = faiss.IndexFlatL2(args.vec_dim)
          index_faiss = faiss.IndexIVFFlat(
              quantizer, args.vec_dim, NLIST
          )
          print(
              "before, index_flat.is_trained: {}".format(
                  index_faiss.is_trained
              )
          )
          index_faiss.train(emb_vecs)
          print(
              "after, index_flat.is_trained: {}".format(
                  index_faiss.is_trained
              )
          )
          index_faiss.add(emb_vecs)
      else:
          print("unsupported index_type input")
          sys.exit()
  
      q_vec = emb_vecs[
          [
              0,
          ]
      ]
      search_sample(index_faiss, q_vec)
  
      print("saving created data and index >>>>>>")
      for file, obj in [(args.dict_path, item_dict)]:
          with open(file, "wb") as f:
              print("saving to: {}".format(file))
              SafePickle.dump(obj, f)
      # https://github.com/matsui528/faiss_tips#io
  
      with open(args.index_save_path, "wb") as f:
          print("saving to: {}".format(args.index_save_path))
          # serialize the index into binary array (np.array).
          # You can save/load it via numpy IO functions.
          chunk = faiss.serialize_index(index_faiss)
          SafePickle.dump(chunk, f)
  
      if args.verbose:
          # sample
          print("start loading data for testing >>>>>>")
          print(emb_vecs.shape)
          print(emb_vecs[0, :10])
          print(item_dict.shape)
          print(item_dict[:10])
  
          print(f"loading {args.dict_path}")
  
          with open(args.dict_path, "rb") as f:
              item_dict_loaded = SafePickle.load(f)
          print(item_dict_loaded.shape)
          print(item_dict_loaded[:10])
  
          q_vec = emb_vecs[
              [
                  0,
              ]
          ]
          print(q_vec.shape)
  
          with open(args.index_save_path, "rb") as f:
              index_loaded = faiss.deserialize_index(SafePickle.load(f))
              search_sample(index_loaded, q_vec)
  
file_path: BigDL/python/friesian/example/recall/generate_test_data.py
imports: from bigdl.friesian.utils import SafePickle
  import string
  import numpy
  import math
  import faiss
  import sys
  import random
  import argparse
label_desc: Use image node_id 4 for calling a global function with example usage: generate_data(args) without return types