python - How to speed up the operation of repeating take first n rows for each group after group_by? - Stack Overflow

admin管理员组

文章数量:1122832

The df contains 100 millions of rows, and group_by columns is like 25-30. Is there a way to speed this operation up from here? or this is the best I can get.

import polars as pl
import numpy as np

rows = 100000000
n_cols = 30
df = pl.DataFrame(np.random.randint(0, 100, size=(n_cols, rows)), schema=[str(col_i) for col_i in range(n_cols)])
First_n_rows_list = [1,2,3]

df = df.sort('col_0').group_by([“col_”+str(i) for i in range(1, n_cols)])    
result = pl.concat([df.head(First_n_rows).with_columns(pl.lit(First_n_rows).alias('First_n_rows').cast(pl.Int8)) for First_n_rows in First_n_rows_list])

The df contains 100 millions of rows, and group_by columns is like 25-30. Is there a way to speed this operation up from here? or this is the best I can get.

import polars as pl
import numpy as np

rows = 100000000
n_cols = 30
df = pl.DataFrame(np.random.randint(0, 100, size=(n_cols, rows)), schema=[str(col_i) for col_i in range(n_cols)])
First_n_rows_list = [1,2,3]

df = df.sort('col_0').group_by([“col_”+str(i) for i in range(1, n_cols)])    
result = pl.concat([df.head(First_n_rows).with_columns(pl.lit(First_n_rows).alias('First_n_rows').cast(pl.Int8)) for First_n_rows in First_n_rows_list])

• python
• python-polars

Share Improve this question Follow edited yesterday user28199045 asked yesterday user28199045user28199045 1333 bronze badges 3

• i am thinking another way to do this, just don't know if this will improve the running time. In each group, get the rank for each row, then duplicate each row max(x_list)-rank times. for example, rank 0 duplicate 3 times, rank 1 duplicate 2 time. – user28199045 Commented yesterday
• it's not fully clear what are you trying to achieve here, maybe some more meaningful example with proper (small) input and output would help. But for a start, you can use df.head(max(x_list)) to reduce the size of the dataframe. – roman Commented yesterday
• The goal is to get first n rows within each group, add a column showing first_n_value for each first_n_value in the list. Then concat all the rows from different first_n_value case. I edited my description to make it clear. Thanks! – user28199045 Commented yesterday

Add a comment  | 

2 Answers 2

Sorted by: Reset to default 0

As you said, you can take head(max(x_list)) and then repeat each row appropriate number of times:

x = head(max(x_list))

(
    df.head(x)
    .with_columns(
        pl.int_range(pl.len() + 1, 1, step=-1)
        .over([str(x) for x in range(1,n_cols)]).alias("x")
    )
    .with_columns(pl.exclude("x").repeat_by("x"))
    .explode(pl.exclude("x"))
)

import polars as pl
import numpy as no

n = 50
df = pl.DataFrame(np.random.randint(0, 100, size = (4, n)), schema= ['A', 'B', 'C', 'D'])
x_list = [1, 2, 3]

grouped = df.group_by(['A', 'B', 'C'])
result = pl.concat([grouped.head(x).with_columns(pl.lit(x).alias('x').cast(pl.Int8)) for x in x_list])

本文标签：