From 9591b6a10c5792d29d277bcf748a50dd7b5ee549 Mon Sep 17 00:00:00 2001 From: Amine Lemaizi <40530011+aminelemaizi@users.noreply.github.com> Date: Sun, 24 Aug 2025 16:59:03 +0100 Subject: [PATCH 1/2] wip --- coded-flows/coded_flows/utils/media.py | 34 +++++++++++---- coded-flows/tests/test_data_save.py | 57 ++++++++++++++++++++++++++ 2 files changed, 82 insertions(+), 9 deletions(-) diff --git a/coded-flows/coded_flows/utils/media.py b/coded-flows/coded_flows/utils/media.py index a1f0770..575895e 100644 --- a/coded-flows/coded_flows/utils/media.py +++ b/coded-flows/coded_flows/utils/media.py @@ -147,27 +147,29 @@ def save_data_to_parquet( filename=None, ) -> str: + random_filename = f"cfdata_{filename if filename else uuid.uuid4().hex}.parquet" + temp_dir = os.path.join(tempfile.gettempdir(), "coded-flows-media") + os.makedirs(temp_dir, exist_ok=True) + file_path = os.path.join(temp_dir, random_filename) + if ( isinstance(data, pd.DataFrame) + or isinstance(data, pd.Series) or isinstance(data, pa.Table) or ( isinstance(data, list) and all(isinstance(item, dict) for item in data[:50]) ) ): - random_filename = f"cfdata_{filename if filename else uuid.uuid4().hex}.parquet" - temp_dir = os.path.join(tempfile.gettempdir(), "coded-flows-media") - os.makedirs(temp_dir, exist_ok=True) - file_path = os.path.join(temp_dir, random_filename) - try: - - # Handle different data types if isinstance(data, pd.DataFrame): data.to_parquet( file_path, row_group_size=50000, index=False, engine="pyarrow" ) - + elif isinstance(data, pd.Series): + data.to_frame().to_parquet( + file_path, row_group_size=50000, index=False, engine="pyarrow" + ) elif isinstance(data, pa.Table): pq.write_table(data, file_path, row_group_size=50000) @@ -180,4 +182,18 @@ def save_data_to_parquet( except Exception as e: raise Exception(f"❌ Error saving data to parquet: {str(e)}") else: - raise TypeError(f"Unsupported data type: {type(data)}") + + if isinstance(data, np.ndarray): + if data.ndim != 1: + raise ValueError(f"NumPy array must be one-dimensional.") + pd.DataFrame(data, columns=["value"]).to_parquet( + file_path, row_group_size=50000, index=False, engine="pyarrow" + ) + elif isinstance(data, list): + pd.DataFrame(data, columns=["value"]).to_parquet( + file_path, row_group_size=50000, index=False, engine="pyarrow" + ) + else: + raise TypeError(f"Unsupported data type: {type(data)}") + + return file_path diff --git a/coded-flows/tests/test_data_save.py b/coded-flows/tests/test_data_save.py index 2732ab7..f26ebc3 100644 --- a/coded-flows/tests/test_data_save.py +++ b/coded-flows/tests/test_data_save.py @@ -253,3 +253,60 @@ def test_arrow_table_to_parquet(): assert table.equals(table_loaded), "Saved and loaded Arrow tables are not equal." os.remove(file_path) + + +def test_numerical_list_to_parquet(): + numerical_list = [1, 2, 3, 4, 5] + file_path = save_data_to_parquet(numerical_list) + + # Check file existence + assert os.path.exists(file_path), "Output file does not exist." + + # Read back from Parquet into DataFrame + df_loaded = pd.read_parquet(file_path, engine="pyarrow") + + # Convert original list to DataFrame for comparison + expected_df = pd.DataFrame( + {"value": numerical_list} + ) # Assuming single column with default name + + pdt.assert_frame_equal(expected_df.sort_index(axis=1), df_loaded.sort_index(axis=1)) + os.remove(file_path) + + +def test_numpy_array_to_parquet(): + import numpy as np + + numpy_array = np.array([10, 20, 30, 40, 50]) + file_path = save_data_to_parquet(numpy_array) + + # Check file existence + assert os.path.exists(file_path), "Output file does not exist." + + # Read back from Parquet into DataFrame + df_loaded = pd.read_parquet(file_path, engine="pyarrow") + + # Convert original numpy array to DataFrame for comparison + expected_df = pd.DataFrame( + {"value": numpy_array} + ) # Assuming single column with default name + + pdt.assert_frame_equal(expected_df.sort_index(axis=1), df_loaded.sort_index(axis=1)) + os.remove(file_path) + + +def test_pandas_series_to_parquet(): + pandas_series = pd.Series([100, 200, 300, 400, 500], name="values") + file_path = save_data_to_parquet(pandas_series) + + # Check file existence + assert os.path.exists(file_path), "Output file does not exist." + + # Read back from Parquet into DataFrame + df_loaded = pd.read_parquet(file_path, engine="pyarrow") + + # Convert original Series to DataFrame for comparison + expected_df = pandas_series.to_frame() + + pdt.assert_frame_equal(expected_df.sort_index(axis=1), df_loaded.sort_index(axis=1)) + os.remove(file_path) From f82d741d42fa2b393c44045e4a980b2c737fbec4 Mon Sep 17 00:00:00 2001 From: Amine Lemaizi <40530011+aminelemaizi@users.noreply.github.com> Date: Sun, 24 Aug 2025 23:07:27 +0100 Subject: [PATCH 2/2] fixing parquet saves & including polars to Datafrale and DataSeries --- coded-flows/coded_flows/types/__init__.py | 4 +- coded-flows/coded_flows/types/extra.py | 78 ++++++-- coded-flows/coded_flows/utils/converters.py | 56 +++++- coded-flows/poetry.lock | 46 ++++- coded-flows/pyproject.toml | 3 +- coded-flows/tests/test_converters.py | 189 +++++++++++++++++++- coded-flows/tests/test_extras.py | 40 +++++ 7 files changed, 387 insertions(+), 29 deletions(-) diff --git a/coded-flows/coded_flows/types/__init__.py b/coded-flows/coded_flows/types/__init__.py index 0871a67..5149964 100644 --- a/coded-flows/coded_flows/types/__init__.py +++ b/coded-flows/coded_flows/types/__init__.py @@ -488,8 +488,8 @@ def is_supported_type(element_type): "Any", "Null", # Data - "DataSeries", # <-- works as a Helper - "DataFrame", # <-- works as a Helper + "DataSeries", + "DataFrame", "ArrowTable", "NDArray", "DataDict", diff --git a/coded-flows/coded_flows/types/extra.py b/coded-flows/coded_flows/types/extra.py index 3694b9f..1de0e97 100644 --- a/coded-flows/coded_flows/types/extra.py +++ b/coded-flows/coded_flows/types/extra.py @@ -2,38 +2,86 @@ import base64 import pandas as pd import pyarrow as pa +import polars as pl from numpy import ndarray from pydantic import GetCoreSchemaHandler from pydantic_core import core_schema -from typing import Any, Type +from typing import Any, Type, Union, List from PIL import Image -class DataSeries(pd.Series): +class DataSeriesMeta(type): + + def __instancecheck__(cls, instance): + return isinstance(instance, (pd.Series, pl.Series)) + + +def serialize_series(series: Union[pd.Series, pl.Series]) -> List[Any]: + return series.to_list() + + +class DataSeries(metaclass=DataSeriesMeta): @classmethod def __get_pydantic_core_schema__( - cls, _source: Type[Any], _handler: GetCoreSchemaHandler + cls, + source: Type[Any], + handler: GetCoreSchemaHandler, ) -> core_schema.CoreSchema: + pandas_schema = core_schema.is_instance_schema(pd.Series) + polars_schema = core_schema.is_instance_schema(pl.Series) - return core_schema.is_instance_schema( - pd.Series, - serialization=core_schema.plain_serializer_function_ser_schema( - lambda instance: list(instance) - ), + union_schema = core_schema.union_schema([pandas_schema, polars_schema]) + + serialization = core_schema.plain_serializer_function_ser_schema( + serialize_series, when_used="json" + ) + + return core_schema.json_or_python_schema( + json_schema=union_schema, + python_schema=union_schema, + serialization=serialization, ) -class DataFrame(pd.DataFrame): +class DataFrameMeta(type): + def __instancecheck__(cls, instance): + return isinstance(instance, (pd.DataFrame, pl.DataFrame, pl.LazyFrame)) + + +def serialize_dataframe( + df: Union[pd.DataFrame, pl.DataFrame, pl.LazyFrame], +) -> list[dict[str, Any]]: + if isinstance(df, pd.DataFrame): + return df.to_dict(orient="records") + if isinstance(df, pl.DataFrame): + return df.to_dicts() + if isinstance(df, pl.LazyFrame): + return df.collect().to_dicts() + raise TypeError(f"Unsupported dataframe type: {type(df)}") + + +class DataFrame(metaclass=DataFrameMeta): @classmethod def __get_pydantic_core_schema__( - cls, _source: Type[Any], _handler: GetCoreSchemaHandler + cls, + _source: Type[Any], + _handler: GetCoreSchemaHandler, ) -> core_schema.CoreSchema: + pandas_schema = core_schema.is_instance_schema(pd.DataFrame) + polars_eager_schema = core_schema.is_instance_schema(pl.DataFrame) + polars_lazy_schema = core_schema.is_instance_schema(pl.LazyFrame) + union_schema = core_schema.union_schema( + [pandas_schema, polars_eager_schema, polars_lazy_schema] + ) - return core_schema.is_instance_schema( - pd.DataFrame, - serialization=core_schema.plain_serializer_function_ser_schema( - lambda instance: instance.to_dict(orient="records") - ), + serialization = core_schema.plain_serializer_function_ser_schema( + serialize_dataframe, when_used="json" + ) + + return core_schema.json_or_python_schema( + json_schema=union_schema, + python_schema=union_schema, + serialization=serialization, ) diff --git a/coded-flows/coded_flows/utils/converters.py b/coded-flows/coded_flows/utils/converters.py index 115b645..de3a2b2 100644 --- a/coded-flows/coded_flows/utils/converters.py +++ b/coded-flows/coded_flows/utils/converters.py @@ -7,6 +7,7 @@ from typing import Any, Callable, Union from pydantic_core import MultiHostUrl import pandas as pd +import polars as pl import pyarrow as pa from ..types import ( AnyUrl, @@ -260,28 +261,65 @@ def dataseries_to_type(output_type: str) -> Callable: elif output_type == "Set": return lambda x: set(x.to_list()) elif output_type == "Json": - return lambda x: x.to_json() + return lambda x: ( + json.dumps(x.to_list()) + if isinstance(x, pl.Series) + else x.to_json(orient="records") + ) elif output_type == "NDArray": return lambda x: x.to_numpy() elif output_type == "ArrowTable": - return lambda x: pa.Table.from_pandas(x.to_frame()) + return lambda x: ( + pa.Table.from_arrays([x.to_arrow()], names=[x.name if x.name else "value"]) + if isinstance(x, pl.Series) + else pa.Table.from_pandas(x.to_frame()) + ) + + +def _collect_if_lazy(df: pl.DataFrame | pl.LazyFrame | pd.DataFrame) -> pl.DataFrame: + return df.collect() if isinstance(df, pl.LazyFrame) else df def dataframe_to_type(output_type: str) -> Callable: + if output_type == "DataRecords": - return lambda x: x.to_dict("records") + return lambda x: ( + _collect_if_lazy(x).to_dicts() + if isinstance(_collect_if_lazy(x), pl.DataFrame) + else x.to_dict("records") + ) elif output_type == "List": - return lambda x: x.to_dict("records") + return lambda x: ( + _collect_if_lazy(x).to_dicts() + if isinstance(_collect_if_lazy(x), pl.DataFrame) + else x.to_dict("records") + ) elif output_type == "Dict": - return lambda x: x.to_dict("list") + return lambda x: ( + _collect_if_lazy(x).to_dict(as_series=False) + if isinstance(_collect_if_lazy(x), pl.DataFrame) + else x.to_dict("list") + ) elif output_type == "DataDict": - return lambda x: x.to_dict("list") + return lambda x: ( + _collect_if_lazy(x).to_dict(as_series=False) + if isinstance(_collect_if_lazy(x), pl.DataFrame) + else x.to_dict("list") + ) elif output_type == "Json": - return lambda x: x.to_json(orient="records") + return lambda x: ( + _collect_if_lazy(x).write_json() + if isinstance(_collect_if_lazy(x), pl.DataFrame) + else x.to_json(orient="records") + ) elif output_type == "NDArray": - return lambda x: x.to_numpy() + return lambda x: _collect_if_lazy(x).to_numpy() elif output_type == "ArrowTable": - return lambda x: pa.Table.from_pandas(x) + return lambda x: ( + _collect_if_lazy(x).to_arrow() + if isinstance(_collect_if_lazy(x), pl.DataFrame) + else pa.Table.from_pandas(x) + ) def arrow_to_type(output_type: str) -> Callable: diff --git a/coded-flows/poetry.lock b/coded-flows/poetry.lock index 4abe654..ce5c4cc 100644 --- a/coded-flows/poetry.lock +++ b/coded-flows/poetry.lock @@ -501,6 +501,50 @@ files = [ dev = ["pre-commit", "tox"] testing = ["coverage", "pytest", "pytest-benchmark"] +[[package]] +name = "polars" +version = "1.32.3" +description = "Blazingly fast DataFrame library" +optional = false +python-versions = ">=3.9" +groups = ["main"] +files = [ + {file = "polars-1.32.3-cp39-abi3-macosx_10_12_x86_64.whl", hash = "sha256:c7c472ea1d50a5104079cb64e34f78f85774bcc69b875ba8daf21233f4c70d42"}, + {file = "polars-1.32.3-cp39-abi3-macosx_11_0_arm64.whl", hash = "sha256:fd87275f0cc795e72a2030b58293198cfa748d4b009cf52218e27db5397ed07f"}, + {file = "polars-1.32.3-cp39-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:c9a9b9668ef310e5a77a7e7daa9c753874779c8da52e93f654bfd7953eb4b60b"}, + {file = "polars-1.32.3-cp39-abi3-manylinux_2_24_aarch64.whl", hash = "sha256:c8f5d2f43b80b68e39bfaa2948ce632563633466576f12e74e8560d6481f5851"}, + {file = "polars-1.32.3-cp39-abi3-win_amd64.whl", hash = "sha256:db56a7cb4898e173d62634e182f74bdff744c62be5470e0fe20df8d10f659af7"}, + {file = "polars-1.32.3-cp39-abi3-win_arm64.whl", hash = "sha256:a2e3f87c60f54eefe67b1bebd3105918d84df0fd6d59cc6b870c2f16d2d26ca1"}, + {file = "polars-1.32.3.tar.gz", hash = "sha256:57c500dc1b5cba49b0589034478db031815f3d57a20cb830b05ecee1a9ba56b1"}, +] + +[package.extras] +adbc = ["adbc-driver-manager[dbapi]", "adbc-driver-sqlite[dbapi]"] +all = ["polars[async,cloudpickle,database,deltalake,excel,fsspec,graph,iceberg,numpy,pandas,plot,pyarrow,pydantic,style,timezone]"] +async = ["gevent"] +calamine = ["fastexcel (>=0.9)"] +cloudpickle = ["cloudpickle"] +connectorx = ["connectorx (>=0.3.2)"] +database = ["polars[adbc,connectorx,sqlalchemy]"] +deltalake = ["deltalake (>=1.0.0)"] +excel = ["polars[calamine,openpyxl,xlsx2csv,xlsxwriter]"] +fsspec = ["fsspec"] +gpu = ["cudf-polars-cu12"] +graph = ["matplotlib"] +iceberg = ["pyiceberg (>=0.7.1)"] +numpy = ["numpy (>=1.16.0)"] +openpyxl = ["openpyxl (>=3.0.0)"] +pandas = ["pandas", "polars[pyarrow]"] +plot = ["altair (>=5.4.0)"] +polars-cloud = ["polars-cloud (>=0.0.1a1)"] +pyarrow = ["pyarrow (>=7.0.0)"] +pydantic = ["pydantic"] +sqlalchemy = ["polars[pandas]", "sqlalchemy"] +style = ["great-tables (>=0.8.0)"] +timezone = ["tzdata ; platform_system == \"Windows\""] +xlsx2csv = ["xlsx2csv (>=0.8.0)"] +xlsxwriter = ["xlsxwriter"] + [[package]] name = "pyarrow" version = "18.1.0" @@ -915,4 +959,4 @@ files = [ [metadata] lock-version = "2.1" python-versions = "^3.10" -content-hash = "9cb7231fb4ac47dd40a57373f28463b3641104ce24944a3980d1d7d1ae0647d6" +content-hash = "f6f5f6bc0bc8655ebe70bc538137c2997ff74f0bf9c5dbb420b41e8e3df00bf6" diff --git a/coded-flows/pyproject.toml b/coded-flows/pyproject.toml index ce61756..65f1db3 100644 --- a/coded-flows/pyproject.toml +++ b/coded-flows/pyproject.toml @@ -1,6 +1,6 @@ [tool.poetry] name = "coded-flows" -version = "0.6.1" +version = "0.6.2" description = "Various utilities for Coded Flows" authors = ["COLOR CODED CODES "] readme = "README.md" @@ -14,6 +14,7 @@ email-validator = "^2.0.0" pycountry = "^24.6.1" pillow = "^10.1.0" pydantic-extra-types = "^2.6.0" +polars = "^1.0.0" [tool.poetry.group.dev.dependencies] pytest = "^8.1.1" diff --git a/coded-flows/tests/test_converters.py b/coded-flows/tests/test_converters.py index a342b84..998177a 100644 --- a/coded-flows/tests/test_converters.py +++ b/coded-flows/tests/test_converters.py @@ -1,5 +1,6 @@ import pytest import pandas as pd +import polars as pl import pyarrow as pa import numpy as np from collections import deque @@ -56,6 +57,24 @@ def sample_series(): return pd.Series([1, 2, 3, 4, 5]) +@pytest.fixture +def sample_polars_dataframe(): + """Fixture to create a sample Polars DataFrame.""" + return pl.DataFrame({"A": [1, 2, 3], "B": [4, 5, 6], "C": [7, 8, 9]}) + + +@pytest.fixture +def sample_polars_series(): + """Fixture to create a sample Polars Series.""" + return pl.Series([1, 2, 3, 4, 5]) + + +@pytest.fixture +def sample_polars_lazyframe(): + """Fixture to create a sample Polars LazyFrame.""" + return pl.DataFrame({"A": [1, 2, 3], "B": [4, 5, 6], "C": [7, 8, 9]}).lazy() + + @pytest.fixture def sample_numpy_array(): """Fixture to create a sample numpy array.""" @@ -154,7 +173,7 @@ def test_dataseries_to_list(sample_series): assert isinstance(output, list) -def test_dataseries_to_list(sample_series): +def test_dataseries_to_tuple(sample_series): try: output = convert_type(sample_series, "DataSeries", "Tuple") except TypeError as e: @@ -194,6 +213,62 @@ def test_dataseries_to_arrow(sample_series): assert isinstance(output, pa.Table) +def test_pl_dataseries_to_dataframe(sample_polars_series): + try: + output = convert_type(sample_polars_series, "DataSeries", "DataFrame") + except TypeError as e: + pytest.fail(f"TypeError encountered for DataSeries -> Dataframe: {e}") + assert isinstance(output, pl.DataFrame) + + +def test_pl_dataseries_to_list(sample_polars_series): + try: + output = convert_type(sample_polars_series, "DataSeries", "List") + except TypeError as e: + pytest.fail(f"TypeError encountered for DataSeries -> List: {e}") + assert isinstance(output, list) + + +def test_pl_dataseries_to_tuple(sample_polars_series): + try: + output = convert_type(sample_polars_series, "DataSeries", "Tuple") + except TypeError as e: + pytest.fail(f"TypeError encountered for DataSeries -> Tuple: {e}") + assert isinstance(output, tuple) + + +def test_pl_dataseries_to_set(sample_polars_series): + try: + output = convert_type(sample_polars_series, "DataSeries", "Set") + except TypeError as e: + pytest.fail(f"TypeError encountered for DataSeries -> Set: {e}") + assert isinstance(output, set) + + +def test_pl_dataseries_to_json(sample_polars_series): + try: + output = convert_type(sample_polars_series, "DataSeries", "Json") + except TypeError as e: + pytest.fail(f"TypeError encountered for DataSeries -> Json: {e}") + assert isinstance(output, str) + + +def test_pl_dataseries_to_numpy(sample_polars_series): + try: + output = convert_type(sample_polars_series, "DataSeries", "NDArray") + except TypeError as e: + pytest.fail(f"TypeError encountered for DataSeries -> NDArray: {e}") + assert isinstance(output, np.ndarray) + + +def test_pl_dataseries_to_arrow(sample_polars_series): + try: + output = convert_type(sample_polars_series, "DataSeries", "ArrowTable") + except TypeError as e: + pytest.fail(f"TypeError encountered for DataSeries -> ArrowTable: {e}") + assert isinstance(output, pa.Table) + + def test_dataframe_to_datarecords(sample_dataframe): try: output = convert_type(sample_dataframe, "DataFrame", "DataRecords") @@ -250,6 +325,118 @@ def test_dataframe_to_arrow(sample_dataframe): assert isinstance(output, pa.Table) +def test_pl_dataframe_to_datarecords(sample_polars_dataframe): + try: + output = convert_type(sample_polars_dataframe, "DataFrame", "DataRecords") + except TypeError as e: + pytest.fail(f"TypeError encountered for DataFrame -> DataRecords: {e}") + assert isinstance(output, list) and all(isinstance(item, dict) for item in output) + + +def test_pl_dataframe_to_list(sample_polars_dataframe): + try: + output = convert_type(sample_polars_dataframe, "DataFrame", "List") + except TypeError as e: + pytest.fail(f"TypeError encountered for DataFrame -> List: {e}") + assert isinstance(output, list) + + +def test_pl_dataframe_to_dict(sample_polars_dataframe): + try: + output = convert_type(sample_polars_dataframe, "DataFrame", "Dict") + except TypeError as e: + pytest.fail(f"TypeError encountered for DataFrame -> Dict: {e}") + assert isinstance(output, dict) + + +def test_pl_dataframe_to_datadict(sample_polars_dataframe): + try: + output = convert_type(sample_polars_dataframe, "DataFrame", "DataDict") + except TypeError as e: + pytest.fail(f"TypeError encountered for DataFrame -> DataDict: {e}") + assert all(isinstance(k, str) and isinstance(v, list) for k, v in output.items()) + + +def test_pl_dataframe_to_json(sample_polars_dataframe): + try: + output = convert_type(sample_polars_dataframe, "DataFrame", "Json") + except TypeError as e: + pytest.fail(f"TypeError encountered for DataFrame -> Json: {e}") + assert isinstance(output, str) + + +def test_pl_dataframe_to_numpy(sample_polars_dataframe): + try: + output = convert_type(sample_polars_dataframe, "DataFrame", "NDArray") + except TypeError as e: + pytest.fail(f"TypeError encountered for DataFrame -> NDArray: {e}") + assert isinstance(output, np.ndarray) + + +def test_pl_dataframe_to_arrow(sample_polars_dataframe): + try: + output = convert_type(sample_polars_dataframe, "DataFrame", "ArrowTable") + except TypeError as e: + pytest.fail(f"TypeError encountered for DataFrame -> ArrowTable: {e}") + assert isinstance(output, pa.Table) + + +def test_pl_lazy_dataframe_to_datarecords(sample_polars_lazyframe): + try: + output = convert_type(sample_polars_lazyframe, "DataFrame", "DataRecords") + except TypeError as e: + pytest.fail(f"TypeError encountered for DataFrame -> DataRecords: {e}") + assert isinstance(output, list) and all(isinstance(item, dict) for item in output) + + +def test_pl_lazy_dataframe_to_list(sample_polars_lazyframe): + try: + output = convert_type(sample_polars_lazyframe, "DataFrame", "List") + except TypeError as e: + pytest.fail(f"TypeError encountered for DataFrame -> List: {e}") + assert isinstance(output, list) + + +def test_pl_lazy_dataframe_to_dict(sample_polars_lazyframe): + try: + output = convert_type(sample_polars_lazyframe, "DataFrame", "Dict") + except TypeError as e: + pytest.fail(f"TypeError encountered for DataFrame -> Dict: {e}") + assert isinstance(output, dict) + + +def test_pl_lazy_dataframe_to_datadict(sample_polars_lazyframe): + try: + output = convert_type(sample_polars_lazyframe, "DataFrame", "DataDict") + except TypeError as e: + pytest.fail(f"TypeError encountered for DataFrame -> DataDict: {e}") + assert all(isinstance(k, str) and isinstance(v, list) for k, v in output.items()) + + +def test_pl_lazy_dataframe_to_json(sample_polars_lazyframe): + try: + output = convert_type(sample_polars_lazyframe, "DataFrame", "Json") + except TypeError as e: + pytest.fail(f"TypeError encountered for DataFrame -> Json: {e}") + assert isinstance(output, str) + + +def test_pl_lazy_dataframe_to_numpy(sample_polars_lazyframe): + try: + output = convert_type(sample_polars_lazyframe, "DataFrame", "NDArray") + except TypeError as e: + pytest.fail(f"TypeError encountered for DataFrame -> NDArray: {e}") + assert isinstance(output, np.ndarray) + + +def test_pl_lazy_dataframe_to_arrow(sample_polars_lazyframe): + try: + output = convert_type(sample_polars_lazyframe, "DataFrame", "ArrowTable") + except TypeError as e: + pytest.fail(f"TypeError encountered for DataFrame -> ArrowTable: {e}") + assert isinstance(output, pa.Table) + + def test_arrow_to_datarecords(sample_arrow_table): try: output = convert_type(sample_arrow_table, "ArrowTable", "DataRecords") diff --git a/coded-flows/tests/test_extras.py b/coded-flows/tests/test_extras.py index 0104393..1183448 100644 --- a/coded-flows/tests/test_extras.py +++ b/coded-flows/tests/test_extras.py @@ -1,5 +1,6 @@ import numpy as np import pandas as pd +import polars as pl import pytest from PIL import Image import pyarrow as pa @@ -26,6 +27,24 @@ def sample_series(): return pd.Series([1, 2, 3, 4, 5]) +@pytest.fixture +def sample_polars_dataframe(): + """Fixture to create a sample Polars DataFrame.""" + return pl.DataFrame({"A": [1, 2, 3], "B": [4, 5, 6], "C": [7, 8, 9]}) + + +@pytest.fixture +def sample_polars_series(): + """Fixture to create a sample Polars Series.""" + return pl.Series([1, 2, 3, 4, 5]) + + +@pytest.fixture +def sample_polars_lazyframe(): + """Fixture to create a sample Polars LazyFrame.""" + return pl.DataFrame({"A": [1, 2, 3], "B": [4, 5, 6], "C": [7, 8, 9]}).lazy() + + @pytest.fixture def sample_numpy_array(): """Fixture to create a sample numpy array.""" @@ -57,6 +76,20 @@ def test_valid_dataframe_type(sample_dataframe): pytest.fail("ValidationError was raised for a valid type match of a DataFrame") +def test_valid_pl_dataframe_type(sample_polars_dataframe): + try: + is_valid_value_type(sample_polars_dataframe, DataFrame) + except ValidationError: + pytest.fail("ValidationError was raised for a valid type match of a DataFrame") + + +def test_valid_pl_lazy_dataframe_type(sample_polars_lazyframe): + try: + is_valid_value_type(sample_polars_lazyframe, DataFrame) + except ValidationError: + pytest.fail("ValidationError was raised for a valid type match of a DataFrame") + + def test_invalid_dataframe_type(): with pytest.raises(ValidationError): is_valid_value_type(3, DataFrame) @@ -70,6 +103,13 @@ def test_valid_series_type(sample_series): pytest.fail("ValidationError was raised for a valid type match of a DataSeries") +def test_valid_pl_series_type(sample_polars_series): + try: + is_valid_value_type(sample_polars_series, DataSeries) + except ValidationError: + pytest.fail("ValidationError was raised for a valid type match of a DataSeries") + + def test_invalid_series_type(): with pytest.raises(ValidationError): is_valid_value_type(3, DataSeries)