title: Efficient nested loops with bootstrap resampling on DataFrame in PySpark
body: <p>I have a DataFrame with the following schema:</p>
  <pre><code>+--------------------+--------+-------+-------+-----------+--------------------+
  |              userid|datadate|  runid|variant|device_type|          prediction|
  +--------------------+--------+-------+-------+-----------+--------------------+
  |0001d15b-e2da-4f4...|20220111|1196752|      1|     Mobile|  0.8827571312010658|
  |00021723-2a0d-497...|20220111|1196752|      1|     Mobile| 0.30763173370229735|
  |00021723-2a0d-497...|20220111|1196752|      0|     Mobile|  0.5336206154783815|
  </code></pre>
  <p>I would like to perform the following operation:
  I want to do for each &quot;runid&quot;, for each &quot;device_type&quot;, some calculations with variant==1 and variant==0, including a resampling loop.
  The ultimate goal is to store these calculations in another DF.</p>
  <p>So in a naive approach the code would look like that:</p>
  <pre><code>for runid in df.select('runid').distinct().rdd.flatMap(list).collect():
      for device in [&quot;Mobile&quot;,&quot;Desktop&quot;]:
          a_variant = df.filter((df.runid  == runid) &amp; (df.device_type  == device) &amp; (df.variant  == 0))
          b_variant = df.filter((df.runid  == runid) &amp; (df.device_type  == device) &amp; (df.variant  == 1))
  
          ## do some more calculations here
  
          # bootstrap loop:
          for samp in range(100):
              sampled_vector_a = a_variant.select(&quot;prediction&quot;).sample(withReplacement = True, fraction = 1.0, seed = 123)
              sampled_vector_b = b_variant.select(&quot;prediction&quot;).sample(withReplacement = True, fraction = 1.0, seed = 123)
  
              ## do some more calculations here
  
        ## do some more calculations here
        ## store calculations in a new DataFrame
  
  </code></pre>
  <p>Currently the process is too slow.
  How can I optimize this process by utilizing spark in the best way?
  Thanks!</p>