title: Replacing a pivot with a lazy groupby operation
body: <p>I'm pivoting a rather large dataframe of shape <code>(10_000_000, 678)</code> into one of approx. shape <code>(770_000, 8_789)</code> to create a dataset for an ML algorithm. It's a relatively slow operation taking about half an hour on a high-ram cluster I am using, and I'm wondering if there is a way to speed it up. Here is a minimum example, with a larger one below:</p>
  <pre class="lang-py prettyprint-override"><code>import polars as pl
  import numpy as np
  
  data = {
      &quot;id&quot;: [1,1,1,2,2,2,3,3,3],
      &quot;rank&quot;: [1,2,3,1,2,3,1,2,3], # rank is always repeating 1-3 (or 0-12 in large example)
      &quot;A&quot;: np.random.random((9)),
      &quot;B&quot;: np.random.random((9)),
  }
  
  df = pl.DataFrame(data)
  df_pivot = df.pivot(values=[&quot;A&quot;, &quot;B&quot;], index=&quot;id&quot;, columns=&quot;rank&quot;)
  
  # Now rename columns, since they are currently:
  # df_pivot.columns
  # ['id', '1', '2', '3', '1', '2', '3']
  ranks = [1,2,3]
  renamed_columns = df_pivot.columns[:1]
  for col in df.columns[2:]:
      for rank in ranks:
          renamed_columns.append(f&quot;{col}_{rank}&quot;)
  df_pivot.columns = renamed_columns
  </code></pre>
  <pre><code># df_pivot
  shape: (3, 7)
  ┌─────┬──────────┬──────────┬──────────┬──────────┬──────────┬──────────┐
  │ id  ┆ A_1      ┆ A_2      ┆ A_3      ┆ B_1      ┆ B_2      ┆ B_3      │
  │ --- ┆ ---      ┆ ---      ┆ ---      ┆ ---      ┆ ---      ┆ ---      │
  │ i64 ┆ f64      ┆ f64      ┆ f64      ┆ f64      ┆ f64      ┆ f64      │
  ╞═════╪══════════╪══════════╪══════════╪══════════╪══════════╪══════════╡
  │ 1   ┆ 0.867957 ┆ 0.854234 ┆ 0.408062 ┆ 0.076254 ┆ 0.899092 ┆ 0.059019 │
  ├╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌┤
  │ 2   ┆ 0.642296 ┆ 0.670476 ┆ 0.480494 ┆ 0.4254   ┆ 0.536173 ┆ 0.492312 │
  ├╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌┤
  │ 3   ┆ 0.778481 ┆ 0.151697 ┆ 0.330138 ┆ 0.6661   ┆ 0.4086   ┆ 0.992057 │
  └─────┴──────────┴──────────┴──────────┴──────────┴──────────┴──────────┘
  </code></pre>
  <p>The <a href="https://github.com/pola-rs/polars/blob/8bc22e054c9bf7b8288a61ea8e54fd267d092f95/polars/polars-lazy/src/frame/pivot.rs#L1" rel="nofollow noreferrer">polars pivot code states</a> that in a comment:</p>
  <blockquote>
  <p>Polars lazy does not implement a pivot because it is impossible to know the schema without materializing the whole dataset. This makes a pivot quite a terrible operation for performant workflows. An optimization can never be pushed down passed a pivot.</p>
  </blockquote>
  <p>And in the <a href="https://github.com/pola-rs/polars/blob/8bc22e054c9bf7b8288a61ea8e54fd267d092f95/polars/polars-core/src/frame/groupby/pivot.rs#L43" rel="nofollow noreferrer">groupby.pivot</a> code:</p>
  <blockquote>
  <p>Polars'/arrow memory is not ideal for transposing operations like pivots. If you have a relatively large table, consider using a groupby over a pivot.</p>
  </blockquote>
  <p>Some questions:</p>
  <ul>
  <li><p>Is it possible to replace the above pivot example by a (preferably lazy) combination of groupby and something else? <a href="https://stackoverflow.com/a/56676656/1018861">This SO post</a> about pandas suggests an equivalency of groupby + &quot;unstack&quot; with pivot. Polars does not implement an unstack function, afaik.</p>
  </li>
  <li><p>Is the above suggestion more performant than the current pivot implementation? (See the larger example below).</p>
  </li>
  <li><p>I actually <em>do</em> know the schema ahead-of-schedule, since in my situation <code>rank</code> is a known series ([1, 2, 3] in the example). If implemented, would a lazy pivot where one can supply the schema be more performant than the eager one?</p>
  </li>
  <li><p>Should I be implementing it differently?</p>
  </li>
  </ul>
  <pre class="lang-py prettyprint-override"><code># Much larger example, but with 10_000 rows instead of 10_000_000
  # 10_000 runs in 3 seconds, 100_000 runs in 40 seconds (M1 macbook)
  
  from string import ascii_lowercase
  import polars as pl
  import numpy as np
  
  ranks = np.arange(13)
  N_ROWS = 10_000 # this could be ~10_000_000
  
  df = (pl.DataFrame({&quot;ID&quot;: np.arange(N_ROWS)})).join(
      pl.DataFrame({&quot;rank&quot;: ranks}), how=&quot;cross&quot;
  )
  
  # create 26**2 dummy column names
  column_names = []
  for letter1 in ascii_lowercase:
      for letter2 in ascii_lowercase:
          column_names.append(letter1 + letter2)
  
  # stack frames to create: ID, ranks, aa, ab, ..., zz
  df = df.hstack(
      pl.DataFrame({letter: np.random.random(len(df)) for letter in column_names})
      )
  
  
  df_pivot = df.pivot(values=df.columns[2:], index=&quot;ID&quot;, columns=&quot;rank&quot;)
  
  renamed_columns = df_pivot.columns[:1]
  for col in df.columns[2:]:
      for rank in ranks:
          renamed_columns.append(f&quot;{col}_{rank}&quot;)
  df_pivot.columns = renamed_columns
  </code></pre>