Add get_feature_names_out() to DataFrameMapper

sklearn_pandas/dataframe_mapper.py

@@ -45,10 +45,12 @@ def _get_feature_names(estimator):
     """
     Attempt to extract feature names based on a given estimator
     """
-    if hasattr(estimator, 'classes_'):
-        return estimator.classes_
+    if hasattr(estimator, 'get_feature_names_out'):
+        return estimator.get_feature_names_out()
     elif hasattr(estimator, 'get_feature_names'):
         return estimator.get_feature_names()
+    elif hasattr(estimator, 'classes_'):
+        return estimator.classes_
     return None
 
 
@@ -236,14 +238,35 @@ def fit(self, X, y=None):
         """
         self._build(X=X)
 
+        transformed_names_ = []
         for columns, transformers, options in self.built_features:
             t1 = datetime.now()
             input_df = options.get('input_df', self.input_df)
+            Xt = self._get_col_subset(X, columns, input_df)
+
+            alias = options.get('alias')
+            prefix = options.get('prefix', '')
+            suffix = options.get('suffix', '')
 
             if transformers is not None:
                 with add_column_names_to_exception(columns):
-                    Xt = self._get_col_subset(X, columns, input_df)
                     _call_fit(transformers.fit, Xt, y)
+                # compute num_cols
+                dummy = np.zeros((1, Xt.shape[1]))
+                if input_df:
+                    if isinstance(Xt, pd.Series):
+                        dummy = pd.Series(dummy.flatten(), index=Xt.index)
+                    elif isinstance(Xt, pd.DataFrame):
+                        dummy = pd.DataFrame(dummy, columns=Xt.columns, index=Xt.index[:1])
+                Xt_dummy = transformers.transform(dummy)
+                if hasattr(Xt_dummy, 'shape'):
+                    num_cols = Xt_dummy.shape[1] if len(Xt_dummy.shape) > 1 else 1
+                else:
+                    num_cols = 1
+            else:
+                num_cols = Xt.shape[1] if len(Xt.shape) > 1 else 1
+
+            transformed_names_ += self.get_names(columns, transformers, num_cols, alias, prefix, suffix)
             logger.info(f"[FIT] {columns}: {_elapsed_secs(t1)} secs")
 
         # handle features not explicitly selected
@@ -252,56 +275,75 @@ def fit(self, X, y=None):
             with add_column_names_to_exception(unsel_cols):
                 Xt = self._get_col_subset(X, unsel_cols, self.input_df)
                 _call_fit(self.built_default.fit, Xt, y)
+            # compute num_cols
+            dummy = np.zeros((1, Xt.shape[1]))
+            if self.input_df:
+                if isinstance(Xt, pd.Series):
+                    dummy = pd.Series(dummy.flatten(), index=Xt.index)
+                elif isinstance(Xt, pd.DataFrame):
+                    dummy = pd.DataFrame(dummy, columns=Xt.columns, index=Xt.index[:1])
+            Xt_dummy = self.built_default.transform(dummy)
+            num_cols = Xt_dummy.shape[1] if len(Xt_dummy.shape) > 1 else 1
+            transformed_names_ += self.get_names(unsel_cols, self.built_default, num_cols)
+        elif self.built_default is None:
+            unsel_cols = self._unselected_columns(X)
+            transformed_names_ += unsel_cols
+
+        self.transformed_names_ = transformed_names_
         return self
 
-    def get_names(self, columns, transformer, x, alias=None, prefix='',
+    def get_names(self, columns, transformer, num_cols, alias=None, prefix='',
                   suffix=''):
         """
         Return verbose names for the transformed columns.
 
         columns       name (or list of names) of the original column(s)
         transformer   transformer - can be a TransformerPipeline
-        x             transformed columns (numpy.ndarray)
+        num_cols      number of output columns
         alias         base name to use for the selected columns
+
+        Rules for output feature names:
+        - Base name: alias if provided, else joined columns with '_', else columns.
+        - Passthrough (transformer=None): output names are the column names.
+        - Transforms: single output -> base name; multi output -> base_name + '_' + feature_name or index.
+        - Prefix/suffix applied to all output names.
         """
         if alias is not None:
-            name = alias
+            base_name = alias
         elif isinstance(columns, list):
-            name = '_'.join(map(str, columns))
+            base_name = '_'.join(map(str, columns))
         else:
-            name = columns
-        num_cols = x.shape[1] if len(x.shape) > 1 else 1
-
-        output = []
-
-        if num_cols > 1:
-            # If there are as many columns as classes in the transformer,
-            # infer column names from classes names.
-
-            # If we are dealing with multiple transformers for these columns
-            # attempt to extract the names from each of them, starting from the
-            # last one
-            if isinstance(transformer, TransformerPipeline):
-                inverse_steps = transformer.steps[::-1]
-                estimators = (estimator for name, estimator in inverse_steps)
-                names_steps = (_get_feature_names(e) for e in estimators)
-                names = next((n for n in names_steps if n is not None), None)
-            # Otherwise use the only estimator present
-            else:
-                names = _get_feature_names(transformer)
+            base_name = columns
 
-            if names is not None and len(names) == num_cols:
-                output = [f"{name}_{o}" for o in names]
-                # otherwise, return name concatenated with '_1', '_2', etc.
+        if transformer is None:
+            # Passthrough: use column names directly
+            if isinstance(columns, list):
+                output = columns
             else:
-                output = [name + '_' + str(o) for o in range(num_cols)]
+                output = [columns]
         else:
-            output = [name]
+            # Transformed
+            if num_cols == 1:
+                output = [base_name]
+            else:
+                # Try to get names from transformer
+                if isinstance(transformer, TransformerPipeline):
+                    inverse_steps = transformer.steps[::-1]
+                    estimators = (estimator for name, estimator in inverse_steps)
+                    names_steps = (_get_feature_names(e) for e in estimators)
+                    names = next((n for n in names_steps if n is not None), None)
+                else:
+                    names = _get_feature_names(transformer)
 
-        if prefix == suffix == "":
-            return output
+                if names is not None and len(names) == num_cols:
+                    output = [f"{base_name}_{o}" for o in names]
+                else:
+                    output = [f"{base_name}_{o}" for o in range(num_cols)]
 
-        return ['{}{}{}'.format(prefix, x, suffix) for x in output]
+        if prefix or suffix:
+            output = [f"{prefix}{x}{suffix}" for x in output]
+
+        return output
 
     def get_dtypes(self, extracted):
         dtypes_features = [self.get_dtype(ex) for ex in extracted]
@@ -326,7 +368,6 @@ def _transform(self, X, y=None, do_fit=False):
             self._build(X=X)
 
         extracted = []
-        transformed_names_ = []
         for columns, transformers, options in self.built_features:
             input_df = options.get('input_df', self.input_df)
 
@@ -352,15 +393,16 @@ def _transform(self, X, y=None, do_fit=False):
                         Xt = transformers.transform(Xt)
                         logger.info(f"[TRANSFORM] {columns}: {_elapsed_secs(t1)} secs")  # NOQA
 
-            extracted.append(_handle_feature(Xt))
-
             alias = options.get('alias')
 
             prefix = options.get('prefix', '')
             suffix = options.get('suffix', '')
 
-            transformed_names_ += self.get_names(
-                columns, transformers, Xt, alias, prefix, suffix)
+            num_cols = Xt.shape[1] if len(Xt.shape) > 1 else 1
+            names = self.get_names(
+                columns, transformers, num_cols, alias, prefix, suffix)
+
+            extracted.append((_handle_feature(Xt), names))
 
         # handle features not explicitly selected
         if self.built_default is not False:
@@ -374,32 +416,39 @@ def _transform(self, X, y=None, do_fit=False):
                         if do_fit:
                             _call_fit(self.built_default.fit, Xt, y)
                         Xt = self.built_default.transform(Xt)
-                transformed_names_ += self.get_names(
-                    unsel_cols, self.built_default, Xt)
+                num_cols = Xt.shape[1] if len(Xt.shape) > 1 else 1
+                names = self.get_names(
+                    unsel_cols, self.built_default, num_cols)
             else:
                 # if not applying a default transformer,
                 # keep column names unmodified
-                transformed_names_ += unsel_cols
+                names = unsel_cols
 
-            extracted.append(_handle_feature(Xt))
+            extracted.append((_handle_feature(Xt), names))
 
+        # Build the combined transformed names
+        transformed_names_ = []
+        for _, names in extracted:
+            transformed_names_ += names
         self.transformed_names_ = transformed_names_
 
         # combine the feature outputs into one array.
         # at this point we lose track of which features
         # were created from which input columns, so it's
         # assumed that that doesn't matter to the model.
 
+        extracted_arrays = [arr for arr, _ in extracted]
+
         # If any of the extracted features is sparse, combine sparsely.
         # Otherwise, combine as normal arrays.
-        if any(sparse.issparse(fea) for fea in extracted):
-            stacked = sparse.hstack(extracted).tocsr()
+        if any(sparse.issparse(fea) for fea in extracted_arrays):
+            stacked = sparse.hstack(extracted_arrays).tocsr()
             # return a sparse matrix only if the mapper was initialized
             # with sparse=True
             if not self.sparse:
                 stacked = stacked.toarray()
         else:
-            stacked = np.hstack(extracted)
+            stacked = np.hstack(extracted_arrays)
 
         if self.df_out:
             # if no rows were dropped preserve the original index,
@@ -411,7 +460,7 @@ def _transform(self, X, y=None, do_fit=False):
                 index = None
 
             # output different data types, if appropriate
-            dtypes = self.get_dtypes(extracted)
+            dtypes = self.get_dtypes(extracted_arrays)
             df_out = pd.DataFrame(
                 stacked,
                 columns=self.transformed_names_,
@@ -441,3 +490,27 @@ def fit_transform(self, X, y=None):
         y       the target vector relative to X, optional
         """
         return self._transform(X, y, True)
+
+    def get_feature_names_out(self, input_features=None):
+        """
+        Get output feature names for transformation.
+
+        Parameters
+        ----------
+        input_features : array-like of str or None, default=None
+            Not used, present here for API consistency by convention.
+
+        Returns
+        -------
+        feature_names_out : ndarray of str objects
+            Feature names for output features.
+
+        Raises
+        ------
+        ValueError
+            If the mapper has not been fitted yet.
+        """
+        if not self.transformed_names_:
+            raise ValueError("This DataFrameMapper instance is not fitted yet. "
+                             "Call 'fit' with appropriate arguments before using this method.")
+        return np.array(self.transformed_names_)

tests/test_dataframe_mapper.py

@@ -230,6 +230,69 @@ def test_transformed_names_complex_alias(complex_dataframe):
     assert mapper.transformed_names_ == ['new_a', 'new_b', 'new_c']
 
 
+def test_get_feature_names_out_not_fitted():
+    """
+    get_feature_names_out raises ValueError if called before fitting
+    """
+    mapper = DataFrameMapper([('a', StandardScaler())])
+    with pytest.raises(ValueError, match="not fitted yet"):
+        mapper.get_feature_names_out()
+
+
+def test_get_feature_names_out_single_column(simple_dataframe):
+    """
+    get_feature_names_out for single column transform
+    """
+    df = simple_dataframe
+    mapper = DataFrameMapper([('a', StandardScaler())])
+    mapper.fit(df)
+    names = mapper.get_feature_names_out()
+    assert list(names) == ['a']
+
+
+def test_get_feature_names_out_multi_column(complex_dataframe):
+    """
+    get_feature_names_out for multi-column transform
+    """
+    df = complex_dataframe
+    mapper = DataFrameMapper([('target', LabelBinarizer())])
+    mapper.fit(df)
+    names = mapper.get_feature_names_out()
+    assert list(names) == ['target_a', 'target_b', 'target_c']
+
+
+def test_get_feature_names_out_passthrough(simple_dataframe):
+    """
+    get_feature_names_out for passthrough columns
+    """
+    df = simple_dataframe
+    mapper = DataFrameMapper([('a', None)])
+    mapper.fit(df)
+    names = mapper.get_feature_names_out()
+    assert list(names) == ['a']
+
+
+def test_get_feature_names_out_mixed(simple_dataframe):
+    """
+    get_feature_names_out for mixed transforms and passthrough
+    """
+    df = DataFrame({'a': [1, 2, 3], 'b': [4, 5, 6], 'c': [7, 8, 9]})
+    mapper = DataFrameMapper([
+        ('a', StandardScaler()),
+        ('b', None),
+        ('c', LabelBinarizer())  # assuming single unique values, but anyway
+    ])
+    mapper.fit(df)
+    names = mapper.get_feature_names_out()
+    # a: scaled -> 'a'
+    # b: passthrough -> 'b'
+    # c: binarizer -> 'c_7', 'c_8', 'c_9' or similar
+    # But since LabelBinarizer on [7,8,9], it will create columns for each unique
+    assert len(names) >= 2  # at least a and b
+    assert 'a' in names
+    assert 'b' in names
+
+
 def test_exception_column_context_transform(simple_dataframe):
     """
     If an exception is raised when transforming a column,