test_asof.py 6.5 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197
  1. import numpy as np
  2. import pytest
  3. from pandas._libs.tslibs import IncompatibleFrequency
  4. from pandas import (
  5. DataFrame,
  6. Period,
  7. Series,
  8. Timestamp,
  9. date_range,
  10. period_range,
  11. to_datetime,
  12. )
  13. import pandas._testing as tm
  14. @pytest.fixture
  15. def date_range_frame():
  16. """
  17. Fixture for DataFrame of ints with date_range index
  18. Columns are ['A', 'B'].
  19. """
  20. N = 50
  21. rng = date_range("1/1/1990", periods=N, freq="53s")
  22. return DataFrame({"A": np.arange(N), "B": np.arange(N)}, index=rng)
  23. class TestFrameAsof:
  24. def test_basic(self, date_range_frame):
  25. # Explicitly cast to float to avoid implicit cast when setting np.nan
  26. df = date_range_frame.astype({"A": "float"})
  27. N = 50
  28. df.loc[df.index[15:30], "A"] = np.nan
  29. dates = date_range("1/1/1990", periods=N * 3, freq="25s")
  30. result = df.asof(dates)
  31. assert result.notna().all(1).all()
  32. lb = df.index[14]
  33. ub = df.index[30]
  34. dates = list(dates)
  35. result = df.asof(dates)
  36. assert result.notna().all(1).all()
  37. mask = (result.index >= lb) & (result.index < ub)
  38. rs = result[mask]
  39. assert (rs == 14).all(1).all()
  40. def test_subset(self, date_range_frame):
  41. N = 10
  42. # explicitly cast to float to avoid implicit upcast when setting to np.nan
  43. df = date_range_frame.iloc[:N].copy().astype({"A": "float"})
  44. df.loc[df.index[4:8], "A"] = np.nan
  45. dates = date_range("1/1/1990", periods=N * 3, freq="25s")
  46. # with a subset of A should be the same
  47. result = df.asof(dates, subset="A")
  48. expected = df.asof(dates)
  49. tm.assert_frame_equal(result, expected)
  50. # same with A/B
  51. result = df.asof(dates, subset=["A", "B"])
  52. expected = df.asof(dates)
  53. tm.assert_frame_equal(result, expected)
  54. # B gives df.asof
  55. result = df.asof(dates, subset="B")
  56. expected = df.resample("25s", closed="right").ffill().reindex(dates)
  57. expected.iloc[20:] = 9
  58. # no "missing", so "B" can retain int dtype (df["A"].dtype platform-dependent)
  59. expected["B"] = expected["B"].astype(df["B"].dtype)
  60. tm.assert_frame_equal(result, expected)
  61. def test_missing(self, date_range_frame):
  62. # GH 15118
  63. # no match found - `where` value before earliest date in index
  64. N = 10
  65. df = date_range_frame.iloc[:N].copy()
  66. result = df.asof("1989-12-31")
  67. expected = Series(
  68. index=["A", "B"], name=Timestamp("1989-12-31"), dtype=np.float64
  69. )
  70. tm.assert_series_equal(result, expected)
  71. result = df.asof(to_datetime(["1989-12-31"]))
  72. expected = DataFrame(
  73. index=to_datetime(["1989-12-31"]), columns=["A", "B"], dtype="float64"
  74. )
  75. tm.assert_frame_equal(result, expected)
  76. # Check that we handle PeriodIndex correctly, dont end up with
  77. # period.ordinal for series name
  78. df = df.to_period("D")
  79. result = df.asof("1989-12-31")
  80. assert isinstance(result.name, Period)
  81. def test_asof_all_nans(self, frame_or_series):
  82. # GH 15713
  83. # DataFrame/Series is all nans
  84. result = frame_or_series([np.nan]).asof([0])
  85. expected = frame_or_series([np.nan])
  86. tm.assert_equal(result, expected)
  87. def test_all_nans(self, date_range_frame):
  88. # GH 15713
  89. # DataFrame is all nans
  90. # testing non-default indexes, multiple inputs
  91. N = 150
  92. rng = date_range_frame.index
  93. dates = date_range("1/1/1990", periods=N, freq="25s")
  94. result = DataFrame(np.nan, index=rng, columns=["A"]).asof(dates)
  95. expected = DataFrame(np.nan, index=dates, columns=["A"])
  96. tm.assert_frame_equal(result, expected)
  97. # testing multiple columns
  98. dates = date_range("1/1/1990", periods=N, freq="25s")
  99. result = DataFrame(np.nan, index=rng, columns=["A", "B", "C"]).asof(dates)
  100. expected = DataFrame(np.nan, index=dates, columns=["A", "B", "C"])
  101. tm.assert_frame_equal(result, expected)
  102. # testing scalar input
  103. result = DataFrame(np.nan, index=[1, 2], columns=["A", "B"]).asof([3])
  104. expected = DataFrame(np.nan, index=[3], columns=["A", "B"])
  105. tm.assert_frame_equal(result, expected)
  106. result = DataFrame(np.nan, index=[1, 2], columns=["A", "B"]).asof(3)
  107. expected = Series(np.nan, index=["A", "B"], name=3)
  108. tm.assert_series_equal(result, expected)
  109. @pytest.mark.parametrize(
  110. "stamp,expected",
  111. [
  112. (
  113. Timestamp("2018-01-01 23:22:43.325+00:00"),
  114. Series(2, name=Timestamp("2018-01-01 23:22:43.325+00:00")),
  115. ),
  116. (
  117. Timestamp("2018-01-01 22:33:20.682+01:00"),
  118. Series(1, name=Timestamp("2018-01-01 22:33:20.682+01:00")),
  119. ),
  120. ],
  121. )
  122. def test_time_zone_aware_index(self, stamp, expected):
  123. # GH21194
  124. # Testing awareness of DataFrame index considering different
  125. # UTC and timezone
  126. df = DataFrame(
  127. data=[1, 2],
  128. index=[
  129. Timestamp("2018-01-01 21:00:05.001+00:00"),
  130. Timestamp("2018-01-01 22:35:10.550+00:00"),
  131. ],
  132. )
  133. result = df.asof(stamp)
  134. tm.assert_series_equal(result, expected)
  135. def test_is_copy(self, date_range_frame):
  136. # GH-27357, GH-30784: ensure the result of asof is an actual copy and
  137. # doesn't track the parent dataframe / doesn't give SettingWithCopy warnings
  138. df = date_range_frame.astype({"A": "float"})
  139. N = 50
  140. df.loc[df.index[15:30], "A"] = np.nan
  141. dates = date_range("1/1/1990", periods=N * 3, freq="25s")
  142. result = df.asof(dates)
  143. with tm.assert_produces_warning(None):
  144. result["C"] = 1
  145. def test_asof_periodindex_mismatched_freq(self):
  146. N = 50
  147. rng = period_range("1/1/1990", periods=N, freq="H")
  148. df = DataFrame(np.random.randn(N), index=rng)
  149. # Mismatched freq
  150. msg = "Input has different freq"
  151. with pytest.raises(IncompatibleFrequency, match=msg):
  152. df.asof(rng.asfreq("D"))
  153. def test_asof_preserves_bool_dtype(self):
  154. # GH#16063 was casting bools to floats
  155. dti = date_range("2017-01-01", freq="MS", periods=4)
  156. ser = Series([True, False, True], index=dti[:-1])
  157. ts = dti[-1]
  158. res = ser.asof([ts])
  159. expected = Series([True], index=[ts])
  160. tm.assert_series_equal(res, expected)