test_nth.py 25 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824
  1. import numpy as np
  2. import pytest
  3. import pandas as pd
  4. from pandas import (
  5. DataFrame,
  6. Index,
  7. MultiIndex,
  8. Series,
  9. Timestamp,
  10. isna,
  11. )
  12. import pandas._testing as tm
  13. def test_first_last_nth(df):
  14. # tests for first / last / nth
  15. grouped = df.groupby("A")
  16. first = grouped.first()
  17. expected = df.loc[[1, 0], ["B", "C", "D"]]
  18. expected.index = Index(["bar", "foo"], name="A")
  19. expected = expected.sort_index()
  20. tm.assert_frame_equal(first, expected)
  21. nth = grouped.nth(0)
  22. expected = df.loc[[0, 1]]
  23. tm.assert_frame_equal(nth, expected)
  24. last = grouped.last()
  25. expected = df.loc[[5, 7], ["B", "C", "D"]]
  26. expected.index = Index(["bar", "foo"], name="A")
  27. tm.assert_frame_equal(last, expected)
  28. nth = grouped.nth(-1)
  29. expected = df.iloc[[5, 7]]
  30. tm.assert_frame_equal(nth, expected)
  31. nth = grouped.nth(1)
  32. expected = df.iloc[[2, 3]]
  33. tm.assert_frame_equal(nth, expected)
  34. # it works!
  35. grouped["B"].first()
  36. grouped["B"].last()
  37. grouped["B"].nth(0)
  38. df.loc[df["A"] == "foo", "B"] = np.nan
  39. assert isna(grouped["B"].first()["foo"])
  40. assert isna(grouped["B"].last()["foo"])
  41. assert isna(grouped["B"].nth(0).iloc[0])
  42. # v0.14.0 whatsnew
  43. df = DataFrame([[1, np.nan], [1, 4], [5, 6]], columns=["A", "B"])
  44. g = df.groupby("A")
  45. result = g.first()
  46. expected = df.iloc[[1, 2]].set_index("A")
  47. tm.assert_frame_equal(result, expected)
  48. expected = df.iloc[[1, 2]]
  49. result = g.nth(0, dropna="any")
  50. tm.assert_frame_equal(result, expected)
  51. @pytest.mark.parametrize("method", ["first", "last"])
  52. def test_first_last_with_na_object(method, nulls_fixture):
  53. # https://github.com/pandas-dev/pandas/issues/32123
  54. groups = DataFrame({"a": [1, 1, 2, 2], "b": [1, 2, 3, nulls_fixture]}).groupby("a")
  55. result = getattr(groups, method)()
  56. if method == "first":
  57. values = [1, 3]
  58. else:
  59. values = [2, 3]
  60. values = np.array(values, dtype=result["b"].dtype)
  61. idx = Index([1, 2], name="a")
  62. expected = DataFrame({"b": values}, index=idx)
  63. tm.assert_frame_equal(result, expected)
  64. @pytest.mark.parametrize("index", [0, -1])
  65. def test_nth_with_na_object(index, nulls_fixture):
  66. # https://github.com/pandas-dev/pandas/issues/32123
  67. df = DataFrame({"a": [1, 1, 2, 2], "b": [1, 2, 3, nulls_fixture]})
  68. groups = df.groupby("a")
  69. result = groups.nth(index)
  70. expected = df.iloc[[0, 2]] if index == 0 else df.iloc[[1, 3]]
  71. tm.assert_frame_equal(result, expected)
  72. @pytest.mark.parametrize("method", ["first", "last"])
  73. def test_first_last_with_None(method):
  74. # https://github.com/pandas-dev/pandas/issues/32800
  75. # None should be preserved as object dtype
  76. df = DataFrame.from_dict({"id": ["a"], "value": [None]})
  77. groups = df.groupby("id", as_index=False)
  78. result = getattr(groups, method)()
  79. tm.assert_frame_equal(result, df)
  80. @pytest.mark.parametrize("method", ["first", "last"])
  81. @pytest.mark.parametrize(
  82. "df, expected",
  83. [
  84. (
  85. DataFrame({"id": "a", "value": [None, "foo", np.nan]}),
  86. DataFrame({"value": ["foo"]}, index=Index(["a"], name="id")),
  87. ),
  88. (
  89. DataFrame({"id": "a", "value": [np.nan]}, dtype=object),
  90. DataFrame({"value": [None]}, index=Index(["a"], name="id")),
  91. ),
  92. ],
  93. )
  94. def test_first_last_with_None_expanded(method, df, expected):
  95. # GH 32800, 38286
  96. result = getattr(df.groupby("id"), method)()
  97. tm.assert_frame_equal(result, expected)
  98. def test_first_last_nth_dtypes(df_mixed_floats):
  99. df = df_mixed_floats.copy()
  100. df["E"] = True
  101. df["F"] = 1
  102. # tests for first / last / nth
  103. grouped = df.groupby("A")
  104. first = grouped.first()
  105. expected = df.loc[[1, 0], ["B", "C", "D", "E", "F"]]
  106. expected.index = Index(["bar", "foo"], name="A")
  107. expected = expected.sort_index()
  108. tm.assert_frame_equal(first, expected)
  109. last = grouped.last()
  110. expected = df.loc[[5, 7], ["B", "C", "D", "E", "F"]]
  111. expected.index = Index(["bar", "foo"], name="A")
  112. expected = expected.sort_index()
  113. tm.assert_frame_equal(last, expected)
  114. nth = grouped.nth(1)
  115. expected = df.iloc[[2, 3]]
  116. tm.assert_frame_equal(nth, expected)
  117. # GH 2763, first/last shifting dtypes
  118. idx = list(range(10))
  119. idx.append(9)
  120. s = Series(data=range(11), index=idx, name="IntCol")
  121. assert s.dtype == "int64"
  122. f = s.groupby(level=0).first()
  123. assert f.dtype == "int64"
  124. def test_first_last_nth_nan_dtype():
  125. # GH 33591
  126. df = DataFrame({"data": ["A"], "nans": Series([np.nan], dtype=object)})
  127. grouped = df.groupby("data")
  128. expected = df.set_index("data").nans
  129. tm.assert_series_equal(grouped.nans.first(), expected)
  130. tm.assert_series_equal(grouped.nans.last(), expected)
  131. expected = df.nans
  132. tm.assert_series_equal(grouped.nans.nth(-1), expected)
  133. tm.assert_series_equal(grouped.nans.nth(0), expected)
  134. def test_first_strings_timestamps():
  135. # GH 11244
  136. test = DataFrame(
  137. {
  138. Timestamp("2012-01-01 00:00:00"): ["a", "b"],
  139. Timestamp("2012-01-02 00:00:00"): ["c", "d"],
  140. "name": ["e", "e"],
  141. "aaaa": ["f", "g"],
  142. }
  143. )
  144. result = test.groupby("name").first()
  145. expected = DataFrame(
  146. [["a", "c", "f"]],
  147. columns=Index([Timestamp("2012-01-01"), Timestamp("2012-01-02"), "aaaa"]),
  148. index=Index(["e"], name="name"),
  149. )
  150. tm.assert_frame_equal(result, expected)
  151. def test_nth():
  152. df = DataFrame([[1, np.nan], [1, 4], [5, 6]], columns=["A", "B"])
  153. g = df.groupby("A")
  154. tm.assert_frame_equal(g.nth(0), df.iloc[[0, 2]])
  155. tm.assert_frame_equal(g.nth(1), df.iloc[[1]])
  156. tm.assert_frame_equal(g.nth(2), df.loc[[]])
  157. tm.assert_frame_equal(g.nth(-1), df.iloc[[1, 2]])
  158. tm.assert_frame_equal(g.nth(-2), df.iloc[[0]])
  159. tm.assert_frame_equal(g.nth(-3), df.loc[[]])
  160. tm.assert_series_equal(g.B.nth(0), df.B.iloc[[0, 2]])
  161. tm.assert_series_equal(g.B.nth(1), df.B.iloc[[1]])
  162. tm.assert_frame_equal(g[["B"]].nth(0), df[["B"]].iloc[[0, 2]])
  163. tm.assert_frame_equal(g.nth(0, dropna="any"), df.iloc[[1, 2]])
  164. tm.assert_frame_equal(g.nth(-1, dropna="any"), df.iloc[[1, 2]])
  165. tm.assert_frame_equal(g.nth(7, dropna="any"), df.iloc[:0])
  166. tm.assert_frame_equal(g.nth(2, dropna="any"), df.iloc[:0])
  167. # out of bounds, regression from 0.13.1
  168. # GH 6621
  169. df = DataFrame(
  170. {
  171. "color": {0: "green", 1: "green", 2: "red", 3: "red", 4: "red"},
  172. "food": {0: "ham", 1: "eggs", 2: "eggs", 3: "ham", 4: "pork"},
  173. "two": {
  174. 0: 1.5456590000000001,
  175. 1: -0.070345000000000005,
  176. 2: -2.4004539999999999,
  177. 3: 0.46206000000000003,
  178. 4: 0.52350799999999997,
  179. },
  180. "one": {
  181. 0: 0.56573799999999996,
  182. 1: -0.9742360000000001,
  183. 2: 1.033801,
  184. 3: -0.78543499999999999,
  185. 4: 0.70422799999999997,
  186. },
  187. }
  188. ).set_index(["color", "food"])
  189. result = df.groupby(level=0, as_index=False).nth(2)
  190. expected = df.iloc[[-1]]
  191. tm.assert_frame_equal(result, expected)
  192. result = df.groupby(level=0, as_index=False).nth(3)
  193. expected = df.loc[[]]
  194. tm.assert_frame_equal(result, expected)
  195. # GH 7559
  196. # from the vbench
  197. df = DataFrame(np.random.randint(1, 10, (100, 2)), dtype="int64")
  198. s = df[1]
  199. g = df[0]
  200. expected = s.groupby(g).first()
  201. expected2 = s.groupby(g).apply(lambda x: x.iloc[0])
  202. tm.assert_series_equal(expected2, expected, check_names=False)
  203. assert expected.name == 1
  204. assert expected2.name == 1
  205. # validate first
  206. v = s[g == 1].iloc[0]
  207. assert expected.iloc[0] == v
  208. assert expected2.iloc[0] == v
  209. with pytest.raises(ValueError, match="For a DataFrame"):
  210. s.groupby(g, sort=False).nth(0, dropna=True)
  211. # doc example
  212. df = DataFrame([[1, np.nan], [1, 4], [5, 6]], columns=["A", "B"])
  213. g = df.groupby("A")
  214. result = g.B.nth(0, dropna="all")
  215. expected = df.B.iloc[[1, 2]]
  216. tm.assert_series_equal(result, expected)
  217. # test multiple nth values
  218. df = DataFrame([[1, np.nan], [1, 3], [1, 4], [5, 6], [5, 7]], columns=["A", "B"])
  219. g = df.groupby("A")
  220. tm.assert_frame_equal(g.nth(0), df.iloc[[0, 3]])
  221. tm.assert_frame_equal(g.nth([0]), df.iloc[[0, 3]])
  222. tm.assert_frame_equal(g.nth([0, 1]), df.iloc[[0, 1, 3, 4]])
  223. tm.assert_frame_equal(g.nth([0, -1]), df.iloc[[0, 2, 3, 4]])
  224. tm.assert_frame_equal(g.nth([0, 1, 2]), df.iloc[[0, 1, 2, 3, 4]])
  225. tm.assert_frame_equal(g.nth([0, 1, -1]), df.iloc[[0, 1, 2, 3, 4]])
  226. tm.assert_frame_equal(g.nth([2]), df.iloc[[2]])
  227. tm.assert_frame_equal(g.nth([3, 4]), df.loc[[]])
  228. business_dates = pd.date_range(start="4/1/2014", end="6/30/2014", freq="B")
  229. df = DataFrame(1, index=business_dates, columns=["a", "b"])
  230. # get the first, fourth and last two business days for each month
  231. key = [df.index.year, df.index.month]
  232. result = df.groupby(key, as_index=False).nth([0, 3, -2, -1])
  233. expected_dates = pd.to_datetime(
  234. [
  235. "2014/4/1",
  236. "2014/4/4",
  237. "2014/4/29",
  238. "2014/4/30",
  239. "2014/5/1",
  240. "2014/5/6",
  241. "2014/5/29",
  242. "2014/5/30",
  243. "2014/6/2",
  244. "2014/6/5",
  245. "2014/6/27",
  246. "2014/6/30",
  247. ]
  248. )
  249. expected = DataFrame(1, columns=["a", "b"], index=expected_dates)
  250. tm.assert_frame_equal(result, expected)
  251. def test_nth_multi_grouper(three_group):
  252. # PR 9090, related to issue 8979
  253. # test nth on multiple groupers
  254. grouped = three_group.groupby(["A", "B"])
  255. result = grouped.nth(0)
  256. expected = three_group.iloc[[0, 3, 4, 7]]
  257. tm.assert_frame_equal(result, expected)
  258. @pytest.mark.parametrize(
  259. "data, expected_first, expected_last",
  260. [
  261. (
  262. {
  263. "id": ["A"],
  264. "time": Timestamp("2012-02-01 14:00:00", tz="US/Central"),
  265. "foo": [1],
  266. },
  267. {
  268. "id": ["A"],
  269. "time": Timestamp("2012-02-01 14:00:00", tz="US/Central"),
  270. "foo": [1],
  271. },
  272. {
  273. "id": ["A"],
  274. "time": Timestamp("2012-02-01 14:00:00", tz="US/Central"),
  275. "foo": [1],
  276. },
  277. ),
  278. (
  279. {
  280. "id": ["A", "B", "A"],
  281. "time": [
  282. Timestamp("2012-01-01 13:00:00", tz="America/New_York"),
  283. Timestamp("2012-02-01 14:00:00", tz="US/Central"),
  284. Timestamp("2012-03-01 12:00:00", tz="Europe/London"),
  285. ],
  286. "foo": [1, 2, 3],
  287. },
  288. {
  289. "id": ["A", "B"],
  290. "time": [
  291. Timestamp("2012-01-01 13:00:00", tz="America/New_York"),
  292. Timestamp("2012-02-01 14:00:00", tz="US/Central"),
  293. ],
  294. "foo": [1, 2],
  295. },
  296. {
  297. "id": ["A", "B"],
  298. "time": [
  299. Timestamp("2012-03-01 12:00:00", tz="Europe/London"),
  300. Timestamp("2012-02-01 14:00:00", tz="US/Central"),
  301. ],
  302. "foo": [3, 2],
  303. },
  304. ),
  305. ],
  306. )
  307. def test_first_last_tz(data, expected_first, expected_last):
  308. # GH15884
  309. # Test that the timezone is retained when calling first
  310. # or last on groupby with as_index=False
  311. df = DataFrame(data)
  312. result = df.groupby("id", as_index=False).first()
  313. expected = DataFrame(expected_first)
  314. cols = ["id", "time", "foo"]
  315. tm.assert_frame_equal(result[cols], expected[cols])
  316. result = df.groupby("id", as_index=False)["time"].first()
  317. tm.assert_frame_equal(result, expected[["id", "time"]])
  318. result = df.groupby("id", as_index=False).last()
  319. expected = DataFrame(expected_last)
  320. cols = ["id", "time", "foo"]
  321. tm.assert_frame_equal(result[cols], expected[cols])
  322. result = df.groupby("id", as_index=False)["time"].last()
  323. tm.assert_frame_equal(result, expected[["id", "time"]])
  324. @pytest.mark.parametrize(
  325. "method, ts, alpha",
  326. [
  327. ["first", Timestamp("2013-01-01", tz="US/Eastern"), "a"],
  328. ["last", Timestamp("2013-01-02", tz="US/Eastern"), "b"],
  329. ],
  330. )
  331. def test_first_last_tz_multi_column(method, ts, alpha):
  332. # GH 21603
  333. category_string = Series(list("abc")).astype("category")
  334. df = DataFrame(
  335. {
  336. "group": [1, 1, 2],
  337. "category_string": category_string,
  338. "datetimetz": pd.date_range("20130101", periods=3, tz="US/Eastern"),
  339. }
  340. )
  341. result = getattr(df.groupby("group"), method)()
  342. expected = DataFrame(
  343. {
  344. "category_string": pd.Categorical(
  345. [alpha, "c"], dtype=category_string.dtype
  346. ),
  347. "datetimetz": [ts, Timestamp("2013-01-03", tz="US/Eastern")],
  348. },
  349. index=Index([1, 2], name="group"),
  350. )
  351. tm.assert_frame_equal(result, expected)
  352. @pytest.mark.parametrize(
  353. "values",
  354. [
  355. pd.array([True, False], dtype="boolean"),
  356. pd.array([1, 2], dtype="Int64"),
  357. pd.to_datetime(["2020-01-01", "2020-02-01"]),
  358. pd.to_timedelta([1, 2], unit="D"),
  359. ],
  360. )
  361. @pytest.mark.parametrize("function", ["first", "last", "min", "max"])
  362. def test_first_last_extension_array_keeps_dtype(values, function):
  363. # https://github.com/pandas-dev/pandas/issues/33071
  364. # https://github.com/pandas-dev/pandas/issues/32194
  365. df = DataFrame({"a": [1, 2], "b": values})
  366. grouped = df.groupby("a")
  367. idx = Index([1, 2], name="a")
  368. expected_series = Series(values, name="b", index=idx)
  369. expected_frame = DataFrame({"b": values}, index=idx)
  370. result_series = getattr(grouped["b"], function)()
  371. tm.assert_series_equal(result_series, expected_series)
  372. result_frame = grouped.agg({"b": function})
  373. tm.assert_frame_equal(result_frame, expected_frame)
  374. def test_nth_multi_index_as_expected():
  375. # PR 9090, related to issue 8979
  376. # test nth on MultiIndex
  377. three_group = DataFrame(
  378. {
  379. "A": [
  380. "foo",
  381. "foo",
  382. "foo",
  383. "foo",
  384. "bar",
  385. "bar",
  386. "bar",
  387. "bar",
  388. "foo",
  389. "foo",
  390. "foo",
  391. ],
  392. "B": [
  393. "one",
  394. "one",
  395. "one",
  396. "two",
  397. "one",
  398. "one",
  399. "one",
  400. "two",
  401. "two",
  402. "two",
  403. "one",
  404. ],
  405. "C": [
  406. "dull",
  407. "dull",
  408. "shiny",
  409. "dull",
  410. "dull",
  411. "shiny",
  412. "shiny",
  413. "dull",
  414. "shiny",
  415. "shiny",
  416. "shiny",
  417. ],
  418. }
  419. )
  420. grouped = three_group.groupby(["A", "B"])
  421. result = grouped.nth(0)
  422. expected = three_group.iloc[[0, 3, 4, 7]]
  423. tm.assert_frame_equal(result, expected)
  424. @pytest.mark.parametrize(
  425. "op, n, expected_rows",
  426. [
  427. ("head", -1, [0]),
  428. ("head", 0, []),
  429. ("head", 1, [0, 2]),
  430. ("head", 7, [0, 1, 2]),
  431. ("tail", -1, [1]),
  432. ("tail", 0, []),
  433. ("tail", 1, [1, 2]),
  434. ("tail", 7, [0, 1, 2]),
  435. ],
  436. )
  437. @pytest.mark.parametrize("columns", [None, [], ["A"], ["B"], ["A", "B"]])
  438. @pytest.mark.parametrize("as_index", [True, False])
  439. def test_groupby_head_tail(op, n, expected_rows, columns, as_index):
  440. df = DataFrame([[1, 2], [1, 4], [5, 6]], columns=["A", "B"])
  441. g = df.groupby("A", as_index=as_index)
  442. expected = df.iloc[expected_rows]
  443. if columns is not None:
  444. g = g[columns]
  445. expected = expected[columns]
  446. result = getattr(g, op)(n)
  447. tm.assert_frame_equal(result, expected)
  448. @pytest.mark.parametrize(
  449. "op, n, expected_cols",
  450. [
  451. ("head", -1, [0]),
  452. ("head", 0, []),
  453. ("head", 1, [0, 2]),
  454. ("head", 7, [0, 1, 2]),
  455. ("tail", -1, [1]),
  456. ("tail", 0, []),
  457. ("tail", 1, [1, 2]),
  458. ("tail", 7, [0, 1, 2]),
  459. ],
  460. )
  461. def test_groupby_head_tail_axis_1(op, n, expected_cols):
  462. # GH 9772
  463. df = DataFrame(
  464. [[1, 2, 3], [1, 4, 5], [2, 6, 7], [3, 8, 9]], columns=["A", "B", "C"]
  465. )
  466. g = df.groupby([0, 0, 1], axis=1)
  467. expected = df.iloc[:, expected_cols]
  468. result = getattr(g, op)(n)
  469. tm.assert_frame_equal(result, expected)
  470. def test_group_selection_cache():
  471. # GH 12839 nth, head, and tail should return same result consistently
  472. df = DataFrame([[1, 2], [1, 4], [5, 6]], columns=["A", "B"])
  473. expected = df.iloc[[0, 2]]
  474. g = df.groupby("A")
  475. result1 = g.head(n=2)
  476. result2 = g.nth(0)
  477. tm.assert_frame_equal(result1, df)
  478. tm.assert_frame_equal(result2, expected)
  479. g = df.groupby("A")
  480. result1 = g.tail(n=2)
  481. result2 = g.nth(0)
  482. tm.assert_frame_equal(result1, df)
  483. tm.assert_frame_equal(result2, expected)
  484. g = df.groupby("A")
  485. result1 = g.nth(0)
  486. result2 = g.head(n=2)
  487. tm.assert_frame_equal(result1, expected)
  488. tm.assert_frame_equal(result2, df)
  489. g = df.groupby("A")
  490. result1 = g.nth(0)
  491. result2 = g.tail(n=2)
  492. tm.assert_frame_equal(result1, expected)
  493. tm.assert_frame_equal(result2, df)
  494. def test_nth_empty():
  495. # GH 16064
  496. df = DataFrame(index=[0], columns=["a", "b", "c"])
  497. result = df.groupby("a").nth(10)
  498. expected = df.iloc[:0]
  499. tm.assert_frame_equal(result, expected)
  500. result = df.groupby(["a", "b"]).nth(10)
  501. expected = df.iloc[:0]
  502. tm.assert_frame_equal(result, expected)
  503. def test_nth_column_order():
  504. # GH 20760
  505. # Check that nth preserves column order
  506. df = DataFrame(
  507. [[1, "b", 100], [1, "a", 50], [1, "a", np.nan], [2, "c", 200], [2, "d", 150]],
  508. columns=["A", "C", "B"],
  509. )
  510. result = df.groupby("A").nth(0)
  511. expected = df.iloc[[0, 3]]
  512. tm.assert_frame_equal(result, expected)
  513. result = df.groupby("A").nth(-1, dropna="any")
  514. expected = df.iloc[[1, 4]]
  515. tm.assert_frame_equal(result, expected)
  516. @pytest.mark.parametrize("dropna", [None, "any", "all"])
  517. def test_nth_nan_in_grouper(dropna):
  518. # GH 26011
  519. df = DataFrame(
  520. {
  521. "a": [np.nan, "a", np.nan, "b", np.nan],
  522. "b": [0, 2, 4, 6, 8],
  523. "c": [1, 3, 5, 7, 9],
  524. }
  525. )
  526. result = df.groupby("a").nth(0, dropna=dropna)
  527. expected = df.iloc[[1, 3]]
  528. tm.assert_frame_equal(result, expected)
  529. @pytest.mark.parametrize("dropna", [None, "any", "all"])
  530. def test_nth_nan_in_grouper_series(dropna):
  531. # GH 26454
  532. df = DataFrame(
  533. {
  534. "a": [np.nan, "a", np.nan, "b", np.nan],
  535. "b": [0, 2, 4, 6, 8],
  536. }
  537. )
  538. result = df.groupby("a")["b"].nth(0, dropna=dropna)
  539. expected = df["b"].iloc[[1, 3]]
  540. tm.assert_series_equal(result, expected)
  541. def test_first_categorical_and_datetime_data_nat():
  542. # GH 20520
  543. df = DataFrame(
  544. {
  545. "group": ["first", "first", "second", "third", "third"],
  546. "time": 5 * [np.datetime64("NaT")],
  547. "categories": Series(["a", "b", "c", "a", "b"], dtype="category"),
  548. }
  549. )
  550. result = df.groupby("group").first()
  551. expected = DataFrame(
  552. {
  553. "time": 3 * [np.datetime64("NaT")],
  554. "categories": Series(["a", "c", "a"]).astype(
  555. pd.CategoricalDtype(["a", "b", "c"])
  556. ),
  557. }
  558. )
  559. expected.index = Index(["first", "second", "third"], name="group")
  560. tm.assert_frame_equal(result, expected)
  561. def test_first_multi_key_groupby_categorical():
  562. # GH 22512
  563. df = DataFrame(
  564. {
  565. "A": [1, 1, 1, 2, 2],
  566. "B": [100, 100, 200, 100, 100],
  567. "C": ["apple", "orange", "mango", "mango", "orange"],
  568. "D": ["jupiter", "mercury", "mars", "venus", "venus"],
  569. }
  570. )
  571. df = df.astype({"D": "category"})
  572. result = df.groupby(by=["A", "B"]).first()
  573. expected = DataFrame(
  574. {
  575. "C": ["apple", "mango", "mango"],
  576. "D": Series(["jupiter", "mars", "venus"]).astype(
  577. pd.CategoricalDtype(["jupiter", "mars", "mercury", "venus"])
  578. ),
  579. }
  580. )
  581. expected.index = MultiIndex.from_tuples(
  582. [(1, 100), (1, 200), (2, 100)], names=["A", "B"]
  583. )
  584. tm.assert_frame_equal(result, expected)
  585. @pytest.mark.parametrize("method", ["first", "last", "nth"])
  586. def test_groupby_last_first_nth_with_none(method, nulls_fixture):
  587. # GH29645
  588. expected = Series(["y"])
  589. data = Series(
  590. [nulls_fixture, nulls_fixture, nulls_fixture, "y", nulls_fixture],
  591. index=[0, 0, 0, 0, 0],
  592. ).groupby(level=0)
  593. if method == "nth":
  594. result = getattr(data, method)(3)
  595. else:
  596. result = getattr(data, method)()
  597. tm.assert_series_equal(result, expected)
  598. @pytest.mark.parametrize(
  599. "arg, expected_rows",
  600. [
  601. [slice(None, 3, 2), [0, 1, 4, 5]],
  602. [slice(None, -2), [0, 2, 5]],
  603. [[slice(None, 2), slice(-2, None)], [0, 1, 2, 3, 4, 6, 7]],
  604. [[0, 1, slice(-2, None)], [0, 1, 2, 3, 4, 6, 7]],
  605. ],
  606. )
  607. def test_slice(slice_test_df, slice_test_grouped, arg, expected_rows):
  608. # Test slices GH #42947
  609. result = slice_test_grouped.nth[arg]
  610. equivalent = slice_test_grouped.nth(arg)
  611. expected = slice_test_df.iloc[expected_rows]
  612. tm.assert_frame_equal(result, expected)
  613. tm.assert_frame_equal(equivalent, expected)
  614. def test_nth_indexed(slice_test_df, slice_test_grouped):
  615. # Test index notation GH #44688
  616. result = slice_test_grouped.nth[0, 1, -2:]
  617. equivalent = slice_test_grouped.nth([0, 1, slice(-2, None)])
  618. expected = slice_test_df.iloc[[0, 1, 2, 3, 4, 6, 7]]
  619. tm.assert_frame_equal(result, expected)
  620. tm.assert_frame_equal(equivalent, expected)
  621. def test_invalid_argument(slice_test_grouped):
  622. # Test for error on invalid argument
  623. with pytest.raises(TypeError, match="Invalid index"):
  624. slice_test_grouped.nth(3.14)
  625. def test_negative_step(slice_test_grouped):
  626. # Test for error on negative slice step
  627. with pytest.raises(ValueError, match="Invalid step"):
  628. slice_test_grouped.nth(slice(None, None, -1))
  629. def test_np_ints(slice_test_df, slice_test_grouped):
  630. # Test np ints work
  631. result = slice_test_grouped.nth(np.array([0, 1]))
  632. expected = slice_test_df.iloc[[0, 1, 2, 3, 4]]
  633. tm.assert_frame_equal(result, expected)
  634. def test_groupby_nth_with_column_axis():
  635. # GH43926
  636. df = DataFrame(
  637. [
  638. [4, 5, 6],
  639. [8, 8, 7],
  640. ],
  641. index=["z", "y"],
  642. columns=["C", "B", "A"],
  643. )
  644. result = df.groupby(df.iloc[1], axis=1).nth(0)
  645. expected = df.iloc[:, [0, 2]]
  646. tm.assert_frame_equal(result, expected)
  647. @pytest.mark.parametrize(
  648. "start, stop, expected_values, expected_columns",
  649. [
  650. (None, None, [0, 1, 2, 3, 4], list("ABCDE")),
  651. (None, 1, [0, 3], list("AD")),
  652. (None, 9, [0, 1, 2, 3, 4], list("ABCDE")),
  653. (None, -1, [0, 1, 3], list("ABD")),
  654. (1, None, [1, 2, 4], list("BCE")),
  655. (1, -1, [1], list("B")),
  656. (-1, None, [2, 4], list("CE")),
  657. (-1, 2, [4], list("E")),
  658. ],
  659. )
  660. @pytest.mark.parametrize("method", ["call", "index"])
  661. def test_nth_slices_with_column_axis(
  662. start, stop, expected_values, expected_columns, method
  663. ):
  664. df = DataFrame([range(5)], columns=[list("ABCDE")])
  665. gb = df.groupby([5, 5, 5, 6, 6], axis=1)
  666. result = {
  667. "call": lambda start, stop: gb.nth(slice(start, stop)),
  668. "index": lambda start, stop: gb.nth[start:stop],
  669. }[method](start, stop)
  670. expected = DataFrame([expected_values], columns=[expected_columns])
  671. tm.assert_frame_equal(result, expected)
  672. @pytest.mark.filterwarnings(
  673. "ignore:invalid value encountered in remainder:RuntimeWarning"
  674. )
  675. def test_head_tail_dropna_true():
  676. # GH#45089
  677. df = DataFrame(
  678. [["a", "z"], ["b", np.nan], ["c", np.nan], ["c", np.nan]], columns=["X", "Y"]
  679. )
  680. expected = DataFrame([["a", "z"]], columns=["X", "Y"])
  681. result = df.groupby(["X", "Y"]).head(n=1)
  682. tm.assert_frame_equal(result, expected)
  683. result = df.groupby(["X", "Y"]).tail(n=1)
  684. tm.assert_frame_equal(result, expected)
  685. result = df.groupby(["X", "Y"]).nth(n=0)
  686. tm.assert_frame_equal(result, expected)
  687. def test_head_tail_dropna_false():
  688. # GH#45089
  689. df = DataFrame([["a", "z"], ["b", np.nan], ["c", np.nan]], columns=["X", "Y"])
  690. expected = DataFrame([["a", "z"], ["b", np.nan], ["c", np.nan]], columns=["X", "Y"])
  691. result = df.groupby(["X", "Y"], dropna=False).head(n=1)
  692. tm.assert_frame_equal(result, expected)
  693. result = df.groupby(["X", "Y"], dropna=False).tail(n=1)
  694. tm.assert_frame_equal(result, expected)
  695. result = df.groupby(["X", "Y"], dropna=False).nth(n=0)
  696. tm.assert_frame_equal(result, expected)