141 lines
3.6 KiB
Python
141 lines
3.6 KiB
Python
from typing import Any, Dict
|
|
import pandas as pd
|
|
from .utils import column_labels_string_to_list
|
|
|
|
class PandasPivot:
|
|
"""
|
|
Computes a pivot table from a pandas DataFrame.
|
|
|
|
## Note on a pivot table
|
|
A pivot table is easier to understand if we start off from the table itself:
|
|
In the below DataFrame:
|
|
|
|
```
|
|
Column labels
|
|
X Y Z
|
|
Index
|
|
A 1 2 3
|
|
B 4 5 6
|
|
C 7 8 9
|
|
```
|
|
|
|
We have index `["A", "B", "C"]` and column labels `["X", "Y", "Z"]`.
|
|
Values in this DataFrame can be accessed by specifying the index and the column label.
|
|
For example:
|
|
|
|
```python
|
|
df.loc["A", "X"] # 1
|
|
df.loc["A", "Y"] # 2
|
|
```
|
|
|
|
Now, let's restructure the table by converting column labels into row values:
|
|
|
|
```
|
|
A X 1
|
|
A Y 2
|
|
A Z 3
|
|
B X 4
|
|
B Y 5
|
|
B Z 6
|
|
C X 7
|
|
C Y 8
|
|
C Z 9
|
|
```
|
|
|
|
Let's add a column label to each of these columns:
|
|
|
|
```
|
|
R C V
|
|
--------
|
|
A X 1
|
|
A Y 2
|
|
A Z 3
|
|
B X 4
|
|
B Y 5
|
|
B Z 6
|
|
C X 7
|
|
C Y 8
|
|
C Z 9
|
|
```
|
|
|
|
Now, the first two columns uniquely identify the value in the third column.
|
|
|
|
**Pivoting** is the reverse of this operation.
|
|
- You start with the table above and produce the original one.
|
|
- In pandas, this means telling it to use:
|
|
- R column: `A`, `B`, `C` as the index (row labels).
|
|
- C column: `X`, `Y`, `Z` as the column labels.
|
|
- V column: The numbers (`1, 2, 3, ...`) as the values for each cell indexed by `[A, B, or C]` and `[X, Y, or Z]`.
|
|
|
|
This can be done using:
|
|
|
|
```python
|
|
df.pivot(index="R", columns="C", values="V")
|
|
```
|
|
|
|
The result:
|
|
|
|
```
|
|
X Y Z
|
|
A 1 2 3
|
|
B 4 5 6
|
|
C 7 8 9
|
|
```
|
|
|
|
Now the data is back to its original, structured form.
|
|
|
|
|
|
category: Transformation
|
|
"""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls) -> Dict[str, Any]:
|
|
"""
|
|
Defines the input types for the `select_columns` function.
|
|
|
|
Returns:
|
|
Dict[str, Any]: A dictionary specifying required input types.
|
|
"""
|
|
return {
|
|
"required": {
|
|
"dataframe": ("DATAFRAME", {}),
|
|
"index": ("STRING", {"default": ""}),
|
|
"columns": ("STRING", {"default": ""}),
|
|
"values": ("STRING", {"default": ""}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES: tuple = ("DATAFRAME",)
|
|
FUNCTION: str = "f"
|
|
CATEGORY: str = "Data Analysis"
|
|
|
|
def f(self, dataframe: pd.DataFrame,
|
|
index: str,
|
|
columns: str,
|
|
values: str) -> tuple:
|
|
"""
|
|
Selects specific columns from a pandas DataFrame.
|
|
|
|
Args:
|
|
dataframe (DataFrame): The DataFrame.
|
|
index (str): A comma-separated string of column names to use as index in the pivot table.
|
|
columns (str): A comma-separated string of column names to use as columns in the pivot table.
|
|
values (str): A comma-separated string of column names to use as values in the pivot table.
|
|
|
|
Returns:
|
|
tuple: A tuple containing the DataFrame.
|
|
"""
|
|
ind = column_labels_string_to_list(dataframe, index)
|
|
col = column_labels_string_to_list(dataframe, columns)
|
|
val = column_labels_string_to_list(dataframe, values)
|
|
|
|
if len(ind) == 1:
|
|
ind = ind[0]
|
|
if len(col) == 1:
|
|
col = col
|
|
if len(val) == 1:
|
|
val = val[0]
|
|
|
|
df2 = dataframe.pivot(index=ind, columns=col, values=val)
|
|
return (df2,)
|