code: def test_manual_build_faiss_index(setup):
      d = 8
      n = 50
      n_test = 10
      x = np.random.RandomState(0).rand(n, d).astype(np.float32)
      y = np.random.RandomState(0).rand(n_test, d).astype(np.float32)
  
      nn = NearestNeighbors(algorithm="kd_tree")
      nn.fit(x)
      _, expected_indices = nn.kneighbors(y, 5)
  
      # test brute-force search
      X = mt.tensor(x, chunk_size=10)
      index = build_faiss_index(
          X, "Flat", None, random_state=0, same_distribution=True
      )
      faiss_index = index.execute().fetch()
  
      index_shards = faiss.IndexShards(d)
      for ind in faiss_index:
          shard = _load_index(ind, -1)
          index_shards.add_shard(shard)
      faiss_index = index_shards
  
      faiss_index.nprob = 10
      _, indices = faiss_index.search(y, k=5)
  
      np.testing.assert_array_equal(indices, expected_indices.fetch())
  
      # test one chunk, brute force
      X = mt.tensor(x, chunk_size=50)
      index = build_faiss_index(
          X, "Flat", None, random_state=0, same_distribution=True
      )
      faiss_index = _load_index(index.execute().fetch(), -1)
  
      faiss_index.nprob = 10
      _, indices = faiss_index.search(y, k=5)
  
      np.testing.assert_array_equal(indices, expected_indices.fetch())
  
      # test train, same distribution
      X = mt.tensor(x, chunk_size=10)
      index = build_faiss_index(
          X, "IVF30,Flat", 30, random_state=0, same_distribution=True
      )
      faiss_index = _load_index(index.execute().fetch(), -1)
  
      assert isinstance(faiss_index, faiss.IndexIVFFlat)
      assert faiss_index.ntotal == n
      assert len(tile(index).chunks) == 1
  
      # test train, distributions are variant
      X = mt.tensor(x, chunk_size=10)
      index = build_faiss_index(
          X, "IVF10,Flat", None, random_state=0, same_distribution=False
      )
      faiss_index = index.execute().fetch()
  
      assert len(faiss_index) == 5
      for ind in faiss_index:
          ind = _load_index(ind, -1)
          assert isinstance(ind, faiss.IndexIVFFlat)
          assert ind.ntotal == 10
  
      # test more index type
      index = build_faiss_index(
          X, "PCAR6,IVF8_HNSW32,SQ8", 10, random_state=0
      )
      faiss_index = index.execute().fetch()
  
      assert len(faiss_index) == 5
      for ind in faiss_index:
          ind = _load_index(ind, -1)
          assert isinstance(ind, faiss.IndexPreTransform)
          assert ind.ntotal == 10
  
      # test one chunk, train
      X = mt.tensor(x, chunk_size=50)
      index = build_faiss_index(
          X, "IVF30,Flat", 30, random_state=0, same_distribution=True
      )
      faiss_index = _load_index(index.execute().fetch(), -1)
  
      assert isinstance(faiss_index, faiss.IndexIVFFlat)
      assert faiss_index.ntotal == n
  
      # test wrong index
      with pytest.raises(ValueError):
          build_faiss_index(X, "unknown_index", None)
  
      # test unknown metric
      with pytest.raises(ValueError):
          build_faiss_index(X, "Flat", None, metric="unknown_metric")
  
file_path: mars/mars/learn/neighbors/tests/test_faiss.py
imports: import numpy
  import pytest
  from ....None import tensor
  from ....core import tile
  from ....session import execute, fetch
  from ..None import NearestNeighbors
  from .._faiss import build_faiss_index, _load_index, faiss_query, _gen_index_string_and_sample_count
label_desc: Use image node_id 1 for calling a global function with example usage: test_manual_build_faiss_index(setup) without return types