diff --git a/mesa_frames/abstract/mixin.py b/mesa_frames/abstract/mixin.py index cc843d08..a5153715 100644 --- a/mesa_frames/abstract/mixin.py +++ b/mesa_frames/abstract/mixin.py @@ -152,6 +152,15 @@ class DataFrameMixin(ABC): def _df_remove(self, df: DataFrame, mask: Mask, index_cols: str) -> DataFrame: return self._df_get_masked_df(df, index_cols, mask, negate=True) + @abstractmethod + def _df_and( + self, + df: DataFrame, + other: DataFrame | Sequence[float | int], + axis: Literal["index", "columns"] = "index", + index_cols: str | list[str] | None = None, + ) -> DataFrame: ... + @abstractmethod def _df_add( self, @@ -258,6 +267,15 @@ def _df_drop_duplicates( keep: Literal["first", "last", False] = "first", ) -> DataFrame: ... + @abstractmethod + def _df_ge( + self, + df: DataFrame, + other: DataFrame | Sequence[float | int], + axis: Literal["index", "columns"] = "index", + index_cols: str | list[str] | None = None, + ) -> DataFrame: ... + @abstractmethod def _df_get_bool_mask( self, @@ -282,6 +300,9 @@ def _df_groupby_cumcount( self, df: DataFrame, by: str | list[str], name: str = "cum_count" ) -> Series: ... + @abstractmethod + def _df_index(self, df: DataFrame, index_name: str) -> Index: ... + @abstractmethod def _df_iterator(self, df: DataFrame) -> Iterator[dict[str, Any]]: ... @@ -302,6 +323,24 @@ def _df_join( suffix="_right", ) -> DataFrame: ... + @abstractmethod + def _df_lt( + self, + df: DataFrame, + other: DataFrame | Sequence[float | int], + axis: Literal["index", "columns"] = "index", + index_cols: str | list[str] | None = None, + ) -> DataFrame: ... + + @abstractmethod + def _df_mod( + self, + df: DataFrame, + other: DataFrame | Sequence[float | int], + axis: Literal["index", "columns"] = "index", + index_cols: str | list[str] | None = None, + ) -> DataFrame: ... + @abstractmethod def _df_mul( self, @@ -337,6 +376,23 @@ def _df_norm( include_cols: bool = False, ) -> Series | DataFrame: ... + @abstractmethod + def _df_or( + self, + df: DataFrame, + other: DataFrame | Sequence[float | int], + axis: Literal["index", "columns"] = "index", + index_cols: str | list[str] | None = None, + ) -> DataFrame: ... + + @abstractmethod + def _df_reindex( + self, + df: DataFrame, + other: Sequence[Hashable] | DataFrame, + index_cols: str | list[str], + ) -> DataFrame: ... + @abstractmethod def _df_rename_columns( self, diff --git a/mesa_frames/abstract/space.py b/mesa_frames/abstract/space.py index 73aec5cd..6032ff3c 100644 --- a/mesa_frames/abstract/space.py +++ b/mesa_frames/abstract/space.py @@ -1341,13 +1341,15 @@ def get_neighborhood( neighbors_df = self._df_drop_duplicates(neighbors_df, self._pos_col_names) # Filter out-of-bound neighbors - neighbors_df = self._df_get_masked_df( - neighbors_df, - mask=self._df_all( - (neighbors_df[self._pos_col_names] < self._dimensions) - & (neighbors_df >= 0) - ), + mask = self._df_all( + self._df_and( + self._df_lt( + neighbors_df[self._pos_col_names], self._dimensions, axis="columns" + ), + neighbors_df >= 0, + ) ) + neighbors_df = self._df_get_masked_df(neighbors_df, mask=mask) if include_center: center_df = self._df_rename_columns( @@ -1409,7 +1411,15 @@ def out_of_bounds(self, pos: GridCoordinate | GridCoordinates) -> DataFrame: raise ValueError("This method is only valid for non-torus grids") pos_df = self._get_df_coords(pos, check_bounds=False) out_of_bounds = self._df_all( - (pos_df < 0) | (pos_df >= self._dimensions), + self._df_or( + pos_df < 0, + self._df_ge( + pos_df, + self._dimensions, + axis="columns", + index_cols=self._pos_col_names, + ), + ), name="out_of_bounds", ) return self._df_concat(objs=[pos_df, out_of_bounds], how="horizontal") @@ -1452,7 +1462,7 @@ def torus_adj(self, pos: GridCoordinate | GridCoordinates) -> DataFrame: The adjusted coordinates """ df_coords = self._get_df_coords(pos) - df_coords = df_coords % self._dimensions + df_coords = self._df_mod(df_coords, self._dimensions, axis="columns") return df_coords def _calculate_differences( @@ -1603,13 +1613,11 @@ def _get_df_coords( if agents.n_unique() != len(agents): raise ValueError("Some agents are present multiple times") if agents is not None: - return self._df_reset_index( - self._df_get_masked_df( - self._agents, index_cols="agent_id", mask=agents - ), - index_cols="agent_id", - drop=True, + df = self._df_get_masked_df( + self._agents, index_cols="agent_id", mask=agents ) + df = self._df_reindex(df, agents, "agent_id") + return self._df_reset_index(df, index_cols="agent_id", drop=True) if isinstance(pos, DataFrame): return pos[self._pos_col_names] elif ( diff --git a/mesa_frames/concrete/pandas/mixin.py b/mesa_frames/concrete/pandas/mixin.py index 514338db..91aec0e0 100644 --- a/mesa_frames/concrete/pandas/mixin.py +++ b/mesa_frames/concrete/pandas/mixin.py @@ -1,4 +1,4 @@ -from collections.abc import Collection, Hashable, Iterator, Sequence +from collections.abc import Callable, Collection, Hashable, Iterator, Sequence from typing import Literal import numpy as np @@ -19,6 +19,21 @@ def _df_add( ) -> pd.DataFrame: return df.add(other=other, axis=axis) + def _df_and( + self, + df: pd.DataFrame, + other: pd.DataFrame | Sequence[float | int], + axis: Literal["index"] | Literal["columns"] = "index", + index_cols: str | list[str] | None = None, + ) -> pd.DataFrame: + return self._df_logical_operation( + df=df, + other=other, + operation=lambda x, y: x & y, + axis=axis, + index_cols=index_cols, + ) + def _df_all( self, df: pd.DataFrame, @@ -137,6 +152,15 @@ def _df_drop_duplicates( ) -> pd.DataFrame: return df.drop_duplicates(subset=subset, keep=keep) + def _df_ge( + self, + df: pd.DataFrame, + other: pd.DataFrame | Sequence[float | int], + axis: Literal["index", "columns"] = "index", + index_cols: str | list[str] | None = None, + ) -> pd.DataFrame: + return df.ge(other, axis=axis) + def _df_get_bool_mask( self, df: pd.DataFrame, @@ -192,6 +216,16 @@ def _df_groupby_cumcount( ) -> pd.Series: return df.groupby(by).cumcount().rename(name) + 1 + def _df_index(self, df: pd.DataFrame, index_col: str | list[str]) -> pd.Index: + if ( + index_col is None + or df.index.name == index_col + or df.index.names == index_col + ): + return df.index + else: + return df.set_index(index_col).index + def _df_iterator(self, df: pd.DataFrame) -> Iterator[dict[str, Any]]: for index, row in df.iterrows(): row_dict = row.to_dict() @@ -246,6 +280,50 @@ def _df_join( else: return df + def _df_lt( + self, + df: pd.DataFrame, + other: pd.DataFrame | Sequence[float | int], + axis: Literal["index", "columns"] = "index", + index_cols: str | list[str] | None = None, + ) -> pd.DataFrame: + return df.lt(other, axis=axis) + + def _df_logical_operation( + self, + df: pd.DataFrame, + other: pd.DataFrame | Sequence[bool], + operation: Callable[ + [pd.DataFrame, Sequence[bool] | pd.DataFrame], pd.DataFrame + ], + axis: Literal["index"] | Literal["columns"] = "index", + index_cols: str | list[str] | None = None, + ) -> pd.DataFrame: + if isinstance(other, pd.DataFrame): + if index_cols is not None: + if df.index.name != index_cols: + df = df.set_index(index_cols) + if other.index.name != index_cols: + other = other.set_index(index_cols) + other = other.reindex(df.index, fill_value=np.nan) + return operation(df, other) + else: # Sequence[bool] + other = pd.Series(other) + if axis == "index": + other.index = df.index + return operation(df, other.values[:, None]).astype(bool) + else: + return operation(df, other.values[None, :]).astype(bool) + + def _df_mod( + self, + df: pd.DataFrame, + other: pd.DataFrame | Sequence[float | int], + axis: Literal["index", "columns"] = "index", + index_cols: str | list[str] | None = None, + ) -> pd.DataFrame: + return df.mod(other, axis=axis) + def _df_mul( self, df: pd.DataFrame, @@ -285,6 +363,31 @@ def _df_norm( else: return srs + def _df_or( + self, + df: pd.DataFrame, + other: pd.DataFrame | Sequence[bool], + axis: Literal["index"] | Literal["columns"] = "index", + index_cols: str | list[str] | None = None, + ) -> pd.DataFrame: + return self._df_logical_operation( + df=df, + other=other, + operation=lambda x, y: x | y, + axis=axis, + index_cols=index_cols, + ) + + def _df_reindex( + self, + df: pd.DataFrame, + other: Sequence[Hashable] | pd.DataFrame, + index_cols: str | list[str], + ) -> pd.DataFrame: + df = df.reindex(other) + df.index.name = index_cols + return df + def _df_rename_columns( self, df: pd.DataFrame, diff --git a/mesa_frames/concrete/polars/mixin.py b/mesa_frames/concrete/polars/mixin.py index ab08973a..862e0685 100644 --- a/mesa_frames/concrete/polars/mixin.py +++ b/mesa_frames/concrete/polars/mixin.py @@ -37,6 +37,21 @@ def _df_all( return pl.Series(name, df.select(pl.col("*").all()).row(0)) return df.with_columns(pl.all_horizontal("*").alias(name))[name] + def _df_and( + self, + df: pl.DataFrame, + other: pl.DataFrame | Sequence[float | int], + axis: Literal["index"] | Literal["columns"] = "index", + index_cols: str | list[str] | None = None, + ) -> pl.DataFrame: + return self._df_operation( + df=df, + other=other, + operation=lambda x, y: x & y, + axis=axis, + index_cols=index_cols, + ) + def _df_column_names(self, df: pl.DataFrame) -> list[str]: return df.columns @@ -186,6 +201,21 @@ def _df_drop_duplicates( .select(original_col_order) ) + def _df_ge( + self, + df: pl.DataFrame, + other: pl.DataFrame | Sequence[float | int], + axis: Literal["index", "columns"] = "index", + index_cols: str | list[str] | None = None, + ) -> pl.DataFrame: + return self._df_operation( + df=df, + other=other, + operation=lambda x, y: x >= y, + axis=axis, + index_cols=index_cols, + ) + def _df_get_bool_mask( self, df: pl.DataFrame, @@ -254,6 +284,9 @@ def _df_groupby_cumcount( ) -> pl.Series: return df.with_columns(pl.cum_count(by).over(by).alias(name))[name] + def _df_index(self, df: pl.DataFrame, index_col: str | list[str]) -> pl.Series: + return df[index_col] + def _df_iterator(self, df: pl.DataFrame) -> Iterator[dict[str, Any]]: return iter(df.iter_rows(named=True)) @@ -279,12 +312,37 @@ def _df_join( left_on, right_on = right_on, left_on how = "left" return left.join( - right, - on=on, - left_on=left_on, - right_on=right_on, - how=how, - suffix=suffix, + right, on=on, left_on=left_on, right_on=right_on, how=how, suffix=suffix + ) + + def _df_lt( + self, + df: pl.DataFrame, + other: pl.DataFrame | Sequence[float | int], + axis: Literal["index", "columns"] = "index", + index_cols: str | list[str] | None = None, + ) -> pl.DataFrame: + return self._df_operation( + df=df, + other=other, + operation=lambda x, y: x < y, + axis=axis, + index_cols=index_cols, + ) + + def _df_mod( + self, + df: pl.DataFrame, + other: pl.DataFrame | Sequence[float | int], + axis: Literal["index"] | Literal["columns"] = "index", + index_cols: str | list[str] | None = None, + ) -> pl.DataFrame: + return self._df_operation( + df=df, + other=other, + operation=lambda x, y: x % y, + axis=axis, + index_cols=index_cols, ) def _df_mul( @@ -377,6 +435,38 @@ def _df_operation( else: raise ValueError("other must be a DataFrame or a Sequence") + def _df_or( + self, + df: pl.DataFrame, + other: pl.DataFrame | Sequence[float | int], + axis: Literal["index"] | Literal["columns"] = "index", + index_cols: str | list[str] | None = None, + ) -> pl.DataFrame: + return self._df_operation( + df=df, + other=other, + operation=lambda x, y: x | y, + axis=axis, + index_cols=index_cols, + ) + + def _df_reindex( + self, + df: pl.DataFrame, + other: Sequence[Hashable] | pl.DataFrame, + index_cols: str | list[str], + ) -> pl.DataFrame: + # If other is a DataFrame, extract the index columns + if isinstance(other, pl.DataFrame): + other = other.select(index_cols) + else: + # If other is a sequence, create a DataFrame with it + other = pl.Series(name=index_cols, values=other).to_frame() + + # Perform a left join to reindex + result = other.join(df, on=index_cols, how="left") + return result + def _df_rename_columns( self, df: pl.DataFrame, old_columns: list[str], new_columns: list[str] ) -> pl.DataFrame: diff --git a/tests/pandas/test_mixin_pandas.py b/tests/pandas/test_mixin_pandas.py new file mode 100644 index 00000000..c7cf5b7c --- /dev/null +++ b/tests/pandas/test_mixin_pandas.py @@ -0,0 +1,62 @@ +import pandas as pd +import pytest + +from mesa_frames.concrete.pandas.mixin import PandasMixin + + +@pytest.fixture +def df_or(): + return PandasMixin()._df_or + + +@pytest.fixture +def df_0(): + return pd.DataFrame( + { + "unique_id": ["x", "y", "z"], + "A": [1, 0, 1], + "B": ["a", "b", "c"], + "C": [True, False, True], + "D": [0, 1, 1], + } + ).set_index("unique_id") + + +@pytest.fixture +def df_1(): + return pd.DataFrame( + { + "unique_id": ["z", "a", "b"], + "A": [0, 1, 0], + "B": ["d", "e", "f"], + "C": [False, True, False], + "E": [1, 0, 1], + } + ).set_index("unique_id") + + +def test_df_or(df_or: df_or, df_0: pd.DataFrame, df_1: pd.DataFrame): + # Test comparing the DataFrame with a sequence element-wise along the rows (axis='index') + df_0["F"] = [True, True, False] + df_1["F"] = [False, False, True] + result = df_or(df_0[["C", "F"]], df_1["F"], axis="index") + assert isinstance(result, pd.DataFrame) + assert result["C"].tolist() == [True, False, True] + assert result["F"].tolist() == [True, True, True] + + # Test comparing the DataFrame with a sequence element-wise along the columns (axis='columns') + result = df_or(df_0[["C", "F"]], [True, False], axis="columns") + assert isinstance(result, pd.DataFrame) + assert result["C"].tolist() == [True, True, True] + assert result["F"].tolist() == [True, True, False] + + # Test comparing DataFrames with index-column alignment + result = df_or( + df_0[["C", "F"]], + df_1[["C", "F"]], + axis="index", + index_cols="unique_id", + ) + assert isinstance(result, pd.DataFrame) + assert result["C"].tolist() == [True, False, True] + assert result["F"].tolist() == [True, True, False] diff --git a/tests/polars/test_mixin_polars.py b/tests/polars/test_mixin_polars.py index 2db79759..8e6761ba 100644 --- a/tests/polars/test_mixin_polars.py +++ b/tests/polars/test_mixin_polars.py @@ -82,6 +82,32 @@ def test_df_all(self, mixin: PolarsMixin): assert result.name == "all" assert result.to_list() == [False, True] + def test_df_and(self, mixin: PolarsMixin, df_0: pl.DataFrame, df_1: pl.DataFrame): + # Test comparing the DataFrame with a sequence element-wise along the rows (axis='index') + df_0 = df_0.with_columns(F=pl.Series([True, True, False])) + df_1 = df_1.with_columns(F=pl.Series([False, False, True])) + result = mixin._df_and(df_0[["C", "F"]], df_1["F"], axis="index") + assert isinstance(result, pl.DataFrame) + assert result["C"].to_list() == [False, False, True] + assert result["F"].to_list() == [False, False, False] + + # Test comparing the DataFrame with a sequence element-wise along the columns (axis='columns') + result = mixin._df_and(df_0[["C", "F"]], [True, False], axis="columns") + assert isinstance(result, pl.DataFrame) + assert result["C"].to_list() == [True, False, True] + assert result["F"].to_list() == [False, False, False] + + # Test comparing DataFrames with index-column alignment + result = mixin._df_and( + df_0[["unique_id", "C", "F"]], + df_1[["unique_id", "C", "F"]], + axis="index", + index_cols="unique_id", + ) + assert isinstance(result, pl.DataFrame) + assert result["C"].to_list() == [None, False, False] + assert result["F"].to_list() == [None, None, False] + def test_df_column_names(self, mixin: PolarsMixin, df_0: pl.DataFrame): cols = mixin._df_column_names(df_0) assert isinstance(cols, list) @@ -354,6 +380,31 @@ def test_df_drop_duplicates(self, mixin: PolarsMixin, df_0: pl.DataFrame): assert dropped.columns == ["unique_id", "A", "B", "C", "D"] assert dropped["B"].to_list() == ["b", "c", "e", "f"] + def test_df_ge(self, mixin: PolarsMixin, df_0: pl.DataFrame, df_1: pl.DataFrame): + # Test comparing the DataFrame with a sequence element-wise along the rows (axis='index') + result = mixin._df_ge(df_0[["A", "D"]], df_1["A"], axis="index") + assert isinstance(result, pl.DataFrame) + assert result["A"].to_list() == [False, False, False] + assert result["D"].to_list() == [False, False, False] + + # Test comparing the DataFrame with a sequence element-wise along the columns (axis='columns') + result = mixin._df_ge(df_0[["A", "D"]], [1, 2], axis="columns") + assert isinstance(result, pl.DataFrame) + assert result["A"].to_list() == [True, True, True] + assert result["D"].to_list() == [False, True, True] + + # Test comparing DataFrames with index-column alignment + df_1 = df_1.with_columns(D=pl.col("E")) + result = mixin._df_ge( + df_0[["unique_id", "A", "D"]], + df_1[["unique_id", "A", "D"]], + axis="index", + index_cols="unique_id", + ) + assert isinstance(result, pl.DataFrame) + assert result["A"].to_list() == [None, None, False] + assert result["D"].to_list() == [None, None, True] + def test_df_get_bool_mask(self, mixin: PolarsMixin, df_0: pl.DataFrame): # Test with pl.Series[bool] mask = mixin._df_get_bool_mask(df_0, "A", pl.Series([True, False, True])) @@ -411,6 +462,11 @@ def test_df_groupby_cumcount(self, df_0: pl.DataFrame, mixin: PolarsMixin): result = mixin._df_groupby_cumcount(df_0, "C") assert result.to_list() == [1, 1, 2] + def test_df_index(self, mixin: PolarsMixin, df_0: pl.DataFrame): + index = mixin._df_index(df_0, "unique_id") + assert isinstance(index, pl.Series) + assert index.to_list() == ["x", "y", "z"] + def test_df_iterator(self, mixin: PolarsMixin, df_0: pl.DataFrame): iterator = mixin._df_iterator(df_0) first_item = next(iterator) @@ -465,6 +521,56 @@ def test_df_join(self, mixin: PolarsMixin): assert joined.row(2) == (2, "b", 1, "x") assert joined.row(3) == (2, "b", 3, "y") + def test_df_lt(self, mixin: PolarsMixin, df_0: pl.DataFrame, df_1: pl.DataFrame): + # Test comparing the DataFrame with a sequence element-wise along the rows (axis='index') + result = mixin._df_lt(df_0[["A", "D"]], df_1["A"], axis="index") + assert isinstance(result, pl.DataFrame) + assert result["A"].to_list() == [True, True, True] + assert result["D"].to_list() == [True, True, True] + + # Test comparing the DataFrame with a sequence element-wise along the columns (axis='columns') + result = mixin._df_lt(df_0[["A", "D"]], [2, 3], axis="columns") + assert isinstance(result, pl.DataFrame) + assert result["A"].to_list() == [True, False, False] + assert result["D"].to_list() == [True, True, False] + + # Test comparing DataFrames with index-column alignment + df_1 = df_1.with_columns(D=pl.col("E")) + result = mixin._df_lt( + df_0[["unique_id", "A", "D"]], + df_1[["unique_id", "A", "D"]], + axis="index", + index_cols="unique_id", + ) + assert isinstance(result, pl.DataFrame) + assert result["A"].to_list() == [None, None, True] + assert result["D"].to_list() == [None, None, False] + + def test_df_mod(self, mixin: PolarsMixin, df_0: pl.DataFrame, df_1: pl.DataFrame): + # Test taking the modulo of the DataFrame by a sequence element-wise along the rows (axis='index') + result = mixin._df_mod(df_0[["A", "D"]], df_1["A"], axis="index") + assert isinstance(result, pl.DataFrame) + assert result["A"].to_list() == [1, 2, 3] + assert result["D"].to_list() == [1, 2, 3] + + # Test taking the modulo of the DataFrame by a sequence element-wise along the columns (axis='columns') + result = mixin._df_mod(df_0[["A", "D"]], [1, 2], axis="columns") + assert isinstance(result, pl.DataFrame) + assert result["A"].to_list() == [0, 0, 0] + assert result["D"].to_list() == [1, 0, 1] + + # Test taking the modulo of DataFrames with index-column alignment + df_1 = df_1.with_columns(D=pl.col("E")) + result = mixin._df_mod( + df_0[["unique_id", "A", "D"]], + df_1[["unique_id", "A", "D"]], + axis="index", + index_cols="unique_id", + ) + assert isinstance(result, pl.DataFrame) + assert result["A"].to_list() == [None, None, 3] + assert result["D"].to_list() == [None, None, 0] + def test_df_mul(self, mixin: PolarsMixin, df_0: pl.DataFrame, df_1: pl.DataFrame): # Test multiplying the DataFrame by a sequence element-wise along the rows (axis='index') result = mixin._df_mul(df_0[["A", "D"]], df_1["A"], axis="index") @@ -507,6 +613,53 @@ def test_df_norm(self, mixin: PolarsMixin): assert norm.row(0, named=True)["norm"] == 5 assert norm.row(1, named=True)["norm"] == 5 + def test_df_or(self, mixin: PolarsMixin, df_0: pl.DataFrame, df_1: pl.DataFrame): + # Test comparing the DataFrame with a sequence element-wise along the rows (axis='index') + df_0 = df_0.with_columns(F=pl.Series([True, True, False])) + df_1 = df_1.with_columns(F=pl.Series([False, False, True])) + result = mixin._df_or(df_0[["C", "F"]], df_1["F"], axis="index") + assert isinstance(result, pl.DataFrame) + assert result["C"].to_list() == [True, False, True] + assert result["F"].to_list() == [True, True, True] + + # Test comparing the DataFrame with a sequence element-wise along the columns (axis='columns') + result = mixin._df_or(df_0[["C", "F"]], [True, False], axis="columns") + assert isinstance(result, pl.DataFrame) + assert result["C"].to_list() == [True, True, True] + assert result["F"].to_list() == [True, True, False] + + # Test comparing DataFrames with index-column alignment + result = mixin._df_or( + df_0[["unique_id", "C", "F"]], + df_1[["unique_id", "C", "F"]], + axis="index", + index_cols="unique_id", + ) + assert isinstance(result, pl.DataFrame) + assert result["C"].to_list() == [True, None, True] + assert result["F"].to_list() == [True, True, False] + + def test_df_reindex( + self, mixin: PolarsMixin, df_0: pl.DataFrame, df_1: pl.DataFrame + ): + # Test with DataFrame + reindexed = mixin._df_reindex(df_0, df_1, "unique_id") + assert isinstance(reindexed, pl.DataFrame) + assert reindexed["unique_id"].to_list() == ["z", "a", "b"] + assert reindexed["A"].to_list() == [3, None, None] + assert reindexed["B"].to_list() == ["c", None, None] + assert reindexed["C"].to_list() == [True, None, None] + assert reindexed["D"].to_list() == [3, None, None] + + # Test with list + reindexed = mixin._df_reindex(df_0, ["z", "a", "b"], "unique_id") + assert isinstance(reindexed, pl.DataFrame) + assert reindexed["unique_id"].to_list() == ["z", "a", "b"] + assert reindexed["A"].to_list() == [3, None, None] + assert reindexed["B"].to_list() == ["c", None, None] + assert reindexed["C"].to_list() == [True, None, None] + assert reindexed["D"].to_list() == [3, None, None] + def test_df_rename_columns(self, mixin: PolarsMixin, df_0: pl.DataFrame): renamed = mixin._df_rename_columns(df_0, ["A", "B"], ["X", "Y"]) assert renamed.columns == ["unique_id", "X", "Y", "C", "D"]