如何在 pandas DataFrame 中随机提取 n 个连续行？

Question

I would like to randomly select n contiguous rows from a DataFrame. This is the only working code I could come up with:

random_row = df.sample(n=1)
start = random_row.index
end = start + n - 1
n_rows = df.iloc[int(start.values):int(end.values)]

But I feel it's bad code as in hacky and not very pythonic. For some reason I could not use the Int64Index es, and that feels very weird. I would expect to be able to index a DataFrame by its... Indexes, but it throws some errors.

Can anyone advise how to make my code better or recommend a better way to do what I'm doing?

Answer 1

I can see 2 problems with your current code:

• It assumes that the rows are sequentially indexed with integers (0, 1, 2, 3, ...)
• There may not be enough rows going down from start until the end of the data frame to meet your requirements

A slightly improved version:

def random_block(frame, k, wrap=True):
    "Return a random block of k contiguous rows from the DataFrame"
    n = len(frame)
    if k == 0:
        raise ValueError('k must be >0')
    elif k > n:
        raise ValueError('k must not be longer than the DataFrame')
    elif k == n:
        return frame
    
    start = np.random.randint(0, n - (0 if wrap else k - 1))
    end = start + k
    return frame.iloc[np.arange(n).take(range(start,end), mode='wrap')]
    
# Usage
result = random_block(df, 10)