Convert SQL result from self-join to square pandas dataframe
<p>Here's a tl;dr version of what I'm after; the details are below:
A SQL query gives me a table with fields [person 1 id], [person 2 id], and [number of times they were in a group together]. I want to convert to a pandas dataframe that's square -- one row per person and one column per person, with the value of each element being number of times they were in a group together. I'm looking for a more elegant way to do that than going through the rows of my result and filling up the dataframe one element at a time.</p>

<hr>

<p>I have a database with a table of assignments, which has a column for the person_id and a column for the assignment_id. It has other stuff too, but for our purposes here, this is what matters:</p>

<pre><code>SELECT person_id, assignment_id FROM assignments;
</code></pre>

<pre>
person_id | assignment_id
----------+--------------
      385 |           42
      163 |           29
       51 |           42
      385 |           37
      163 |           37
         ...
</pre>

<p>I want to see how often any two people have been on the same assignment. So I do:</p>

<pre><code>SELECT a1.person_id AS p1_id, a2.person_id AS p2_id, COUNT(*)
FROM assignments AS a1
INNER JOIN assignments AS a2 ON a1.assignment_id = a2.assignment_id AND a1.person_id &lt; a2.person_id
GROUP BY a1.person_id, a2.person_id
</code></pre>

<p>Which gives output like:</p>

<pre>
p1_id | p2_id | count
------+-------+------
   51 |   385 |    1
  163 |   385 |    1
         ...
</pre>

<p>Now I'm building a Python script to access the data and want to dump it into a dataframe with a row for each person, a column for each person, and the cell having the number of times they shared an assignment. So the output would be something like this (I don't care what goes in the * cells -- could reasonably be 0 or the number of assignments the person did -- and don't really care whether how the first row and column are formatted):</p>

<pre>
 p1_id |  p_51  |  p_163 |  p_385
-------+--------+--------+--------
   51  |    *   |    0   |    1
  163  |    0   |    *   |    1
  385  |    1   |    1   |    *
</pre>

<p>I'll only have about 20 people, so it wouldn't hurt performance measurably to just set the values one by one, but I'm trying to learn good practice for when I have larger data sets. What's the right way to do something like this?</p>

<p>(I'm open to modifying the SQL query, if that's the best way to handle it.)</p>

<p>You could use <a href="http://pandas.pydata.org/pandas-docs/stable/generated/pandas.DataFrame.groupby.html" rel="nofollow"><code>groupby</code></a> after converting the required columns to type <code>str</code> and aggregate by joining them as well as taking their counts.</p>

<pre><code>df[['person_id', 'assignment_id']] = df[['person_id', 'assignment_id']].astype(str)

df = df.groupby(['assignment_id'], as_index=False, sort=False)['person_id']   \
       .agg({'col':','.join})['col']                                         \
       .str.split(',').apply(lambda x: sorted(x, reverse=True))              \
       .apply(pd.Series).add_prefix('p_id_')                                 \
       .set_index('p_id_0', drop=False) 
</code></pre>

<p>You could simplify further by using <a href="http://pandas.pydata.org/pandas-docs/stable/generated/pandas.get_dummies.html" rel="nofollow"><code>get_dummies</code></a> to obtain the indicator variables for the index, <code>p_id_0</code> as shown:</p>

<pre><code>df1 = pd.get_dummies(df['p_id_1']).add_prefix('p_')
print (df1)

        p_163  p_385
p_id_0              
51        0.0    1.0
163       0.0    0.0
385       1.0    0.0

df2 = pd.get_dummies(df['p_id_0']).add_prefix('p_')
print (df2)

        p_163  p_385  p_51
p_id_0                    
51        0.0    0.0   1.0
163       1.0    0.0   0.0
385       0.0    1.0   0.0
</code></pre>

<p>Then, concatenating these individual <code>dataframes</code> after mapping all values of the indexed frame to 0's followed by grouping the same column named columns together:</p>

<pre><code>df_final = pd.concat([df1, df2.applymap(lambda x: 0)], axis=1).add_prefix('p_')
print (df_final.groupby(df.columns, axis=1).sum())

        p_163  p_385  p_51
p_id_0                    
51        0.0    1.0   0.0
163       0.0    0.0   0.0
385       1.0    0.0   0.0
</code></pre>

python postgresql pandas