code: def test_reservoir_sample_map_partitions_correctness():
      N, k = 20, 10
      seq = list(range(N))
      distribution = [0 for _ in range(N)]
      expected_distribution = [0 for _ in range(N)]
      reps = 2000
      for _ in range(reps):
          picks, _ = random._sample_map_partitions(seq, k)
          for pick in picks:
              distribution[pick] += 1
          for pick in rnd.sample(seq, k=k):
              expected_distribution[pick] += 1
  
      # convert to probabilities
      distribution = [c / (reps * k) for c in distribution]
      expected_distribution = [
          c / (reps * k) for c in expected_distribution
      ]
  
      # use bhattacharyya distance to asses the similarity of distributions
      assert math.isclose(
          0.0,
          bhattacharyya(distribution, expected_distribution),
          abs_tol=1e-2,
      )
  
imports: from __future__ import annotations
  import math
  import random
  import pytest
  import dask.bag
  from dask.bag import random
label_desc: Use image node_id 17 for calling a global function with example usage: test_reservoir_sample_map_partitions_correctness() without return types