code: def gen_data(doc_count, query_count, dimension, dtype=np.float32):
      if dtype == np.float32:
          rs = np.random.RandomState(0)
          doc = pd.DataFrame(
              rs.rand(doc_count, dimension).astype(dtype)
          )
          query = rs.rand(query_count, dimension).astype(dtype)
      elif dtype == np.int32:
          rs = np.random.RandomState(0)
          doc = pd.DataFrame(
              (rs.rand(doc_count, dimension) * 1000).astype(dtype)
          )
          query = (rs.rand(query_count, dimension) * 1000).astype(dtype)
      else:
          raise ValueError(f"Unsupported dtype {dtype}")
      return doc, query
  
file_path: mars/mars/learn/proxima/simple_index/tests/test_simple_index.py
imports: import os
  import tempfile
  import numpy
  import pandas
  import pytest
  from .....None import dataframe
  from .....None import tensor
  from ...core import proxima
  from ..None import build_index, search_index, recall
label_desc: Use image node_id 2 for calling a global function with example usage: gen_data(doc_count, query_count, dimension, dtype) and returns: doc,  query