Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.11/site-packages/pandas/core/groupby/indexing.py: 26%

1from __future__ import annotations 

2 

3from typing import ( 

4 TYPE_CHECKING, 

5 Literal, 

6 cast, 

7) 

8import warnings 

9 

10import numpy as np 

11 

12from pandas.errors import Pandas4Warning 

13from pandas.util._decorators import cache_readonly 

14from pandas.util._exceptions import find_stack_level 

15 

16from pandas.core.dtypes.common import ( 

17 is_integer, 

18 is_list_like, 

19) 

20 

21if TYPE_CHECKING: 

22 from collections.abc import Iterable 

23 

24 from pandas._typing import PositionalIndexer 

25 

26 from pandas import ( 

27 DataFrame, 

28 Series, 

29 ) 

30 from pandas.core.groupby import groupby 

31 

32 

33class GroupByIndexingMixin: 

34 """ 

35 Mixin for adding ._positional_selector to GroupBy. 

36 """ 

37 

38 @cache_readonly 

39 def _positional_selector(self) -> GroupByPositionalSelector: 

40 """ 

41 Return positional selection for each group. 

42 

43 ``groupby._positional_selector[i:j]`` is similar to 

44 ``groupby.apply(lambda x: x.iloc[i:j])`` 

45 but much faster and preserves the original index and order. 

46 

47 ``_positional_selector[]`` is compatible with and extends :meth:`~GroupBy.head` 

48 and :meth:`~GroupBy.tail`. For example: 

49 

50 - ``head(5)`` 

51 - ``_positional_selector[5:-5]`` 

52 - ``tail(5)`` 

53 

54 together return all the rows. 

55 

56 Allowed inputs for the index are: 

57 

58 - An integer valued iterable, e.g. ``range(2, 4)``. 

59 - A comma separated list of integers and slices, e.g. ``5``, ``2, 4``, ``2:4``. 

60 

61 The output format is the same as :meth:`~GroupBy.head` and 

62 :meth:`~GroupBy.tail`, namely 

63 a subset of the ``DataFrame`` or ``Series`` with the index and order preserved. 

64 

65 Returns 

66 ------- 

67 Series 

68 The filtered subset of the original Series. 

69 DataFrame 

70 The filtered subset of the original DataFrame. 

71 

72 See Also 

73 -------- 

74 DataFrame.iloc : Purely integer-location based indexing for selection by 

75 position. 

76 GroupBy.head : Return first n rows of each group. 

77 GroupBy.tail : Return last n rows of each group. 

78 GroupBy.nth : Take the nth row from each group if n is an int, or a 

79 subset of rows, if n is a list of ints. 

80 

81 Notes 

82 ----- 

83 - The slice step cannot be negative. 

84 - If the index specification results in overlaps, the item is not duplicated. 

85 - If the index specification changes the order of items, then 

86 they are returned in their original order. 

87 By contrast, ``DataFrame.iloc`` can change the row order. 

88 - ``groupby()`` parameters such as as_index and dropna are ignored. 

89 

90 The differences between ``_positional_selector[]`` and :meth:`~GroupBy.nth` 

91 with ``as_index=False`` are: 

92 

93 - Input to ``_positional_selector`` can include 

94 one or more slices whereas ``nth`` 

95 just handles an integer or a list of integers. 

96 - ``_positional_selector`` can accept a slice relative to the 

97 last row of each group. 

98 - ``_positional_selector`` does not have an equivalent to the 

99 ``nth()`` ``dropna`` parameter. 

100 

101 Examples 

102 -------- 

103 >>> df = pd.DataFrame( 

104 ... [["a", 1], ["a", 2], ["a", 3], ["b", 4], ["b", 5]], columns=["A", "B"] 

105 ... ) 

106 >>> df.groupby("A")._positional_selector[1:2] 

107 A B 

108 1 a 2 

109 4 b 5 

110 

111 >>> df.groupby("A")._positional_selector[1, -1] 

112 A B 

113 1 a 2 

114 2 a 3 

115 4 b 5 

116 """ 

117 if TYPE_CHECKING: 

118 groupby_self = cast("groupby.GroupBy", self) 

119 else: 

120 groupby_self = self 

121 

122 return GroupByPositionalSelector(groupby_self) 

123 

124 def _make_mask_from_positional_indexer( 

125 self, 

126 arg: PositionalIndexer | tuple, 

127 ) -> np.ndarray: 

128 if is_list_like(arg): 

129 if all(is_integer(i) for i in cast("Iterable", arg)): 

130 mask = self._make_mask_from_list(cast("Iterable[int]", arg)) 

131 else: 

132 mask = self._make_mask_from_tuple(cast("tuple", arg)) 

133 

134 elif isinstance(arg, slice): 

135 mask = self._make_mask_from_slice(arg) 

136 elif is_integer(arg): 

137 mask = self._make_mask_from_int(cast("int", arg)) 

138 else: 

139 raise TypeError( 

140 f"Invalid index {type(arg)}. " 

141 "Must be integer, list-like, slice or a tuple of " 

142 "integers and slices" 

143 ) 

144 

145 if isinstance(mask, bool): 

146 if mask: 

147 mask = self._ascending_count >= 0 

148 else: 

149 mask = self._ascending_count < 0 

150 

151 return cast("np.ndarray", mask) 

152 

153 def _make_mask_from_int(self, arg: int) -> np.ndarray: 

154 if arg >= 0: 

155 return self._ascending_count == arg 

156 else: 

157 return self._descending_count == (-arg - 1) 

158 

159 def _make_mask_from_list(self, args: Iterable[int]) -> bool | np.ndarray: 

160 positive = [arg for arg in args if arg >= 0] 

161 negative = [-arg - 1 for arg in args if arg < 0] 

162 

163 mask: bool | np.ndarray = False 

164 

165 if positive: 

166 mask |= np.isin(self._ascending_count, positive) 

167 

168 if negative: 

169 mask |= np.isin(self._descending_count, negative) 

170 

171 return mask 

172 

173 def _make_mask_from_tuple(self, args: tuple) -> bool | np.ndarray: 

174 mask: bool | np.ndarray = False 

175 

176 for arg in args: 

177 if is_integer(arg): 

178 mask |= self._make_mask_from_int(cast("int", arg)) 

179 elif isinstance(arg, slice): 

180 mask |= self._make_mask_from_slice(arg) 

181 else: 

182 raise ValueError( 

183 f"Invalid argument {type(arg)}. Should be int or slice." 

184 ) 

185 

186 return mask 

187 

188 def _make_mask_from_slice(self, arg: slice) -> bool | np.ndarray: 

189 start = arg.start 

190 stop = arg.stop 

191 step = arg.step 

192 

193 if step is not None and step < 0: 

194 raise ValueError(f"Invalid step {step}. Must be non-negative") 

195 

196 mask: bool | np.ndarray = True 

197 

198 if step is None: 

199 step = 1 

200 

201 if start is None: 

202 if step > 1: 

203 mask &= self._ascending_count % step == 0 

204 

205 elif start >= 0: 

206 mask &= self._ascending_count >= start 

207 

208 if step > 1: 

209 mask &= (self._ascending_count - start) % step == 0 

210 

211 else: 

212 mask &= self._descending_count < -start 

213 

214 offset_array = self._descending_count + start + 1 

215 limit_array = ( 

216 self._ascending_count + self._descending_count + (start + 1) 

217 ) < 0 

218 offset_array = np.where(limit_array, self._ascending_count, offset_array) 

219 

220 mask &= offset_array % step == 0 

221 

222 if stop is not None: 

223 if stop >= 0: 

224 mask &= self._ascending_count < stop 

225 else: 

226 mask &= self._descending_count >= -stop 

227 

228 return mask 

229 

230 @cache_readonly 

231 def _ascending_count(self) -> np.ndarray: 

232 if TYPE_CHECKING: 

233 groupby_self = cast("groupby.GroupBy", self) 

234 else: 

235 groupby_self = self 

236 

237 return groupby_self._cumcount_array() 

238 

239 @cache_readonly 

240 def _descending_count(self) -> np.ndarray: 

241 if TYPE_CHECKING: 

242 groupby_self = cast("groupby.GroupBy", self) 

243 else: 

244 groupby_self = self 

245 

246 return groupby_self._cumcount_array(ascending=False) 

247 

248 

249class GroupByPositionalSelector: 

250 """ 

251 Return positional selection for each group. 

252 

253 ``groupby._positional_selector[i:j]`` is similar to 

254 ``groupby.apply(lambda x: x.iloc[i:j])`` 

255 but much faster and preserves the original index and order. 

256 

257 ``_positional_selector[]`` is compatible with and extends :meth:`~GroupBy.head` 

258 and :meth:`~GroupBy.tail`. For example: 

259 

260 - ``head(5)`` 

261 - ``_positional_selector[5:-5]`` 

262 - ``tail(5)`` 

263 

264 together return all the rows. 

265 

266 Allowed inputs for the index are: 

267 

268 - An integer valued iterable, e.g. ``range(2, 4)``. 

269 - A comma separated list of integers and slices, e.g. ``5``, ``2, 4``, ``2:4``. 

270 

271 The output format is the same as :meth:`~GroupBy.head` and 

272 :meth:`~GroupBy.tail`, namely 

273 a subset of the ``DataFrame`` or ``Series`` with the index and order preserved. 

274 

275 Returns 

276 ------- 

277 Series 

278 The filtered subset of the original Series. 

279 DataFrame 

280 The filtered subset of the original DataFrame. 

281 

282 See Also 

283 -------- 

284 DataFrame.iloc : Purely integer-location based indexing for selection by 

285 position. 

286 GroupBy.head : Return first n rows of each group. 

287 GroupBy.tail : Return last n rows of each group. 

288 GroupBy.nth : Take the nth row from each group if n is an int, or a 

289 subset of rows, if n is a list of ints. 

290 

291 Notes 

292 ----- 

293 - The slice step cannot be negative. 

294 - If the index specification results in overlaps, the item is not duplicated. 

295 - If the index specification changes the order of items, then 

296 they are returned in their original order. 

297 By contrast, ``DataFrame.iloc`` can change the row order. 

298 - ``groupby()`` parameters such as as_index and dropna are ignored. 

299 

300 The differences between ``_positional_selector[]`` and :meth:`~GroupBy.nth` 

301 with ``as_index=False`` are: 

302 

303 - Input to ``_positional_selector`` can include 

304 one or more slices whereas ``nth`` 

305 just handles an integer or a list of integers. 

306 - ``_positional_selector`` can accept a slice relative to the 

307 last row of each group. 

308 - ``_positional_selector`` does not have an equivalent to the 

309 ``nth()`` ``dropna`` parameter. 

310 

311 Examples 

312 -------- 

313 >>> df = pd.DataFrame( 

314 ... [["a", 1], ["a", 2], ["a", 3], ["b", 4], ["b", 5]], columns=["A", "B"] 

315 ... ) 

316 >>> df.groupby("A")._positional_selector[1:2] 

317 A B 

318 1 a 2 

319 4 b 5 

320 

321 >>> df.groupby("A")._positional_selector[1, -1] 

322 A B 

323 1 a 2 

324 2 a 3 

325 4 b 5 

326 """ 

327 

328 def __init__(self, groupby_object: groupby.GroupBy) -> None: 

329 self.groupby_object = groupby_object 

330 

331 def __getitem__(self, arg: PositionalIndexer | tuple) -> DataFrame | Series: 

332 """ 

333 Select by positional index per group. 

334 

335 Implements GroupBy._positional_selector 

336 

337 Parameters 

338 ---------- 

339 arg : PositionalIndexer | tuple 

340 Allowed values are: 

341 - int 

342 - int valued iterable such as list or range 

343 - slice with step either None or positive 

344 - tuple of integers and slices 

345 

346 Returns 

347 ------- 

348 Series 

349 The filtered subset of the original groupby Series. 

350 DataFrame 

351 The filtered subset of the original groupby DataFrame. 

352 

353 See Also 

354 -------- 

355 DataFrame.iloc : Integer-location based indexing for selection by position. 

356 GroupBy.head : Return first n rows of each group. 

357 GroupBy.tail : Return last n rows of each group. 

358 GroupBy._positional_selector : Return positional selection for each group. 

359 GroupBy.nth : Take the nth row from each group if n is an int, or a 

360 subset of rows, if n is a list of ints. 

361 """ 

362 mask = self.groupby_object._make_mask_from_positional_indexer(arg) 

363 return self.groupby_object._mask_selected_obj(mask) 

364 

365 

366class GroupByNthSelector: 

367 """ 

368 Dynamically substituted for GroupBy.nth to enable both call and index 

369 """ 

370 

371 def __init__(self, groupby_object: groupby.GroupBy) -> None: 

372 self.groupby_object = groupby_object 

373 

374 def __call__( 

375 self, 

376 n: PositionalIndexer | tuple, 

377 dropna: Literal["any", "all"] | None = None, 

378 ) -> DataFrame | Series: 

379 return self.groupby_object._nth(n, dropna) 

380 

381 def __getitem__(self, n: PositionalIndexer | tuple) -> DataFrame | Series: 

382 warnings.warn( 

383 "GroupBy.nth[...] is deprecated and will be removed in a future " 

384 "version of pandas. Use GroupBy.nth(...) instead.", 

385 Pandas4Warning, 

386 stacklevel=find_stack_level(), 

387 ) 

388 return self.groupby_object._nth(n)