From 9ad40ebabdd92157cd98dbdd0a523d843889b7ae Mon Sep 17 00:00:00 2001 From: Amine Lemaizi <40530011+aminelemaizi@users.noreply.github.com> Date: Wed, 10 Sep 2025 19:38:32 +0100 Subject: [PATCH] including polars for json and parquet temp data --- coded-flows/coded_flows/utils/media.py | 64 +++++++++++++- coded-flows/pyproject.toml | 2 +- coded-flows/tests/test_data_save.py | 115 +++++++++++++++++++++++-- 3 files changed, 170 insertions(+), 11 deletions(-) diff --git a/coded-flows/coded_flows/utils/media.py b/coded-flows/coded_flows/utils/media.py index 954830e..871b499 100644 --- a/coded-flows/coded_flows/utils/media.py +++ b/coded-flows/coded_flows/utils/media.py @@ -6,6 +6,7 @@ import pyarrow.parquet as pq import numpy as np import pandas as pd +import polars as pl from typing import Union, List, Dict, Any from io import BytesIO from PIL import Image @@ -58,19 +59,36 @@ def _save_arrow_table_to_json(table: pa.Table, filename: str = None) -> str: records = table.to_pylist() with open(file_path, "w") as f: json.dump(records, f, separators=(",", ":")) + return file_path + +def _save_polars_to_json(df: pl.DataFrame, filename: str = None) -> str: + random_filename = f"cfdata_{filename if filename else uuid.uuid4().hex}.json" + 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) + df.write_json(file_path) return file_path -# List, DataSeries, NDArray, DataRecords, DataFrame, Arrow +# List, DataSeries, NDArray, DataRecords, DataFrame, Arrow, Polars def save_data_to_json( *data_args: Union[ - pd.DataFrame, pd.Series, pa.Table, np.ndarray, List[Dict[str, Any]], List[Any] + pd.DataFrame, + pd.Series, + pa.Table, + np.ndarray, + List[Dict[str, Any]], + List[Any], + pl.DataFrame, + pl.Series, + pl.LazyFrame, ], labels: List[str] = [], is_table: bool = False, filename: str = None, ) -> str: + labels = ["values"] if is_table else labels if not is_table and len(data_args) != len(labels): @@ -86,6 +104,8 @@ def save_data_to_json( if is_table and ( isinstance(data, pd.DataFrame) or isinstance(data, pa.Table) + or isinstance(data, pl.DataFrame) + or isinstance(data, pl.LazyFrame) or ( isinstance(data, list) and all(isinstance(item, dict) for item in data[:50]) ) @@ -95,6 +115,11 @@ def save_data_to_json( table_df = data elif isinstance(data, pa.Table): return _save_arrow_table_to_json(data, filename) + elif isinstance(data, pl.DataFrame): + return _save_polars_to_json(data, filename) + elif isinstance(data, pl.LazyFrame): + # Collect LazyFrame to DataFrame first + return _save_polars_to_json(data.collect(), filename) else: table_df = pd.DataFrame.from_records(data) @@ -112,6 +137,17 @@ def save_data_to_json( if label not in data.column_names: raise ValueError(f"Label '{label}' not found in Arrow table columns.") col_data = data.column(label).to_pylist() + elif isinstance(data, pl.DataFrame): + if label not in data.columns: + raise ValueError(f"Label '{label}' not found in DataFrame columns.") + col_data = data[label].to_list() + elif isinstance(data, pl.LazyFrame): + collected_data = data.collect() + if label not in collected_data.columns: + raise ValueError(f"Label '{label}' not found in DataFrame columns.") + col_data = collected_data[label].to_list() + elif isinstance(data, pl.Series): + col_data = data.to_list() elif isinstance(data, pd.Series): col_data = data.values elif isinstance(data, np.ndarray): @@ -142,7 +178,15 @@ def save_data_to_json( def save_data_to_parquet( data: Union[ - pd.DataFrame, pd.Series, pa.Table, np.ndarray, List[Dict[str, Any]], List[Any] + pd.DataFrame, + pd.Series, + pl.DataFrame, + pl.Series, + pl.LazyFrame, + pa.Table, + np.ndarray, + List[Dict[str, Any]], + List[Any], ], filename=None, ) -> str: @@ -155,6 +199,9 @@ def save_data_to_parquet( if ( isinstance(data, pd.DataFrame) or isinstance(data, pd.Series) + or isinstance(data, pl.DataFrame) + or isinstance(data, pl.LazyFrame) + or isinstance(data, pl.Series) or isinstance(data, pa.Table) or ( isinstance(data, list) and all(isinstance(item, dict) for item in data[:50]) @@ -169,9 +216,18 @@ def save_data_to_parquet( data.to_frame().to_parquet( file_path, row_group_size=50000, index=False, engine="pyarrow" ) + elif isinstance(data, pl.DataFrame): + data.write_parquet(file_path, row_group_size=50000, use_pyarrow=True) + elif isinstance(data, pl.LazyFrame): + data.collect().write_parquet( + file_path, row_group_size=50000, use_pyarrow=True + ) + elif isinstance(data, pl.Series): + data.to_frame().write_parquet( + file_path, row_group_size=50000, use_pyarrow=True + ) elif isinstance(data, pa.Table): pq.write_table(data, file_path, row_group_size=50000) - else: table = pa.Table.from_pylist(data) pq.write_table(table, file_path, row_group_size=50000) diff --git a/coded-flows/pyproject.toml b/coded-flows/pyproject.toml index 9571afb..450e805 100644 --- a/coded-flows/pyproject.toml +++ b/coded-flows/pyproject.toml @@ -1,6 +1,6 @@ [tool.poetry] name = "coded-flows" -version = "0.6.4" +version = "0.6.5" description = "Various utilities for Coded Flows" authors = ["COLOR CODED CODES "] readme = "README.md" diff --git a/coded-flows/tests/test_data_save.py b/coded-flows/tests/test_data_save.py index f33e86d..c08e647 100644 --- a/coded-flows/tests/test_data_save.py +++ b/coded-flows/tests/test_data_save.py @@ -1,6 +1,7 @@ import os import pytest import pandas as pd +import polars as pl import pandas.testing as pdt import pyarrow as pa import pyarrow.parquet as pq @@ -14,10 +15,11 @@ def test_valid_inputs(): arr = np.array([20, 21, 22]) data_records = [{"key1": 30}, {"key1": 31}, {"key1": 32}] arrow_table = pa.table({"col1": [1, 2, 3], "col2": ["a", "b", "c"]}) + dfp = pl.DataFrame({"uu": [1, 2, 3]}) - labels = ["x", "y", "z", "w", "col1"] + labels = ["x", "y", "z", "w", "col1", "uu"] file_path = save_data_to_json( - df, series, arr, data_records, arrow_table, labels=labels + df, series, arr, data_records, arrow_table, dfp, labels=labels ) # Check file existence @@ -27,7 +29,7 @@ def test_valid_inputs(): with open(file_path, "r") as f: json_content = f.read() - expected_content = '[{"x":1,"y":100,"z":20,"w":null,"col1":1},{"x":2,"y":200,"z":21,"w":null,"col1":2},{"x":3,"y":300,"z":22,"w":null,"col1":3}]' + expected_content = '[{"x":1,"y":100,"z":20,"w":null,"col1":1,"uu":1},{"x":2,"y":200,"z":21,"w":null,"col1":2,"uu":2},{"x":3,"y":300,"z":22,"w":null,"col1":3,"uu":3}]' assert json_content.strip() == expected_content, "JSON content mismatch." os.remove(file_path) @@ -51,6 +53,23 @@ def test_missing_column_in_dataframe(): save_data_to_json(df, labels=labels) +def test_missing_column_in_polars_dataframe(): + df = pl.DataFrame({"x": [1, 2, 3]}) + labels = ["z"] + + with pytest.raises(ValueError, match="Label 'z' not found in DataFrame columns."): + save_data_to_json(df, labels=labels) + + +def test_missing_column_in_polars_lazyframe(): + df = pl.DataFrame({"x": [1, 2, 3]}) + df_lazy = df.lazy() + labels = ["z"] + + with pytest.raises(ValueError, match="Label 'z' not found in DataFrame columns."): + save_data_to_json(df_lazy, labels=labels) + + def test_invalid_numpy_array(): arr = np.array([[1, 2], [3, 4]]) # 2D array, invalid labels = ["x"] @@ -90,15 +109,18 @@ def test_variable_lengths(): def test_series_input(): series1 = pd.Series([1, 2, 3], name="x") series2 = pd.Series([10, 20, 30], name="y") - labels = ["x", "y"] + series3 = pl.Series("z", [100, 200, 300]) + labels = ["x", "y", "z"] - file_path = save_data_to_json(series1, series2, labels=labels) + file_path = save_data_to_json(series1, series2, series3, labels=labels) # Check JSON content with open(file_path, "r") as f: json_content = f.read() - expected_content = '[{"x":1,"y":10},{"x":2,"y":20},{"x":3,"y":30}]' + expected_content = ( + '[{"x":1,"y":10,"z":100},{"x":2,"y":20,"z":200},{"x":3,"y":30,"z":300}]' + ) assert ( json_content.strip() == expected_content ), "JSON content mismatch for Series inputs." @@ -124,6 +146,25 @@ def test_save_data_to_json_is_table_true_with_dataframe(): os.remove(json_path) +def test_save_data_to_json_is_table_true_with_polars_dataframe(): + # Create a sample DataFrame + df = pl.DataFrame({"col1": [1, 2, 3], "col2": ["a", "b", "c"]}) + + # Call the function with `is_table=True` + json_path = save_data_to_json(df, is_table=True) + + # Assert the file exists + assert os.path.exists(json_path), "JSON file was not created." + + # Assert the file content matches the DataFrame content + with open(json_path, "r") as f: + saved_data = pl.read_json(f) + assert df.equals(saved_data) + + # Clean up + os.remove(json_path) + + def test_save_data_to_json_is_table_true_with_arrow_table(): # Create a sample Arrow Table table = pa.table({"col1": [1, 2, 3], "col2": ["a", "b", "c"]}) @@ -226,6 +267,20 @@ def test_save_data_to_json_is_table_true_series(): os.remove(json_path) +def test_save_data_to_json_is_table_true_polars_series(): + # Series input + series = pl.Series("col1", [1, 2, 3]) + json_path = save_data_to_json(series, is_table=True) + assert os.path.exists(json_path) + with open(json_path, "r") as f: + saved_data = pl.read_json(f) + expected_df = pl.DataFrame({"values": series}) + assert expected_df.equals(saved_data) + + # Clean up + os.remove(json_path) + + def test_dataframe_to_parquet(): df = pd.DataFrame({"x": [1, 2, 3], "z": [7, 8, 9]}) @@ -241,6 +296,37 @@ def test_dataframe_to_parquet(): os.remove(file_path) +def test_polars_dataframe_to_parquet(): + df = pl.DataFrame({"x": [1, 2, 3], "z": [7, 8, 9]}) + + file_path = save_data_to_parquet(df) + + # Check file existence + assert os.path.exists(file_path), "Output file does not exist." + + # Read back from Parquet into DataFrame + df_loaded = pl.read_parquet(file_path, use_pyarrow=True) + + assert df.equals(df_loaded) + os.remove(file_path) + + +def test_polars_lazyframe_to_parquet(): + df = pl.DataFrame({"x": [1, 2, 3], "z": [7, 8, 9]}) + dfl = df.lazy() + + file_path = save_data_to_parquet(dfl) + + # Check file existence + assert os.path.exists(file_path), "Output file does not exist." + + # Read back from Parquet into DataFrame + df_loaded = pl.read_parquet(file_path, use_pyarrow=True) + + assert df.equals(df_loaded) + os.remove(file_path) + + def test_arrow_table_to_parquet(): table = pa.table({"x": [1, 2, 3], "z": [7, 8, 9]}) @@ -310,3 +396,20 @@ def test_pandas_series_to_parquet(): pdt.assert_frame_equal(expected_df.sort_index(axis=1), df_loaded.sort_index(axis=1)) os.remove(file_path) + + +def test_polars_series_to_parquet(): + polars_series = pl.Series("values", [100, 200, 300, 400, 500]) + file_path = save_data_to_parquet(polars_series) + + # Check file existence + assert os.path.exists(file_path), "Output file does not exist." + + # Read back from Parquet into DataFrame + df_loaded = pl.read_parquet(file_path, use_pyarrow=True) + + # Convert original Series to DataFrame for comparison + expected_df = polars_series.to_frame() + + assert expected_df.equals(df_loaded) + os.remove(file_path)