From 1a057571c9695ece272a07377a1751e650893786 Mon Sep 17 00:00:00 2001 From: Gabriel Luiz Freitas Almeida Date: Tue, 26 Nov 2024 10:04:44 -0300 Subject: [PATCH] feat: introduces DataSet class to improve the experience with lists of Data (#4834) * Add DataSet class to handle conversion between Data objects and DataFrame * feat: enhance DataSet class with detailed docstrings and examples for better usability * feat: add custom constructor property to DataSet for improved DataFrame compatibility * test: add unit tests for DataSet class methods and functionality --- src/backend/base/langflow/schema/data_set.py | 103 ++++++++++++ .../tests/unit/schema/test_schema_data_set.py | 159 ++++++++++++++++++ 2 files changed, 262 insertions(+) create mode 100644 src/backend/base/langflow/schema/data_set.py create mode 100644 src/backend/tests/unit/schema/test_schema_data_set.py diff --git a/src/backend/base/langflow/schema/data_set.py b/src/backend/base/langflow/schema/data_set.py new file mode 100644 index 000000000..449ceead0 --- /dev/null +++ b/src/backend/base/langflow/schema/data_set.py @@ -0,0 +1,103 @@ +import pandas as pd + +from langflow.schema.data import Data + + +class DataSet(pd.DataFrame): + """A pandas DataFrame subclass specialized for handling collections of Data objects. + + This class extends pandas.DataFrame to provide seamless integration between + Langflow's Data objects and pandas' powerful data manipulation capabilities. + + Key Features: + - Direct initialization from a list of Data objects + - Maintains all pandas DataFrame functionality + - Conversion back to Data objects when needed + + Notes: + - Nested dictionaries within Data objects are preserved in their column representation + - All pandas DataFrame operations (groupby, merge, concat, etc.) remain available + - Column dtypes are inferred from the Data objects' contents + + Examples: + >>> data_objects = [ + ... Data(data={"name": "John", "age": 30}), + ... Data(data={"name": "Jane", "age": 25}) + ... ] + >>> dataset = DataSet.from_data_list(data_objects) + >>> dataset['age'].mean() + 27.5 + >>> original_data = dataset.to_data_list() + + Inheritance: + This class inherits all functionality from pandas.DataFrame, meaning any + operation that works on a DataFrame will work on a DataSet: + - Filtering: dataset[dataset['age'] > 25] + - Aggregation: dataset.groupby('category').mean() + - Statistical operations: dataset.describe() + - etc. + """ + + @classmethod + def from_data_list(cls, data_list: list[Data]) -> "DataSet": + """Creates a DataSet from a list of Data objects. + + This method converts a list of Data objects into a DataFrame structure, + preserving all data from the original Data objects. + + Args: + data_list (list[Data]): A list of Data objects to convert into a DataFrame. + Each Data object's internal dictionary becomes a row in the DataFrame. + + Returns: + DataSet: A new DataSet instance containing all data from the input list. + + Examples: + >>> data_objects = [ + ... Data(data={"name": "John", "age": 30}), + ... Data(data={"name": "Jane", "age": 25}) + ... ] + >>> dataset = DataSet.from_data_list(data_objects) + >>> print(dataset.columns) + Index(['name', 'age'], dtype='object') + + Notes: + - Column names are derived from the keys in the Data objects + - If Data objects have different keys, the resulting DataFrame will have + NaN values for missing data + - The original structure of nested data is preserved in the DataFrame + """ + data_dicts = [d.data for d in data_list] + return cls(data_dicts) + + def to_data_list(self) -> list[Data]: + """Converts the DataSet back to a list of Data objects. + + This method transforms each row of the DataFrame back into a Data object, + reconstructing the original data structure. + + Returns: + list[Data]: A list of Data objects, where each object corresponds to + a row in the DataFrame. + + Examples: + >>> dataset = DataSet({'name': ['John'], 'age': [30]}) + >>> data_objects = dataset.to_data_list() + >>> print(data_objects[0].data) + {'name': 'John', 'age': 30} + + Notes: + - Each row is converted to a dictionary using to_dict() + - The resulting Data objects will contain all columns as keys in their + internal dictionary + - Any modifications made to the DataFrame will be reflected in the + resulting Data objects + """ + return [Data(data=row.to_dict()) for _, row in self.iterrows()] + + @property + def _constructor(self): + def _c(*args, **kwargs): + return DataSet(*args, **kwargs).__finalize__(self) + + return _c diff --git a/src/backend/tests/unit/schema/test_schema_data_set.py b/src/backend/tests/unit/schema/test_schema_data_set.py new file mode 100644 index 000000000..a40eee69a --- /dev/null +++ b/src/backend/tests/unit/schema/test_schema_data_set.py @@ -0,0 +1,159 @@ +import pandas as pd +import pytest +from langflow.schema.data import Data +from langflow.schema.data_set import DataSet + + +@pytest.fixture +def sample_data_objects() -> list[Data]: + """Fixture providing a list of sample Data objects.""" + return [ + Data(data={"name": "John", "age": 30, "city": "New York"}), + Data(data={"name": "Jane", "age": 25, "city": "Boston"}), + Data(data={"name": "Bob", "age": 35, "city": "Chicago"}), + ] + + +@pytest.fixture +def sample_dataset(sample_data_objects) -> DataSet: + """Fixture providing a sample DataSet instance.""" + return DataSet.from_data_list(sample_data_objects) + + +def test_from_data_list_basic(): + """Test basic functionality of from_data_list.""" + data_objects = [Data(data={"name": "John", "age": 30}), Data(data={"name": "Jane", "age": 25})] + dataset = DataSet.from_data_list(data_objects) + + assert isinstance(dataset, DataSet) + assert isinstance(dataset, pd.DataFrame) + assert len(dataset) == 2 + assert list(dataset.columns) == ["name", "age"] + assert dataset.iloc[0]["name"] == "John" + assert dataset.iloc[1]["age"] == 25 + + +def test_from_data_list_empty(): + """Test from_data_list with empty input.""" + dataset = DataSet.from_data_list([]) + assert isinstance(dataset, DataSet) + assert len(dataset) == 0 + + +def test_from_data_list_missing_fields(): + """Test from_data_list with inconsistent data fields.""" + data_objects = [ + Data(data={"name": "John", "age": 30}), + Data(data={"name": "Jane", "city": "Boston"}), # Missing age + ] + dataset = DataSet.from_data_list(data_objects) + + assert isinstance(dataset, DataSet) + assert set(dataset.columns) == {"name", "age", "city"} + assert pd.isna(dataset.iloc[1]["age"]) + assert pd.isna(dataset.iloc[0]["city"]) + + +def test_from_data_list_nested_data(): + """Test from_data_list with nested dictionary data.""" + data_objects = [ + Data(data={"name": "John", "address": {"city": "New York", "zip": "10001"}}), + Data(data={"name": "Jane", "address": {"city": "Boston", "zip": "02108"}}), + ] + dataset = DataSet.from_data_list(data_objects) + + assert isinstance(dataset, DataSet) + assert isinstance(dataset["address"][0], dict) + assert dataset["address"][0]["city"] == "New York" + + +def test_to_data_list_basic(sample_dataset, sample_data_objects): + """Test basic functionality of to_data_list.""" + result = sample_dataset.to_data_list() + + assert isinstance(result, list) + assert all(isinstance(item, Data) for item in result) + assert len(result) == len(sample_data_objects) + + # Check if data is preserved + for original, converted in zip(sample_data_objects, result, strict=False): + assert original.data == converted.data + + +def test_to_data_list_empty(): + """Test to_data_list with empty DataFrame.""" + empty_dataset = DataSet() + result = empty_dataset.to_data_list() + assert isinstance(result, list) + assert len(result) == 0 + + +def test_to_data_list_modified_data(sample_dataset): + """Test to_data_list after DataFrame modifications.""" + # Modify the dataset + sample_dataset["new_column"] = [1, 2, 3] + sample_dataset.iloc[0, sample_dataset.columns.get_loc("age")] = 31 + + result = sample_dataset.to_data_list() + + assert isinstance(result, list) + assert all(isinstance(item, Data) for item in result) + assert result[0].data["new_column"] == 1 + assert result[0].data["age"] == 31 + + +def test_dataset_pandas_operations(sample_dataset): + """Test that pandas operations work correctly on DataSet.""" + # Test filtering + filtered = sample_dataset[sample_dataset["age"] > 30] + assert isinstance(filtered, DataSet), f"Expected DataSet, got {type(filtered)}" + assert len(filtered) == 1 + assert filtered.iloc[0]["name"] == "Bob" + + # Test aggregation + mean_age = sample_dataset["age"].mean() + assert mean_age == 30 + + # Test groupby + grouped = sample_dataset.groupby("city").agg({"age": "mean"}) + assert isinstance(grouped, pd.DataFrame) + assert len(grouped) == 3 + + +def test_dataset_with_null_values(): + """Test handling of null values in DataSet.""" + data_objects = [Data(data={"name": "John", "age": None}), Data(data={"name": None, "age": 25})] + dataset = DataSet.from_data_list(data_objects) + + assert pd.isna(dataset.iloc[0]["age"]) + assert pd.isna(dataset.iloc[1]["name"]) + + # Test that null values are preserved when converting back + result = dataset.to_data_list() + assert pd.isna(result[0].data["age"]), f"Expected NaN, got {result[0].data['age']}" + assert pd.isna(result[1].data["name"]), f"Expected NaN, got {result[1].data['name']}" + + +def test_dataset_type_preservation(): + """Test that data types are preserved through conversion.""" + data_objects = [ + Data( + data={ + "int_val": 1, + "float_val": 1.5, + "str_val": "test", + "bool_val": True, + "list_val": [1, 2, 3], + "dict_val": {"key": "value"}, + } + ) + ] + dataset = DataSet.from_data_list(data_objects) + result = dataset.to_data_list() + + assert isinstance(result[0].data["int_val"], int) + assert isinstance(result[0].data["float_val"], float) + assert isinstance(result[0].data["str_val"], str) + assert isinstance(result[0].data["bool_val"], bool) + assert isinstance(result[0].data["list_val"], list) + assert isinstance(result[0].data["dict_val"], dict)