| 664 | |
| 665 | |
| 666 | class PivotTable(ApplyConcatApply): |
| 667 | _parameters = ["frame", "columns", "index", "values", "aggfunc"] |
| 668 | _defaults = {"columns": None, "index": None, "values": None, "aggfunc": "mean"} |
| 669 | |
| 670 | @functools.cached_property |
| 671 | def _meta(self): |
| 672 | df = self.frame._meta |
| 673 | columns = self.operand("columns") |
| 674 | values = self.operand("values") |
| 675 | index = self.operand("index") |
| 676 | columns_contents = pd.CategoricalIndex(df[columns].cat.categories, name=columns) |
| 677 | |
| 678 | if is_scalar(values): |
| 679 | new_columns = columns_contents |
| 680 | else: |
| 681 | new_columns = pd.MultiIndex.from_product( |
| 682 | (sorted(values), columns_contents), names=[None, columns] |
| 683 | ) |
| 684 | |
| 685 | if self.operand("aggfunc") in ["first", "last"]: |
| 686 | # Infer datatype as non-numeric values are allowed |
| 687 | if is_scalar(values): |
| 688 | meta = pd.DataFrame( |
| 689 | columns=new_columns, |
| 690 | dtype=df[values].dtype, |
| 691 | index=pd.Index(df[index]), |
| 692 | ) |
| 693 | else: |
| 694 | meta = pd.DataFrame( |
| 695 | columns=new_columns, |
| 696 | index=pd.Index(df[index]), |
| 697 | ) |
| 698 | for value_col in values: |
| 699 | meta[value_col] = meta[value_col].astype( |
| 700 | df[values].dtypes[value_col] |
| 701 | ) |
| 702 | else: |
| 703 | # Use float64 as other aggregate functions require numerical data |
| 704 | meta = pd.DataFrame( |
| 705 | columns=new_columns, dtype=np.float64, index=pd.Index(df[index]) |
| 706 | ) |
| 707 | return meta |
| 708 | |
| 709 | def _lower(self): |
| 710 | args = [ |
| 711 | self.frame, |
| 712 | self.operand("columns"), |
| 713 | self.operand("index"), |
| 714 | self.operand("values"), |
| 715 | ] |
| 716 | if self.aggfunc == "sum": |
| 717 | return PivotTableSum(*args) |
| 718 | elif self.aggfunc == "mean": |
| 719 | return PivotTableSum(*args) / PivotTableCount(*args) |
| 720 | elif self.aggfunc == "count": |
| 721 | return PivotTableCount(*args) |
| 722 | elif self.aggfunc == "first": |
| 723 | return PivotTableFirst(*args) |