Perform groupby with a standard numerical aggregation function (e.g. mean) with Numba.
(
self,
func: Callable,
dtype_mapping: dict[np.dtype, Any],
engine_kwargs: dict[str, bool] | None,
**aggregator_kwargs,
)
| 1359 | ) |
| 1360 | |
| 1361 | def _numba_agg_general( |
| 1362 | self, |
| 1363 | func: Callable, |
| 1364 | dtype_mapping: dict[np.dtype, Any], |
| 1365 | engine_kwargs: dict[str, bool] | None, |
| 1366 | **aggregator_kwargs, |
| 1367 | ): |
| 1368 | """ |
| 1369 | Perform groupby with a standard numerical aggregation function (e.g. mean) |
| 1370 | with Numba. |
| 1371 | """ |
| 1372 | if not self.as_index: |
| 1373 | raise NotImplementedError( |
| 1374 | "as_index=False is not supported. Use .reset_index() instead." |
| 1375 | ) |
| 1376 | |
| 1377 | data = self._obj_with_exclusions |
| 1378 | df = data if data.ndim == 2 else data.to_frame() |
| 1379 | |
| 1380 | aggregator = executor.generate_shared_aggregator( |
| 1381 | func, |
| 1382 | dtype_mapping, |
| 1383 | True, # is_grouped_kernel |
| 1384 | **get_jit_arguments(engine_kwargs), |
| 1385 | ) |
| 1386 | # Pass group ids to kernel directly if it can handle it |
| 1387 | # (This is faster since it doesn't require a sort) |
| 1388 | ids = self._grouper.ids |
| 1389 | ngroups = self._grouper.ngroups |
| 1390 | |
| 1391 | res_mgr = df._mgr.apply( |
| 1392 | aggregator, labels=ids, ngroups=ngroups, **aggregator_kwargs |
| 1393 | ) |
| 1394 | res_mgr.axes[1] = self._grouper.result_index |
| 1395 | result = df._constructor_from_mgr(res_mgr, axes=res_mgr.axes) |
| 1396 | |
| 1397 | if data.ndim == 1: |
| 1398 | result = result.squeeze("columns") |
| 1399 | result.name = data.name |
| 1400 | else: |
| 1401 | result.columns = data.columns |
| 1402 | return result |
| 1403 | |
| 1404 | @final |
| 1405 | def _transform_with_numba(self, func, *args, engine_kwargs=None, **kwargs): |