Added support for integer column label in Pandas Select Columns node.

modules/pandas_wrapper/pandas_select_columns.py

@@ -1,6 +1,6 @@
 from typing import Any, Dict
-from io import StringIO
 import pandas as pd
+from .utils import column_labels_string_to_list
 
 class PandasSelectColumns:
     """
@@ -8,7 +8,7 @@ class PandasSelectColumns:
 
     category: Data subset selection
     """
-    
+
     @classmethod
     def INPUT_TYPES(cls) -> Dict[str, Any]:
         """
@@ -25,10 +25,10 @@ class PandasSelectColumns:
         }
 
     RETURN_TYPES: tuple = ("DATAFRAME",)
-    FUNCTION: str = "select_columns"
+    FUNCTION: str = "f"
     CATEGORY: str = "Data Analysis"
 
-    def select_columns(self, dataframe: pd.DataFrame, column_names: str) -> tuple:
+    def f(self, dataframe: pd.DataFrame, column_names: str) -> tuple:
         """
         Selects specific columns from a pandas DataFrame.
 
@@ -39,6 +39,6 @@ class PandasSelectColumns:
         Returns:
             tuple: A tuple containing the DataFrame with selected columns.
         """
-        selected_columns = [col.strip() for col in column_names.split(",")]
+        selected_columns = column_labels_string_to_list(dataframe, column_names)
         df2 = dataframe[selected_columns]
         return (df2,)

modules/pandas_wrapper/utils.py

@@ -1,11 +1,51 @@
+from typing import List, Union
 import json
 import numpy as np
 import pandas as pd
+import pandas as pd
+
+
+def column_labels_string_to_list(dataframe: pd.DataFrame, s: str) -> List[Union[str, int]]:
+    """
+    Converts a comma-separated string of column labels into a list of valid column names for a Pandas DataFrame.
+
+    This function ensures that each label in the input string corresponds to an existing column in the DataFrame.
+    If a label is a numeric string and corresponds to an integer column name, it will be converted to an integer.
+    If a column name does not exist in the DataFrame, a ValueError is raised.
+
+    Args:
+        dataframe (pd.DataFrame): The DataFrame to validate column names against.
+        s (str): A comma-separated string representing column labels.
+
+    Returns:
+        List[Union[str, int]]: A list of valid column labels.
+
+    Raises:
+        ValueError: If any label in `s` does not match an existing column name.
+    """
+    columns = [col.strip() for col in s.split(",")]
+    valid_columns = set(dataframe.columns)  # Use a set for faster lookups
+    output_columns = []
+
+    for c in columns:
+        if c in valid_columns:
+            output_columns.append(c)
+        else:
+            try:
+                numeric_c = int(c)
+                if numeric_c in valid_columns:
+                    output_columns.append(numeric_c)
+                else:
+                    raise ValueError(f"Column '{c}' not found in the DataFrame.")
+            except ValueError:
+                raise ValueError(f"Column '{c}' not found in the DataFrame.")
+
+    return output_columns
 
 
 def series_to_jsons(series: pd.Series) -> str:
     """
-    Convert a pandas Series to a JSON string.
+    Converts a pandas Series to a JSON string.
 
     Args:
         series (pd.Series): The pandas Series to convert.