|
12 | 12 |
|
13 | 13 | from pandas._libs.tslibs import BaseOffset, to_offset |
14 | 14 | import pandas._libs.window.aggregations as window_aggregations |
15 | | -from pandas._typing import ArrayLike, Axis, FrameOrSeries, Scalar |
| 15 | +from pandas._typing import ArrayLike, Axis, FrameOrSeries, Label |
16 | 16 | from pandas.compat._optional import import_optional_dependency |
17 | 17 | from pandas.compat.numpy import function as nv |
18 | 18 | from pandas.util._decorators import Appender, Substitution, cache_readonly, doc |
@@ -381,21 +381,31 @@ def _wrap_result(self, result, block=None, obj=None): |
381 | 381 | return type(obj)(result, index=index, columns=block.columns) |
382 | 382 | return result |
383 | 383 |
|
384 | | - def _wrap_results(self, results, blocks, obj, exclude=None) -> FrameOrSeries: |
| 384 | + def _wrap_results(self, results, obj, skipped: List[int]) -> FrameOrSeries: |
385 | 385 | """ |
386 | 386 | Wrap the results. |
387 | 387 |
|
388 | 388 | Parameters |
389 | 389 | ---------- |
390 | 390 | results : list of ndarrays |
391 | | - blocks : list of blocks |
392 | 391 | obj : conformed data (may be resampled) |
393 | | - exclude: list of columns to exclude, default to None |
| 392 | + skipped: List[int] |
| 393 | + Indices of blocks that are skipped. |
394 | 394 | """ |
395 | 395 | from pandas import Series, concat |
396 | 396 |
|
| 397 | + exclude: List[Label] = [] |
| 398 | + if obj.ndim == 2: |
| 399 | + orig_blocks = list(obj._to_dict_of_blocks(copy=False).values()) |
| 400 | + for i in skipped: |
| 401 | + exclude.extend(orig_blocks[i].columns) |
| 402 | + else: |
| 403 | + orig_blocks = [obj] |
| 404 | + |
| 405 | + kept_blocks = [blk for i, blk in enumerate(orig_blocks) if i not in skipped] |
| 406 | + |
397 | 407 | final = [] |
398 | | - for result, block in zip(results, blocks): |
| 408 | + for result, block in zip(results, kept_blocks): |
399 | 409 |
|
400 | 410 | result = self._wrap_result(result, block=block, obj=obj) |
401 | 411 | if result.ndim == 1: |
@@ -491,24 +501,21 @@ def _apply_blockwise( |
491 | 501 |
|
492 | 502 | skipped: List[int] = [] |
493 | 503 | results: List[ArrayLike] = [] |
494 | | - exclude: List[Scalar] = [] |
495 | 504 | for i, b in enumerate(blocks): |
496 | 505 | try: |
497 | 506 | values = self._prep_values(b.values) |
498 | 507 |
|
499 | 508 | except (TypeError, NotImplementedError) as err: |
500 | 509 | if isinstance(obj, ABCDataFrame): |
501 | 510 | skipped.append(i) |
502 | | - exclude.extend(b.columns) |
503 | 511 | continue |
504 | 512 | else: |
505 | 513 | raise DataError("No numeric types to aggregate") from err |
506 | 514 |
|
507 | 515 | result = homogeneous_func(values) |
508 | 516 | results.append(result) |
509 | 517 |
|
510 | | - block_list = [blk for i, blk in enumerate(blocks) if i not in skipped] |
511 | | - return self._wrap_results(results, block_list, obj, exclude) |
| 518 | + return self._wrap_results(results, obj, skipped) |
512 | 519 |
|
513 | 520 | def _apply( |
514 | 521 | self, |
@@ -1283,7 +1290,7 @@ def count(self): |
1283 | 1290 | ).sum() |
1284 | 1291 | results.append(result) |
1285 | 1292 |
|
1286 | | - return self._wrap_results(results, blocks, obj) |
| 1293 | + return self._wrap_results(results, obj, skipped=[]) |
1287 | 1294 |
|
1288 | 1295 | _shared_docs["apply"] = dedent( |
1289 | 1296 | r""" |
|
0 commit comments