feat: Add DataFrameInput and refactor DataSet references (#4898)

* refactor: rename DataSet to DataFrame

* feat: add DataFrameInput type and update InputTypes definition

* feat: add DataFrame type to CUSTOM_COMPONENT_SUPPORTED_TYPES

* refactor: replace DataSet with DataFrame in schema and tests

src/backend/base/langflow/field_typing/constants.py

@@ -19,6 +19,7 @@ from langchain_core.vectorstores import VectorStore, VectorStoreRetriever
 from langchain_text_splitters import TextSplitter
 
 from langflow.schema.data import Data
+from langflow.schema.dataframe import DataFrame
 from langflow.schema.message import Message
 
 NestedDict: TypeAlias = dict[str, str | dict]
@@ -77,4 +78,5 @@ CUSTOM_COMPONENT_SUPPORTED_TYPES = {
     "Callable": Callable,
     "LanguageModel": LanguageModel,
     "Retriever": Retriever,
+    "DataFrame": DataFrame,
 }

src/backend/base/langflow/inputs/inputs.py

@@ -1,6 +1,6 @@
 import warnings
 from collections.abc import AsyncIterator, Iterator
-from typing import Any, get_args
+from typing import Any, TypeAlias, get_args
 
 from pandas import DataFrame
 from pydantic import Field, field_validator
@@ -78,6 +78,10 @@ class DataInput(HandleInput, InputTraceMixin, ListableInputMixin, ToolModeMixin)
     input_types: list[str] = ["Data"]
 
 
+class DataFrameInput(HandleInput, InputTraceMixin, ListableInputMixin, ToolModeMixin):
+    input_types: list[str] = ["DataFrame"]
+
+
 class PromptInput(BaseInputMixin, ListableInputMixin, InputTraceMixin, ToolModeMixin):
     field_type: SerializableFieldTypes = FieldTypes.PROMPT
 
@@ -492,7 +496,7 @@ class DefaultPromptField(Input):
     value: Any = ""  # Set the value to empty string
 
 
-InputTypes = (
+InputTypes: TypeAlias = (
     Input
     | DefaultPromptField
     | BoolInput
@@ -516,6 +520,7 @@ InputTypes = (
     | TableInput
     | LinkInput
     | SliderInput
+    | DataFrameInput
 )
 
 InputTypesMap: dict[str, type[InputTypes]] = {t.__name__: t for t in get_args(InputTypes)}

src/backend/base/langflow/schema/__init__.py

@@ -1,6 +1,6 @@
 from .data import Data
-from .data_set import DataSet
+from .dataframe import DataFrame
 from .dotdict import dotdict
 from .message import Message
 
-__all__ = ["Data", "dotdict", "Message", "DataSet"]
+__all__ = ["Data", "dotdict", "Message", "DataFrame"]

src/backend/base/langflow/schema/dataframe.py

@@ -1,11 +1,12 @@
 from typing import cast
 
 import pandas as pd
+from pandas import DataFrame as pandas_DataFrame
 
 from langflow.schema.data import Data
 
 
-class DataSet(pd.DataFrame):
+class DataFrame(pandas_DataFrame):
     """A pandas DataFrame subclass specialized for handling collections of Data objects.
 
     This class extends pandas.DataFrame to provide seamless integration between
@@ -22,13 +23,13 @@ class DataSet(pd.DataFrame):
 
     Examples:
         >>> # From Data objects
-        >>> dataset = DataSet([Data(data={"name": "John"}), Data(data={"name": "Jane"})])
+        >>> dataset = DataFrame([Data(data={"name": "John"}), Data(data={"name": "Jane"})])
 
         >>> # From dictionaries
-        >>> dataset = DataSet([{"name": "John"}, {"name": "Jane"}])
+        >>> dataset = DataFrame([{"name": "John"}, {"name": "Jane"}])
 
         >>> # From dictionary of lists
-        >>> dataset = DataSet({"name": ["John", "Jane"], "age": [30, 25]})
+        >>> dataset = DataFrame({"name": ["John", "Jane"], "age": [30, 25]})
     """
 
     def __init__(self, data: None | list[dict | Data] | dict | pd.DataFrame = None, **kwargs):
@@ -49,36 +50,36 @@ class DataSet(pd.DataFrame):
         super().__init__(**kwargs)
 
     def to_data_list(self) -> list[Data]:
-        """Converts the DataSet back to a list of Data objects."""
+        """Converts the DataFrame back to a list of Data objects."""
         list_of_dicts = self.to_dict(orient="records")
         return [Data(data=row) for row in list_of_dicts]
 
-    def add_row(self, data: dict | Data) -> "DataSet":
+    def add_row(self, data: dict | Data) -> "DataFrame":
         """Adds a single row to the dataset.
 
         Args:
             data: Either a Data object or a dictionary to add as a new row
 
         Returns:
-            DataSet: A new DataSet with the added row
+            DataFrame: A new DataFrame with the added row
 
         Example:
-            >>> dataset = DataSet([{"name": "John"}])
+            >>> dataset = DataFrame([{"name": "John"}])
             >>> dataset = dataset.add_row({"name": "Jane"})
         """
         if isinstance(data, Data):
             data = data.data
         new_df = self._constructor([data])
-        return cast(DataSet, pd.concat([self, new_df], ignore_index=True))
+        return cast(DataFrame, pd.concat([self, new_df], ignore_index=True))
 
-    def add_rows(self, data: list[dict | Data]) -> "DataSet":
+    def add_rows(self, data: list[dict | Data]) -> "DataFrame":
         """Adds multiple rows to the dataset.
 
         Args:
             data: List of Data objects or dictionaries to add as new rows
 
         Returns:
-            DataSet: A new DataSet with the added rows
+            DataFrame: A new DataFrame with the added rows
         """
         processed_data = []
         for item in data:
@@ -87,11 +88,11 @@ class DataSet(pd.DataFrame):
             else:
                 processed_data.append(item)
         new_df = self._constructor(processed_data)
-        return cast(DataSet, pd.concat([self, new_df], ignore_index=True))
+        return cast(DataFrame, pd.concat([self, new_df], ignore_index=True))
 
     @property
     def _constructor(self):
         def _c(*args, **kwargs):
-            return DataSet(*args, **kwargs).__finalize__(self)
+            return DataFrame(*args, **kwargs).__finalize__(self)
 
         return _c