[ENH] Enhancing polars support by introducing set_output

skpro/datatypes/_adapter/polars.py

@@ -1,5 +1,8 @@
 # copyright: sktime developers, BSD-3-Clause License (see LICENSE file)
 """Common utilities for polars based data containers."""
+import numpy as np
+import pandas as pd
+
 from skpro.datatypes._common import _req
 from skpro.datatypes._common import _ret as ret
 
@@ -93,6 +96,12 @@ def convert_polars_to_pandas_with_index(obj):
             pd_df = pd_df.set_index(col, drop=True)
             pd_df.index.name = col.split("__index__")[1]
 
+    # check to see if we need to convert single melted polars columns
+    # back to multi_index
+    if all([True if col.startswith("__") else False for col in pd_df.columns]):
+        multi_index_columns = transform_single_column_to_multiindex_columns(pd_df)
+        pd_df.columns = pd.MultiIndex.from_arrays(multi_index_columns)
+
     return pd_df
 
 
@@ -140,6 +149,11 @@ def convert_pandas_to_polars_with_index(
         else:
             obj = obj.rename(columns={"index": "__index__"})
 
+    n_column_levels = check_n_level_of_dataframe(obj)
+
+    if n_column_levels > 1:
+        obj.columns = transform_pandas_multiindex_columns_to_single_column(obj)
+
     pl_df = from_pandas(
         data=obj,
         schema_overrides=schema_overrides,
@@ -151,3 +165,83 @@ def convert_pandas_to_polars_with_index(
         pl_df = pl_df.lazy()
 
     return pl_df
+
+
+def transform_pandas_multiindex_columns_to_single_column(X_input: pd.DataFrame):
+    """Convert function to return a list containing melted columns.
+
+    Assumes a multi-index column pandas DataFrame
+    Parameters
+    ----------
+    X : pandas DataFrame
+        pandas DataFrame containing a multi-index column (nlevels > 1)
+
+    Returns
+    -------
+    df_cols : a list object containing strings of all of the melted columns
+    """
+    df_cols = []
+    for col in X_input.columns:
+        df_cols.append("__" + "__".join(str(x) for x in col if x != "") + "__")
+    # in case "__index__" is in one of the tuples inside X_input
+    df_cols = [col.replace("____", "__") for col in df_cols]
+
+    return df_cols
+
+
+def transform_single_column_to_multiindex_columns(obj):
+    """Convert function to return a list containing un-melted columns."""
+    obj_columns = obj.columns
+
+    df_cols = []
+    for col in obj_columns:
+        items = col.split("__")
+        items = [item for item in items if item]
+        df_cols.append(items)
+
+    # take the transpose of the list of lists
+    df_cols = np.array(df_cols).T.tolist()
+
+    for multi_index_array in df_cols:
+        # try to convert item to float if the string is a supposed float value
+        for i in range(len(multi_index_array)):
+            try:
+                multi_index_array[i] = float(multi_index_array[i])
+            except ValueError:
+                pass
+    return df_cols
+
+
+def check_n_level_of_dataframe(X_input, axis=1):
+    """Convert function to check the number of levels inside a pd/pl frame.
+
+    Parameters
+    ----------
+    X_input : polars or pandas DataFrame
+        A given polars or pandas DataFrame. Note that the polars portion of this
+        code requires the soft dependencies polars and pyarrow to be installed
+    axis : [0,1]
+        Specify the index or columns of a pandas DataFrame. If 0, uses the index
+        If 1, uses the columns. This parameter is ignored if X_input is not
+        a pandas DataFrame.
+
+    Returns
+    -------
+    levels : int
+        An integer specifying the number of levels given a DataFrame
+    """
+    import polars as pl
+
+    if axis not in [0, 1]:
+        raise ValueError(f"axis must be in [0,1] " f"found {axis}.")
+    levels = None
+    if isinstance(X_input, pd.DataFrame):
+        if axis == 0:
+            levels = X_input.index.nlevels
+        elif axis == 1:
+            levels = X_input.columns.nlevels
+
+    if isinstance(X_input, pl.DataFrame) or isinstance(X_input, pl.LazyFrame):
+        levels = 1
+
+    return levels

skpro/datatypes/tests/test_polars.py

@@ -12,6 +12,14 @@
     import polars as pl
 
     from skpro.datatypes import check_is_mtype, convert
+    from skpro.datatypes._adapter.polars import (
+        check_n_level_of_dataframe,
+        transform_pandas_multiindex_columns_to_single_column,
+    )
+    from skpro.datatypes._table._convert import (
+        convert_pandas_to_polars_eager,
+        convert_polars_to_pandas,
+    )
 
 TEST_ALPHAS = [0.05, 0.1, 0.25]
 
@@ -46,6 +54,44 @@ def _pd_to_pl(df):
     return convert(df, from_type="pd_DataFrame_Table", to_type="polars_eager_table")
 
 
+@pytest.fixture
+def load_pandas_multi_index_column_fixture():
+    arrays = [
+        ["A", "A", "A", "A"],
+        ["Foo", "Foo", "Bar", "Bar"],
+        ["One", "Two", "One", "Two"],
+    ]
+    columns = pd.MultiIndex.from_arrays(arrays)
+
+    # Create the DataFrame
+    data = [[1, 2, 3, 4], [5, 6, 7, 8], [9, 10, 11, 12]]
+    pd_multi_column_fixture = pd.DataFrame(data, columns=columns)
+
+    return pd_multi_column_fixture
+
+
+@pytest.fixture
+def load_pandas_simple_column_fixture():
+    data = {"test_target": [10, 20, 30]}
+
+    # Create the DataFrame with a custom index
+    pd_simple_column_fixture = pd.DataFrame(
+        data, index=pd.Index(["row1", "row2", "row3"])
+    )
+
+    return pd_simple_column_fixture
+
+
+@pytest.fixture
+def load_polars_simple_fixture():
+    data = {"column1": [1, 2, 3], "column2": [4, 5, 6], "column3": [7, 8, 9]}
+
+    # Create the DataFrame
+    pl_simple_fixture = pl.DataFrame(data)
+
+    return pl_simple_fixture
+
+
 @pytest.fixture
 def polars_load_diabetes_polars(polars_load_diabetes_pandas):
     X_train, X_test, y_train = polars_load_diabetes_pandas
@@ -167,3 +213,122 @@ def test_polars_eager_regressor_in_predict_quantiles(
     assert y_pred_quantile.columns[0] == ("target", 0.05)
     assert y_pred_quantile.columns[1] == ("target", 0.1)
     assert y_pred_quantile.columns[2] == ("target", 0.25)
+
+
+@pytest.mark.skipif(
+    not run_test_module_changed("skpro.datatypes")
+    or not _check_soft_dependencies(["polars", "pyarrow"], severity="none"),
+    reason="skip test if polars/pyarrow is not installed in environment",
+)
+def test_check_column_level_of_dataframe_pandas(
+    load_pandas_multi_index_column_fixture,
+    load_pandas_simple_column_fixture,
+):
+    pd_multi_column_fixture = load_pandas_multi_index_column_fixture
+    pd_simple_column_fixture = load_pandas_simple_column_fixture
+
+    n_levels_multi_pd = check_n_level_of_dataframe(pd_multi_column_fixture)
+    n_levels_simple_pd = check_n_level_of_dataframe(pd_simple_column_fixture)
+    n_levels_simple_pd_index = check_n_level_of_dataframe(
+        pd_simple_column_fixture, axis=0
+    )
+
+    assert n_levels_multi_pd == 3
+    assert n_levels_simple_pd == 1
+    assert n_levels_simple_pd_index == 1
+
+
+@pytest.mark.skipif(
+    not _check_soft_dependencies(["polars", "pyarrow"], severity="none"),
+    reason="skip test if polars/pyarrow is not installed in environment",
+)
+def test_check_column_level_of_dataframe_polars(
+    load_polars_simple_fixture,
+):
+    pl_simple_column_fixture = load_polars_simple_fixture
+    n_levels_simple_pl = check_n_level_of_dataframe(pl_simple_column_fixture)
+    assert n_levels_simple_pl == 1
+
+
+@pytest.mark.skipif(
+    not run_test_module_changed("skpro.datatypes")
+    or not _check_soft_dependencies(["polars", "pyarrow"], severity="none"),
+    reason="skip test if polars/pyarrow is not installed in environment",
+)
+def test_convert_multiindex_columns_to_single_column(
+    load_pandas_multi_index_column_fixture,
+):
+    pd_multi_column_fixture1 = load_pandas_multi_index_column_fixture
+    df_list1 = transform_pandas_multiindex_columns_to_single_column(
+        pd_multi_column_fixture1
+    )
+    assert df_list1 == [
+        "__A__Foo__One__",
+        "__A__Foo__Two__",
+        "__A__Bar__One__",
+        "__A__Bar__Two__",
+    ]
+
+    pd_multi_column_fixture2 = load_pandas_multi_index_column_fixture
+    df_list2 = transform_pandas_multiindex_columns_to_single_column(
+        pd_multi_column_fixture2
+    )
+    assert df_list2 == [
+        "__A__Foo__One__",
+        "__A__Foo__Two__",
+        "__A__Bar__One__",
+        "__A__Bar__Two__",
+    ]
+
+
+@pytest.mark.skipif(
+    not run_test_module_changed("skpro.datatypes")
+    or not _check_soft_dependencies(["polars", "pyarrow"], severity="none"),
+    reason="skip test if polars/pyarrow is not installed in environment",
+)
+def test_convert_single_column_to_multiindex_column(
+    estimator,
+    polars_load_diabetes_pandas,
+):
+    X_train, X_test, y_train = polars_load_diabetes_pandas
+    estimator.fit(X_train, y_train)
+
+    # test for predict
+    y_pred = estimator.predict(X_test)
+    assert isinstance(y_pred, pd.DataFrame)
+
+    y_pred_pl = convert_pandas_to_polars_eager(y_pred)
+    y_pred_pd = convert_polars_to_pandas(y_pred_pl)
+    assert all(y_pred_pd.columns == y_pred.columns)
+    assert all(y_pred_pd.index == y_pred.index)
+    assert all(y_pred_pd.values == y_pred.values)
+
+    # test for interval
+    y_pred_interval = estimator.predict_interval(X_test)
+    assert isinstance(y_pred_interval, pd.DataFrame)
+
+    y_pred_interval_pl = convert_pandas_to_polars_eager(y_pred_interval)
+    y_pred_interval_pd = convert_polars_to_pandas(y_pred_interval_pl)
+    assert all(y_pred_interval_pd.columns == y_pred_interval.columns)
+    assert all(y_pred_interval_pd.index == y_pred_interval.index)
+    assert y_pred_interval_pd.equals(y_pred_interval)
+
+    # test for quantile
+    y_pred_quantile = estimator.predict_quantiles(X_test, alpha=TEST_ALPHAS)
+    assert isinstance(y_pred_quantile, pd.DataFrame)
+
+    y_pred_quantile_pl = convert_pandas_to_polars_eager(y_pred_quantile)
+    y_pred_quantile_pd = convert_polars_to_pandas(y_pred_quantile_pl)
+    assert all(y_pred_quantile_pd.columns == y_pred_quantile.columns)
+    assert all(y_pred_quantile_pd.index == y_pred_quantile.index)
+    assert y_pred_quantile_pd.equals(y_pred_quantile)
+
+    # test for var
+    y_pred_var = estimator.predict_interval(X_test)
+    assert isinstance(y_pred_var, pd.DataFrame)
+
+    y_pred_var_pl = convert_pandas_to_polars_eager(y_pred_var)
+    y_pred_var_pd = convert_polars_to_pandas(y_pred_var_pl)
+    assert all(y_pred_var_pd.columns == y_pred_var.columns)
+    assert all(y_pred_var_pd.index == y_pred_var.index)
+    assert y_pred_var_pd.equals(y_pred_var)