From 060da1d42f9720e7c579cb0759bcef3fdb7b6c1d Mon Sep 17 00:00:00 2001 From: FBruzzesi Date: Sun, 2 Feb 2025 15:17:15 +0100 Subject: [PATCH] parse names only --- narwhals/_arrow/dataframe.py | 19 ++++++++++++++----- narwhals/_dask/dataframe.py | 4 ++-- narwhals/_duckdb/dataframe.py | 19 ++++++++++++++----- narwhals/_pandas_like/dataframe.py | 8 ++++---- narwhals/_polars/dataframe.py | 8 ++++---- narwhals/_spark_like/dataframe.py | 4 ++-- narwhals/dataframe.py | 13 ++----------- 7 files changed, 42 insertions(+), 33 deletions(-) diff --git a/narwhals/_arrow/dataframe.py b/narwhals/_arrow/dataframe.py index f43b86004..2797758ef 100644 --- a/narwhals/_arrow/dataframe.py +++ b/narwhals/_arrow/dataframe.py @@ -764,15 +764,24 @@ def sample( def unpivot( self: Self, - on: list[str] | None, - index: list[str], + on: str | list[str] | None, + index: str | list[str] | None, variable_name: str, value_name: str, ) -> Self: native_frame = self._native_frame n_rows = len(self) - on_ = [c for c in self.columns if c not in index] if on is None else on + index_: list[str] = ( + [] if index is None else [index] if isinstance(index, str) else index + ) + on_: list[str] = ( + [c for c in self.columns if c not in index_] + if on is None + else [on] + if isinstance(on, str) + else on + ) promote_kwargs = ( {"promote_options": "permissive"} @@ -784,11 +793,11 @@ def unpivot( [ pa.Table.from_arrays( [ - *[native_frame.column(idx_col) for idx_col in index], + *[native_frame.column(idx_col) for idx_col in index_], pa.array([on_col] * n_rows, pa.string()), native_frame.column(on_col), ], - names=[*index, variable_name, value_name], + names=[*index_, variable_name, value_name], ) for on_col in on_ ], diff --git a/narwhals/_dask/dataframe.py b/narwhals/_dask/dataframe.py index 844942fcc..4faa431bb 100644 --- a/narwhals/_dask/dataframe.py +++ b/narwhals/_dask/dataframe.py @@ -417,8 +417,8 @@ def gather_every(self: Self, n: int, offset: int) -> Self: def unpivot( self: Self, - on: list[str] | None, - index: list[str], + on: str | list[str] | None, + index: str | list[str] | None, variable_name: str, value_name: str, ) -> Self: diff --git a/narwhals/_duckdb/dataframe.py b/narwhals/_duckdb/dataframe.py index 7c618f789..ad724f883 100644 --- a/narwhals/_duckdb/dataframe.py +++ b/narwhals/_duckdb/dataframe.py @@ -429,11 +429,22 @@ def drop_nulls(self: Self, subset: list[str] | None) -> Self: def unpivot( self: Self, - on: list[str] | None, - index: list[str], + on: str | list[str] | None, + index: str | list[str] | None, variable_name: str, value_name: str, ) -> Self: + index_: list[str] = ( + [] if index is None else [index] if isinstance(index, str) else index + ) + on_: list[str] = ( + [c for c in self.columns if c not in index_] + if on is None + else [on] + if isinstance(on, str) + else on + ) + if variable_name == "": msg = "`variable_name` cannot be empty string for duckdb backend." raise NotImplementedError(msg) @@ -442,10 +453,8 @@ def unpivot( msg = "`value_name` cannot be empty string for duckdb backend." raise NotImplementedError(msg) - on_ = [c for c in self.columns if c not in index] if on is None else on - cols_to_select = ", ".join( - f'"{col}"' for col in [*index, variable_name, value_name] + f'"{col}"' for col in [*index_, variable_name, value_name] ) unpivot_on = ", ".join(f'"{col}"' for col in on_) diff --git a/narwhals/_pandas_like/dataframe.py b/narwhals/_pandas_like/dataframe.py index 393ef8db0..f3e8c73a9 100644 --- a/narwhals/_pandas_like/dataframe.py +++ b/narwhals/_pandas_like/dataframe.py @@ -1057,8 +1057,8 @@ def sample( def unpivot( self: Self, - on: list[str] | None, - index: list[str], + on: str | list[str] | None, + index: str | list[str] | None, variable_name: str, value_name: str, ) -> Self: @@ -1066,8 +1066,8 @@ def unpivot( self._native_frame.melt( id_vars=index, value_vars=on, - var_name=variable_name if variable_name is not None else "variable", - value_name=value_name if value_name is not None else "value", + var_name=variable_name, + value_name=value_name, ) ) diff --git a/narwhals/_polars/dataframe.py b/narwhals/_polars/dataframe.py index 4b3b45366..0d81ac1ab 100644 --- a/narwhals/_polars/dataframe.py +++ b/narwhals/_polars/dataframe.py @@ -310,8 +310,8 @@ def drop(self: Self, columns: list[str], strict: bool) -> Self: # noqa: FBT001 def unpivot( self: Self, - on: list[str] | None, - index: list[str], + on: str | list[str] | None, + index: str | list[str] | None, variable_name: str, value_name: str, ) -> Self: @@ -506,8 +506,8 @@ def drop(self: Self, columns: list[str], strict: bool) -> Self: # noqa: FBT001 def unpivot( self: Self, - on: list[str] | None, - index: list[str], + on: str | list[str] | None, + index: str | list[str] | None, variable_name: str, value_name: str, ) -> Self: diff --git a/narwhals/_spark_like/dataframe.py b/narwhals/_spark_like/dataframe.py index 6e4e02c29..c537287c4 100644 --- a/narwhals/_spark_like/dataframe.py +++ b/narwhals/_spark_like/dataframe.py @@ -412,8 +412,8 @@ def explode(self: Self, columns: str | Sequence[str], *more_columns: str) -> Sel def unpivot( self: Self, - on: list[str] | None, - index: list[str], + on: str | list[str] | None, + index: str | list[str] | None, variable_name: str, value_name: str, ) -> Self: diff --git a/narwhals/dataframe.py b/narwhals/dataframe.py index 3fb8f57a0..dd8f39ec5 100644 --- a/narwhals/dataframe.py +++ b/narwhals/dataframe.py @@ -330,22 +330,13 @@ def unpivot( variable_name: str | None, value_name: str | None, ) -> Self: - on_ = [on] if isinstance(on, str) else on - index_ = ( - [index] - if isinstance(index, str) - else index - if isinstance(index, list) - else [] - ) - variable_name = variable_name if variable_name is not None else "variable" value_name = value_name if value_name is not None else "value" return self._from_compliant_dataframe( self._compliant_frame.unpivot( - on=on_, - index=index_, + on=on, + index=index, variable_name=variable_name, value_name=value_name, )