scikit-learn
diff --git a/‎sklearn/compose/_column_transformer.py
Lines changed: 11 additions & 2 deletions b/‎sklearn/compose/_column_transformer.py
Lines changed: 11 additions & 2 deletions
diff --git a/‎sklearn/compose/tests/test_column_transformer.py
Lines changed: 40 additions & 0 deletions b/‎sklearn/compose/tests/test_column_transformer.py
Lines changed: 40 additions & 0 deletions
@@ -61,13 +61,15 @@ class ColumnTransformer(_BaseComposition, TransformerMixin):
             strings 'drop' and 'passthrough' are accepted as well, to
             indicate to drop the columns or to pass them through untransformed,
             respectively.
-        column(s) : string or int, array-like of string or int, slice or \
-boolean mask array
+        column(s) : string or int, array-like of string or int, slice, \
+boolean mask array or callable
             Indexes the data on its second axis. Integers are interpreted as
             positional columns, while strings can reference DataFrame columns
             by name.  A scalar string or int should be used where
             ``transformer`` expects X to be a 1d array-like (vector),
             otherwise a 2d array will be passed to the transformer.
+            A callable is passed the input data `X` and can return any of the
+            above.
 
     remainder : {'passthrough', 'drop'} or estimator, default 'passthrough'
         By default, all remaining columns that were not specified in
@@ -499,6 +501,7 @@ def _get_column(X, key):
     Supported key types (key):
     - scalar: output is 1D
     - lists, slices, boolean masks: output is 2D
+    - callable that returns any of the above
 
     Supported key data types:
 
@@ -510,6 +513,9 @@ def _get_column(X, key):
           can use any hashable object as key).
 
     """
+    if callable(key):
+        key = key(X)
+
     # check whether we have string column names or integers
     if _check_key_type(key, int):
         column_names = False
@@ -551,6 +557,9 @@ def _get_column_indices(X, key):
     """
     n_columns = X.shape[1]
 
+    if callable(key):
+        key = key(X)
+
     if _check_key_type(key, int):
         if isinstance(key, int):
             return [key]
 
@@ -99,6 +99,12 @@ def test_column_transformer():
         assert_array_equal(ct.fit_transform(X_array), res)
         assert_array_equal(ct.fit(X_array).transform(X_array), res)
 
+        # callable that returns any of the allowed specifiers
+        ct = ColumnTransformer([('trans', Trans(), lambda x: selection)],
+                               remainder='drop')
+        assert_array_equal(ct.fit_transform(X_array), res)
+        assert_array_equal(ct.fit(X_array).transform(X_array), res)
+
     ct = ColumnTransformer([('trans1', Trans(), [0]),
                             ('trans2', Trans(), [1])])
     assert_array_equal(ct.fit_transform(X_array), X_res_both)
@@ -166,6 +172,12 @@ def test_column_transformer_dataframe():
         assert_array_equal(ct.fit_transform(X_df), res)
         assert_array_equal(ct.fit(X_df).transform(X_df), res)
 
+        # callable that returns any of the allowed specifiers
+        ct = ColumnTransformer([('trans', Trans(), lambda X: selection)],
+                               remainder='drop')
+        assert_array_equal(ct.fit_transform(X_df), res)
+        assert_array_equal(ct.fit(X_df).transform(X_df), res)
+
     ct = ColumnTransformer([('trans1', Trans(), ['first']),
                             ('trans2', Trans(), ['second'])])
     assert_array_equal(ct.fit_transform(X_df), X_res_both)
@@ -777,3 +789,31 @@ def test_column_transformer_no_estimators():
     assert len(ct.transformers_) == 1
     assert ct.transformers_[-1][0] == 'remainder'
     assert ct.transformers_[-1][2] == [0, 1, 2]
+
+
+def test_column_transformer_callable_specifier():
+    # assert that function gets the full array / dataframe
+    X_array = np.array([[0, 1, 2], [2, 4, 6]]).T
+    X_res_first = np.array([[0, 1, 2]]).T
+
+    def func(X):
+        assert_array_equal(X, X_array)
+        return [0]
+
+    ct = ColumnTransformer([('trans', Trans(), func)],
+                           remainder='drop')
+    assert_array_equal(ct.fit_transform(X_array), X_res_first)
+    assert_array_equal(ct.fit(X_array).transform(X_array), X_res_first)
+
+    pd = pytest.importorskip('pandas')
+    X_df = pd.DataFrame(X_array, columns=['first', 'second'])
+
+    def func(X):
+        assert_array_equal(X.columns, X_df.columns)
+        assert_array_equal(X.values, X_df.values)
+        return ['first']
+
+    ct = ColumnTransformer([('trans', Trans(), func)],
+                           remainder='drop')
+    assert_array_equal(ct.fit_transform(X_df), X_res_first)
+    assert_array_equal(ct.fit(X_df).transform(X_df), X_res_first)