74 lines
2.8 KiB
Python
74 lines
2.8 KiB
Python
from typing import Any, Dict
|
|
import pandas as pd
|
|
|
|
class PandasLocRowMultiIndexDataFrame:
|
|
"""
|
|
Pandas Loc Row MultiIndex DataFrame:
|
|
Selects a row or rows from a pandas DataFrame with MultiIndex and output as a DataFrame.
|
|
|
|
Specify comma-separated MultiIndex labels (row index) and the data type for the MultiIndex to select a row.
|
|
|
|
Below screenshots illustrate these options:
|
|
|
|
|
|
The top DataFrame has an integer-type MultiIndex, while the bottom DataFrame has a string-type MultiIndex. In both cases, two index values are specified to select a single row while the correct MultiIndex data type is also selected. Note that mixed data types for MultiIndex are not currently supported.
|
|
|
|
category: Data subset selection
|
|
"""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls) -> Dict[str, Any]:
|
|
"""
|
|
Defines the input types for the function.
|
|
|
|
Returns:
|
|
Dict[str, Any]: A dictionary specifying required input types.
|
|
"""
|
|
return {
|
|
"required": {
|
|
"dataframe": ("DATAFRAME", {}),
|
|
"row_index": ("STRING", {"default": ""}),
|
|
"row_index_type":(("string", "int"),)
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES: tuple = ("DATAFRAME",)
|
|
FUNCTION: str = "f"
|
|
CATEGORY: str = "Data Analysis"
|
|
|
|
def f(self, dataframe: pd.DataFrame,
|
|
row_index: str,
|
|
row_index_type: str) -> tuple:
|
|
"""
|
|
Selects a specific row using a comma-separated MultiIndex labels (row index) from a pandas DataFrame, and returns the DataFrame.
|
|
|
|
Args:
|
|
dataframe (DataFrame): The DataFrame.
|
|
row_index (str): A comma-separated MultiIndex label for the row.
|
|
row_index_type (str): The data type of the row index.
|
|
|
|
Returns:
|
|
tuple: A tuple containing the Pandas DataFrame.
|
|
"""
|
|
# We need to pass a tuple of MultiIndex.
|
|
# e.g. if the first index is A, second is B, we need to pass:
|
|
# ("A", "B")
|
|
ind_list = [ind.strip() for ind in row_index.split(",")]
|
|
|
|
# This will throw an exception if the MultiIndex is not an int.
|
|
if row_index_type == "int":
|
|
ind_list = [int(ind) for ind in ind_list]
|
|
|
|
ind_tuple = tuple(ind_list) # Convert the list to a tuple
|
|
# Ensure partial indexing is supported by adding slice(None) for missing levels
|
|
if len(ind_tuple) < dataframe.index.nlevels:
|
|
filler_count = dataframe.index.nlevels - len(ind_tuple)
|
|
for _ in range(filler_count):
|
|
ind_tuple += (slice(None),)
|
|
|
|
result = dataframe.loc[ind_tuple]
|
|
if isinstance(result, pd.Series): # loc can return Series so convert to DataFrame
|
|
result = result.to_frame().T
|
|
|
|
return (result,)
|