Splitting a sorted lazy frame into smaller chunks, preserving whole groups

user_1018861 0
Viewed 129

I need to split a multi-GB file that I am intending to add more ML-features to (adding lead and lag features). Without splitting the file, I am running out of ram, even in lazy mode.

The data is sorted by the first column, which is an id column that always contains 13 consecutive entries of the same value. It is important to avoid splitting the group of ids into separate files. This means we could split the file at every 13th, 26th, 39th (etc) position and preserve the whole group of ids.

Ideally, I'd like each file to end up being ~100MB in size. The actual df.shape is (5965869, 193), and assuming each column has a size of 8 byte, that means that I would need to split it into approximately files of 5000*13 rows.

I could get the first chunk using df.limit(5000*13), but I'm unsure on how to proceed after that. I also don't know how I would deal with the potential case where each file would contain N rows, but the last file contains fewer rows, since the total row count may not evenly divide into N, but leave a remainder.

The following is a minimal example of 1320 = 260 rows. How could I split this into 3 files - where the first two files have 138 rows each, and the last takes the remaining 13*4 rows?

import polars as pl
import numpy as np

df = pl.DataFrame({"id":13*list(range(20)), "value":np.arange(13*20)}).sort(by="id")
ldf = df.lazy()

# Either via groupby somehow
ldf.groupby("id")...take(...) # how to do this in a loop?

# or by reading the file in multiples of 13 rows
ldf...take_in_row_chunks(chunksize=some_multiple_of_13)
python-polars
1 Answers

I dont know if i understand correctly the question but you want to split it each n rows into a new dataframe.

A possible approach would be to create a new id each 13th column and then split into the dataframes into a dictionary, for simplicity i will use a split each n numbers in order for it to be reproducible. I apologize for not knowing how to insert code here in advance.


    import math
    
    df = pl.DataFrame(
        {
            "foo": ["A", "A", "B", "B", "C","C","a"],
            "N": [1, 2, 2, 4, 2,3,3],
            "bar": ["k", "l", "m", "m", "l","o","v"],
        }
    )
    
    a = df.height # number of rows
    n = 3  # split each rows
    b = (np.linspace(1, math.ceil(a/n), num=math.ceil(a/n)))
    c = (np.repeat(b, n))[:-( n - (a%n))]
    df = df.lazy().with_column(pl.lit(c).alias('new_index')).collect()
    m = df.partition_by(groups="new_index", maintain_order=True, as_dict=True)

result will be a dictionary where key is the n number of the split.

user_187456310
Related