ENH: Add split method to DataFrame for flexible row-based partitioning

[gclopton]

Feature Type

• Adding new functionality to pandas

• Changing existing functionality in pandas

• Removing existing functionality in pandas

Problem Description

Currently, pandas does not provide a direct, built-in method to split a DataFrame into multiple smaller DataFrames based on a number of specified rows. Users seeking to partition a DataFrame into chunks either have to relay on workarounds using loops and manual indexing. Adding such a feature can save users time in scenarios requiring a lot of processing of data into smaller segments, such as batch processing, cross-validation in machine learning, or dividing data for parallel processing.

Feature Description

The .split() method that I am proposing for pandas DataFrames would allow users to divide a DataFrame into smaller DataFrames based on a specified number of rows, with flexible handling of any remainder rows. The method can be described in detail as follows:

DataFrame.split(n, remainder='first' or 'last' or None)

Parameters:

1. n (integer) Would represent number of rows each resulting dataframe should contain
2. remainder (string, optional): Specifies how to handle remainder rows that do not fit evenly into the split. It accepts the following values:
   2.1. 'first': Include the remainder rows in the first split DataFrame.
   2.2. 'last': Include the remainder rows in the last split DataFrame.
   2.3. 'None': If the DataFrame cannot be evenly split, raise an error. This is the default behavior.

Pseudocode Description:

def split_dataframe(df, n, remainder=None):
1.) Get total number of rows in the dataframe
2.) Check for divisibility of remainder is None:
2.1) If length is not perfectly divisible by n, raise an error.
3.) If the remainder is 'first':
3.1) Calculate the number of rows for the first split to include the remainder.
3.2) Split the DataFrame accordingly, adjusting subsequent splits to have n rows.
4.) If the remainder is 'last':
4.1) Split the DataFrame into partitions of n rows, except for the last partition
4.2) Calculate and append the last partition separately to include any remainder.
5.) Else, if the dataframe is perfectly divisible ( or if the remainder is None and the divisibility check has passed):
5.1.) Split the DataFrame into equal parts of n rows.
6) Return the list of split DataFrames.

Example of usage:

Say that we have a DataFrame consisting of 100 rows. Then we could split this DataFrame into a list of 10 dataframes:

split_dfs = df.split(10)

In another case, we have a dataframe consisting of 99 rows. Then we could split the dataframe into 10 dataframes, where the first DataFrame consists of 9 rows and the rest of the DataFrames have 10 rows.

split_dfs = df.split(9, remainder='first')

In the other case, we could split the DataFrame into a list of 10 DataFrames, where the last DataFrame consists of the remaining 9 rows.

Alternative Solutions

One approach could involve splitting a DataFrame into smaller chunks using numpy.array_split. This function can divide an array or DataFrame into a specified number of parts and handle any remainders by evenly distributing them across splits. However there are limitations to this approach:

1.) The result is a list of numpy arrays, not DataFrames. So you lose the DataFrame context and metadata (like column names).
2.) It requires an additional step to convert these arrays back into DataFrames, adding more complexity.

We could also do manual looping and slicing. The problem is that it requires boilerplate code, which is error prone and inefficient especially when working with larger groups of people. This approach also lacks the simplicity and ease of use that a built-in method would provide.

There are also third-party libraries and packages for splitting DataFrames, such as dask.DataFrame or more_itertools.chunked. Though dask.dataframe allows for processing data into chunks, you still would have too many modifications to implement the functionality that I've described for df.split. more_itertools.chunked can be used to split an iterable into smaller iterables of a specified size, which can be applied to DataFrames. The problem is you need to convert chunks back into DataFrames. So it would be much simpler to have a built in df.split method.

Additional Context

Practical use Cases: I've encountered several scenarios in data preprocessing for machine learning where batch processing of DataFrame chunks was necessary. Implementing a custom solution each time has not been ideal. I've also run into several situations where I've needed to split DataFrame into individual DataFrames for subsequent calculations on the partitioned data.

[gclopton]

I would like to contribute to the implementation of this feature, pending approval of the feature request.

In addition to the initial proposal, I believe enhancing the function to include an option to split along either axis would significantly increase its utility.

DataFrame.split(n, axis=1, remainder='first', or 'last' or None) where axis=0 (split along rows) is the default.

[rhshadrach]

Thanks for the request - it seems to me pandas provides a number of ways to split up existing DataFrames that can be utilized, I don't see this feature provided any more utility that isn't already possible. For example, if chunker is any iterable (e.g. [0, 0, 0, 1, 1, 1]), then one can do:

[chunk for _, chunk in df.groupby(chunker)]

In addition, since pandas does not support multiple processes or distributed computing, it seems to me most common operations would benefit from not splitting up a DataFrame.

[gclopton]

Thank you for the feedback, especially regarding your suggestion on using df.groupby(chunker) as a solution.

I want to clarify the intent behind proposing this method for splitting DataFrames into smaller chunks.

The primary motivation is to enhance pandas with additional flexibility for handling large datasets, which would include scenarios such as:

1.) stepwise processing or algorithms requiring partitioned input, where splitting a DataFrame into smaller segments is necessary.
2.) Chunk-wise data export/import, simplifying the management of large datasets.
3.) Memory constraints, where processing an entire dataset at once is not feasible.
4.) External parallelization, where despite pandas not supporting distributed computing, processing smaller, independent chunks in parallel using external tools can significantly speed up operations.

The suggested method of using df.groupby(chunker) is a great solution for segmenting a DataFrame based on predefined criteria but differs from a bit from what I was suggesting:

1.) Uniformity and simplicity: The df.split() method aims to uniformly split a DataFrame into chunks based on a specified size, offering a straightforward solution when division based on size rather than data-driven groupings is important. This especially useful for tasks requiring equal-sized partions and for users seeking a simple, direct way to divide their dataframe.
2.) Ease of use: For those aiming to divide DataFrame into smaller parts without complex grouping criteria, the 'df.split()' method simplifies the process by requiring only the desired chunk size as input and optionally how to handle remainders.

I am thinking of something roughly along the lines of this:

import pandas as pd
import numpy as np

def split_dataframe(df, n, remainder=None):
    # Get number of rows in DataFrame
    total_rows = len(df)
    
    # Initialize list to hold the split DataFrames
    dfs = []
    
    # Calculate the number of DataFrames
    num_dfs = total_rows // n
    extra_rows = total_rows % n
    
    if remainder == 'first' and extra_rows > 0:
        # Size of first chunk is equal to remainder. Remaining chunks are same size.
        dfs.append(df.iloc[:n + extra_rows])
        start_idx = n + extra_rows
        for _ in range(1, num_dfs):
            dfs.append(df.iloc[start_idx:start_idx + n])
            start_idx += n
    elif remainder == 'last' and extra_rows > 0:
        # Last chunk consists of remaining rows. The other chunks are the same size.
        for i in range(num_dfs):
            dfs.append(df.iloc[i*n:(i+1)*n])
        # Add extra rows to the last chunk
        dfs.append(df.iloc[num_dfs*n:])
    elif remainder == 'spread':
        # Evenly spread extra rows across the first few chunks
        for i in range(num_dfs + (1 if extra_rows > 0 else 0)):
            size = n + (1 if i < extra_rows else 0)
            dfs.append(df.iloc[i*size:min((i+1)*size, total_rows)])
    else:
        if extra_rows > 0:
            raise ValueError(f"DataFrame length is not perfectly divisible by {n}. Please specify the 'remainder' parameter.")
        # If remainder is None and the DataFrame is perfectly divisible
        for i in range(num_dfs):
            dfs.append(df.iloc[i*n:(i+1)*n])
    
    return dfs


# Example usage
df1 = pd.DataFrame(np.arange(100), columns=['Column'])
df2 = pd.DataFrame(np.arange(99), columns=['Column'])
dfs_equal = split_dataframe(df1, 10)  # Uses the default 'equal' distribution
dfs_first = split_dataframe(df2, 10, remainder='first')
dfs_last = split_dataframe(df2, 10, remainder='last')
dfs_spread = split_dataframe(df2, 10, remainder='spread')