# Buggy code
def count(self, axis=0, level=None, numeric_only=False):
    """Get the count of non-null objects in the DataFrame.

    Arguments:
        axis: 0 or 'index' for row-wise, 1 or 'columns' for column-wise.
        level: If the axis is a MultiIndex (hierarchical), count along a
            particular level, collapsing into a DataFrame.
        numeric_only: Include only float, int, boolean data

    Returns:
        The count, in a Series (or DataFrame if level is specified).
    """

    axis = self._get_axis_number(axis) if axis is not None else 0

    if level is not None:
        if not isinstance(self.axes[axis], pandas.MultiIndex):
            # error thrown by pandas
            raise TypeError("Can only count levels on hierarchical columns.")

        if isinstance(level, str):
            level = self.axes[axis].names.index(level)

        new_names = dict(
            zip(
                range(len(self.axes[axis].levels[level])),
                self.axes[axis].levels[level],
            )
        )
        return (
            self.groupby(self.axes[axis].codes[level], axis=axis)
            .count()
            .rename(new_names, axis=axis)
            .rename_axis(self.axes[axis].names[level], axis=axis)
        )

    return self._reduce_dimension(
        self._query_compiler.count(axis=axis, level=level, numeric_only=numeric_only)
    )


# Traceback (KeyError)
p0 Index(['data0', 'data1', 'data2'], dtype='object')
---------------------------------------------------------------------------
KeyError                                  Traceback (most recent call last)
<ipython-input-6-704c3a79c97b> in <module>()
----> 1 df_g=csv2.groupby('p0').get_group('A')

~/.local/lib/python3.7/site-packages/modin/pandas/dataframe.py in groupby(self, by, axis, level, as_index, sort, group_keys, squeeze, observed, **kwargs)
353         elif isinstance(by, string_types):
354             idx_name = by
--> 355             by = self.__getitem__(by)._query_compiler
356         elif is_list_like(by):
357             if isinstance(by, Series):

~/.local/lib/python3.7/site-packages/modin/pandas/base.py in __getitem__(self, key)
3220             return self._getitem_slice(indexer)
3221         else:
-> 3222             return self._getitem(key)
3223
3224     def _getitem_slice(self, key):

~/.local/lib/python3.7/site-packages/modin/pandas/dataframe.py in _getitem(self, key)
1835         else:
1836             print(key, self.columns)
-> 1837             return self._getitem_column(key)
1838
1839     def _getitem_column(self, key):

~/.local/lib/python3.7/site-packages/modin/pandas/dataframe.py in _getitem_column(self, key)
1839     def _getitem_column(self, key):
1840         if key not in self.keys():
-> 1841             raise KeyError("{}".format(key))
1842         s = self._reduce_dimension(self._query_compiler.getitem_column_array([key]))
1843         s._parent = self

KeyError: 'p0'

# Fixed code
def count(self, axis=0, level=None, numeric_only=False):
    """Get the count of non-null objects in the DataFrame.

    Arguments:
        axis: 0 or 'index' for row-wise, 1 or 'columns' for column-wise.
        level: If the axis is a MultiIndex (hierarchical), count along a
            particular level, collapsing into a DataFrame.
        numeric_only: Include only float, int, boolean data

    Returns:
        The count, in a Series (or DataFrame if level is specified).
    """

    axis = self._get_axis_number(axis) if axis is not None else 0

    if level is not None:
        if not isinstance(self.axes[axis], pandas.MultiIndex):
            # error thrown by pandas
            raise TypeError("Can only count levels on hierarchical columns.")

        if isinstance(level, str):
            level = self.axes[axis].names.index(level)
        return self.groupby(level=level, axis=axis).count()

    return self._reduce_dimension(
        self._query_compiler.count(axis=axis, level=level, numeric_only=numeric_only)
    )