add drop_axis option to columnselector

Author: rasbt

docs/sources/CHANGELOG.md

@@ -25,6 +25,7 @@ The CHANGELOG for the current development version is available at
 -   New function implementing the 5x2cv paired t-test procedure (`paired_ttest_5x2cv`) proposed by Dieterrich (1998)
     to compare the performance of two models. ([#325](https://github.com/rasbt/mlxtend/issues/325))
 - A `refit` parameter was added to stacking classes (similar to the `refit` parameter in the `EnsembleVoteClassifier`), to support classifiers and regressors that follow the scikit-learn API but are not compatible with scikit-learn's `clone` function ([#325](https://github.com/rasbt/mlxtend/issues/324))
+- The `ColumnSelector` now has a `drop_axis` argument to use it in pipelines with `CountVectorizers` ([#333](https://github.com/rasbt/mlxtend/pull/333)
 
 ##### Changes
 

docs/sources/user_guide/feature_selection/ColumnSelector.ipynb

@@ -3,9 +3,7 @@
   {
    "cell_type": "code",
    "execution_count": 1,
-   "metadata": {
-    "collapsed": true
-   },
+   "metadata": {},
    "outputs": [],
    "source": [
     "%matplotlib inline"
@@ -72,9 +70,7 @@
   {
    "cell_type": "code",
    "execution_count": 2,
-   "metadata": {
-    "collapsed": true
-   },
+   "metadata": {},
    "outputs": [],
    "source": [
     "from sklearn.datasets import load_iris\n",
@@ -183,9 +179,7 @@
   {
    "cell_type": "code",
    "execution_count": 5,
-   "metadata": {
-    "collapsed": true
-   },
+   "metadata": {},
    "outputs": [],
    "source": [
     "from sklearn.datasets import load_iris\n",
@@ -283,7 +277,7 @@
      "text": [
       "## ColumnSelector\n",
       "\n",
-      "*ColumnSelector(cols=None)*\n",
+      "*ColumnSelector(cols=None, drop_axis=False)*\n",
       "\n",
       "Base class for all estimators in scikit-learn\n",
       "\n",
@@ -402,15 +396,6 @@
     "    s = f.read() + '<br><br>'\n",
     "print(s)"
    ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {
-    "collapsed": true
-   },
-   "outputs": [],
-   "source": []
   }
  ],
  "metadata": {
@@ -430,7 +415,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.6.1"
+   "version": "3.6.4"
   }
  },
  "nbformat": 4,

mlxtend/feature_selection/column_selector.py

@@ -13,7 +13,7 @@
 
 class ColumnSelector(BaseEstimator):
 
-    def __init__(self, cols=None):
+    def __init__(self, cols=None, drop_axis=False):
         """Object for selecting specific columns from a data set.
 
         Parameters
@@ -22,8 +22,18 @@ def __init__(self, cols=None):
             A list specifying the feature indices to be selected. For example,
             [1, 4, 5] to select the 2nd, 5th, and 6th feature columns.
             If None, returns all columns in the array.
+
+        drop_axis : bool (default=False)
+            Drops last axis if True and the only one column is selected. This
+            is useful, e.g., when the ColumnSelector is used for selecting
+            only one column and the resulting array should be fed to e.g.,
+            a scikit-learn column selector. E.g., instead of returning an
+            array with shape (n_samples, 1), drop_axis=True will return an
+            aray with shape (n_samples,).
+
         """
         self.cols = cols
+        self.drop_axis = drop_axis
 
     def fit_transform(self, X, y=None):
         """ Return a slice of the input array.
@@ -60,7 +70,7 @@ def transform(self, X, y=None):
 
         """
         t = X[:, self.cols]
-        if len(t.shape) == 1:
+        if len(t.shape) == 1 and not self.drop_axis:
             t = t[:, np.newaxis]
         return t
 

mlxtend/feature_selection/tests/test_column_selector.py

@@ -17,7 +17,16 @@
 def test_ColumnSelector():
     X1_in = np.ones((4, 8))
     X1_out = ColumnSelector(cols=(1, 3)).transform(X1_in)
-    assert(X1_out.shape == (4, 2))
+    assert X1_out.shape == (4, 2)
+
+
+def test_ColumnSelector_drop_axis():
+    X1_in = np.ones((4, 8))
+    X1_out = ColumnSelector(cols=(1), drop_axis=True).transform(X1_in)
+    assert X1_out.shape == (4,)
+
+    X1_out = ColumnSelector(cols=(1)).transform(X1_in)
+    assert X1_out.shape == (4, 1)
 
 
 def test_ColumnSelector_in_gridsearch():