test_setops.py 20 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597
  1. from datetime import datetime
  2. import numpy as np
  3. import pytest
  4. import pytz
  5. import pandas.util._test_decorators as td
  6. import pandas as pd
  7. from pandas import (
  8. DataFrame,
  9. DatetimeIndex,
  10. Index,
  11. Series,
  12. bdate_range,
  13. date_range,
  14. )
  15. import pandas._testing as tm
  16. from pandas.tseries.offsets import (
  17. BMonthEnd,
  18. Minute,
  19. MonthEnd,
  20. )
  21. START, END = datetime(2009, 1, 1), datetime(2010, 1, 1)
  22. class TestDatetimeIndexSetOps:
  23. tz = [
  24. None,
  25. "UTC",
  26. "Asia/Tokyo",
  27. "US/Eastern",
  28. "dateutil/Asia/Singapore",
  29. "dateutil/US/Pacific",
  30. ]
  31. # TODO: moved from test_datetimelike; dedup with version below
  32. def test_union2(self, sort):
  33. everything = tm.makeDateIndex(10)
  34. first = everything[:5]
  35. second = everything[5:]
  36. union = first.union(second, sort=sort)
  37. tm.assert_index_equal(union, everything)
  38. @pytest.mark.parametrize("box", [np.array, Series, list])
  39. def test_union3(self, sort, box):
  40. everything = tm.makeDateIndex(10)
  41. first = everything[:5]
  42. second = everything[5:]
  43. # GH 10149 support listlike inputs other than Index objects
  44. expected = first.union(second, sort=sort)
  45. case = box(second.values)
  46. result = first.union(case, sort=sort)
  47. tm.assert_index_equal(result, expected)
  48. @pytest.mark.parametrize("tz", tz)
  49. def test_union(self, tz, sort):
  50. rng1 = date_range("1/1/2000", freq="D", periods=5, tz=tz)
  51. other1 = date_range("1/6/2000", freq="D", periods=5, tz=tz)
  52. expected1 = date_range("1/1/2000", freq="D", periods=10, tz=tz)
  53. expected1_notsorted = DatetimeIndex(list(other1) + list(rng1))
  54. rng2 = date_range("1/1/2000", freq="D", periods=5, tz=tz)
  55. other2 = date_range("1/4/2000", freq="D", periods=5, tz=tz)
  56. expected2 = date_range("1/1/2000", freq="D", periods=8, tz=tz)
  57. expected2_notsorted = DatetimeIndex(list(other2) + list(rng2[:3]))
  58. rng3 = date_range("1/1/2000", freq="D", periods=5, tz=tz)
  59. other3 = DatetimeIndex([], tz=tz)
  60. expected3 = date_range("1/1/2000", freq="D", periods=5, tz=tz)
  61. expected3_notsorted = rng3
  62. for rng, other, exp, exp_notsorted in [
  63. (rng1, other1, expected1, expected1_notsorted),
  64. (rng2, other2, expected2, expected2_notsorted),
  65. (rng3, other3, expected3, expected3_notsorted),
  66. ]:
  67. result_union = rng.union(other, sort=sort)
  68. tm.assert_index_equal(result_union, exp)
  69. result_union = other.union(rng, sort=sort)
  70. if sort is None:
  71. tm.assert_index_equal(result_union, exp)
  72. else:
  73. tm.assert_index_equal(result_union, exp_notsorted)
  74. def test_union_coverage(self, sort):
  75. idx = DatetimeIndex(["2000-01-03", "2000-01-01", "2000-01-02"])
  76. ordered = DatetimeIndex(idx.sort_values(), freq="infer")
  77. result = ordered.union(idx, sort=sort)
  78. tm.assert_index_equal(result, ordered)
  79. result = ordered[:0].union(ordered, sort=sort)
  80. tm.assert_index_equal(result, ordered)
  81. assert result.freq == ordered.freq
  82. def test_union_bug_1730(self, sort):
  83. rng_a = date_range("1/1/2012", periods=4, freq="3H")
  84. rng_b = date_range("1/1/2012", periods=4, freq="4H")
  85. result = rng_a.union(rng_b, sort=sort)
  86. exp = list(rng_a) + list(rng_b[1:])
  87. if sort is None:
  88. exp = DatetimeIndex(sorted(exp))
  89. else:
  90. exp = DatetimeIndex(exp)
  91. tm.assert_index_equal(result, exp)
  92. def test_union_bug_1745(self, sort):
  93. left = DatetimeIndex(["2012-05-11 15:19:49.695000"])
  94. right = DatetimeIndex(
  95. [
  96. "2012-05-29 13:04:21.322000",
  97. "2012-05-11 15:27:24.873000",
  98. "2012-05-11 15:31:05.350000",
  99. ]
  100. )
  101. result = left.union(right, sort=sort)
  102. exp = DatetimeIndex(
  103. [
  104. "2012-05-11 15:19:49.695000",
  105. "2012-05-29 13:04:21.322000",
  106. "2012-05-11 15:27:24.873000",
  107. "2012-05-11 15:31:05.350000",
  108. ]
  109. )
  110. if sort is None:
  111. exp = exp.sort_values()
  112. tm.assert_index_equal(result, exp)
  113. def test_union_bug_4564(self, sort):
  114. from pandas import DateOffset
  115. left = date_range("2013-01-01", "2013-02-01")
  116. right = left + DateOffset(minutes=15)
  117. result = left.union(right, sort=sort)
  118. exp = list(left) + list(right)
  119. if sort is None:
  120. exp = DatetimeIndex(sorted(exp))
  121. else:
  122. exp = DatetimeIndex(exp)
  123. tm.assert_index_equal(result, exp)
  124. def test_union_freq_both_none(self, sort):
  125. # GH11086
  126. expected = bdate_range("20150101", periods=10)
  127. expected._data.freq = None
  128. result = expected.union(expected, sort=sort)
  129. tm.assert_index_equal(result, expected)
  130. assert result.freq is None
  131. def test_union_freq_infer(self):
  132. # When taking the union of two DatetimeIndexes, we infer
  133. # a freq even if the arguments don't have freq. This matches
  134. # TimedeltaIndex behavior.
  135. dti = date_range("2016-01-01", periods=5)
  136. left = dti[[0, 1, 3, 4]]
  137. right = dti[[2, 3, 1]]
  138. assert left.freq is None
  139. assert right.freq is None
  140. result = left.union(right)
  141. tm.assert_index_equal(result, dti)
  142. assert result.freq == "D"
  143. def test_union_dataframe_index(self):
  144. rng1 = date_range("1/1/1999", "1/1/2012", freq="MS")
  145. s1 = Series(np.random.randn(len(rng1)), rng1)
  146. rng2 = date_range("1/1/1980", "12/1/2001", freq="MS")
  147. s2 = Series(np.random.randn(len(rng2)), rng2)
  148. df = DataFrame({"s1": s1, "s2": s2})
  149. exp = date_range("1/1/1980", "1/1/2012", freq="MS")
  150. tm.assert_index_equal(df.index, exp)
  151. def test_union_with_DatetimeIndex(self, sort):
  152. i1 = Index(np.arange(0, 20, 2, dtype=np.int64))
  153. i2 = date_range(start="2012-01-03 00:00:00", periods=10, freq="D")
  154. # Works
  155. i1.union(i2, sort=sort)
  156. # Fails with "AttributeError: can't set attribute"
  157. i2.union(i1, sort=sort)
  158. # TODO: moved from test_datetimelike; de-duplicate with version below
  159. def test_intersection2(self):
  160. first = tm.makeDateIndex(10)
  161. second = first[5:]
  162. intersect = first.intersection(second)
  163. assert tm.equalContents(intersect, second)
  164. # GH 10149
  165. cases = [klass(second.values) for klass in [np.array, Series, list]]
  166. for case in cases:
  167. result = first.intersection(case)
  168. assert tm.equalContents(result, second)
  169. third = Index(["a", "b", "c"])
  170. result = first.intersection(third)
  171. expected = Index([], dtype=object)
  172. tm.assert_index_equal(result, expected)
  173. @pytest.mark.parametrize(
  174. "tz", [None, "Asia/Tokyo", "US/Eastern", "dateutil/US/Pacific"]
  175. )
  176. def test_intersection(self, tz, sort):
  177. # GH 4690 (with tz)
  178. base = date_range("6/1/2000", "6/30/2000", freq="D", name="idx")
  179. # if target has the same name, it is preserved
  180. rng2 = date_range("5/15/2000", "6/20/2000", freq="D", name="idx")
  181. expected2 = date_range("6/1/2000", "6/20/2000", freq="D", name="idx")
  182. # if target name is different, it will be reset
  183. rng3 = date_range("5/15/2000", "6/20/2000", freq="D", name="other")
  184. expected3 = date_range("6/1/2000", "6/20/2000", freq="D", name=None)
  185. rng4 = date_range("7/1/2000", "7/31/2000", freq="D", name="idx")
  186. expected4 = DatetimeIndex([], freq="D", name="idx")
  187. for rng, expected in [
  188. (rng2, expected2),
  189. (rng3, expected3),
  190. (rng4, expected4),
  191. ]:
  192. result = base.intersection(rng)
  193. tm.assert_index_equal(result, expected)
  194. assert result.freq == expected.freq
  195. # non-monotonic
  196. base = DatetimeIndex(
  197. ["2011-01-05", "2011-01-04", "2011-01-02", "2011-01-03"], tz=tz, name="idx"
  198. )
  199. rng2 = DatetimeIndex(
  200. ["2011-01-04", "2011-01-02", "2011-02-02", "2011-02-03"], tz=tz, name="idx"
  201. )
  202. expected2 = DatetimeIndex(["2011-01-04", "2011-01-02"], tz=tz, name="idx")
  203. rng3 = DatetimeIndex(
  204. ["2011-01-04", "2011-01-02", "2011-02-02", "2011-02-03"],
  205. tz=tz,
  206. name="other",
  207. )
  208. expected3 = DatetimeIndex(["2011-01-04", "2011-01-02"], tz=tz, name=None)
  209. # GH 7880
  210. rng4 = date_range("7/1/2000", "7/31/2000", freq="D", tz=tz, name="idx")
  211. expected4 = DatetimeIndex([], tz=tz, name="idx")
  212. assert expected4.freq is None
  213. for rng, expected in [
  214. (rng2, expected2),
  215. (rng3, expected3),
  216. (rng4, expected4),
  217. ]:
  218. result = base.intersection(rng, sort=sort)
  219. if sort is None:
  220. expected = expected.sort_values()
  221. tm.assert_index_equal(result, expected)
  222. assert result.freq == expected.freq
  223. # parametrize over both anchored and non-anchored freqs, as they
  224. # have different code paths
  225. @pytest.mark.parametrize("freq", ["T", "B"])
  226. def test_intersection_empty(self, tz_aware_fixture, freq):
  227. # empty same freq GH2129
  228. tz = tz_aware_fixture
  229. rng = date_range("6/1/2000", "6/15/2000", freq=freq, tz=tz)
  230. result = rng[0:0].intersection(rng)
  231. assert len(result) == 0
  232. assert result.freq == rng.freq
  233. result = rng.intersection(rng[0:0])
  234. assert len(result) == 0
  235. assert result.freq == rng.freq
  236. # no overlap GH#33604
  237. check_freq = freq != "T" # We don't preserve freq on non-anchored offsets
  238. result = rng[:3].intersection(rng[-3:])
  239. tm.assert_index_equal(result, rng[:0])
  240. if check_freq:
  241. # We don't preserve freq on non-anchored offsets
  242. assert result.freq == rng.freq
  243. # swapped left and right
  244. result = rng[-3:].intersection(rng[:3])
  245. tm.assert_index_equal(result, rng[:0])
  246. if check_freq:
  247. # We don't preserve freq on non-anchored offsets
  248. assert result.freq == rng.freq
  249. def test_intersection_bug_1708(self):
  250. from pandas import DateOffset
  251. index_1 = date_range("1/1/2012", periods=4, freq="12H")
  252. index_2 = index_1 + DateOffset(hours=1)
  253. result = index_1.intersection(index_2)
  254. assert len(result) == 0
  255. @pytest.mark.parametrize("tz", tz)
  256. def test_difference(self, tz, sort):
  257. rng_dates = ["1/2/2000", "1/3/2000", "1/1/2000", "1/4/2000", "1/5/2000"]
  258. rng1 = DatetimeIndex(rng_dates, tz=tz)
  259. other1 = date_range("1/6/2000", freq="D", periods=5, tz=tz)
  260. expected1 = DatetimeIndex(rng_dates, tz=tz)
  261. rng2 = DatetimeIndex(rng_dates, tz=tz)
  262. other2 = date_range("1/4/2000", freq="D", periods=5, tz=tz)
  263. expected2 = DatetimeIndex(rng_dates[:3], tz=tz)
  264. rng3 = DatetimeIndex(rng_dates, tz=tz)
  265. other3 = DatetimeIndex([], tz=tz)
  266. expected3 = DatetimeIndex(rng_dates, tz=tz)
  267. for rng, other, expected in [
  268. (rng1, other1, expected1),
  269. (rng2, other2, expected2),
  270. (rng3, other3, expected3),
  271. ]:
  272. result_diff = rng.difference(other, sort)
  273. if sort is None and len(other):
  274. # We dont sort (yet?) when empty GH#24959
  275. expected = expected.sort_values()
  276. tm.assert_index_equal(result_diff, expected)
  277. def test_difference_freq(self, sort):
  278. # GH14323: difference of DatetimeIndex should not preserve frequency
  279. index = date_range("20160920", "20160925", freq="D")
  280. other = date_range("20160921", "20160924", freq="D")
  281. expected = DatetimeIndex(["20160920", "20160925"], freq=None)
  282. idx_diff = index.difference(other, sort)
  283. tm.assert_index_equal(idx_diff, expected)
  284. tm.assert_attr_equal("freq", idx_diff, expected)
  285. other = date_range("20160922", "20160925", freq="D")
  286. idx_diff = index.difference(other, sort)
  287. expected = DatetimeIndex(["20160920", "20160921"], freq=None)
  288. tm.assert_index_equal(idx_diff, expected)
  289. tm.assert_attr_equal("freq", idx_diff, expected)
  290. def test_datetimeindex_diff(self, sort):
  291. dti1 = date_range(freq="Q-JAN", start=datetime(1997, 12, 31), periods=100)
  292. dti2 = date_range(freq="Q-JAN", start=datetime(1997, 12, 31), periods=98)
  293. assert len(dti1.difference(dti2, sort)) == 2
  294. @pytest.mark.parametrize("tz", [None, "Asia/Tokyo", "US/Eastern"])
  295. def test_setops_preserve_freq(self, tz):
  296. rng = date_range("1/1/2000", "1/1/2002", name="idx", tz=tz)
  297. result = rng[:50].union(rng[50:100])
  298. assert result.name == rng.name
  299. assert result.freq == rng.freq
  300. assert result.tz == rng.tz
  301. result = rng[:50].union(rng[30:100])
  302. assert result.name == rng.name
  303. assert result.freq == rng.freq
  304. assert result.tz == rng.tz
  305. result = rng[:50].union(rng[60:100])
  306. assert result.name == rng.name
  307. assert result.freq is None
  308. assert result.tz == rng.tz
  309. result = rng[:50].intersection(rng[25:75])
  310. assert result.name == rng.name
  311. assert result.freqstr == "D"
  312. assert result.tz == rng.tz
  313. nofreq = DatetimeIndex(list(rng[25:75]), name="other")
  314. result = rng[:50].union(nofreq)
  315. assert result.name is None
  316. assert result.freq == rng.freq
  317. assert result.tz == rng.tz
  318. result = rng[:50].intersection(nofreq)
  319. assert result.name is None
  320. assert result.freq == rng.freq
  321. assert result.tz == rng.tz
  322. def test_intersection_non_tick_no_fastpath(self):
  323. # GH#42104
  324. dti = DatetimeIndex(
  325. [
  326. "2018-12-31",
  327. "2019-03-31",
  328. "2019-06-30",
  329. "2019-09-30",
  330. "2019-12-31",
  331. "2020-03-31",
  332. ],
  333. freq="Q-DEC",
  334. )
  335. result = dti[::2].intersection(dti[1::2])
  336. expected = dti[:0]
  337. tm.assert_index_equal(result, expected)
  338. class TestBusinessDatetimeIndex:
  339. def test_union(self, sort):
  340. rng = bdate_range(START, END)
  341. # overlapping
  342. left = rng[:10]
  343. right = rng[5:10]
  344. the_union = left.union(right, sort=sort)
  345. assert isinstance(the_union, DatetimeIndex)
  346. # non-overlapping, gap in middle
  347. left = rng[:5]
  348. right = rng[10:]
  349. the_union = left.union(right, sort=sort)
  350. assert isinstance(the_union, Index)
  351. # non-overlapping, no gap
  352. left = rng[:5]
  353. right = rng[5:10]
  354. the_union = left.union(right, sort=sort)
  355. assert isinstance(the_union, DatetimeIndex)
  356. # order does not matter
  357. if sort is None:
  358. tm.assert_index_equal(right.union(left, sort=sort), the_union)
  359. else:
  360. expected = DatetimeIndex(list(right) + list(left))
  361. tm.assert_index_equal(right.union(left, sort=sort), expected)
  362. # overlapping, but different offset
  363. rng = date_range(START, END, freq=BMonthEnd())
  364. the_union = rng.union(rng, sort=sort)
  365. assert isinstance(the_union, DatetimeIndex)
  366. def test_union_not_cacheable(self, sort):
  367. rng = date_range("1/1/2000", periods=50, freq=Minute())
  368. rng1 = rng[10:]
  369. rng2 = rng[:25]
  370. the_union = rng1.union(rng2, sort=sort)
  371. if sort is None:
  372. tm.assert_index_equal(the_union, rng)
  373. else:
  374. expected = DatetimeIndex(list(rng[10:]) + list(rng[:10]))
  375. tm.assert_index_equal(the_union, expected)
  376. rng1 = rng[10:]
  377. rng2 = rng[15:35]
  378. the_union = rng1.union(rng2, sort=sort)
  379. expected = rng[10:]
  380. tm.assert_index_equal(the_union, expected)
  381. def test_intersection(self):
  382. rng = date_range("1/1/2000", periods=50, freq=Minute())
  383. rng1 = rng[10:]
  384. rng2 = rng[:25]
  385. the_int = rng1.intersection(rng2)
  386. expected = rng[10:25]
  387. tm.assert_index_equal(the_int, expected)
  388. assert isinstance(the_int, DatetimeIndex)
  389. assert the_int.freq == rng.freq
  390. the_int = rng1.intersection(rng2.view(DatetimeIndex))
  391. tm.assert_index_equal(the_int, expected)
  392. # non-overlapping
  393. the_int = rng[:10].intersection(rng[10:])
  394. expected = DatetimeIndex([])
  395. tm.assert_index_equal(the_int, expected)
  396. def test_intersection_bug(self):
  397. # GH #771
  398. a = bdate_range("11/30/2011", "12/31/2011")
  399. b = bdate_range("12/10/2011", "12/20/2011")
  400. result = a.intersection(b)
  401. tm.assert_index_equal(result, b)
  402. assert result.freq == b.freq
  403. def test_intersection_list(self):
  404. # GH#35876
  405. # values is not an Index -> no name -> retain "a"
  406. values = [pd.Timestamp("2020-01-01"), pd.Timestamp("2020-02-01")]
  407. idx = DatetimeIndex(values, name="a")
  408. res = idx.intersection(values)
  409. tm.assert_index_equal(res, idx)
  410. def test_month_range_union_tz_pytz(self, sort):
  411. from pytz import timezone
  412. tz = timezone("US/Eastern")
  413. early_start = datetime(2011, 1, 1)
  414. early_end = datetime(2011, 3, 1)
  415. late_start = datetime(2011, 3, 1)
  416. late_end = datetime(2011, 5, 1)
  417. early_dr = date_range(start=early_start, end=early_end, tz=tz, freq=MonthEnd())
  418. late_dr = date_range(start=late_start, end=late_end, tz=tz, freq=MonthEnd())
  419. early_dr.union(late_dr, sort=sort)
  420. @td.skip_if_windows
  421. def test_month_range_union_tz_dateutil(self, sort):
  422. from pandas._libs.tslibs.timezones import dateutil_gettz
  423. tz = dateutil_gettz("US/Eastern")
  424. early_start = datetime(2011, 1, 1)
  425. early_end = datetime(2011, 3, 1)
  426. late_start = datetime(2011, 3, 1)
  427. late_end = datetime(2011, 5, 1)
  428. early_dr = date_range(start=early_start, end=early_end, tz=tz, freq=MonthEnd())
  429. late_dr = date_range(start=late_start, end=late_end, tz=tz, freq=MonthEnd())
  430. early_dr.union(late_dr, sort=sort)
  431. @pytest.mark.parametrize("sort", [False, None])
  432. def test_intersection_duplicates(self, sort):
  433. # GH#38196
  434. idx1 = Index(
  435. [
  436. pd.Timestamp("2019-12-13"),
  437. pd.Timestamp("2019-12-12"),
  438. pd.Timestamp("2019-12-12"),
  439. ]
  440. )
  441. result = idx1.intersection(idx1, sort=sort)
  442. expected = Index([pd.Timestamp("2019-12-13"), pd.Timestamp("2019-12-12")])
  443. tm.assert_index_equal(result, expected)
  444. class TestCustomDatetimeIndex:
  445. def test_union(self, sort):
  446. # overlapping
  447. rng = bdate_range(START, END, freq="C")
  448. left = rng[:10]
  449. right = rng[5:10]
  450. the_union = left.union(right, sort=sort)
  451. assert isinstance(the_union, DatetimeIndex)
  452. # non-overlapping, gap in middle
  453. left = rng[:5]
  454. right = rng[10:]
  455. the_union = left.union(right, sort)
  456. assert isinstance(the_union, Index)
  457. # non-overlapping, no gap
  458. left = rng[:5]
  459. right = rng[5:10]
  460. the_union = left.union(right, sort=sort)
  461. assert isinstance(the_union, DatetimeIndex)
  462. # order does not matter
  463. if sort is None:
  464. tm.assert_index_equal(right.union(left, sort=sort), the_union)
  465. # overlapping, but different offset
  466. rng = date_range(START, END, freq=BMonthEnd())
  467. the_union = rng.union(rng, sort=sort)
  468. assert isinstance(the_union, DatetimeIndex)
  469. def test_intersection_bug(self):
  470. # GH #771
  471. a = bdate_range("11/30/2011", "12/31/2011", freq="C")
  472. b = bdate_range("12/10/2011", "12/20/2011", freq="C")
  473. result = a.intersection(b)
  474. tm.assert_index_equal(result, b)
  475. assert result.freq == b.freq
  476. @pytest.mark.parametrize(
  477. "tz", [None, "UTC", "Europe/Berlin", pytz.FixedOffset(-60)]
  478. )
  479. def test_intersection_dst_transition(self, tz):
  480. # GH 46702: Europe/Berlin has DST transition
  481. idx1 = date_range("2020-03-27", periods=5, freq="D", tz=tz)
  482. idx2 = date_range("2020-03-30", periods=5, freq="D", tz=tz)
  483. result = idx1.intersection(idx2)
  484. expected = date_range("2020-03-30", periods=2, freq="D", tz=tz)
  485. tm.assert_index_equal(result, expected)