How do I access the whole DataFrame in map/agg function?
<p>I want to add a new column to each row of my DataFrame and need to do some calculation over the whole DataFrame for this. Is this possible?</p>

<h2>Context</h2>

<p>I have a DataFrame that contains a three dimensional index (to locate a cell in a grid) and a value per row:</p>

<pre><code>+---+-----+----+---+
|  t|    x|   y|  n|
+---+-----+----+---+
|  6|-4008|7321|  5|
|  7|-4007|7317|  1|
|  6|-4008|7323|  2|
|  6|-4008|7324| 17|
|  7|-4007|7326| 16|
+---+-----+----+---+
…
</code></pre>

<p>Now I want to add a fifth column (w) that holds the <code>n</code> values of the neighboring cells for each row.<br>
Example:</p>

<pre><code>+---+-----+----+---+---+
|  t|    x|   y|  n|  w|
+---+-----+----+---+---+
|  6|-4008|7318|  5|  1|
|  7|-4007|7317|  1|  5|
+---+-----+----+---+---+
…
</code></pre>

<p>For this I have to locate the current cell in the grid (easily done via <code>t, x, y</code>), find the cells surrounding it and <del>concatenate</del> count and sum up their <code>n</code> values. The result of this would be what I want to insert into the <code>w</code> column.</p>

<p>Now how do I perform this calculation for every single row, when I need access to the whole DataFrame for it?</p>

<hr>

<h2>Update 1</h2>

<p>My initial idea was to something like:</p>

<pre class="lang-scala prettyprint-override"><code>// df is my DataFrame that holds t, x, y and n
df.map(row =&gt; Row(row(0), row(1), row(2), row(3), calculateW(df, row)))
</code></pre>

<p>But I get an NPE every time I try to access <code>df</code> inside <code>calculateW()</code></p>

<hr>

<h2>Update 2, containing "solution"</h2>

<p>Since I am new to <em>thinking in Spark</em> I came up with the following solution. It works but is terribly slow.</p>

<pre class="lang-scala prettyprint-override"><code>def udfCalcZ = udf(...)
def udfCalcP = udf(...)

val dfWithPandZ = df.as("a")
  .join(
    df.as("b"),
    $"b.t".geq($"a.t" - 1) &amp;&amp; $"b.t".leq($"a.t" + 1)
      &amp;&amp; $"b.y".geq($"a.y" - 1) &amp;&amp; $"b.y".leq($"a.y" + 1)
      &amp;&amp; $"b.x".geq($"a.x" - 1) &amp;&amp; $"b.x".leq($"a.x" + 1)
  )
  .groupBy($"a.x", $"a.y", $"a.t", $"a.n").agg(count($"b.n").as("wLength"), sum($"b.n").as("wSum"))
  .withColumn("zscore", udfCalcZ($"wLength", $"wSum"))
  .withColumn("pvalue", udfCalcP($"zscore"))
</code></pre>