pandas
dataframe
data manipulation
python
data processing

Split a large pandas dataframe

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design

Introduction

Splitting a large pandas DataFrame is useful when you need to process data in chunks, write batch files, or keep downstream operations within memory limits. The best split strategy depends on what "large" means for your workflow. Sometimes you want equal row chunks, sometimes logical groups, and sometimes a lazy iterator so you do not create many intermediate frames at once.

Split by row count with numpy.array_split

If you want a fixed number of roughly equal parts, numpy.array_split is a convenient option.

python
1import numpy as np
2import pandas as pd
3
4df = pd.DataFrame({"value": range(10)})
5parts = np.array_split(df, 3)
6
7for i, part in enumerate(parts, start=1):
8    print(f"Part {i}")
9    print(part)

This returns a list of smaller DataFrames. It is simple and readable when you know the number of output chunks ahead of time.

Split by chunk size with iloc

If your workflow cares about a fixed number of rows per chunk, iterate with iloc.

python
1import pandas as pd
2
3df = pd.DataFrame({"value": range(12)})
4chunk_size = 5
5
6for start in range(0, len(df), chunk_size):
7    chunk = df.iloc[start:start + chunk_size]
8    print(chunk)

This pattern is often better than array_split because it gives direct control over the maximum chunk size. It is especially useful when sending batches to another system or writing output files with predictable row counts.

Use a generator for memory-friendly processing

If the DataFrame is already large, creating a list of all chunks may be unnecessary. A generator yields one chunk at a time.

python
1import pandas as pd
2
3def chunk_dataframe(df, chunk_size):
4    for start in range(0, len(df), chunk_size):
5        yield df.iloc[start:start + chunk_size]
6
7
8df = pd.DataFrame({"value": range(12)})
9
10for chunk in chunk_dataframe(df, 4):
11    print(chunk)

This keeps the chunking logic reusable and avoids materializing the entire chunk list at once.

Split by logical groups instead of row counts

Sometimes equal sizes are not what you want. If each customer, date, or category should stay together, split by grouping.

python
1import pandas as pd
2
3df = pd.DataFrame(
4    {
5        "team": ["A", "A", "B", "B", "C"],
6        "score": [10, 12, 9, 15, 11],
7    }
8)
9
10grouped = {name: group for name, group in df.groupby("team")}
11
12for name, group in grouped.items():
13    print(name)
14    print(group)

This is a semantic split rather than a size-based one. It is helpful when downstream logic requires each chunk to be internally consistent by key.

Splitting does not solve every memory problem

A common misconception is that splitting a DataFrame always reduces memory usage. If the full DataFrame is already loaded into memory, splitting it after the fact does not magically undo that cost. It only helps with how you process or export the data afterward.

If memory is the real problem, the better design may be reading the source in chunks from the start, such as using pd.read_csv(..., chunksize=...).

python
1import pandas as pd
2
3for chunk in pd.read_csv("large.csv", chunksize=10000):
4    print(chunk.shape)

That pattern is more scalable than loading everything and splitting later.

Common Pitfalls

The biggest mistake is building a huge list of chunk DataFrames when a generator would do. That can waste memory and defeat the purpose of chunked processing.

Another issue is assuming array_split guarantees identical sizes. It produces roughly equal parts, not perfectly equal ones in every case.

Developers also split by row count when the data really should be split by a business key such as customer id or date. That can break downstream joins or aggregations.

Finally, remember that chunking after the load does not reduce the memory required to create the original DataFrame. If memory is tight, stream the input source in chunks from the beginning.

Summary

  • Use np.array_split when you want a fixed number of roughly equal parts.
  • Use iloc when you want a specific chunk size in rows.
  • Prefer generators when you want chunked processing without building a full chunk list.
  • Split by groupby when logical group boundaries matter more than equal sizes.
  • If memory is the real issue, read the source data in chunks instead of splitting after load.

Related reading
Free course
Beginner
7 lessons
2 hours
Tackling System Design Interview Problems

A short course that equips you with the skills to approach system design interviews methodically.

Start the free course
Track what you have practised

A free account saves your progress, solutions and study plan across every problem on Codemia.

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design

All Rights Reserved.