Skip to content

Commit 4d1e7b4

Browse files
wjsiwjsi
committed
Fix tests
1 parent 7666f99 commit 4d1e7b4

File tree

6 files changed

+99
-68
lines changed

6 files changed

+99
-68
lines changed

‎mars/dataframe/groupby/aggregation.py

Lines changed: 36 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -701,34 +701,36 @@ def _do_custom_agg(op, custom_reduction, *input_objs):
701701
out = op.outputs[0]
702702
for group_key in input_objs[0].groups.keys():
703703
group_objs = [o.get_group(group_key) for o in input_objs]
704+
704705
agg_done = False
705706
if op.stage == OperandStage.map:
706-
result = custom_reduction.pre(group_objs[0])
707+
res_tuple = custom_reduction.pre(group_objs[0])
707708
agg_done = custom_reduction.pre_with_agg
708-
if not isinstance(result, tuple):
709-
result = (result,)
709+
if not isinstance(res_tuple, tuple):
710+
res_tuple = (res_tuple,)
710711
else:
711-
result = group_objs
712+
res_tuple = group_objs
712713

713714
if not agg_done:
714-
result = custom_reduction.agg(*result)
715-
if not isinstance(result, tuple):
716-
result = (result,)
715+
res_tuple = custom_reduction.agg(*res_tuple)
716+
if not isinstance(res_tuple, tuple):
717+
res_tuple = (res_tuple,)
717718

718719
if op.stage == OperandStage.agg:
719-
result = custom_reduction.post(*result)
720-
if not isinstance(result, tuple):
721-
result = (result,)
722-
723-
if out.ndim == 2:
724-
if result[0].ndim == 1:
725-
result = tuple(r.to_frame().T for r in result)
726-
if op.stage == OperandStage.agg:
727-
result = tuple(r.astype(out.dtypes) for r in result)
728-
else:
729-
result = tuple(xdf.Series(r) for r in result)
720+
res_tuple = custom_reduction.post(*res_tuple)
721+
if not isinstance(res_tuple, tuple):
722+
res_tuple = (res_tuple,)
723+
724+
new_res_list = []
725+
for r in res_tuple:
726+
if out.ndim == 2 and r.ndim == 1:
727+
r = r.to_frame().T
728+
elif out.ndim < 2:
729+
if getattr(r, "ndim", 0) == 2:
730+
r = r.iloc[0, :]
731+
else:
732+
r = xdf.Series(r)
730733

731-
for r in result:
732734
if len(input_objs[0].grouper.names) == 1:
733735
r.index = xdf.Index(
734736
[group_key], name=input_objs[0].grouper.names[0]
@@ -737,7 +739,21 @@ def _do_custom_agg(op, custom_reduction, *input_objs):
737739
r.index = xdf.MultiIndex.from_tuples(
738740
[group_key], names=input_objs[0].grouper.names
739741
)
740-
results.append(result)
742+
743+
if op.groupby_params.get("selection"):
744+
# correct columns for groupby-selection-agg paradigms
745+
selection = op.groupby_params["selection"]
746+
r.columns = [selection] if input_objs[0].ndim == 1 else selection
747+
748+
if out.ndim == 2 and op.stage == OperandStage.agg:
749+
dtype_cols = set(out.dtypes.index) & set(r.columns)
750+
conv_dtypes = {
751+
k: v for k, v in out.dtypes.items() if k in dtype_cols
752+
}
753+
r = r.astype(conv_dtypes)
754+
new_res_list.append(r)
755+
756+
results.append(tuple(new_res_list))
741757
if not results and op.stage == OperandStage.agg:
742758
empty_df = pd.DataFrame(
743759
[], columns=out.dtypes.index, index=out.index_value.to_pandas()[:0]

‎mars/dataframe/groupby/tests/test_groupby.py

Lines changed: 33 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -396,28 +396,28 @@ def test_groupby_fill():
396396
)
397397
mdf = md.DataFrame(df1, chunk_size=3)
398398

399-
r = tile(getattr(mdf.groupby(["one", "two"]), "ffill")())
399+
r = tile(mdf.groupby(["one", "two"]).ffill())
400400
assert r.op.output_types[0] == OutputType.dataframe
401401
assert r.shape == (len(df1), 1)
402402
assert len(r.chunks) == 3
403403
assert r.chunks[0].shape == (np.nan, 1)
404404
assert r.dtypes.index.tolist() == ["three"]
405405

406-
r = tile(getattr(mdf.groupby(["two"]), "bfill")())
406+
r = tile(mdf.groupby(["two"]).bfill())
407407
assert r.op.output_types[0] == OutputType.dataframe
408408
assert r.shape == (len(df1), 2)
409409
assert len(r.chunks) == 3
410410
assert r.chunks[0].shape == (np.nan, 2)
411411
assert r.dtypes.index.tolist() == ["one", "three"]
412412

413-
r = tile(getattr(mdf.groupby(["two"]), "backfill")())
413+
r = tile(mdf.groupby(["two"]).backfill())
414414
assert r.op.output_types[0] == OutputType.dataframe
415415
assert r.shape == (len(df1), 2)
416416
assert len(r.chunks) == 3
417417
assert r.chunks[0].shape == (np.nan, 2)
418418
assert r.dtypes.index.tolist() == ["one", "three"]
419419

420-
r = tile(getattr(mdf.groupby(["one"]), "fillna")(5))
420+
r = tile(mdf.groupby(["one"]).fillna(5))
421421
assert r.op.output_types[0] == OutputType.dataframe
422422
assert r.shape == (len(df1), 2)
423423
assert len(r.chunks) == 3
@@ -426,25 +426,25 @@ def test_groupby_fill():
426426

427427
s1 = pd.Series([4, 3, 9, np.nan, np.nan, 7, 10, 8, 1, 6])
428428
ms1 = md.Series(s1, chunk_size=3)
429-
r = tile(getattr(ms1.groupby(lambda x: x % 2), "ffill")())
429+
r = tile(ms1.groupby(lambda x: x % 2).ffill())
430430
assert r.op.output_types[0] == OutputType.series
431431
assert len(r.chunks) == 4
432432
assert r.shape == (len(s1),)
433433
assert r.chunks[0].shape == (np.nan,)
434434

435-
r = tile(getattr(ms1.groupby(lambda x: x % 2), "bfill")())
435+
r = tile(ms1.groupby(lambda x: x % 2).bfill())
436436
assert r.op.output_types[0] == OutputType.series
437437
assert len(r.chunks) == 4
438438
assert r.shape == (len(s1),)
439439
assert r.chunks[0].shape == (np.nan,)
440440

441-
r = tile(getattr(ms1.groupby(lambda x: x % 2), "backfill")())
441+
r = tile(ms1.groupby(lambda x: x % 2).backfill())
442442
assert r.op.output_types[0] == OutputType.series
443443
assert len(r.chunks) == 4
444444
assert r.shape == (len(s1),)
445445
assert r.chunks[0].shape == (np.nan,)
446446

447-
r = tile(getattr(ms1.groupby(lambda x: x % 2), "fillna")(5))
447+
r = tile(ms1.groupby(lambda x: x % 2).fillna(5))
448448
assert r.op.output_types[0] == OutputType.series
449449
assert len(r.chunks) == 4
450450
assert r.shape == (len(s1),)
@@ -453,26 +453,47 @@ def test_groupby_fill():
453453
s1 = pd.Series([4, 3, 9, np.nan, np.nan, 7, 10, 8, 1, 6])
454454
ms1 = md.Series(s1, chunk_size=3)
455455

456-
r = tile(getattr(ms1.groupby(lambda x: x % 2), "ffill")())
456+
r = tile(ms1.groupby(lambda x: x % 2).ffill())
457457
assert r.op.output_types[0] == OutputType.series
458458
assert len(r.chunks) == 4
459459
assert r.shape == (len(s1),)
460460
assert r.chunks[0].shape == (np.nan,)
461461

462-
r = tile(getattr(ms1.groupby(lambda x: x % 2), "bfill")())
462+
r = tile(ms1.groupby(lambda x: x % 2).bfill())
463463
assert r.op.output_types[0] == OutputType.series
464464
assert len(r.chunks) == 4
465465
assert r.shape == (len(s1),)
466466
assert r.chunks[0].shape == (np.nan,)
467467

468-
r = tile(getattr(ms1.groupby(lambda x: x % 2), "backfill")())
468+
r = tile(ms1.groupby(lambda x: x % 2).backfill())
469469
assert r.op.output_types[0] == OutputType.series
470470
assert len(r.chunks) == 4
471471
assert r.shape == (len(s1),)
472472
assert r.chunks[0].shape == (np.nan,)
473473

474-
r = tile(getattr(ms1.groupby(lambda x: x % 2), "fillna")(5))
474+
r = tile(ms1.groupby(lambda x: x % 2).fillna(5))
475475
assert r.op.output_types[0] == OutputType.series
476476
assert len(r.chunks) == 4
477477
assert r.shape == (len(s1),)
478478
assert r.chunks[0].shape == (np.nan,)
479+
480+
481+
def test_groupby_nunique():
482+
df1 = pd.DataFrame(
483+
[
484+
[1, 1, 10],
485+
[1, 1, np.nan],
486+
[1, 1, np.nan],
487+
[1, 2, np.nan],
488+
[1, 2, 20],
489+
[1, 2, np.nan],
490+
[1, 3, np.nan],
491+
[1, 3, np.nan],
492+
],
493+
columns=["one", "two", "three"],
494+
)
495+
mdf = md.DataFrame(df1, chunk_size=3)
496+
497+
r = tile(mdf.groupby(["one", "two"]).nunique())
498+
assert len(r.chunks) == 1
499+
assert isinstance(r.chunks[0].op, DataFrameGroupByAgg)

‎mars/dataframe/groupby/tests/test_groupby_execution.py

Lines changed: 11 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -1241,13 +1241,16 @@ def test_groupby_nunique(setup):
12411241
# test with as_index=False
12421242
mdf = md.DataFrame(df1, chunk_size=13)
12431243
if _agg_size_as_frame:
1244+
res = mdf.groupby("b", as_index=False)["a"].nunique().execute().fetch()
1245+
expected = df1.groupby("b", as_index=False)["a"].nunique()
12441246
pd.testing.assert_frame_equal(
1245-
mdf.groupby("b", as_index=False)["a"]
1246-
.nunique()
1247-
.execute()
1248-
.fetch()
1249-
.sort_values(by="b", ignore_index=True),
1250-
df1.groupby("b", as_index=False)["a"]
1251-
.nunique()
1252-
.sort_values(by="b", ignore_index=True),
1247+
res.sort_values(by="b", ignore_index=True),
1248+
expected.sort_values(by="b", ignore_index=True),
1249+
)
1250+
1251+
res = mdf.groupby("b", as_index=False)[["a", "c"]].nunique().execute().fetch()
1252+
expected = df1.groupby("b", as_index=False)[["a", "c"]].nunique()
1253+
pd.testing.assert_frame_equal(
1254+
res.sort_values(by="b", ignore_index=True),
1255+
expected.sort_values(by="b", ignore_index=True),
12531256
)

‎mars/dataframe/reduction/nunique.py

Lines changed: 18 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -41,61 +41,53 @@ def __init__(
4141
self._dropna = dropna
4242
self._use_arrow_dtype = use_arrow_dtype
4343

44-
@staticmethod
45-
def _drop_duplicates_to_arrow(v, explode=False):
44+
def _drop_duplicates(self, xdf, value, explode=False):
4645
if explode:
47-
v = v.explode()
48-
try:
49-
return ArrowListArray([v.drop_duplicates().to_numpy()])
50-
except pa.ArrowInvalid:
51-
# fallback due to diverse dtypes
52-
return [v.drop_duplicates().to_list()]
46+
value = value.explode()
47+
48+
if not self._use_arrow_dtype or xdf is cudf:
49+
return [value.drop_duplicates().to_list()]
50+
else:
51+
try:
52+
return ArrowListArray([value.drop_duplicates().to_numpy()])
53+
except pa.ArrowInvalid:
54+
# fallback due to diverse dtypes
55+
return [value.drop_duplicates().to_list()]
5356

5457
def pre(self, in_data): # noqa: W0221 # pylint: disable=arguments-differ
5558
xdf = cudf if self.is_gpu() else pd
5659
if isinstance(in_data, xdf.Series):
57-
unique_values = in_data.drop_duplicates()
60+
unique_values = self._drop_duplicates(xdf, in_data)
5861
return xdf.Series(unique_values, name=in_data.name)
5962
else:
6063
if self._axis == 0:
6164
data = dict()
6265
for d, v in in_data.iteritems():
63-
if not self._use_arrow_dtype or xdf is cudf:
64-
data[d] = [v.drop_duplicates().to_list()]
65-
else:
66-
data[d] = self._drop_duplicates_to_arrow(v)
66+
data[d] = self._drop_duplicates(xdf, v)
6767
df = xdf.DataFrame(data)
6868
else:
6969
df = xdf.DataFrame(columns=[0])
7070
for d, v in in_data.iterrows():
71-
if not self._use_arrow_dtype or xdf is cudf:
72-
df.loc[d] = [v.drop_duplicates().to_list()]
73-
else:
74-
df.loc[d] = self._drop_duplicates_to_arrow(v)
71+
df.loc[d] = self._drop_duplicates(xdf, v)
7572
return df
7673

7774
def agg(self, in_data): # noqa: W0221 # pylint: disable=arguments-differ
7875
xdf = cudf if self.is_gpu() else pd
7976
if isinstance(in_data, xdf.Series):
80-
unique_values = in_data.explode().drop_duplicates()
77+
unique_values = self._drop_duplicates(xdf, in_data, explode=True)
8178
return xdf.Series(unique_values, name=in_data.name)
8279
else:
8380
if self._axis == 0:
8481
data = dict()
8582
for d, v in in_data.iteritems():
86-
if not self._use_arrow_dtype or xdf is cudf:
87-
data[d] = [v.explode().drop_duplicates().to_list()]
88-
else:
83+
if self._use_arrow_dtype and xdf is not cudf:
8984
v = pd.Series(v.to_numpy())
90-
data[d] = self._drop_duplicates_to_arrow(v, explode=True)
85+
data[d] = self._drop_duplicates(xdf, v, explode=True)
9186
df = xdf.DataFrame(data)
9287
else:
9388
df = xdf.DataFrame(columns=[0])
9489
for d, v in in_data.iterrows():
95-
if not self._use_arrow_dtype or xdf is cudf:
96-
df.loc[d] = [v.explode().drop_duplicates().to_list()]
97-
else:
98-
df.loc[d] = self._drop_duplicates_to_arrow(v, explode=True)
90+
df.loc[d] = self._drop_duplicates(xdf, v, explode=True)
9991
return df
10092

10193
def post(self, in_data): # noqa: W0221 # pylint: disable=arguments-differ

‎mars/dataframe/reduction/tests/test_reduction_execution.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -679,7 +679,7 @@ def test_nunique(setup, check_ref_counts):
679679

680680

681681
@pytest.mark.skipif(pa is None, reason="pyarrow not installed")
682-
def test_use_arrow_dtype_n_unique(setup, check_ref_counts):
682+
def test_use_arrow_dtype_nunique(setup, check_ref_counts):
683683
with option_context({"dataframe.use_arrow_dtype": True, "combine_size": 2}):
684684
rs = np.random.RandomState(0)
685685
data1 = pd.DataFrame(

‎mars/oscar/backends/message.pyx

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -556,7 +556,6 @@ cpdef reset_random_seed():
556556
global _rnd_is_seed_set
557557

558558
seed_bytes = getrandbits(64).to_bytes(8, "little")
559-
# memcpy(&seed, <char *>seed_bytes, 8)
560559
_rnd_gen.seed((<uint_fast64_t *><char *>seed_bytes)[0])
561560
_rnd_is_seed_set = True
562561

0 commit comments

Comments
 (0)