Skip to content

Commit 1d8b188

Browse files
authored
BUG/ENH: groupby.quantile support non-nano (#50978)
1 parent b39c78b commit 1d8b188

File tree

3 files changed

+25
-11
lines changed

3 files changed

+25
-11
lines changed

pandas/core/groupby/groupby.py

+2-2
Original file line numberDiff line numberDiff line change
@@ -3188,10 +3188,10 @@ def pre_processor(vals: ArrayLike) -> tuple[np.ndarray, Dtype | None]:
31883188
elif is_bool_dtype(vals.dtype) and isinstance(vals, ExtensionArray):
31893189
out = vals.to_numpy(dtype=float, na_value=np.nan)
31903190
elif is_datetime64_dtype(vals.dtype):
3191-
inference = np.dtype("datetime64[ns]")
3191+
inference = vals.dtype
31923192
out = np.asarray(vals).astype(float)
31933193
elif is_timedelta64_dtype(vals.dtype):
3194-
inference = np.dtype("timedelta64[ns]")
3194+
inference = vals.dtype
31953195
out = np.asarray(vals).astype(float)
31963196
elif isinstance(vals, ExtensionArray) and is_float_dtype(vals):
31973197
inference = np.dtype(np.float64)

pandas/tests/groupby/test_quantile.py

+18-6
Original file line numberDiff line numberDiff line change
@@ -26,33 +26,45 @@
2626
([np.nan, 4.0, np.nan, 2.0, np.nan], [np.nan, 4.0, np.nan, 2.0, np.nan]),
2727
# Timestamps
2828
(
29-
list(pd.date_range("1/1/18", freq="D", periods=5)),
30-
list(pd.date_range("1/1/18", freq="D", periods=5))[::-1],
29+
pd.date_range("1/1/18", freq="D", periods=5),
30+
pd.date_range("1/1/18", freq="D", periods=5)[::-1],
31+
),
32+
(
33+
pd.date_range("1/1/18", freq="D", periods=5).as_unit("s"),
34+
pd.date_range("1/1/18", freq="D", periods=5)[::-1].as_unit("s"),
3135
),
3236
# All NA
3337
([np.nan] * 5, [np.nan] * 5),
3438
],
3539
)
3640
@pytest.mark.parametrize("q", [0, 0.25, 0.5, 0.75, 1])
3741
def test_quantile(interpolation, a_vals, b_vals, q, request):
38-
if interpolation == "nearest" and q == 0.5 and b_vals == [4, 3, 2, 1]:
42+
if (
43+
interpolation == "nearest"
44+
and q == 0.5
45+
and isinstance(b_vals, list)
46+
and b_vals == [4, 3, 2, 1]
47+
):
3948
request.node.add_marker(
4049
pytest.mark.xfail(
4150
reason="Unclear numpy expectation for nearest "
4251
"result with equidistant data"
4352
)
4453
)
54+
all_vals = pd.concat([pd.Series(a_vals), pd.Series(b_vals)])
4555

4656
a_expected = pd.Series(a_vals).quantile(q, interpolation=interpolation)
4757
b_expected = pd.Series(b_vals).quantile(q, interpolation=interpolation)
4858

49-
df = DataFrame(
50-
{"key": ["a"] * len(a_vals) + ["b"] * len(b_vals), "val": a_vals + b_vals}
51-
)
59+
df = DataFrame({"key": ["a"] * len(a_vals) + ["b"] * len(b_vals), "val": all_vals})
5260

5361
expected = DataFrame(
5462
[a_expected, b_expected], columns=["val"], index=Index(["a", "b"], name="key")
5563
)
64+
if all_vals.dtype.kind == "M" and expected.dtypes.values[0].kind == "M":
65+
# TODO(non-nano): this should be unnecessary once array_to_datetime
66+
# correctly infers non-nano from Timestamp.unit
67+
expected = expected.astype(all_vals.dtype)
5668
result = df.groupby("key").quantile(q, interpolation=interpolation)
5769

5870
tm.assert_frame_equal(result, expected)

pandas/tests/resample/test_timedelta.py

+5-3
Original file line numberDiff line numberDiff line change
@@ -174,10 +174,12 @@ def test_resample_with_timedelta_yields_no_empty_groups(duplicates):
174174
tm.assert_frame_equal(result, expected)
175175

176176

177-
def test_resample_quantile_timedelta():
177+
@pytest.mark.parametrize("unit", ["s", "ms", "us", "ns"])
178+
def test_resample_quantile_timedelta(unit):
178179
# GH: 29485
180+
dtype = np.dtype(f"m8[{unit}]")
179181
df = DataFrame(
180-
{"value": pd.to_timedelta(np.arange(4), unit="s")},
182+
{"value": pd.to_timedelta(np.arange(4), unit="s").astype(dtype)},
181183
index=pd.date_range("20200101", periods=4, tz="UTC"),
182184
)
183185
result = df.resample("2D").quantile(0.99)
@@ -189,7 +191,7 @@ def test_resample_quantile_timedelta():
189191
]
190192
},
191193
index=pd.date_range("20200101", periods=2, tz="UTC", freq="2D"),
192-
)
194+
).astype(dtype)
193195
tm.assert_frame_equal(result, expected)
194196

195197

0 commit comments

Comments
 (0)