Skip to content

Commit

Permalink
Optimize get_window_ranges function (#489)
Browse files Browse the repository at this point in the history
deque.appendleft is roughly 10x faster than list.insert(0, ...)
  • Loading branch information
gwaramadze authored Sep 9, 2024
1 parent be120a3 commit 1240e61
Show file tree
Hide file tree
Showing 2 changed files with 10 additions and 10 deletions.
9 changes: 5 additions & 4 deletions quixstreams/dataframe/windows/base.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
from typing import Any, Optional, Callable, Tuple, List
from collections import deque
from typing import Any, Deque, Optional, Callable, Tuple

from typing_extensions import TypedDict

Expand All @@ -17,7 +18,7 @@ class WindowResult(TypedDict):

def get_window_ranges(
timestamp_ms: int, duration_ms: int, step_ms: Optional[int] = None
) -> List[Tuple[int, int]]:
) -> Deque[Tuple[int, int]]:
"""
Get a list of window ranges for the given timestamp.
:param timestamp_ms: timestamp in milliseconds
Expand All @@ -28,14 +29,14 @@ def get_window_ranges(
if not step_ms:
step_ms = duration_ms

window_ranges = []
window_ranges = deque()
current_window_start = timestamp_ms - (timestamp_ms % step_ms)

while (
current_window_start > timestamp_ms - duration_ms and current_window_start >= 0
):
window_end = current_window_start + duration_ms
window_ranges.insert(0, (current_window_start, window_end))
window_ranges.appendleft((current_window_start, window_end))
current_window_start -= step_ms

return window_ranges
11 changes: 5 additions & 6 deletions tests/test_quixstreams/test_dataframe/test_windows/test_base.py
Original file line number Diff line number Diff line change
Expand Up @@ -16,12 +16,10 @@ class TestGetWindowRanges:
],
)
def test_get_window_ranges_with_step(self, timestamp, duration, step, expected):
assert (
get_window_ranges(
timestamp_ms=timestamp, duration_ms=duration, step_ms=step
)
== expected
result = get_window_ranges(
timestamp_ms=timestamp, duration_ms=duration, step_ms=step
)
assert list(result) == expected

@pytest.mark.parametrize(
"timestamp, duration, expected",
Expand All @@ -31,4 +29,5 @@ def test_get_window_ranges_with_step(self, timestamp, duration, step, expected):
],
)
def test_get_window_ranges_no_step(self, timestamp, duration, expected):
assert get_window_ranges(timestamp, duration) == expected
result = get_window_ranges(timestamp, duration)
assert list(result) == expected

0 comments on commit 1240e61

Please sign in to comment.