Repository: pyjanitor-devs/pyjanitor (python)

Issue:
transform_columns ignores `column_names` when `new_column_names` is specified
# Brief Description

`transform_columns` ignores the cols specified in `column_names` if `new_column_names` is specified (wrongly). For example, See the MWE.


# Minimally Reproducible Code

```python
>>> import numpy as np
>>> import pandas as pd
>>> import janitor
>>> df = pd.DataFrame({
...     "col1": [5, 10, 15],
...     "col2": [3, 6, 9],
...     "col3": [10, 100, 1_000],
... })
>>> df
    col1  col2  col3
0     5     3    10
1    10     6   100
2    15     9  1000

>>> df.transform_columns(
...     ["col1", "col3"],
...     np.log10,
...     new_column_names={"col1": "transform1", "col2": "transform2"},
... )
   col1  col2  col3  transform1  transform2
0     5     3    10    0.698970    0.477121
1    10     6   100    1.000000    0.778151
2    15     9  1000    1.176091    0.954243
```
Note that I specified that I wanted to transform "col1" and "col3" with log10, but "wrongly" specified the keys in the dictionary of `new_column_names` as "col1" and "col2".
Yet, the transformation is applied on "col1" and "col2" (obvious, because the expected output of `log10(col3)` is `[1, 2, 3]` which is not present in the output).


Pull request description:
[ENH/DOC] Improve transform_columns
# PR Description

Please describe the changes proposed in the pull request:

- Bug fix for `transform_columns` by setting up the `dest_column_names` dictionary using the columns in `column_names` instead of just using the user-entered `new_column_names`. Also added a test to avoid a similar regression in the future. Resolves #1063 .
- Make `transform_column(s)` not mutate the original dataframe.
- Check the column to be transformed actually exists, if not raise the error as early as possible.
- Add MWE for `transform_column(s)`. Related to #972 .



Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 52ca80e..53e9a2c 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -23,6 +23,7 @@
 -   [ENH] Add `xlsx_cells` for reading a spreadsheet as a table of individual cells. #929 @samukweku.
 -   [ENH] Let `filter_string` suit parameters of `Series.str.contains` Issue #1003 and #1047. @Zeroto521
 -   [ENH] `names_glue` in `pivot_wider` now takes a string form, using str.format_map under the hood. `levels_order` is also deprecated. @samukweku
+-   [BUG] Fixed bug in `transform_columns` which ignored the `column_names` specification when `new_column_names` dictionary was provided as an argument, issue #1063. @thatlittleboy
 
 ## [v0.22.0] - 2021-11-21
 
diff --git a/janitor/functions/transform_columns.py b/janitor/functions/transform_columns.py
index 100f7cb..ec75250 100644
--- a/janitor/functions/transform_columns.py
+++ b/janitor/functions/transform_columns.py
@@ -2,7 +2,18 @@ from typing import Callable, Dict, Hashable, List, Optional, Tuple, Union
 import pandas_flavor as pf
 import pandas as pd
 
-from janitor.utils import check, deprecated_alias
+from janitor.utils import check, check_column, deprecated_alias
+
+
+def _get_transform_column_result(
+    series: pd.Series,
+    function: Callable,
+    elementwise: bool,
+) -> pd.Series:
+    """Perform the actual computation for Series transformation."""
+    if elementwise:
+        return series.apply(function)
+    return function(series)
 
 
 @pf.register_dataframe_method
@@ -14,72 +25,60 @@ def transform_column(
     dest_column_name: Optional[str] = None,
     elementwise: bool = True,
 ) -> pd.DataFrame:
-    """Transform the given column in-place using the provided function.
-
-    Functions can be applied one of two ways:
-
-    - Element-wise (default; `elementwise=True``)
-    - Column-wise  (alternative; `elementwise=False``)
-
-    If the function is applied "elementwise",
-    then the first argument of the function signature
-    should be the individual element of each function.
-    This is the default behaviour of `transform_column``,
-    because it is easy to understand.
-    For example:
+    """Transform the given column using the provided function.
 
+    Meant to be the method-chaining equivalent of:
+    ```python
+    df[dest_column_name] = df[column_name].apply(function)
+    ```
 
+    Functions can be applied in one of two ways:
 
-        def elemwise_func(x):
-            modified_x = ... # do stuff here
-            return modified_x
+    - **Element-wise** (default; `elementwise=True`). Then, the individual
+    column elements will be passed in as the first argument of `function`.
+    - **Column-wise** (`elementwise=False`). Then, `function` is expected to
+    take in a pandas Series and return a sequence that is of identical length
+    to the original.
 
-        df.transform_column(column_name="my_column", function=elementwise_func)
-
-    On the other hand, columnwise application of a function
-    behaves as if the function takes in a pandas Series
-    and emits back a sequence that is of identical length to the original.
-    One place where this is desirable
-    is to gain access to `pandas` native string methods,
-    which are super fast!
-
-
-
-        def columnwise_func(s: pd.Series) -> pd.Series:
-            return s.str[0:5]
-
-        df.transform_column(
-            column_name="my_column",
-            lambda s: s.str[0:5],
-            elementwise=False
-        )
+    If `dest_column_name` is provided, then the transformation result is stored
+    in that column. Otherwise, the transformed result is stored under the name
+    of the original column.
 
     This method does not mutate the original DataFrame.
 
-    Let's say we wanted to apply a log10 transform a column of data.
-
-    Originally one would write code like this:
-
-
-
-        # YOU NO LONGER NEED TO WRITE THIS!
-        df[column_name] = df[column_name].apply(np.log10)
-
-    With the method chaining syntax, we can do the following instead:
-
-
-
-        df = (
-            pd.DataFrame(...)
-            .transform_column(column_name, np.log10)
-        )
-
-    With the functional syntax:
-
-
-
-        df = pd.DataFrame(...)
-        df = transform_column(df, column_name, np.log10)
+    Example: Transform a column in-place with an element-wise function.
+
+        >>> import pandas as pd
+        >>> import janitor
+        >>> df = pd.DataFrame({
+        ...     "a": [2, 3, 4],
+        ...     "b": ["area", "pyjanitor", "grapefruit"],
+        ... })
+        >>> df
+           a           b
+        0  2        area
+        1  3   pyjanitor
+        2  4  grapefruit
+        >>> df.transform_column(
+        ...     column_name="a",
+        ...     function=lambda x: x**2 - 1,
+        ... )
+            a           b
+        0   3        area
+        1   8   pyjanitor
+        2  15  grapefruit
+
+    Example: Transform a column in-place with an column-wise function.
+
+        >>> df.transform_column(
+        ...     column_name="b",
+        ...     function=lambda srs: srs.str[:5],
+        ...     elementwise=False,
+        ... )
+           a      b
+        0  2   area
+        1  3  pyjan
+        2  4  grape
 
     :param df: A pandas DataFrame.
     :param column_name: The column to transform.
@@ -89,24 +88,33 @@ def transform_column(
         name being overwritten. If a name is provided here, then a new column
         with the transformed values will be created.
     :param elementwise: Whether to apply the function elementwise or not.
-        If elementwise is True, then the function's first argument
+        If `elementwise` is True, then the function's first argument
         should be the data type of each datum in the column of data,
         and should return a transformed datum.
-        If elementwise is False, then the function's should expect
+        If `elementwise` is False, then the function's should expect
         a pandas Series passed into it, and return a pandas Series.
 
     :returns: A pandas DataFrame with a transformed column.
     """
+    check_column(df, column_name)
+
     if dest_column_name is None:
         dest_column_name = column_name
+    elif dest_column_name != column_name:
+        # If `dest_column_name` is provided and equals `column_name`, then we
+        # assume that the user's intent is to perform an in-place
+        # transformation (Same behaviour as when `dest_column_name` = None).
+        # Otherwise we throw an error if `dest_column_name` already exists in
+        # df.
+        check_column(df, dest_column_name, present=False)
 
-    if elementwise:
-        result = df[column_name].apply(function)
-    else:
-        result = function(df[column_name])
+    result = _get_transform_column_result(
+        df[column_name],
+        function,
+        elementwise,
+    )
 
-    df = df.assign(**{dest_column_name: result})
-    return df
+    return df.assign(**{dest_column_name: result})
 
 
 @pf.register_dataframe_method
@@ -121,95 +129,101 @@ def transform_columns(
 ) -> pd.DataFrame:
     """Transform multiple columns through the same transformation.
 
-    This method mutates the original DataFrame.
+    This method does not mutate the original DataFrame.
 
     Super syntactic sugar!
-
-    Basically wraps `transform_column` and calls it repeatedly over all column
-    names provided.
+    Essentially wraps [`transform_column`][janitor.functions.transform_columns.transform_column]
+    and calls it repeatedly over all column names provided.
 
     User can optionally supply either a suffix to create a new set of columns
     with the specified suffix, or provide a dictionary mapping each original
-    column name to its corresponding new column name. Note that all column
-    names must be strings.
-
-    A few examples below. Firstly, to just log10 transform a list of columns
-    without creating new columns to hold the transformed values:
-
-
-
-        df = (
-            pd.DataFrame(...)
-            .transform_columns(['col1', 'col2', 'col3'], np.log10)
-        )
-
-    Secondly, to add a '_log' suffix when creating a new column, which we think
-    is going to be the most common use case:
-
-
-
-        df = (
-            pd.DataFrame(...)
-            .transform_columns(
-                ['col1', 'col2', 'col3'],
-                np.log10,
-                suffix="_log"
-            )
-        )
-
-    Finally, to provide new names explicitly:
-
-
-
-        df = (
-            pd.DataFrame(...)
-            .transform_column(
-                ['col1', 'col2', 'col3'],
-                np.log10,
-                new_column_names={
-                    'col1': 'transform1',
-                    'col2': 'transform2',
-                    'col3': 'transform3',
-                    }
-                )
-        )
+    column name in `column_names` to its corresponding new column name.
+    Note that all column names must be strings.
+
+    Example: log10 transform a list of columns, replacing original columns.
+
+        >>> import numpy as np
+        >>> import pandas as pd
+        >>> import janitor
+        >>> df = pd.DataFrame({
+        ...     "col1": [5, 10, 15],
+        ...     "col2": [3, 6, 9],
+        ...     "col3": [10, 100, 1_000],
+        ... })
+        >>> df
+           col1  col2  col3
+        0     5     3    10
+        1    10     6   100
+        2    15     9  1000
+        >>> df.transform_columns(["col1", "col2", "col3"], np.log10)
+               col1      col2  col3
+        0  0.698970  0.477121   1.0
+        1  1.000000  0.778151   2.0
+        2  1.176091  0.954243   3.0
+
+    Example: Using the `suffix` parameter to create new columns.
+
+        >>> df.transform_columns(["col1", "col3"], np.log10, suffix="_log")
+           col1  col2  col3  col1_log  col3_log
+        0     5     3    10  0.698970       1.0
+        1    10     6   100  1.000000       2.0
+        2    15     9  1000  1.176091       3.0
+
+    Example: Using the `new_column_names` parameter to create new columns.
+
+        >>> df.transform_columns(
+        ...     ["col1", "col3"],
+        ...     np.log10,
+        ...     new_column_names={"col1": "transform1"},
+        ... )
+           col1  col2  col3  transform1
+        0     5     3   1.0    0.698970
+        1    10     6   2.0    1.000000
+        2    15     9   3.0    1.176091
 
     :param df: A pandas DataFrame.
     :param column_names: An iterable of columns to transform.
     :param function: A function to apply on each column.
-    :param suffix: (optional) Suffix to use when creating new columns to hold
+    :param suffix: Suffix to use when creating new columns to hold
         the transformed values.
     :param elementwise: Passed on to `transform_column`; whether or not
         to apply the transformation function elementwise (True)
         or columnwise (False).
-    :param new_column_names: (optional) An explicit mapping of old column names
-        to new column names.
+    :param new_column_names: An explicit mapping of old column names in
+        `column_names` to new column names. If any column specified in
+        `column_names` is not a key in this dictionary, the transformation
+        will happen in-place for that column.
     :returns: A pandas DataFrame with transformed columns.
-    :raises ValueError: if both `suffix` and `new_column_names` are
-        specified
-    """
-    dest_column_names = dict(zip(column_names, column_names))
-
+    :raises ValueError: If both `suffix` and `new_column_names` are
+        specified.
+    """  # noqa: E501
     check("column_names", column_names, [list, tuple])
+    check_column(df, column_names)
 
     if suffix is not None and new_column_names is not None:
         raise ValueError(
-            "only one of suffix or new_column_names should be specified"
+            "Only one of `suffix` or `new_column_names` should be specified."
         )
 
-    if suffix:  # If suffix is specified...
+    if suffix:
         check("suffix", suffix, [str])
-        for col in column_names:
-            dest_column_names[col] = col + suffix
-
-    if new_column_names:  # If new_column_names is specified...
+        dest_column_names = {col: col + suffix for col in column_names}
+    elif new_column_names:
         check("new_column_names", new_column_names, [dict])
-        dest_column_names = new_column_names
+        dest_column_names = {
+            col: new_column_names.get(col, col) for col in column_names
+        }
+    else:
+        dest_column_names = dict(zip(column_names, column_names))
 
-    # Now, transform columns.
+    results = {}
     for old_col, new_col in dest_column_names.items():
-        df = transform_column(
-            df, old_col, function, new_col, elementwise=elementwise
+        if old_col != new_col:
+            check_column(df, new_col, present=False)
+        results[new_col] = _get_transform_column_result(
+            df[old_col],
+            function,
+            elementwise=elementwise,
         )
 
-    return df
+    return df.assign(**results)


Test patch:
diff --git a/tests/functions/test_transform_column.py b/tests/functions/test_transform_column.py
index 2180120..5a2e192 100644
--- a/tests/functions/test_transform_column.py
+++ b/tests/functions/test_transform_column.py
@@ -7,6 +7,18 @@ import pytest
 from pandas.testing import assert_frame_equal, assert_series_equal
 
 
+@pytest.mark.functions
+def test_transform_nonexisting_column(dataframe):
+    """Checks an error is raised when the column being transformed
+    is non-existent.
+    """
+    with pytest.raises(
+        ValueError,
+        match="_foobar_ not present",
+    ):
+        dataframe.transform_column("_foobar_", np.log10)
+
+
 @pytest.mark.functions
 @pytest.mark.parametrize("elementwise", [True, False])
 def test_transform_column(dataframe, elementwise):
@@ -18,7 +30,6 @@ def test_transform_column(dataframe, elementwise):
     and on a single pandas Series.
     We use np.log10 as an example
     """
-
     df = dataframe.transform_column("a", np.log10, elementwise=elementwise)
     expected = pd.Series(np.log10([1, 2, 3] * 3))
     expected.name = "a"
@@ -28,16 +39,44 @@ def test_transform_column(dataframe, elementwise):
 @pytest.mark.functions
 def test_transform_column_with_dest(dataframe):
     """Test creating a new destination column."""
-
     expected_df = dataframe.assign(a_log10=np.log10(dataframe["a"]))
 
     df = dataframe.copy().transform_column(
-        "a", np.log10, dest_column_name="a_log10"
+        "a",
+        np.log10,
+        dest_column_name="a_log10",
     )
 
     assert_frame_equal(df, expected_df)
 
 
+@pytest.mark.functions
+def test_transform_column_dest_column_already_present(dataframe):
+    """Test behaviour when new destination column is provided and already
+    exists."""
+    # If dest_column_name already exists, throw an error
+    with pytest.raises(
+        ValueError,
+        match="pyjanitor already present in dataframe",
+    ):
+        _ = dataframe.assign(pyjanitor=1).transform_column(
+            "a",
+            np.log10,
+            dest_column_name="pyjanitor",
+        )
+
+    # unless dest_column_name is the same as column_name (the column being
+    # transformed); assume user wants an in-place transformation in this
+    # case.
+    expected_df = dataframe.copy().assign(a=np.log10(dataframe["a"]))
+    result_df = dataframe.transform_column(
+        "a",
+        np.log10,
+        dest_column_name="a",
+    )
+    assert_frame_equal(result_df, expected_df)
+
+
 @pytest.mark.functions
 def test_transform_column_no_mutation(dataframe):
     """Test checking that transform_column doesn't mutate the dataframe."""
diff --git a/tests/functions/test_transform_columns.py b/tests/functions/test_transform_columns.py
index 0c94d2d..91956e5 100644
--- a/tests/functions/test_transform_columns.py
+++ b/tests/functions/test_transform_columns.py
@@ -1,3 +1,4 @@
+"""Tests for transform_columns."""
 import numpy as np
 import pandas as pd
 import pytest
@@ -6,8 +7,7 @@ from pandas.testing import assert_frame_equal
 
 @pytest.mark.functions
 def test_transform_columns(dataframe):
-    # replacing the data of the original column
-
+    """Checks in-place transformation of multiple columns is as expected."""
     df = (
         dataframe.add_column("another", 10)
         .add_column("column", 100)
@@ -21,8 +21,7 @@ def test_transform_columns(dataframe):
 
 @pytest.mark.functions
 def test_transform_column_with_suffix(dataframe):
-    # creating a new destination column
-
+    """Checks `suffix` creates new columns as expected."""
     df = (
         dataframe.add_column("another", 10)
         .add_column("column", 100)
@@ -37,7 +36,7 @@ def test_transform_column_with_suffix(dataframe):
 
 @pytest.mark.functions
 def test_transform_column_with_new_names(dataframe):
-    # creating a new destination column
+    """Checks `new_column_names` creates new columns as expected."""
     df = (
         dataframe.add_column("another", 10)
         .add_column("column", 100)
@@ -54,9 +53,38 @@ def test_transform_column_with_new_names(dataframe):
     assert "column" in df.columns
 
 
+@pytest.mark.functions
+def test_transform_column_with_incomplete_new_names(dataframe):
+    """Use of `new_column_names` with additional columns (not in `column_names`
+    should passthrough silently. Related to bug #1063.
+    """
+    df = (
+        dataframe.add_column("another", 10)
+        .add_column("column", 100)
+        .transform_columns(
+            ["another", "column"],
+            np.log10,
+            new_column_names={
+                "another": "hello",
+                "fakecol": "world",
+            },
+        )
+    )
+
+    assert "another" in df.columns
+    assert "column" in df.columns
+    assert "hello" in df.columns
+    assert "world" not in df.columns
+
+
 @pytest.mark.functions
 def test_suffix_newname_validation(dataframe):
-    with pytest.raises(ValueError):
+    """Check ValueError is raised when both suffix and new_column_names are
+    provided."""
+    with pytest.raises(
+        ValueError,
+        match="Only one of `suffix` or `new_column_names` should be specified",
+    ):
         _ = (
             dataframe.add_column("another", 10)
             .add_column("column", 100)