from numpy.random import permutation
from legate import pandas as lp
from tests.utils import equals, equals_scalar
for index in [
pd.RangeIndex(1, 21, 2),
pd.RangeIndex(21, 1, -2),
pd.Index(permutation(10)),
]:
print(f"Index: {index}")
sr = pd.Series(range(10), index=index)
lsr = lp.Series(range(10), index=index)
assert equals(lsr[:], sr[:])
assert equals_scalar(lsr[index[3]], sr[index[3]])
assert equals(lsr[3:-4], sr[3:-4])
assert equals(lsr[:-3], sr[:-3])
assert equals(lsr[2:], sr[2:])
pd_mask = sr % 3 == 0
lp_mask = mask = lsr % 3 == 0
for mask in [pd_mask, pd_mask.values, pd_mask.to_list(), lp_mask]:
assert equals(lsr[mask], sr[pd_mask])
sr[pd_mask] = sr[pd_mask] + 100
lsr[lp_mask] = lsr[lp_mask] + 100
assert equals(lsr, sr)
from numpy.random import permutation
from legate import pandas as lp
from tests.utils import equals, equals_scalar
for index in [
pd.RangeIndex(1, 21, 2),
pd.RangeIndex(21, 1, -2),
pd.Index(permutation(10)),
]:
print(f"Index: {index}")
sr = pd.Series(range(10), index=index)
lsr = lp.Series(range(10), index=index)
assert equals(lsr.loc[:], sr.loc[:])
assert equals_scalar(lsr.loc[index[0]], sr.loc[index[0]])
assert equals(lsr.loc[index[3]:index[-4]], sr.loc[index[3]:index[-4]])
assert equals(lsr.loc[:index[-3]], sr.loc[:index[-3]])
assert equals(lsr.loc[index[2]:], sr.loc[index[2]:])
pd_mask = sr % 3 == 0
lp_mask = mask = lsr % 3 == 0
for mask in [pd_mask, pd_mask.values, pd_mask.to_list(), lp_mask]:
assert equals(lsr.loc[mask], sr.loc[pd_mask])
sr.loc[pd_mask] = sr.loc[pd_mask] + 100
lsr.loc[lp_mask] = lsr.loc[lp_mask] + 100
assert equals(lsr, sr)
from numpy.random import permutation
from legate import pandas as lp
from tests.utils import equals, equals_scalar
for index in [
pd.RangeIndex(1, 21, 2),
pd.RangeIndex(21, 1, -2),
pd.Index(permutation(10)),
]:
print(f"Index: {index}")
sr = pd.Series(range(10), index=index)
lsr = lp.Series(range(10), index=index)
assert equals(lsr.iloc[:], sr.iloc[:])
assert equals_scalar(lsr.iloc[0], sr.iloc[0])
assert equals(lsr.iloc[3:-4], sr.iloc[3:-4])
assert equals(lsr.iloc[:-3], sr.iloc[:-3])
assert equals(lsr.iloc[2:], sr.iloc[2:])
pd_mask = sr % 3 == 0
lp_mask = mask = lsr % 3 == 0
for mask in [pd_mask, pd_mask.values, pd_mask.to_list(), lp_mask]:
assert equals(lsr.iloc[mask], sr.iloc[pd_mask.to_list()])
sr.iloc[pd_mask.values] = sr.iloc[pd_mask.values] + 100
lsr.iloc[lp_mask] = lsr.iloc[lp_mask] + 100
assert equals(lsr, sr)
for index in [
pd.RangeIndex(10),
pd.RangeIndex(1, 21, 2),
pd.RangeIndex(21, 1, -2),
pd.Index(permutation(10)),
]:
print(f"Index: {index}")
sr = pd.Series(range(10), index=index)
lsr = lp.Series(sr)
for idx in range(3, 8):
print(f"Testing lsr.at[{index[idx]}].__getitem__")
out_pd = sr.at[index[idx]]
out_lp = lsr.at[index[idx]]
assert equals_scalar(out_lp, out_pd)
for idx, val in enumerate([100, 200, 300]):
print(f"Testing lsr.at[{index[idx + 3]}].__setitem__")
sr.at[index[idx + 3]] = val
lsr.at[index[idx + 3]] = val
out_pd = sr.at[index[idx + 3]]
out_lp = lsr.at[index[idx + 3]]
assert equals_scalar(out_lp, out_pd)
# Negative tests
lsr = lp.Series(range(10))
_test(ValueError, lsr, slice(None))
must_fail(ex, _loc)
index = pd.MultiIndex.from_arrays([
[3, 2, 1, 0, 3, 2, 1, 0, 4, 3],
[0, 0, 0, 0, 1, 1, 1, 1, 2, 2],
[0, 1, 2, 3, 4, 0, 1, 2, 3, 4],
])
sr = pd.Series(range(10), index=index)
lsr = lp.Series(sr)
_test(KeyError, lsr, (2, 0, 3))
assert equals_scalar(
lsr.loc[(2, 0, 1), ],
sr.loc[(2, 0, 1), ],
)
assert equals(
lsr.loc[(2, 0), ],
sr.loc[(2, 0), ],
)
assert equals(lsr.loc[2], sr.loc[2])
assert equals(lsr.loc[4], sr.loc[4])
lsr.loc[(3, 0, 0), ] = 50
sr.loc[(3, 0, 0)] = 50
assert equals(lsr, sr)
lsr.loc[(2, 0), ] = 100
sr.loc[(2, 0)] = 100
assert equals(lsr, sr)
)
df["c"] = df["c"].astype(pd.StringDtype())
df["d"] = df["d"].astype("category")
ldf = lp.DataFrame(df)
assert equals(ldf.loc[:, "a"], df.loc[:, "a"])
assert equals(ldf.loc[:, ["b"]], df.loc[:, ["b"]])
assert equals(ldf.loc[:, ["b", "d"]], df.loc[:, ["b", "d"]])
assert equals(
ldf.loc[:, [True, False, True, False]],
df.loc[:, [True, False, True, False]],
)
assert equals(ldf.loc[:, "a":], df.loc[:, "a":])
assert equals(ldf.loc[:, :"c"], df.loc[:, :"c"])
assert equals_scalar(ldf.loc[index[0], "a"], df.loc[index[0], "a"])
assert equals(
ldf.loc[index[3], ["a", "b"]].to_pandas().T.squeeze(),
df.loc[index[3], ["a", "b"]],
)
assert equals(
ldf.loc[index[3]:index[-4], "b":"d"],
df.loc[index[3]:index[-4], "b":"d"],
)
assert equals(ldf.loc[:index[-3], "b":"d"], df.loc[:index[-3], "b":"d"])
assert equals(ldf.loc[index[2]:, "b":"d"], df.loc[index[2]:, "b":"d"])
def _test(ex, df, *args):
def _loc():
df.loc[args]
must_fail(ex, _loc)
index = pd.MultiIndex.from_arrays([
[3, 2, 1, 0, 3, 2, 1, 0, 4, 3],
[0, 0, 0, 0, 1, 1, 1, 1, 2, 2],
[0, 1, 2, 3, 4, 0, 1, 2, 3, 4],
])
df = pd.DataFrame({"a": range(10)}, index=index)
ldf = lp.DataFrame(df)
_test(KeyError, ldf, (2, 0, 3))
assert equals_scalar(ldf.loc[(2, 0, 1), "a"], df.loc[(2, 0, 1), "a"])
assert equals(ldf.loc[(2, 0), "a"], df.loc[(2, 0), "a"])
assert equals(ldf.loc[2, "a"], df.loc[2, "a"])
assert equals(ldf.loc[4, "a"], df.loc[4, "a"])
ldf.loc[(2, 0), "a"] = 100
df.loc[(2, 0), "a"] = 100
assert equals(ldf, df)
ldf.loc[4, "a"] = 200
df.loc[4, "a"] = 200
assert equals(ldf, df)