test_setops.py 33 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909910911912913914915916917918919920921922923924925926927928929930931932933934935936937938939940941942943944945946947948949950951952953954955956957958959960961962963964965966967968969970971972973
  1. """
  2. The tests in this package are to ensure the proper resultant dtypes of
  3. set operations.
  4. """
  5. from datetime import datetime
  6. import operator
  7. import numpy as np
  8. import pytest
  9. from pandas._libs import lib
  10. from pandas.core.dtypes.cast import find_common_type
  11. from pandas import (
  12. CategoricalDtype,
  13. CategoricalIndex,
  14. DatetimeTZDtype,
  15. Index,
  16. MultiIndex,
  17. PeriodDtype,
  18. RangeIndex,
  19. Series,
  20. Timestamp,
  21. )
  22. import pandas._testing as tm
  23. from pandas.api.types import (
  24. is_signed_integer_dtype,
  25. pandas_dtype,
  26. )
  27. def equal_contents(arr1, arr2) -> bool:
  28. """
  29. Checks if the set of unique elements of arr1 and arr2 are equivalent.
  30. """
  31. return frozenset(arr1) == frozenset(arr2)
  32. @pytest.fixture(
  33. params=tm.ALL_REAL_NUMPY_DTYPES
  34. + [
  35. "object",
  36. "category",
  37. "datetime64[ns]",
  38. "timedelta64[ns]",
  39. ]
  40. )
  41. def any_dtype_for_small_pos_integer_indexes(request):
  42. """
  43. Dtypes that can be given to an Index with small positive integers.
  44. This means that for any dtype `x` in the params list, `Index([1, 2, 3], dtype=x)` is
  45. valid and gives the correct Index (sub-)class.
  46. """
  47. return request.param
  48. def test_union_same_types(index):
  49. # Union with a non-unique, non-monotonic index raises error
  50. # Only needed for bool index factory
  51. idx1 = index.sort_values()
  52. idx2 = index.sort_values()
  53. assert idx1.union(idx2).dtype == idx1.dtype
  54. def test_union_different_types(index_flat, index_flat2, request, using_infer_string):
  55. # This test only considers combinations of indices
  56. # GH 23525
  57. idx1 = index_flat
  58. idx2 = index_flat2
  59. if (
  60. not idx1.is_unique
  61. and not idx2.is_unique
  62. and idx1.dtype.kind == "i"
  63. and idx2.dtype.kind == "b"
  64. ) or (
  65. not idx2.is_unique
  66. and not idx1.is_unique
  67. and idx2.dtype.kind == "i"
  68. and idx1.dtype.kind == "b"
  69. ):
  70. # Each condition had idx[1|2].is_monotonic_decreasing
  71. # but failed when e.g.
  72. # idx1 = Index(
  73. # [True, True, True, True, True, True, True, True, False, False], dtype='bool'
  74. # )
  75. # idx2 = Index([0, 0, 1, 1, 2, 2], dtype='int64')
  76. mark = pytest.mark.xfail(
  77. reason="GH#44000 True==1", raises=ValueError, strict=False
  78. )
  79. request.applymarker(mark)
  80. common_dtype = find_common_type([idx1.dtype, idx2.dtype])
  81. if using_infer_string:
  82. if len(idx1) == 0 and (idx1.dtype.kind == "O" or isinstance(idx1, RangeIndex)):
  83. common_dtype = idx2.dtype
  84. elif len(idx2) == 0 and (
  85. idx2.dtype.kind == "O" or isinstance(idx2, RangeIndex)
  86. ):
  87. common_dtype = idx1.dtype
  88. warn = None
  89. msg = "'<' not supported between"
  90. if not len(idx1) or not len(idx2):
  91. pass
  92. elif (idx1.dtype.kind == "c" and (not lib.is_np_dtype(idx2.dtype, "iufc"))) or (
  93. idx2.dtype.kind == "c" and (not lib.is_np_dtype(idx1.dtype, "iufc"))
  94. ):
  95. # complex objects non-sortable
  96. warn = RuntimeWarning
  97. elif (
  98. isinstance(idx1.dtype, PeriodDtype) and isinstance(idx2.dtype, CategoricalDtype)
  99. ) or (
  100. isinstance(idx2.dtype, PeriodDtype) and isinstance(idx1.dtype, CategoricalDtype)
  101. ):
  102. warn = FutureWarning
  103. msg = r"PeriodDtype\[B\] is deprecated"
  104. mark = pytest.mark.xfail(
  105. reason="Warning not produced on all builds",
  106. raises=AssertionError,
  107. strict=False,
  108. )
  109. request.applymarker(mark)
  110. any_uint64 = np.uint64 in (idx1.dtype, idx2.dtype)
  111. idx1_signed = is_signed_integer_dtype(idx1.dtype)
  112. idx2_signed = is_signed_integer_dtype(idx2.dtype)
  113. # Union with a non-unique, non-monotonic index raises error
  114. # This applies to the boolean index
  115. idx1 = idx1.sort_values()
  116. idx2 = idx2.sort_values()
  117. with tm.assert_produces_warning(warn, match=msg):
  118. res1 = idx1.union(idx2)
  119. res2 = idx2.union(idx1)
  120. if any_uint64 and (idx1_signed or idx2_signed):
  121. assert res1.dtype == np.dtype("O")
  122. assert res2.dtype == np.dtype("O")
  123. else:
  124. assert res1.dtype == common_dtype
  125. assert res2.dtype == common_dtype
  126. @pytest.mark.parametrize(
  127. "idx1,idx2",
  128. [
  129. (Index(np.arange(5), dtype=np.int64), RangeIndex(5)),
  130. (Index(np.arange(5), dtype=np.float64), Index(np.arange(5), dtype=np.int64)),
  131. (Index(np.arange(5), dtype=np.float64), RangeIndex(5)),
  132. (Index(np.arange(5), dtype=np.float64), Index(np.arange(5), dtype=np.uint64)),
  133. ],
  134. )
  135. def test_compatible_inconsistent_pairs(idx1, idx2):
  136. # GH 23525
  137. res1 = idx1.union(idx2)
  138. res2 = idx2.union(idx1)
  139. assert res1.dtype in (idx1.dtype, idx2.dtype)
  140. assert res2.dtype in (idx1.dtype, idx2.dtype)
  141. @pytest.mark.parametrize(
  142. "left, right, expected",
  143. [
  144. ("int64", "int64", "int64"),
  145. ("int64", "uint64", "object"),
  146. ("int64", "float64", "float64"),
  147. ("uint64", "float64", "float64"),
  148. ("uint64", "uint64", "uint64"),
  149. ("float64", "float64", "float64"),
  150. ("datetime64[ns]", "int64", "object"),
  151. ("datetime64[ns]", "uint64", "object"),
  152. ("datetime64[ns]", "float64", "object"),
  153. ("datetime64[ns, CET]", "int64", "object"),
  154. ("datetime64[ns, CET]", "uint64", "object"),
  155. ("datetime64[ns, CET]", "float64", "object"),
  156. ("Period[D]", "int64", "object"),
  157. ("Period[D]", "uint64", "object"),
  158. ("Period[D]", "float64", "object"),
  159. ],
  160. )
  161. @pytest.mark.parametrize("names", [("foo", "foo", "foo"), ("foo", "bar", None)])
  162. def test_union_dtypes(left, right, expected, names):
  163. left = pandas_dtype(left)
  164. right = pandas_dtype(right)
  165. a = Index([], dtype=left, name=names[0])
  166. b = Index([], dtype=right, name=names[1])
  167. result = a.union(b)
  168. assert result.dtype == expected
  169. assert result.name == names[2]
  170. # Testing name retention
  171. # TODO: pin down desired dtype; do we want it to be commutative?
  172. result = a.intersection(b)
  173. assert result.name == names[2]
  174. @pytest.mark.parametrize("values", [[1, 2, 2, 3], [3, 3]])
  175. def test_intersection_duplicates(values):
  176. # GH#31326
  177. a = Index(values)
  178. b = Index([3, 3])
  179. result = a.intersection(b)
  180. expected = Index([3])
  181. tm.assert_index_equal(result, expected)
  182. class TestSetOps:
  183. # Set operation tests shared by all indexes in the `index` fixture
  184. @pytest.mark.parametrize("case", [0.5, "xxx"])
  185. @pytest.mark.parametrize(
  186. "method", ["intersection", "union", "difference", "symmetric_difference"]
  187. )
  188. def test_set_ops_error_cases(self, case, method, index):
  189. # non-iterable input
  190. msg = "Input must be Index or array-like"
  191. with pytest.raises(TypeError, match=msg):
  192. getattr(index, method)(case)
  193. @pytest.mark.filterwarnings(r"ignore:PeriodDtype\[B\] is deprecated:FutureWarning")
  194. def test_intersection_base(self, index):
  195. if isinstance(index, CategoricalIndex):
  196. pytest.skip(f"Not relevant for {type(index).__name__}")
  197. first = index[:5].unique()
  198. second = index[:3].unique()
  199. intersect = first.intersection(second)
  200. tm.assert_index_equal(intersect, second)
  201. if isinstance(index.dtype, DatetimeTZDtype):
  202. # The second.values below will drop tz, so the rest of this test
  203. # is not applicable.
  204. return
  205. # GH#10149
  206. cases = [second.to_numpy(), second.to_series(), second.to_list()]
  207. for case in cases:
  208. result = first.intersection(case)
  209. assert equal_contents(result, second)
  210. if isinstance(index, MultiIndex):
  211. msg = "other must be a MultiIndex or a list of tuples"
  212. with pytest.raises(TypeError, match=msg):
  213. first.intersection([1, 2, 3])
  214. @pytest.mark.filterwarnings(r"ignore:PeriodDtype\[B\] is deprecated:FutureWarning")
  215. def test_union_base(self, index):
  216. index = index.unique()
  217. first = index[3:]
  218. second = index[:5]
  219. everything = index
  220. union = first.union(second)
  221. tm.assert_index_equal(union.sort_values(), everything.sort_values())
  222. if isinstance(index.dtype, DatetimeTZDtype):
  223. # The second.values below will drop tz, so the rest of this test
  224. # is not applicable.
  225. return
  226. # GH#10149
  227. cases = [second.to_numpy(), second.to_series(), second.to_list()]
  228. for case in cases:
  229. result = first.union(case)
  230. assert equal_contents(result, everything)
  231. if isinstance(index, MultiIndex):
  232. msg = "other must be a MultiIndex or a list of tuples"
  233. with pytest.raises(TypeError, match=msg):
  234. first.union([1, 2, 3])
  235. @pytest.mark.filterwarnings(r"ignore:PeriodDtype\[B\] is deprecated:FutureWarning")
  236. def test_difference_base(self, sort, index):
  237. first = index[2:]
  238. second = index[:4]
  239. if index.inferred_type == "boolean":
  240. # i think (TODO: be sure) there assumptions baked in about
  241. # the index fixture that don't hold here?
  242. answer = set(first).difference(set(second))
  243. elif isinstance(index, CategoricalIndex):
  244. answer = []
  245. else:
  246. answer = index[4:]
  247. result = first.difference(second, sort)
  248. assert equal_contents(result, answer)
  249. # GH#10149
  250. cases = [second.to_numpy(), second.to_series(), second.to_list()]
  251. for case in cases:
  252. result = first.difference(case, sort)
  253. assert equal_contents(result, answer)
  254. if isinstance(index, MultiIndex):
  255. msg = "other must be a MultiIndex or a list of tuples"
  256. with pytest.raises(TypeError, match=msg):
  257. first.difference([1, 2, 3], sort)
  258. @pytest.mark.filterwarnings(r"ignore:PeriodDtype\[B\] is deprecated:FutureWarning")
  259. def test_symmetric_difference(self, index, using_infer_string, request):
  260. if (
  261. using_infer_string
  262. and index.dtype == "object"
  263. and index.inferred_type == "string"
  264. ):
  265. request.applymarker(pytest.mark.xfail(reason="TODO: infer_string"))
  266. if isinstance(index, CategoricalIndex):
  267. pytest.skip(f"Not relevant for {type(index).__name__}")
  268. if len(index) < 2:
  269. pytest.skip("Too few values for test")
  270. if index[0] in index[1:] or index[-1] in index[:-1]:
  271. # index fixture has e.g. an index of bools that does not satisfy this,
  272. # another with [0, 0, 1, 1, 2, 2]
  273. pytest.skip("Index values no not satisfy test condition.")
  274. first = index[1:]
  275. second = index[:-1]
  276. answer = index[[0, -1]]
  277. result = first.symmetric_difference(second)
  278. tm.assert_index_equal(result.sort_values(), answer.sort_values())
  279. # GH#10149
  280. cases = [second.to_numpy(), second.to_series(), second.to_list()]
  281. for case in cases:
  282. result = first.symmetric_difference(case)
  283. assert equal_contents(result, answer)
  284. if isinstance(index, MultiIndex):
  285. msg = "other must be a MultiIndex or a list of tuples"
  286. with pytest.raises(TypeError, match=msg):
  287. first.symmetric_difference([1, 2, 3])
  288. @pytest.mark.parametrize(
  289. "fname, sname, expected_name",
  290. [
  291. ("A", "A", "A"),
  292. ("A", "B", None),
  293. ("A", None, None),
  294. (None, "B", None),
  295. (None, None, None),
  296. ],
  297. )
  298. def test_corner_union(self, index_flat, fname, sname, expected_name):
  299. # GH#9943, GH#9862
  300. # Test unions with various name combinations
  301. # Do not test MultiIndex or repeats
  302. if not index_flat.is_unique:
  303. index = index_flat.unique()
  304. else:
  305. index = index_flat
  306. # Test copy.union(copy)
  307. first = index.copy().set_names(fname)
  308. second = index.copy().set_names(sname)
  309. union = first.union(second)
  310. expected = index.copy().set_names(expected_name)
  311. tm.assert_index_equal(union, expected)
  312. # Test copy.union(empty)
  313. first = index.copy().set_names(fname)
  314. second = index.drop(index).set_names(sname)
  315. union = first.union(second)
  316. expected = index.copy().set_names(expected_name)
  317. tm.assert_index_equal(union, expected)
  318. # Test empty.union(copy)
  319. first = index.drop(index).set_names(fname)
  320. second = index.copy().set_names(sname)
  321. union = first.union(second)
  322. expected = index.copy().set_names(expected_name)
  323. tm.assert_index_equal(union, expected)
  324. # Test empty.union(empty)
  325. first = index.drop(index).set_names(fname)
  326. second = index.drop(index).set_names(sname)
  327. union = first.union(second)
  328. expected = index.drop(index).set_names(expected_name)
  329. tm.assert_index_equal(union, expected)
  330. @pytest.mark.parametrize(
  331. "fname, sname, expected_name",
  332. [
  333. ("A", "A", "A"),
  334. ("A", "B", None),
  335. ("A", None, None),
  336. (None, "B", None),
  337. (None, None, None),
  338. ],
  339. )
  340. def test_union_unequal(self, index_flat, fname, sname, expected_name):
  341. if not index_flat.is_unique:
  342. index = index_flat.unique()
  343. else:
  344. index = index_flat
  345. # test copy.union(subset) - need sort for unicode and string
  346. first = index.copy().set_names(fname)
  347. second = index[1:].set_names(sname)
  348. union = first.union(second).sort_values()
  349. expected = index.set_names(expected_name).sort_values()
  350. tm.assert_index_equal(union, expected)
  351. @pytest.mark.parametrize(
  352. "fname, sname, expected_name",
  353. [
  354. ("A", "A", "A"),
  355. ("A", "B", None),
  356. ("A", None, None),
  357. (None, "B", None),
  358. (None, None, None),
  359. ],
  360. )
  361. def test_corner_intersect(self, index_flat, fname, sname, expected_name):
  362. # GH#35847
  363. # Test intersections with various name combinations
  364. if not index_flat.is_unique:
  365. index = index_flat.unique()
  366. else:
  367. index = index_flat
  368. # Test copy.intersection(copy)
  369. first = index.copy().set_names(fname)
  370. second = index.copy().set_names(sname)
  371. intersect = first.intersection(second)
  372. expected = index.copy().set_names(expected_name)
  373. tm.assert_index_equal(intersect, expected)
  374. # Test copy.intersection(empty)
  375. first = index.copy().set_names(fname)
  376. second = index.drop(index).set_names(sname)
  377. intersect = first.intersection(second)
  378. expected = index.drop(index).set_names(expected_name)
  379. tm.assert_index_equal(intersect, expected)
  380. # Test empty.intersection(copy)
  381. first = index.drop(index).set_names(fname)
  382. second = index.copy().set_names(sname)
  383. intersect = first.intersection(second)
  384. expected = index.drop(index).set_names(expected_name)
  385. tm.assert_index_equal(intersect, expected)
  386. # Test empty.intersection(empty)
  387. first = index.drop(index).set_names(fname)
  388. second = index.drop(index).set_names(sname)
  389. intersect = first.intersection(second)
  390. expected = index.drop(index).set_names(expected_name)
  391. tm.assert_index_equal(intersect, expected)
  392. @pytest.mark.parametrize(
  393. "fname, sname, expected_name",
  394. [
  395. ("A", "A", "A"),
  396. ("A", "B", None),
  397. ("A", None, None),
  398. (None, "B", None),
  399. (None, None, None),
  400. ],
  401. )
  402. def test_intersect_unequal(self, index_flat, fname, sname, expected_name):
  403. if not index_flat.is_unique:
  404. index = index_flat.unique()
  405. else:
  406. index = index_flat
  407. # test copy.intersection(subset) - need sort for unicode and string
  408. first = index.copy().set_names(fname)
  409. second = index[1:].set_names(sname)
  410. intersect = first.intersection(second).sort_values()
  411. expected = index[1:].set_names(expected_name).sort_values()
  412. tm.assert_index_equal(intersect, expected)
  413. @pytest.mark.filterwarnings(r"ignore:PeriodDtype\[B\] is deprecated:FutureWarning")
  414. def test_intersection_name_retention_with_nameless(self, index):
  415. if isinstance(index, MultiIndex):
  416. index = index.rename(list(range(index.nlevels)))
  417. else:
  418. index = index.rename("foo")
  419. other = np.asarray(index)
  420. result = index.intersection(other)
  421. assert result.name == index.name
  422. # empty other, same dtype
  423. result = index.intersection(other[:0])
  424. assert result.name == index.name
  425. # empty `self`
  426. result = index[:0].intersection(other)
  427. assert result.name == index.name
  428. def test_difference_preserves_type_empty(self, index, sort):
  429. # GH#20040
  430. # If taking difference of a set and itself, it
  431. # needs to preserve the type of the index
  432. if not index.is_unique:
  433. pytest.skip("Not relevant since index is not unique")
  434. result = index.difference(index, sort=sort)
  435. expected = index[:0]
  436. tm.assert_index_equal(result, expected, exact=True)
  437. def test_difference_name_retention_equals(self, index, names):
  438. if isinstance(index, MultiIndex):
  439. names = [[x] * index.nlevels for x in names]
  440. index = index.rename(names[0])
  441. other = index.rename(names[1])
  442. assert index.equals(other)
  443. result = index.difference(other)
  444. expected = index[:0].rename(names[2])
  445. tm.assert_index_equal(result, expected)
  446. def test_intersection_difference_match_empty(self, index, sort):
  447. # GH#20040
  448. # Test that the intersection of an index with an
  449. # empty index produces the same index as the difference
  450. # of an index with itself. Test for all types
  451. if not index.is_unique:
  452. pytest.skip("Not relevant because index is not unique")
  453. inter = index.intersection(index[:0])
  454. diff = index.difference(index, sort=sort)
  455. tm.assert_index_equal(inter, diff, exact=True)
  456. @pytest.mark.filterwarnings("ignore:invalid value encountered in cast:RuntimeWarning")
  457. @pytest.mark.filterwarnings(r"ignore:PeriodDtype\[B\] is deprecated:FutureWarning")
  458. @pytest.mark.parametrize(
  459. "method", ["intersection", "union", "difference", "symmetric_difference"]
  460. )
  461. def test_setop_with_categorical(index_flat, sort, method, using_infer_string):
  462. # MultiIndex tested separately in tests.indexes.multi.test_setops
  463. index = index_flat
  464. other = index.astype("category")
  465. exact = "equiv" if isinstance(index, RangeIndex) else True
  466. result = getattr(index, method)(other, sort=sort)
  467. expected = getattr(index, method)(index, sort=sort)
  468. if (
  469. using_infer_string
  470. and index.empty
  471. and method in ("union", "symmetric_difference")
  472. ):
  473. expected = expected.astype("category")
  474. tm.assert_index_equal(result, expected, exact=exact)
  475. result = getattr(index, method)(other[:5], sort=sort)
  476. expected = getattr(index, method)(index[:5], sort=sort)
  477. if (
  478. using_infer_string
  479. and index.empty
  480. and method in ("union", "symmetric_difference")
  481. ):
  482. expected = expected.astype("category")
  483. tm.assert_index_equal(result, expected, exact=exact)
  484. def test_intersection_duplicates_all_indexes(index):
  485. # GH#38743
  486. if index.empty:
  487. # No duplicates in empty indexes
  488. pytest.skip("Not relevant for empty Index")
  489. idx = index
  490. idx_non_unique = idx[[0, 0, 1, 2]]
  491. assert idx.intersection(idx_non_unique).equals(idx_non_unique.intersection(idx))
  492. assert idx.intersection(idx_non_unique).is_unique
  493. def test_union_duplicate_index_subsets_of_each_other(
  494. any_dtype_for_small_pos_integer_indexes,
  495. ):
  496. # GH#31326
  497. dtype = any_dtype_for_small_pos_integer_indexes
  498. a = Index([1, 2, 2, 3], dtype=dtype)
  499. b = Index([3, 3, 4], dtype=dtype)
  500. expected = Index([1, 2, 2, 3, 3, 4], dtype=dtype)
  501. if isinstance(a, CategoricalIndex):
  502. expected = Index([1, 2, 2, 3, 3, 4])
  503. result = a.union(b)
  504. tm.assert_index_equal(result, expected)
  505. result = a.union(b, sort=False)
  506. tm.assert_index_equal(result, expected)
  507. def test_union_with_duplicate_index_and_non_monotonic(
  508. any_dtype_for_small_pos_integer_indexes,
  509. ):
  510. # GH#36289
  511. dtype = any_dtype_for_small_pos_integer_indexes
  512. a = Index([1, 0, 0], dtype=dtype)
  513. b = Index([0, 1], dtype=dtype)
  514. expected = Index([0, 0, 1], dtype=dtype)
  515. result = a.union(b)
  516. tm.assert_index_equal(result, expected)
  517. result = b.union(a)
  518. tm.assert_index_equal(result, expected)
  519. def test_union_duplicate_index_different_dtypes():
  520. # GH#36289
  521. a = Index([1, 2, 2, 3])
  522. b = Index(["1", "0", "0"])
  523. expected = Index([1, 2, 2, 3, "1", "0", "0"])
  524. result = a.union(b, sort=False)
  525. tm.assert_index_equal(result, expected)
  526. def test_union_same_value_duplicated_in_both():
  527. # GH#36289
  528. a = Index([0, 0, 1])
  529. b = Index([0, 0, 1, 2])
  530. result = a.union(b)
  531. expected = Index([0, 0, 1, 2])
  532. tm.assert_index_equal(result, expected)
  533. @pytest.mark.parametrize("dup", [1, np.nan])
  534. def test_union_nan_in_both(dup):
  535. # GH#36289
  536. a = Index([np.nan, 1, 2, 2])
  537. b = Index([np.nan, dup, 1, 2])
  538. result = a.union(b, sort=False)
  539. expected = Index([np.nan, dup, 1.0, 2.0, 2.0])
  540. tm.assert_index_equal(result, expected)
  541. def test_union_rangeindex_sort_true():
  542. # GH 53490
  543. idx1 = RangeIndex(1, 100, 6)
  544. idx2 = RangeIndex(1, 50, 3)
  545. result = idx1.union(idx2, sort=True)
  546. expected = Index(
  547. [
  548. 1,
  549. 4,
  550. 7,
  551. 10,
  552. 13,
  553. 16,
  554. 19,
  555. 22,
  556. 25,
  557. 28,
  558. 31,
  559. 34,
  560. 37,
  561. 40,
  562. 43,
  563. 46,
  564. 49,
  565. 55,
  566. 61,
  567. 67,
  568. 73,
  569. 79,
  570. 85,
  571. 91,
  572. 97,
  573. ]
  574. )
  575. tm.assert_index_equal(result, expected)
  576. def test_union_with_duplicate_index_not_subset_and_non_monotonic(
  577. any_dtype_for_small_pos_integer_indexes,
  578. ):
  579. # GH#36289
  580. dtype = any_dtype_for_small_pos_integer_indexes
  581. a = Index([1, 0, 2], dtype=dtype)
  582. b = Index([0, 0, 1], dtype=dtype)
  583. expected = Index([0, 0, 1, 2], dtype=dtype)
  584. if isinstance(a, CategoricalIndex):
  585. expected = Index([0, 0, 1, 2])
  586. result = a.union(b)
  587. tm.assert_index_equal(result, expected)
  588. result = b.union(a)
  589. tm.assert_index_equal(result, expected)
  590. def test_union_int_categorical_with_nan():
  591. ci = CategoricalIndex([1, 2, np.nan])
  592. assert ci.categories.dtype.kind == "i"
  593. idx = Index([1, 2])
  594. result = idx.union(ci)
  595. expected = Index([1, 2, np.nan], dtype=np.float64)
  596. tm.assert_index_equal(result, expected)
  597. result = ci.union(idx)
  598. tm.assert_index_equal(result, expected)
  599. class TestSetOpsUnsorted:
  600. # These may eventually belong in a dtype-specific test_setops, or
  601. # parametrized over a more general fixture
  602. def test_intersect_str_dates(self):
  603. dt_dates = [datetime(2012, 2, 9), datetime(2012, 2, 22)]
  604. index1 = Index(dt_dates, dtype=object)
  605. index2 = Index(["aa"], dtype=object)
  606. result = index2.intersection(index1)
  607. expected = Index([], dtype=object)
  608. tm.assert_index_equal(result, expected)
  609. @pytest.mark.parametrize("index", ["string"], indirect=True)
  610. def test_intersection(self, index, sort):
  611. first = index[:20]
  612. second = index[:10]
  613. intersect = first.intersection(second, sort=sort)
  614. if sort in (None, False):
  615. tm.assert_index_equal(intersect.sort_values(), second.sort_values())
  616. else:
  617. tm.assert_index_equal(intersect, second)
  618. # Corner cases
  619. inter = first.intersection(first, sort=sort)
  620. assert inter is first
  621. @pytest.mark.parametrize(
  622. "index2,keeps_name",
  623. [
  624. (Index([3, 4, 5, 6, 7], name="index"), True), # preserve same name
  625. (Index([3, 4, 5, 6, 7], name="other"), False), # drop diff names
  626. (Index([3, 4, 5, 6, 7]), False),
  627. ],
  628. )
  629. def test_intersection_name_preservation(self, index2, keeps_name, sort):
  630. index1 = Index([1, 2, 3, 4, 5], name="index")
  631. expected = Index([3, 4, 5])
  632. result = index1.intersection(index2, sort)
  633. if keeps_name:
  634. expected.name = "index"
  635. assert result.name == expected.name
  636. tm.assert_index_equal(result, expected)
  637. @pytest.mark.parametrize("index", ["string"], indirect=True)
  638. @pytest.mark.parametrize(
  639. "first_name,second_name,expected_name",
  640. [("A", "A", "A"), ("A", "B", None), (None, "B", None)],
  641. )
  642. def test_intersection_name_preservation2(
  643. self, index, first_name, second_name, expected_name, sort
  644. ):
  645. first = index[5:20]
  646. second = index[:10]
  647. first.name = first_name
  648. second.name = second_name
  649. intersect = first.intersection(second, sort=sort)
  650. assert intersect.name == expected_name
  651. def test_chained_union(self, sort):
  652. # Chained unions handles names correctly
  653. i1 = Index([1, 2], name="i1")
  654. i2 = Index([5, 6], name="i2")
  655. i3 = Index([3, 4], name="i3")
  656. union = i1.union(i2.union(i3, sort=sort), sort=sort)
  657. expected = i1.union(i2, sort=sort).union(i3, sort=sort)
  658. tm.assert_index_equal(union, expected)
  659. j1 = Index([1, 2], name="j1")
  660. j2 = Index([], name="j2")
  661. j3 = Index([], name="j3")
  662. union = j1.union(j2.union(j3, sort=sort), sort=sort)
  663. expected = j1.union(j2, sort=sort).union(j3, sort=sort)
  664. tm.assert_index_equal(union, expected)
  665. @pytest.mark.parametrize("index", ["string"], indirect=True)
  666. def test_union(self, index, sort):
  667. first = index[5:20]
  668. second = index[:10]
  669. everything = index[:20]
  670. union = first.union(second, sort=sort)
  671. if sort in (None, False):
  672. tm.assert_index_equal(union.sort_values(), everything.sort_values())
  673. else:
  674. tm.assert_index_equal(union, everything)
  675. @pytest.mark.parametrize("klass", [np.array, Series, list])
  676. @pytest.mark.parametrize("index", ["string"], indirect=True)
  677. def test_union_from_iterables(self, index, klass, sort):
  678. # GH#10149
  679. first = index[5:20]
  680. second = index[:10]
  681. everything = index[:20]
  682. case = klass(second.values)
  683. result = first.union(case, sort=sort)
  684. if sort in (None, False):
  685. tm.assert_index_equal(result.sort_values(), everything.sort_values())
  686. else:
  687. tm.assert_index_equal(result, everything)
  688. @pytest.mark.parametrize("index", ["string"], indirect=True)
  689. def test_union_identity(self, index, sort):
  690. first = index[5:20]
  691. union = first.union(first, sort=sort)
  692. # i.e. identity is not preserved when sort is True
  693. assert (union is first) is (not sort)
  694. # This should no longer be the same object, since [] is not consistent,
  695. # both objects will be recast to dtype('O')
  696. union = first.union(Index([], dtype=first.dtype), sort=sort)
  697. assert (union is first) is (not sort)
  698. union = Index([], dtype=first.dtype).union(first, sort=sort)
  699. assert (union is first) is (not sort)
  700. @pytest.mark.parametrize("index", ["string"], indirect=True)
  701. @pytest.mark.parametrize("second_name,expected", [(None, None), ("name", "name")])
  702. def test_difference_name_preservation(self, index, second_name, expected, sort):
  703. first = index[5:20]
  704. second = index[:10]
  705. answer = index[10:20]
  706. first.name = "name"
  707. second.name = second_name
  708. result = first.difference(second, sort=sort)
  709. if sort is True:
  710. tm.assert_index_equal(result, answer)
  711. else:
  712. answer.name = second_name
  713. tm.assert_index_equal(result.sort_values(), answer.sort_values())
  714. if expected is None:
  715. assert result.name is None
  716. else:
  717. assert result.name == expected
  718. def test_difference_empty_arg(self, index, sort):
  719. first = index.copy()
  720. first = first[5:20]
  721. first.name = "name"
  722. result = first.difference([], sort)
  723. expected = index[5:20].unique()
  724. expected.name = "name"
  725. tm.assert_index_equal(result, expected)
  726. def test_difference_should_not_compare(self):
  727. # GH 55113
  728. left = Index([1, 1])
  729. right = Index([True])
  730. result = left.difference(right)
  731. expected = Index([1])
  732. tm.assert_index_equal(result, expected)
  733. @pytest.mark.parametrize("index", ["string"], indirect=True)
  734. def test_difference_identity(self, index, sort):
  735. first = index[5:20]
  736. first.name = "name"
  737. result = first.difference(first, sort)
  738. assert len(result) == 0
  739. assert result.name == first.name
  740. @pytest.mark.parametrize("index", ["string"], indirect=True)
  741. def test_difference_sort(self, index, sort):
  742. first = index[5:20]
  743. second = index[:10]
  744. result = first.difference(second, sort)
  745. expected = index[10:20]
  746. if sort is None:
  747. expected = expected.sort_values()
  748. tm.assert_index_equal(result, expected)
  749. @pytest.mark.parametrize("opname", ["difference", "symmetric_difference"])
  750. def test_difference_incomparable(self, opname):
  751. a = Index([3, Timestamp("2000"), 1])
  752. b = Index([2, Timestamp("1999"), 1])
  753. op = operator.methodcaller(opname, b)
  754. with tm.assert_produces_warning(RuntimeWarning):
  755. # sort=None, the default
  756. result = op(a)
  757. expected = Index([3, Timestamp("2000"), 2, Timestamp("1999")])
  758. if opname == "difference":
  759. expected = expected[:2]
  760. tm.assert_index_equal(result, expected)
  761. # sort=False
  762. op = operator.methodcaller(opname, b, sort=False)
  763. result = op(a)
  764. tm.assert_index_equal(result, expected)
  765. @pytest.mark.parametrize("opname", ["difference", "symmetric_difference"])
  766. def test_difference_incomparable_true(self, opname):
  767. a = Index([3, Timestamp("2000"), 1])
  768. b = Index([2, Timestamp("1999"), 1])
  769. op = operator.methodcaller(opname, b, sort=True)
  770. msg = "'<' not supported between instances of 'Timestamp' and 'int'"
  771. with pytest.raises(TypeError, match=msg):
  772. op(a)
  773. def test_symmetric_difference_mi(self, sort):
  774. index1 = MultiIndex.from_tuples(zip(["foo", "bar", "baz"], [1, 2, 3]))
  775. index2 = MultiIndex.from_tuples([("foo", 1), ("bar", 3)])
  776. result = index1.symmetric_difference(index2, sort=sort)
  777. expected = MultiIndex.from_tuples([("bar", 2), ("baz", 3), ("bar", 3)])
  778. if sort is None:
  779. expected = expected.sort_values()
  780. tm.assert_index_equal(result, expected)
  781. @pytest.mark.parametrize(
  782. "index2,expected",
  783. [
  784. (Index([0, 1, np.nan]), Index([2.0, 3.0, 0.0])),
  785. (Index([0, 1]), Index([np.nan, 2.0, 3.0, 0.0])),
  786. ],
  787. )
  788. def test_symmetric_difference_missing(self, index2, expected, sort):
  789. # GH#13514 change: {nan} - {nan} == {}
  790. # (GH#6444, sorting of nans, is no longer an issue)
  791. index1 = Index([1, np.nan, 2, 3])
  792. result = index1.symmetric_difference(index2, sort=sort)
  793. if sort is None:
  794. expected = expected.sort_values()
  795. tm.assert_index_equal(result, expected)
  796. def test_symmetric_difference_non_index(self, sort):
  797. index1 = Index([1, 2, 3, 4], name="index1")
  798. index2 = np.array([2, 3, 4, 5])
  799. expected = Index([1, 5], name="index1")
  800. result = index1.symmetric_difference(index2, sort=sort)
  801. if sort in (None, True):
  802. tm.assert_index_equal(result, expected)
  803. else:
  804. tm.assert_index_equal(result.sort_values(), expected)
  805. assert result.name == "index1"
  806. result = index1.symmetric_difference(index2, result_name="new_name", sort=sort)
  807. expected.name = "new_name"
  808. if sort in (None, True):
  809. tm.assert_index_equal(result, expected)
  810. else:
  811. tm.assert_index_equal(result.sort_values(), expected)
  812. assert result.name == "new_name"
  813. def test_union_ea_dtypes(self, any_numeric_ea_and_arrow_dtype):
  814. # GH#51365
  815. idx = Index([1, 2, 3], dtype=any_numeric_ea_and_arrow_dtype)
  816. idx2 = Index([3, 4, 5], dtype=any_numeric_ea_and_arrow_dtype)
  817. result = idx.union(idx2)
  818. expected = Index([1, 2, 3, 4, 5], dtype=any_numeric_ea_and_arrow_dtype)
  819. tm.assert_index_equal(result, expected)
  820. def test_union_string_array(self, any_string_dtype):
  821. idx1 = Index(["a"], dtype=any_string_dtype)
  822. idx2 = Index(["b"], dtype=any_string_dtype)
  823. result = idx1.union(idx2)
  824. expected = Index(["a", "b"], dtype=any_string_dtype)
  825. tm.assert_index_equal(result, expected)