Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

SNOW-1638407 Added some simple tests for TimedeltaIndex indexing #2183

Merged
merged 5 commits into from
Aug 30, 2024
Merged

SNOW-1638407 Added some simple tests for TimedeltaIndex indexing #2183

Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
dd8c606
some simple tests for timedeltaindex indexing
sfc-gh-vbudati Aug 28, 2024
0d26309
Merge branch 'main' into vbudati/SNOW-1638407-test-timedelta-index-in…
sfc-gh-vbudati Aug 28, 2024
e32e3cb
add df enlargement test
sfc-gh-vbudati Aug 28, 2024
91e223e
Merge branch 'main' into vbudati/SNOW-1638407-test-timedelta-index-in…
sfc-gh-vbudati Aug 28, 2024
5c48014
Merge branch 'main' into vbudati/SNOW-1638407-test-timedelta-index-in…
sfc-gh-vbudati Aug 29, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
180 changes: 180 additions & 0 deletions tests/integ/modin/types/test_timedelta_indexing.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -389,3 +389,183 @@ def loc_enlargement(key, item, df):
loc_enlargement(key, item, snow_td.copy()).to_pandas().dtypes,
snow_td.dtypes,
)


@pytest.mark.parametrize(
"key, join_count",
[(2, 2), ([2, 1], 2), (slice(1, None), 0), ([True, False, False, True], 1)],
)
def test_index_get_timedelta(key, join_count):
td_idx = native_pd.TimedeltaIndex(
[
native_pd.Timedelta("1 days 1 hour"),
native_pd.Timedelta("2 days 1 minute"),
native_pd.Timedelta("3 days 1 nanoseconds"),
native_pd.Timedelta("100 nanoseconds"),
]
)
snow_td_idx = pd.TimedeltaIndex(td_idx)

with SqlCounter(query_count=1, join_count=join_count):
if is_scalar(key):
assert snow_td_idx[key] == td_idx[key]
else:
eval_snowpark_pandas_result(snow_td_idx, td_idx, lambda idx: idx[key])


@pytest.mark.parametrize(
"key, api, query_count, join_count",
[
[2, "iat", 1, 2],
[native_pd.Timedelta("1 days 1 hour"), "at", 2, 2],
[[2, 1], "iloc", 1, 2],
[
[
native_pd.Timedelta("1 days 1 hour"),
native_pd.Timedelta("1 days 1 hour"),
],
"loc",
1,
1,
],
[slice(1, None), "iloc", 1, 0],
[[True, False, False, True], "iloc", 1, 1],
[[True, False, False, True], "loc", 1, 1],
],
)
def test_series_with_timedelta_index(key, api, query_count, join_count):
td_idx = native_pd.TimedeltaIndex(
[
native_pd.Timedelta("1 days 1 hour"),
native_pd.Timedelta("2 days 1 minute"),
native_pd.Timedelta("3 days 1 nanoseconds"),
native_pd.Timedelta("100 nanoseconds"),
]
)
snow_td_idx = pd.TimedeltaIndex(td_idx)

data = [1, 2, 3, 4]
native_series = native_pd.Series(data, index=td_idx)
snow_series = pd.Series(data, index=snow_td_idx)

with SqlCounter(query_count=query_count, join_count=join_count):
if is_scalar(key):
assert getattr(snow_series, api)[key] == getattr(native_series, api)[key]
else:
eval_snowpark_pandas_result(
snow_series, native_series, lambda s: getattr(s, api)[key]
)


@pytest.mark.parametrize(
"key, api, query_count, join_count",
[
[2, "iat", 1, 2],
[native_pd.Timedelta("1 days 1 hour"), "at", 2, 2],
[[2, 1], "iloc", 1, 2],
[
[
native_pd.Timedelta("1 days 1 hour"),
native_pd.Timedelta("1 days 1 hour"),
],
"loc",
1,
1,
],
[slice(1, None), "iloc", 1, 0],
[[True, False, False, True], "iloc", 1, 1],
[[True, False, False, True], "loc", 1, 1],
],
)
def test_df_with_timedelta_index(key, api, query_count, join_count):
td_idx = native_pd.TimedeltaIndex(
[
native_pd.Timedelta("1 days 1 hour"),
native_pd.Timedelta("2 days 1 minute"),
native_pd.Timedelta("3 days 1 nanoseconds"),
native_pd.Timedelta("100 nanoseconds"),
]
)
snow_td_idx = pd.TimedeltaIndex(td_idx)

data = [[1, 2, 3, 4], [5, 6, 7, 8], [9, 10, 11, 12], [13, 14, 15, 16]]
native_df = native_pd.DataFrame(data, index=td_idx)
snow_df = pd.DataFrame(data, index=snow_td_idx)

with SqlCounter(query_count=query_count, join_count=join_count):
if is_scalar(key):
assert getattr(snow_df, api)[key, 0] == getattr(native_df, api)[key, 0]
else:
eval_snowpark_pandas_result(
snow_df, native_df, lambda s: getattr(s, api)[key]
)


def test_df_with_timedelta_index_enlargement_during_indexing():
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

would it be possible for this test to be parametrized, rather than a sequence of SqlCounter blocks?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think it would become more dense since some of the tests are run on snow_df["a"]. If you still prefer that I parameterize the test I can try to do that

td_idx = native_pd.TimedeltaIndex(
[
native_pd.Timedelta("1 days 1 hour"),
native_pd.Timedelta("2 days 1 minute"),
native_pd.Timedelta("3 days 1 nanoseconds"),
native_pd.Timedelta("100 nanoseconds"),
]
)
snow_td_idx = pd.TimedeltaIndex(td_idx)

data = [[1, 2, 3, 4], [5, 6, 7, 8], [9, 10, 11, 12], [13, 14, 15, 16]]
cols = ["a", "b", "c", "d"]
native_df = native_pd.DataFrame(data, index=td_idx, columns=cols)
snow_df = pd.DataFrame(data, index=snow_td_idx, columns=cols)

def setitem_enlargement(key, item, df):
df[key] = item
return df

item = 23

key = native_pd.Timedelta("2 days")
with SqlCounter(query_count=1, join_count=0):
eval_snowpark_pandas_result(
snow_df.copy(),
native_df.copy(),
functools.partial(setitem_enlargement, key, item),
)

key = native_pd.Timedelta("2 days 45 minutes")
with SqlCounter(query_count=1, join_count=1):
eval_snowpark_pandas_result(
snow_df["a"].copy(),
native_df["a"].copy(),
functools.partial(setitem_enlargement, key, item),
)

def loc_enlargement(key, item, df):
df.loc[key] = item
return df

key = (slice(None, None, None), "x")

with SqlCounter(query_count=1, join_count=0):
eval_snowpark_pandas_result(
snow_df.copy(),
native_df.copy(),
functools.partial(loc_enlargement, key, item),
)

key = native_pd.Timedelta("2 days 25 minutes")
with SqlCounter(query_count=1, join_count=1):
eval_snowpark_pandas_result(
snow_df["a"].copy(),
native_df["a"].copy(),
functools.partial(loc_enlargement, key, item),
)

# single row
key = (native_pd.Timedelta("2 days 45 minutes"), slice(None, None, None))

with SqlCounter(query_count=1, join_count=1):
eval_snowpark_pandas_result(
snow_df.copy(),
native_df.copy(),
functools.partial(loc_enlargement, key, item),
)
Loading