test_align.py 18 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484
  1. from datetime import timezone
  2. import numpy as np
  3. import pytest
  4. import pandas as pd
  5. from pandas import (
  6. DataFrame,
  7. Index,
  8. Series,
  9. date_range,
  10. )
  11. import pandas._testing as tm
  12. class TestDataFrameAlign:
  13. def test_align_asfreq_method_raises(self):
  14. df = DataFrame({"A": [1, np.nan, 2]})
  15. msg = "Invalid fill method"
  16. msg2 = "The 'method', 'limit', and 'fill_axis' keywords"
  17. with pytest.raises(ValueError, match=msg):
  18. with tm.assert_produces_warning(FutureWarning, match=msg2):
  19. df.align(df.iloc[::-1], method="asfreq")
  20. def test_frame_align_aware(self):
  21. idx1 = date_range("2001", periods=5, freq="h", tz="US/Eastern")
  22. idx2 = date_range("2001", periods=5, freq="2h", tz="US/Eastern")
  23. df1 = DataFrame(np.random.default_rng(2).standard_normal((len(idx1), 3)), idx1)
  24. df2 = DataFrame(np.random.default_rng(2).standard_normal((len(idx2), 3)), idx2)
  25. new1, new2 = df1.align(df2)
  26. assert df1.index.tz == new1.index.tz
  27. assert df2.index.tz == new2.index.tz
  28. # different timezones convert to UTC
  29. # frame with frame
  30. df1_central = df1.tz_convert("US/Central")
  31. new1, new2 = df1.align(df1_central)
  32. assert new1.index.tz is timezone.utc
  33. assert new2.index.tz is timezone.utc
  34. # frame with Series
  35. new1, new2 = df1.align(df1_central[0], axis=0)
  36. assert new1.index.tz is timezone.utc
  37. assert new2.index.tz is timezone.utc
  38. df1[0].align(df1_central, axis=0)
  39. assert new1.index.tz is timezone.utc
  40. assert new2.index.tz is timezone.utc
  41. def test_align_float(self, float_frame, using_copy_on_write):
  42. af, bf = float_frame.align(float_frame)
  43. assert af._mgr is not float_frame._mgr
  44. af, bf = float_frame.align(float_frame, copy=False)
  45. if not using_copy_on_write:
  46. assert af._mgr is float_frame._mgr
  47. else:
  48. assert af._mgr is not float_frame._mgr
  49. # axis = 0
  50. other = float_frame.iloc[:-5, :3]
  51. af, bf = float_frame.align(other, axis=0, fill_value=-1)
  52. tm.assert_index_equal(bf.columns, other.columns)
  53. # test fill value
  54. join_idx = float_frame.index.join(other.index)
  55. diff_a = float_frame.index.difference(join_idx)
  56. diff_a_vals = af.reindex(diff_a).values
  57. assert (diff_a_vals == -1).all()
  58. af, bf = float_frame.align(other, join="right", axis=0)
  59. tm.assert_index_equal(bf.columns, other.columns)
  60. tm.assert_index_equal(bf.index, other.index)
  61. tm.assert_index_equal(af.index, other.index)
  62. # axis = 1
  63. other = float_frame.iloc[:-5, :3].copy()
  64. af, bf = float_frame.align(other, axis=1)
  65. tm.assert_index_equal(bf.columns, float_frame.columns)
  66. tm.assert_index_equal(bf.index, other.index)
  67. # test fill value
  68. join_idx = float_frame.index.join(other.index)
  69. diff_a = float_frame.index.difference(join_idx)
  70. diff_a_vals = af.reindex(diff_a).values
  71. assert (diff_a_vals == -1).all()
  72. af, bf = float_frame.align(other, join="inner", axis=1)
  73. tm.assert_index_equal(bf.columns, other.columns)
  74. msg = (
  75. "The 'method', 'limit', and 'fill_axis' keywords in DataFrame.align "
  76. "are deprecated"
  77. )
  78. with tm.assert_produces_warning(FutureWarning, match=msg):
  79. af, bf = float_frame.align(other, join="inner", axis=1, method="pad")
  80. tm.assert_index_equal(bf.columns, other.columns)
  81. msg = (
  82. "The 'method', 'limit', and 'fill_axis' keywords in DataFrame.align "
  83. "are deprecated"
  84. )
  85. with tm.assert_produces_warning(FutureWarning, match=msg):
  86. af, bf = float_frame.align(
  87. other.iloc[:, 0], join="inner", axis=1, method=None, fill_value=None
  88. )
  89. tm.assert_index_equal(bf.index, Index([]).astype(bf.index.dtype))
  90. msg = (
  91. "The 'method', 'limit', and 'fill_axis' keywords in DataFrame.align "
  92. "are deprecated"
  93. )
  94. with tm.assert_produces_warning(FutureWarning, match=msg):
  95. af, bf = float_frame.align(
  96. other.iloc[:, 0], join="inner", axis=1, method=None, fill_value=0
  97. )
  98. tm.assert_index_equal(bf.index, Index([]).astype(bf.index.dtype))
  99. # Try to align DataFrame to Series along bad axis
  100. msg = "No axis named 2 for object type DataFrame"
  101. with pytest.raises(ValueError, match=msg):
  102. float_frame.align(af.iloc[0, :3], join="inner", axis=2)
  103. def test_align_frame_with_series(self, float_frame):
  104. # align dataframe to series with broadcast or not
  105. idx = float_frame.index
  106. s = Series(range(len(idx)), index=idx)
  107. left, right = float_frame.align(s, axis=0)
  108. tm.assert_index_equal(left.index, float_frame.index)
  109. tm.assert_index_equal(right.index, float_frame.index)
  110. assert isinstance(right, Series)
  111. msg = "The 'broadcast_axis' keyword in DataFrame.align is deprecated"
  112. with tm.assert_produces_warning(FutureWarning, match=msg):
  113. left, right = float_frame.align(s, broadcast_axis=1)
  114. tm.assert_index_equal(left.index, float_frame.index)
  115. expected = {c: s for c in float_frame.columns}
  116. expected = DataFrame(
  117. expected, index=float_frame.index, columns=float_frame.columns
  118. )
  119. tm.assert_frame_equal(right, expected)
  120. def test_align_series_condition(self):
  121. # see gh-9558
  122. df = DataFrame({"a": [1, 2, 3], "b": [4, 5, 6]})
  123. result = df[df["a"] == 2]
  124. expected = DataFrame([[2, 5]], index=[1], columns=["a", "b"])
  125. tm.assert_frame_equal(result, expected)
  126. result = df.where(df["a"] == 2, 0)
  127. expected = DataFrame({"a": [0, 2, 0], "b": [0, 5, 0]})
  128. tm.assert_frame_equal(result, expected)
  129. def test_align_int(self, int_frame):
  130. # test other non-float types
  131. other = DataFrame(index=range(5), columns=["A", "B", "C"])
  132. msg = (
  133. "The 'method', 'limit', and 'fill_axis' keywords in DataFrame.align "
  134. "are deprecated"
  135. )
  136. with tm.assert_produces_warning(FutureWarning, match=msg):
  137. af, bf = int_frame.align(other, join="inner", axis=1, method="pad")
  138. tm.assert_index_equal(bf.columns, other.columns)
  139. def test_align_mixed_type(self, float_string_frame):
  140. msg = (
  141. "The 'method', 'limit', and 'fill_axis' keywords in DataFrame.align "
  142. "are deprecated"
  143. )
  144. with tm.assert_produces_warning(FutureWarning, match=msg):
  145. af, bf = float_string_frame.align(
  146. float_string_frame, join="inner", axis=1, method="pad"
  147. )
  148. tm.assert_index_equal(bf.columns, float_string_frame.columns)
  149. def test_align_mixed_float(self, mixed_float_frame):
  150. # mixed floats/ints
  151. other = DataFrame(index=range(5), columns=["A", "B", "C"])
  152. msg = (
  153. "The 'method', 'limit', and 'fill_axis' keywords in DataFrame.align "
  154. "are deprecated"
  155. )
  156. with tm.assert_produces_warning(FutureWarning, match=msg):
  157. af, bf = mixed_float_frame.align(
  158. other.iloc[:, 0], join="inner", axis=1, method=None, fill_value=0
  159. )
  160. tm.assert_index_equal(bf.index, Index([]))
  161. def test_align_mixed_int(self, mixed_int_frame):
  162. other = DataFrame(index=range(5), columns=["A", "B", "C"])
  163. msg = (
  164. "The 'method', 'limit', and 'fill_axis' keywords in DataFrame.align "
  165. "are deprecated"
  166. )
  167. with tm.assert_produces_warning(FutureWarning, match=msg):
  168. af, bf = mixed_int_frame.align(
  169. other.iloc[:, 0], join="inner", axis=1, method=None, fill_value=0
  170. )
  171. tm.assert_index_equal(bf.index, Index([]))
  172. @pytest.mark.parametrize(
  173. "l_ordered,r_ordered,expected",
  174. [
  175. [True, True, pd.CategoricalIndex],
  176. [True, False, Index],
  177. [False, True, Index],
  178. [False, False, pd.CategoricalIndex],
  179. ],
  180. )
  181. def test_align_categorical(self, l_ordered, r_ordered, expected):
  182. # GH-28397
  183. df_1 = DataFrame(
  184. {
  185. "A": np.arange(6, dtype="int64"),
  186. "B": Series(list("aabbca")).astype(
  187. pd.CategoricalDtype(list("cab"), ordered=l_ordered)
  188. ),
  189. }
  190. ).set_index("B")
  191. df_2 = DataFrame(
  192. {
  193. "A": np.arange(5, dtype="int64"),
  194. "B": Series(list("babca")).astype(
  195. pd.CategoricalDtype(list("cab"), ordered=r_ordered)
  196. ),
  197. }
  198. ).set_index("B")
  199. aligned_1, aligned_2 = df_1.align(df_2)
  200. assert isinstance(aligned_1.index, expected)
  201. assert isinstance(aligned_2.index, expected)
  202. tm.assert_index_equal(aligned_1.index, aligned_2.index)
  203. def test_align_multiindex(self):
  204. # GH#10665
  205. # same test cases as test_align_multiindex in test_series.py
  206. midx = pd.MultiIndex.from_product(
  207. [range(2), range(3), range(2)], names=("a", "b", "c")
  208. )
  209. idx = Index(range(2), name="b")
  210. df1 = DataFrame(np.arange(12, dtype="int64"), index=midx)
  211. df2 = DataFrame(np.arange(2, dtype="int64"), index=idx)
  212. # these must be the same results (but flipped)
  213. res1l, res1r = df1.align(df2, join="left")
  214. res2l, res2r = df2.align(df1, join="right")
  215. expl = df1
  216. tm.assert_frame_equal(expl, res1l)
  217. tm.assert_frame_equal(expl, res2r)
  218. expr = DataFrame([0, 0, 1, 1, np.nan, np.nan] * 2, index=midx)
  219. tm.assert_frame_equal(expr, res1r)
  220. tm.assert_frame_equal(expr, res2l)
  221. res1l, res1r = df1.align(df2, join="right")
  222. res2l, res2r = df2.align(df1, join="left")
  223. exp_idx = pd.MultiIndex.from_product(
  224. [range(2), range(2), range(2)], names=("a", "b", "c")
  225. )
  226. expl = DataFrame([0, 1, 2, 3, 6, 7, 8, 9], index=exp_idx)
  227. tm.assert_frame_equal(expl, res1l)
  228. tm.assert_frame_equal(expl, res2r)
  229. expr = DataFrame([0, 0, 1, 1] * 2, index=exp_idx)
  230. tm.assert_frame_equal(expr, res1r)
  231. tm.assert_frame_equal(expr, res2l)
  232. def test_align_series_combinations(self):
  233. df = DataFrame({"a": [1, 3, 5], "b": [1, 3, 5]}, index=list("ACE"))
  234. s = Series([1, 2, 4], index=list("ABD"), name="x")
  235. # frame + series
  236. res1, res2 = df.align(s, axis=0)
  237. exp1 = DataFrame(
  238. {"a": [1, np.nan, 3, np.nan, 5], "b": [1, np.nan, 3, np.nan, 5]},
  239. index=list("ABCDE"),
  240. )
  241. exp2 = Series([1, 2, np.nan, 4, np.nan], index=list("ABCDE"), name="x")
  242. tm.assert_frame_equal(res1, exp1)
  243. tm.assert_series_equal(res2, exp2)
  244. # series + frame
  245. res1, res2 = s.align(df)
  246. tm.assert_series_equal(res1, exp2)
  247. tm.assert_frame_equal(res2, exp1)
  248. def test_multiindex_align_to_series_with_common_index_level(self):
  249. # GH-46001
  250. foo_index = Index([1, 2, 3], name="foo")
  251. bar_index = Index([1, 2], name="bar")
  252. series = Series([1, 2], index=bar_index, name="foo_series")
  253. df = DataFrame(
  254. {"col": np.arange(6)},
  255. index=pd.MultiIndex.from_product([foo_index, bar_index]),
  256. )
  257. expected_r = Series([1, 2] * 3, index=df.index, name="foo_series")
  258. result_l, result_r = df.align(series, axis=0)
  259. tm.assert_frame_equal(result_l, df)
  260. tm.assert_series_equal(result_r, expected_r)
  261. def test_multiindex_align_to_series_with_common_index_level_missing_in_left(self):
  262. # GH-46001
  263. foo_index = Index([1, 2, 3], name="foo")
  264. bar_index = Index([1, 2], name="bar")
  265. series = Series(
  266. [1, 2, 3, 4], index=Index([1, 2, 3, 4], name="bar"), name="foo_series"
  267. )
  268. df = DataFrame(
  269. {"col": np.arange(6)},
  270. index=pd.MultiIndex.from_product([foo_index, bar_index]),
  271. )
  272. expected_r = Series([1, 2] * 3, index=df.index, name="foo_series")
  273. result_l, result_r = df.align(series, axis=0)
  274. tm.assert_frame_equal(result_l, df)
  275. tm.assert_series_equal(result_r, expected_r)
  276. def test_multiindex_align_to_series_with_common_index_level_missing_in_right(self):
  277. # GH-46001
  278. foo_index = Index([1, 2, 3], name="foo")
  279. bar_index = Index([1, 2, 3, 4], name="bar")
  280. series = Series([1, 2], index=Index([1, 2], name="bar"), name="foo_series")
  281. df = DataFrame(
  282. {"col": np.arange(12)},
  283. index=pd.MultiIndex.from_product([foo_index, bar_index]),
  284. )
  285. expected_r = Series(
  286. [1, 2, np.nan, np.nan] * 3, index=df.index, name="foo_series"
  287. )
  288. result_l, result_r = df.align(series, axis=0)
  289. tm.assert_frame_equal(result_l, df)
  290. tm.assert_series_equal(result_r, expected_r)
  291. def test_multiindex_align_to_series_with_common_index_level_missing_in_both(self):
  292. # GH-46001
  293. foo_index = Index([1, 2, 3], name="foo")
  294. bar_index = Index([1, 3, 4], name="bar")
  295. series = Series(
  296. [1, 2, 3], index=Index([1, 2, 4], name="bar"), name="foo_series"
  297. )
  298. df = DataFrame(
  299. {"col": np.arange(9)},
  300. index=pd.MultiIndex.from_product([foo_index, bar_index]),
  301. )
  302. expected_r = Series([1, np.nan, 3] * 3, index=df.index, name="foo_series")
  303. result_l, result_r = df.align(series, axis=0)
  304. tm.assert_frame_equal(result_l, df)
  305. tm.assert_series_equal(result_r, expected_r)
  306. def test_multiindex_align_to_series_with_common_index_level_non_unique_cols(self):
  307. # GH-46001
  308. foo_index = Index([1, 2, 3], name="foo")
  309. bar_index = Index([1, 2], name="bar")
  310. series = Series([1, 2], index=bar_index, name="foo_series")
  311. df = DataFrame(
  312. np.arange(18).reshape(6, 3),
  313. index=pd.MultiIndex.from_product([foo_index, bar_index]),
  314. )
  315. df.columns = ["cfoo", "cbar", "cfoo"]
  316. expected = Series([1, 2] * 3, index=df.index, name="foo_series")
  317. result_left, result_right = df.align(series, axis=0)
  318. tm.assert_series_equal(result_right, expected)
  319. tm.assert_index_equal(result_left.columns, df.columns)
  320. def test_missing_axis_specification_exception(self):
  321. df = DataFrame(np.arange(50).reshape((10, 5)))
  322. series = Series(np.arange(5))
  323. with pytest.raises(ValueError, match=r"axis=0 or 1"):
  324. df.align(series)
  325. @pytest.mark.parametrize("method", ["pad", "bfill"])
  326. @pytest.mark.parametrize("axis", [0, 1, None])
  327. @pytest.mark.parametrize("fill_axis", [0, 1])
  328. @pytest.mark.parametrize("how", ["inner", "outer", "left", "right"])
  329. @pytest.mark.parametrize(
  330. "left_slice",
  331. [
  332. [slice(4), slice(10)],
  333. [slice(0), slice(0)],
  334. ],
  335. )
  336. @pytest.mark.parametrize(
  337. "right_slice",
  338. [
  339. [slice(2, None), slice(6, None)],
  340. [slice(0), slice(0)],
  341. ],
  342. )
  343. @pytest.mark.parametrize("limit", [1, None])
  344. def test_align_fill_method(
  345. self, how, method, axis, fill_axis, float_frame, left_slice, right_slice, limit
  346. ):
  347. frame = float_frame
  348. left = frame.iloc[left_slice[0], left_slice[1]]
  349. right = frame.iloc[right_slice[0], right_slice[1]]
  350. msg = (
  351. "The 'method', 'limit', and 'fill_axis' keywords in DataFrame.align "
  352. "are deprecated"
  353. )
  354. with tm.assert_produces_warning(FutureWarning, match=msg):
  355. aa, ab = left.align(
  356. right,
  357. axis=axis,
  358. join=how,
  359. method=method,
  360. limit=limit,
  361. fill_axis=fill_axis,
  362. )
  363. join_index, join_columns = None, None
  364. ea, eb = left, right
  365. if axis is None or axis == 0:
  366. join_index = left.index.join(right.index, how=how)
  367. ea = ea.reindex(index=join_index)
  368. eb = eb.reindex(index=join_index)
  369. if axis is None or axis == 1:
  370. join_columns = left.columns.join(right.columns, how=how)
  371. ea = ea.reindex(columns=join_columns)
  372. eb = eb.reindex(columns=join_columns)
  373. msg = "DataFrame.fillna with 'method' is deprecated"
  374. with tm.assert_produces_warning(FutureWarning, match=msg):
  375. ea = ea.fillna(axis=fill_axis, method=method, limit=limit)
  376. eb = eb.fillna(axis=fill_axis, method=method, limit=limit)
  377. tm.assert_frame_equal(aa, ea)
  378. tm.assert_frame_equal(ab, eb)
  379. def test_align_series_check_copy(self):
  380. # GH#
  381. df = DataFrame({0: [1, 2]})
  382. ser = Series([1], name=0)
  383. expected = ser.copy()
  384. result, other = df.align(ser, axis=1)
  385. ser.iloc[0] = 100
  386. tm.assert_series_equal(other, expected)
  387. def test_align_identical_different_object(self):
  388. # GH#51032
  389. df = DataFrame({"a": [1, 2]})
  390. ser = Series([3, 4])
  391. result, result2 = df.align(ser, axis=0)
  392. tm.assert_frame_equal(result, df)
  393. tm.assert_series_equal(result2, ser)
  394. assert df is not result
  395. assert ser is not result2
  396. def test_align_identical_different_object_columns(self):
  397. # GH#51032
  398. df = DataFrame({"a": [1, 2]})
  399. ser = Series([1], index=["a"])
  400. result, result2 = df.align(ser, axis=1)
  401. tm.assert_frame_equal(result, df)
  402. tm.assert_series_equal(result2, ser)
  403. assert df is not result
  404. assert ser is not result2