diff --git a/doc/source/whatsnew/v1.1.0.rst b/doc/source/whatsnew/v1.1.0.rst index 1cd325dad9f07..f1273ab334740 100644 --- a/doc/source/whatsnew/v1.1.0.rst +++ b/doc/source/whatsnew/v1.1.0.rst @@ -141,6 +141,7 @@ Reshaping ^^^^^^^^^ - +- Bug in :meth:`DataFrame.count` not returning subclassed data types. - Bug in :meth:`DataFrame.pivot_table` when only MultiIndexed columns is set (:issue:`17038`) - Fix incorrect error message in :meth:`DataFrame.pivot` when ``columns`` is set to ``None``. (:issue:`30924`) - Bug in :func:`crosstab` when inputs are two Series and have tuple names, the output will keep dummy MultiIndex as columns. (:issue:`18321`) diff --git a/pandas/core/frame.py b/pandas/core/frame.py index 1a49388d81243..5394178511de8 100644 --- a/pandas/core/frame.py +++ b/pandas/core/frame.py @@ -7811,7 +7811,7 @@ def count(self, axis=0, level=None, numeric_only=False): # GH #423 if len(frame._get_axis(axis)) == 0: - result = Series(0, index=frame._get_agg_axis(axis)) + result = frame._constructor_sliced(0, index=frame._get_agg_axis(axis)) else: if frame._is_mixed_type or frame._data.any_extension_types: # the or any_extension_types is really only hit for single- @@ -7821,7 +7821,9 @@ def count(self, axis=0, level=None, numeric_only=False): # GH13407 series_counts = notna(frame).sum(axis=axis) counts = series_counts.values - result = Series(counts, index=frame._get_agg_axis(axis)) + result = frame._constructor_sliced( + counts, index=frame._get_agg_axis(axis) + ) return result.astype("int64") @@ -7860,7 +7862,7 @@ def _count_level(self, level, axis=0, numeric_only=False): level_codes = ensure_int64(count_axis.codes[level]) counts = lib.count_level_2d(mask, level_codes, len(level_index), axis=0) - result = DataFrame(counts, index=level_index, columns=agg_axis) + result = frame._constructor(counts, index=level_index, columns=agg_axis) if axis == 1: # Undo our earlier transpose diff --git a/pandas/tests/frame/test_subclass.py b/pandas/tests/frame/test_subclass.py index 4a436d70dc48f..813b30cb96009 100644 --- a/pandas/tests/frame/test_subclass.py +++ b/pandas/tests/frame/test_subclass.py @@ -557,3 +557,36 @@ def strech(row): result = df.apply(lambda x: [1, 2, 3], axis=1) assert not isinstance(result, tm.SubclassedDataFrame) tm.assert_series_equal(result, expected) + + def test_subclassed_count(self): + # GH 31139 + + df = tm.SubclassedDataFrame( + { + "Person": ["John", "Myla", "Lewis", "John", "Myla"], + "Age": [24.0, np.nan, 21.0, 33, 26], + "Single": [False, True, True, True, False], + } + ) + result = df.count() + assert isinstance(result, tm.SubclassedSeries) + + df = tm.SubclassedDataFrame({"A": [1, 0, 3], "B": [0, 5, 6], "C": [7, 8, 0]}) + result = df.count() + assert isinstance(result, tm.SubclassedSeries) + + df = tm.SubclassedDataFrame( + [[10, 11, 12, 13], [20, 21, 22, 23], [30, 31, 32, 33], [40, 41, 42, 43]], + index=MultiIndex.from_tuples( + list(zip(list("AABB"), list("cdcd"))), names=["aaa", "ccc"] + ), + columns=MultiIndex.from_tuples( + list(zip(list("WWXX"), list("yzyz"))), names=["www", "yyy"] + ), + ) + result = df.count(level=1) + assert isinstance(result, tm.SubclassedDataFrame) + + df = tm.SubclassedDataFrame() + result = df.count() + assert isinstance(result, tm.SubclassedSeries)