code: def test_batch_size(nrows, ncols, n_parts, batch_size, client):
      n_neighbors = 10
      n_clusters = 5
      from cuml.dask.neighbors import NearestNeighbors as daskNN
  
      from sklearn.datasets import make_blobs
  
      nrows = _scale_rows(client, nrows)
  
      X, y = make_blobs(
          n_samples=int(nrows),
          n_features=ncols,
          centers=n_clusters,
          random_state=0,
      )
  
      X = X.astype(np.float32)
  
      X_cudf = _prep_training_data(client, X, n_parts)
  
      cumlModel = daskNN(
          n_neighbors=n_neighbors,
          batch_size=batch_size,
          streams_per_handle=5,
      )
  
      cumlModel.fit(X_cudf)
  
      out_d, out_i = cumlModel.kneighbors(X_cudf)
  
      local_i = out_i.compute().to_numpy()
  
      y_hat, _ = predict(local_i, y, n_neighbors)
  
      assert array_equal(y_hat, y)
  
file_path: cuml/python/cuml/tests/dask/test_dask_nearest_neighbors.py
imports: import platform
  from cuml.testing.utils import array_equal
  from sklearn.neighbors import KNeighborsClassifier
  from cuml.testing.utils import unit_param, quality_param, stress_param
  from cuml.dask.common import utils
  from cuml.common import has_scipy
  from cuml.internals.safe_imports import cpu_only_import
  import pytest
  from cuml.internals.safe_imports import gpu_only_import
label_desc: Use image node_id 5 for calling a global function with example usage: test_batch_size(nrows, ncols, n_parts, batch_size, client) without return types