title: User-defined function on large crossjoined dataframes
body: <p>I need to find the sum <code>dsum</code> of a user-defined function <code>calcd</code> run over a dataframe consisting of one or more columns containing a range of floats.  The dataframe is the Cartesian product of a given range across N columns.</p>
  <p>For example, two-column input with the range {-5.0, -4.9, ..., 5.0} might be:</p>
  <div class="s-table-container">
  <table class="s-table">
  <thead>
  <tr>
  <th>x1</th>
  <th>x2</th>
  </tr>
  </thead>
  <tbody>
  <tr>
  <td>-5.0</td>
  <td>-5.0</td>
  </tr>
  <tr>
  <td>-5.0</td>
  <td>-4.9</td>
  </tr>
  <tr>
  <td>...</td>
  <td>...</td>
  </tr>
  <tr>
  <td>5.0</td>
  <td>4.9</td>
  </tr>
  <tr>
  <td>5.0</td>
  <td>5.0</td>
  </tr>
  </tbody>
  </table>
  </div>
  <p>For each of these, I would generate another column <code>di</code> using a UDF <code>dcalc</code>:</p>
  <div class="s-table-container">
  <table class="s-table">
  <thead>
  <tr>
  <th>x1</th>
  <th>x2</th>
  <th>di</th>
  </tr>
  </thead>
  <tbody>
  <tr>
  <td>-5.0</td>
  <td>-5.0</td>
  <td>?</td>
  </tr>
  <tr>
  <td>-5.0</td>
  <td>-4.9</td>
  <td>?</td>
  </tr>
  <tr>
  <td>...</td>
  <td>...</td>
  <td>...</td>
  </tr>
  <tr>
  <td>5.0</td>
  <td>5.0</td>
  <td>?</td>
  </tr>
  </tbody>
  </table>
  </div>
  <p>Finally, I calculate the sum <code>dsum</code> of all the <code>di</code> values.</p>
  <p>Ultimately, I want to do this over many columns (<code>x1</code>...<code>xN</code>) and ranges producing dataframes on the order of hundreds of millions or billions of rows.  My code to accomplish this is:</p>
  <pre><code>import numpy
  from scipy import stats
  from pyspark.sql import SparkSession
  from pyspark.sql.functions import udf
  
  eps   = 10.1
  iterx = 0.1
  sens  = 1
  scale = 0.7262317494490057
  
  spark       = SparkSession.builder.master('local[*]').appName('test').getOrCreate()
  sparkRange1 = spark.range(0,101,1)
  sparkRange2 = sparkRange1.withColumn('x1',(sparkRange1.id-50)*0.1).drop('id')
  sparkRange3 = sparkRange2.crossJoin(sparkRange2.withColumnRenamed('x1', 'x2'))
  
  @udf(&quot;float&quot;)
  def calcd(x1,x2):
      p1 = stats.norm.pdf(x1,loc=0,scale=scale)*            \
           stats.norm.pdf(x2,loc=0,scale=scale)*iterx**2
      p2 = stats.norm.pdf(x1,loc=sens,scale=scale)*         \
           stats.norm.pdf(x2,loc=sens,scale=scale)*iterx**2
      pl = numpy.log(p1/p2)
      if pl &gt; eps:
          d = p1 - numpy.exp(eps) * p2
      else:
          d = 0
      return float(d)
  
  result = sparkRange3.withColumn(&quot;di&quot;, calcd('x1','x2'))
  result.printSchema()
  result.show()
  
  dsum = result.agg({'di': 'sum'}).collect()[0][0]
  </code></pre>
  <p>Unfortunately, the final <code>agg</code> is very slow.  Four columns and ~100M rows takes over an hour given use of 42 cores, but I can calculate the same result in about 30 seconds using C++ on a single core.  I need this to scale, so (Py)Spark seems like it would be a natural fit for the problem.</p>
  <p>Is there another approach anyone can recommend?  Is the slowness a result of how I'm creating the input dataframe?  Thanks in advance.</p>