feat: Support bool and bytes types in describe(include='all') (#994)

* feat: Support bool and bytes types in describe(include='all') * update aggregation unit tests * fix typo and remove unnecessary helper * remove unnecessary dep * fix wording

Author: sycai

bigframes/dataframe.py

@@ -2303,7 +2303,7 @@ def melt(
  self._block.melt(id_col_ids, val_col_ids, var_name, value_name)
  )
 
- _NUMERICAL_DISCRIBE_AGGS = (
+ _NUMERIC_DESCRIBE_AGGS = (
  "count",
  "mean",
  "std",
@@ -2313,41 +2313,53 @@ def melt(
  "75%",
  "max",
  )
- _NON_NUMERICAL_DESCRIBE_AGGS = ("count", "nunique")
+ _NON_NUMERIC_DESCRIBE_AGGS = ("count", "nunique")
 
  def describe(self, include: None | Literal["all"] = None) -> DataFrame:
+
+ allowed_non_numeric_types = {
+ bigframes.dtypes.STRING_DTYPE,
+ bigframes.dtypes.BOOL_DTYPE,
+ bigframes.dtypes.BYTES_DTYPE,
+ }
+
  if include is None:
  numeric_df = self._drop_non_numeric(permissive=False)
  if len(numeric_df.columns) == 0:
- # Describe eligible non-numerical columns
- result = self._drop_non_string().agg(self._NON_NUMERICAL_DESCRIBE_AGGS)
+ # Describe eligible non-numeric columns
+ result = self.select_dtypes(include=allowed_non_numeric_types).agg(
+ self._NON_NUMERIC_DESCRIBE_AGGS
+ )
  else:
- # Otherwise, only describe numerical columns
- result = numeric_df.agg(self._NUMERICAL_DISCRIBE_AGGS)
+ # Otherwise, only describe numeric columns
+ result = numeric_df.agg(self._NUMERIC_DESCRIBE_AGGS)
  return typing.cast(DataFrame, result)
 
  elif include == "all":
  numeric_result = typing.cast(
  DataFrame,
  self._drop_non_numeric(permissive=False).agg(
- self._NUMERICAL_DISCRIBE_AGGS
+ self._NUMERIC_DESCRIBE_AGGS
  ),
  )
- string_result = typing.cast(
+
+ non_numeric_result = typing.cast(
  DataFrame,
- self._drop_non_string().agg(self._NON_NUMERICAL_DESCRIBE_AGGS),
+ self.select_dtypes(include=allowed_non_numeric_types).agg(
+ self._NON_NUMERIC_DESCRIBE_AGGS
+ ),
  )
 
  if len(numeric_result.columns) == 0:
- return string_result
- elif len(string_result.columns) == 0:
+ return non_numeric_result
+ elif len(non_numeric_result.columns) == 0:
  return numeric_result
  else:
  import bigframes.core.reshape as rs
 
  # Use reindex after join to preserve the original column order.
  return rs.concat(
- [numeric_result, string_result], axis=1
+ [non_numeric_result, numeric_result], axis=1
  )._reindex_columns(self.columns)
 
  else:
@@ -2549,26 +2561,18 @@ def unstack(self, level: LevelsType = -1):
  return DataFrame(pivot_block)
 
  def _drop_non_numeric(self, permissive=True) -> DataFrame:
- numerical_types = (
+ numeric_types = (
  set(bigframes.dtypes.NUMERIC_BIGFRAMES_TYPES_PERMISSIVE)
  if permissive
  else set(bigframes.dtypes.NUMERIC_BIGFRAMES_TYPES_RESTRICTIVE)
  )
  non_numeric_cols = [
  col_id
  for col_id, dtype in zip(self._block.value_columns, self._block.dtypes)
- if dtype not in numerical_types
+ if dtype not in numeric_types
  ]
  return DataFrame(self._block.drop_columns(non_numeric_cols))
 
- def _drop_non_string(self) -> DataFrame:
- string_cols = [
- col_id
- for col_id, dtype in zip(self._block.value_columns, self._block.dtypes)
- if dtype == bigframes.dtypes.STRING_DTYPE
- ]
- return DataFrame(self._block.select_columns(string_cols))
-
  def _drop_non_bool(self) -> DataFrame:
  non_bool_cols = [
  col_id

bigframes/operations/aggregations.py

@@ -568,7 +568,7 @@ def is_agg_op_supported(dtype: dtypes.Dtype, op: AggregateOp) -> bool:
  if dtype in dtypes.NUMERIC_BIGFRAMES_TYPES_PERMISSIVE:
  return True
 
- if dtype == dtypes.STRING_DTYPE:
+ if dtype in (dtypes.STRING_DTYPE, dtypes.BOOL_DTYPE, dtypes.BYTES_DTYPE):
  return isinstance(op, (CountOp, NuniqueOp))
 
  # For all other types, support no aggregation

tests/system/small/test_dataframe.py

@@ -2619,15 +2619,15 @@ def test_df_describe(scalars_dfs):
 
 @skip_legacy_pandas
 @pytest.mark.parametrize("include", [None, "all"])
-def test_df_describe_non_numerical(scalars_dfs, include):
+def test_df_describe_non_numeric(scalars_dfs, include):
  scalars_df, scalars_pandas_df = scalars_dfs
 
- non_numerical_columns = ["string_col"]
+ non_numeric_columns = ["string_col", "bytes_col", "bool_col"]
 
- modified_bf = scalars_df[non_numerical_columns]
+ modified_bf = scalars_df[non_numeric_columns]
  bf_result = modified_bf.describe(include=include).to_pandas()
 
- modified_pd_df = scalars_pandas_df[non_numerical_columns]
+ modified_pd_df = scalars_pandas_df[non_numeric_columns]
  pd_result = modified_pd_df.describe(include=include)
 
  # Reindex results with the specified keys and their order, because
@@ -2639,8 +2639,8 @@ def test_df_describe_non_numerical(scalars_dfs, include):
  ).rename(index={"unique": "nunique"})
 
  pd.testing.assert_frame_equal(
- pd_result[non_numerical_columns].astype("Int64"),
- bf_result[non_numerical_columns],
+ pd_result[non_numeric_columns].astype("Int64"),
+ bf_result[non_numeric_columns],
  check_index_type=False,
  )
 
@@ -2649,12 +2649,12 @@ def test_df_describe_non_numerical(scalars_dfs, include):
 def test_df_describe_mixed_types_include_all(scalars_dfs):
  scalars_df, scalars_pandas_df = scalars_dfs
 
- numerical_columns = [
+ numeric_columns = [
  "int64_col",
  "float64_col",
  ]
- non_numerical_columns = ["string_col"]
- supported_columns = numerical_columns + non_numerical_columns
+ non_numeric_columns = ["string_col"]
+ supported_columns = numeric_columns + non_numeric_columns
 
  modified_bf = scalars_df[supported_columns]
  bf_result = modified_bf.describe(include="all").to_pandas()
@@ -2678,14 +2678,14 @@ def test_df_describe_mixed_types_include_all(scalars_dfs):
  ).rename(index={"unique": "nunique"})
 
  pd.testing.assert_frame_equal(
- pd_result[numerical_columns].astype("Float64"),
- bf_result[numerical_columns],
+ pd_result[numeric_columns].astype("Float64"),
+ bf_result[numeric_columns],
  check_index_type=False,
  )
 
  pd.testing.assert_frame_equal(
- pd_result[non_numerical_columns].astype("Int64"),
- bf_result[non_numerical_columns],
+ pd_result[non_numeric_columns].astype("Int64"),
+ bf_result[non_numeric_columns],
  check_index_type=False,
  )
 

tests/unit/operations/test_aggregations.py

@@ -55,38 +55,29 @@
  first_op,
  ]
 )
-_STRING_SUPPORTED_OPS = set([count_op, nunique_op])
 
 
 @pytest.mark.parametrize("dtype", dtypes.NUMERIC_BIGFRAMES_TYPES_PERMISSIVE)
 @pytest.mark.parametrize("op", _ALL_OPS)
-def test_is_agg_op_supported_numerical_support_all(dtype, op):
+def test_is_agg_op_supported_numeric_support_all(dtype, op):
  assert is_agg_op_supported(dtype, op) is True
 
 
-@pytest.mark.parametrize("dtype", [dtypes.STRING_DTYPE])
-@pytest.mark.parametrize("op", _STRING_SUPPORTED_OPS)
-def test_is_agg_op_supported_string_support_ops(dtype, op):
- assert is_agg_op_supported(dtype, op) is True
-
-
-@pytest.mark.parametrize("dtype", [dtypes.STRING_DTYPE])
-@pytest.mark.parametrize("op", _ALL_OPS - _STRING_SUPPORTED_OPS)
-def test_is_agg_op_supported_string_not_support_ops(dtype, op):
- assert is_agg_op_supported(dtype, op) is False
-
-
 @pytest.mark.parametrize(
- "dtype",
+ ("dtype", "supported_ops"),
  [
- dtypes.BYTES_DTYPE,
- dtypes.DATE_DTYPE,
- dtypes.TIME_DTYPE,
- dtypes.DATETIME_DTYPE,
- dtypes.TIMESTAMP_DTYPE,
- dtypes.GEO_DTYPE,
+ (dtypes.STRING_DTYPE, {count_op, nunique_op}),
+ (dtypes.BYTES_DTYPE, {count_op, nunique_op}),
+ (dtypes.DATE_DTYPE, set()),
+ (dtypes.TIME_DTYPE, set()),
+ (dtypes.DATETIME_DTYPE, set()),
+ (dtypes.TIMESTAMP_DTYPE, set()),
+ (dtypes.GEO_DTYPE, set()),
  ],
 )
-@pytest.mark.parametrize("op", _ALL_OPS)
-def test_is_agg_op_supported_non_numerical_no_support(dtype, op):
- assert is_agg_op_supported(dtype, op) is False
+def test_is_agg_op_supported_non_numeric(dtype, supported_ops):
+ for op in supported_ops:
+ assert is_agg_op_supported(dtype, op) is True
+
+ for op in _ALL_OPS - supported_ops:
+ assert is_agg_op_supported(dtype, op) is False