test_get_level_values.py 3.9 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124
  1. import numpy as np
  2. import pandas as pd
  3. from pandas import (
  4. CategoricalIndex,
  5. Index,
  6. MultiIndex,
  7. Timestamp,
  8. date_range,
  9. )
  10. import pandas._testing as tm
  11. class TestGetLevelValues:
  12. def test_get_level_values_box_datetime64(self):
  13. dates = date_range("1/1/2000", periods=4)
  14. levels = [dates, [0, 1]]
  15. codes = [[0, 0, 1, 1, 2, 2, 3, 3], [0, 1, 0, 1, 0, 1, 0, 1]]
  16. index = MultiIndex(levels=levels, codes=codes)
  17. assert isinstance(index.get_level_values(0)[0], Timestamp)
  18. def test_get_level_values(idx):
  19. result = idx.get_level_values(0)
  20. expected = Index(["foo", "foo", "bar", "baz", "qux", "qux"], name="first")
  21. tm.assert_index_equal(result, expected)
  22. assert result.name == "first"
  23. result = idx.get_level_values("first")
  24. expected = idx.get_level_values(0)
  25. tm.assert_index_equal(result, expected)
  26. # GH 10460
  27. index = MultiIndex(
  28. levels=[CategoricalIndex(["A", "B"]), CategoricalIndex([1, 2, 3])],
  29. codes=[np.array([0, 0, 0, 1, 1, 1]), np.array([0, 1, 2, 0, 1, 2])],
  30. )
  31. exp = CategoricalIndex(["A", "A", "A", "B", "B", "B"])
  32. tm.assert_index_equal(index.get_level_values(0), exp)
  33. exp = CategoricalIndex([1, 2, 3, 1, 2, 3])
  34. tm.assert_index_equal(index.get_level_values(1), exp)
  35. def test_get_level_values_all_na():
  36. # GH#17924 when level entirely consists of nan
  37. arrays = [[np.nan, np.nan, np.nan], ["a", np.nan, 1]]
  38. index = MultiIndex.from_arrays(arrays)
  39. result = index.get_level_values(0)
  40. expected = Index([np.nan, np.nan, np.nan], dtype=np.float64)
  41. tm.assert_index_equal(result, expected)
  42. result = index.get_level_values(1)
  43. expected = Index(["a", np.nan, 1], dtype=object)
  44. tm.assert_index_equal(result, expected)
  45. def test_get_level_values_int_with_na():
  46. # GH#17924
  47. arrays = [["a", "b", "b"], [1, np.nan, 2]]
  48. index = MultiIndex.from_arrays(arrays)
  49. result = index.get_level_values(1)
  50. expected = Index([1, np.nan, 2])
  51. tm.assert_index_equal(result, expected)
  52. arrays = [["a", "b", "b"], [np.nan, np.nan, 2]]
  53. index = MultiIndex.from_arrays(arrays)
  54. result = index.get_level_values(1)
  55. expected = Index([np.nan, np.nan, 2])
  56. tm.assert_index_equal(result, expected)
  57. def test_get_level_values_na():
  58. arrays = [[np.nan, np.nan, np.nan], ["a", np.nan, 1]]
  59. index = MultiIndex.from_arrays(arrays)
  60. result = index.get_level_values(0)
  61. expected = Index([np.nan, np.nan, np.nan])
  62. tm.assert_index_equal(result, expected)
  63. result = index.get_level_values(1)
  64. expected = Index(["a", np.nan, 1])
  65. tm.assert_index_equal(result, expected)
  66. arrays = [["a", "b", "b"], pd.DatetimeIndex([0, 1, pd.NaT])]
  67. index = MultiIndex.from_arrays(arrays)
  68. result = index.get_level_values(1)
  69. expected = pd.DatetimeIndex([0, 1, pd.NaT])
  70. tm.assert_index_equal(result, expected)
  71. arrays = [[], []]
  72. index = MultiIndex.from_arrays(arrays)
  73. result = index.get_level_values(0)
  74. expected = Index([], dtype=object)
  75. tm.assert_index_equal(result, expected)
  76. def test_get_level_values_when_periods():
  77. # GH33131. See also discussion in GH32669.
  78. # This test can probably be removed when PeriodIndex._engine is removed.
  79. from pandas import (
  80. Period,
  81. PeriodIndex,
  82. )
  83. idx = MultiIndex.from_arrays(
  84. [PeriodIndex([Period("2019Q1"), Period("2019Q2")], name="b")]
  85. )
  86. idx2 = MultiIndex.from_arrays(
  87. [idx._get_level_values(level) for level in range(idx.nlevels)]
  88. )
  89. assert all(x.is_monotonic_increasing for x in idx2.levels)
  90. def test_values_loses_freq_of_underlying_index():
  91. # GH#49054
  92. idx = pd.DatetimeIndex(date_range("20200101", periods=3, freq="BM"))
  93. expected = idx.copy(deep=True)
  94. idx2 = Index([1, 2, 3])
  95. midx = MultiIndex(levels=[idx, idx2], codes=[[0, 1, 2], [0, 1, 2]])
  96. midx.values
  97. assert idx.freq is not None
  98. tm.assert_index_equal(idx, expected)