How to Implement row_number Window Function in Python
This code implements a SQL-style ROW_NUMBER() window function in pure Python, partitioning rows by a set of columns and ranking them within each partition by an ordered set of columns.
Python code
33 linesfrom collections import defaultdict
import itertools
def row_number(rows, partition_by, order_by):
partitions = defaultdict(list)
for index, row in enumerate(rows):
key = tuple(row[col] for col in partition_by)
partitions[key].append((index, row))
result = []
for key in partitions:
partition_rows = partitions[key]
partition_rows.sort(key=lambda item: tuple(item[1][col] for col in order_by))
for rank, (index, row) in enumerate(partition_rows, start=1):
result.append((index, rank, row))
result.sort(key=lambda item: item[0])
return [(rank, row) for index, rank, row in result]
if __name__ == "__main__":
data = [
{"dept": "sales", "emp": "Alice", "salary": 5000},
{"dept": "sales", "emp": "Bob", "salary": 6000},
{"dept": "hr", "emp": "Carol", "salary": 4000},
{"dept": "hr", "emp": "Dave", "salary": 4500},
{"dept": "sales", "emp": "Eve", "salary": 5500},
]
ranked = row_number(data, partition_by=["dept"], order_by=["salary"])
for rank, row in ranked:
print(row["dept"], row["emp"], row["salary"], rank)Output
sales Alice 5000 1
sales Eve 5500 2
sales Bob 6000 3
hr Carol 4000 1
hr Dave 4500 2How it works
This implementation groups rows into partitions using a defaultdict keyed by the tuple of partition column values. Within each partition, rows are sorted by the order-by columns using a tuple of values for stable comparison. The enumerate function assigns a rank starting at 1 for each row in the sorted partition. Finally, results are collected and re-sorted by the original index to preserve input order. This approach mirrors SQL's ROW_NUMBER() semantics without any external dependencies, making it suitable for teaching or small-scale data processing.
Common mistakes
- Forgetting to convert partition and order keys to tuples, which causes unhashable list errors.
- Sorting the global result by rank instead of original index, losing input order.
- Assuming stable sort without explicitly sorting by all order-by columns in the right direction.
- Mutating the input rows list when partitioning causes unexpected behavior.
Variations
- Use pandas' groupby and cumcount to compute row_number more succinctly.
- Use itertools.groupby on sorted data to avoid a separate partition dictionary.
Real-world use cases
- Assigning sequential numbers to customer orders per user before pagination or deduplication.
- Ranking top N products per category in an e-commerce recommendation engine.
- Implementing a lightweight analytics query in a Python service that avoids a full SQL engine.
Sponsored
More from Big data & Spark
Keep learning
Related tutorials and quizzes for this topic.