Add any, all, prod to dataframe.

Change-Id: I0d7d53c0a7174f9a1251bb8b6261d855db69fdf2

Author: TrevorBergeron

bigframes/core/__init__.py

@@ -797,15 +797,20 @@ def _ibis_window_from_spec(self, window_spec: WindowSpec, allow_ties: bool = Fal
         )
 
     def transpose_single_row(
-        self, labels, *, index_col_id: str = "index", value_col_id: str = "values"
+        self,
+        labels,
+        *,
+        index_col_id: str = "index",
+        value_col_id: str = "values",
+        dtype=pandas.Float64Dtype(),
     ) -> ArrayValue:
-        """Pivot a single row into a 3 column expression with index, values and offsets. Only works if all values can be cast to float."""
+        """Pivot a single row into a 3 column expression with index, values and offsets. Only works if all values can be cast to a common type."""
         table = self.to_ibis_expr(ordering_mode="unordered")
         sub_expressions = []
         for i, col_id in enumerate(self._column_names.keys()):
             sub_expr = table.select(
                 ibis_types.literal(labels[i]).name(index_col_id),
-                _numeric_to_float(table[col_id]).name(value_col_id),
+                ops.AsTypeOp(dtype)._as_ibis(table[col_id]).name(value_col_id),
                 ibis_types.literal(i).name(ORDER_ID_COLUMN),
             )
             sub_expressions.append(sub_expr)
@@ -998,12 +1003,3 @@ def _as_identity(value: ibis_types.Value):
     if value.type().is_float64() or value.type().is_geospatial():
         return value.cast(ibis_dtypes.str)
     return value
-
-
-def _numeric_to_float(value: ibis_types.Value):
-    if value.type().is_float64():
-        return value
-    if value.type().is_boolean():
-        return value.cast(ibis_dtypes.int64).cast(ibis_dtypes.float64)
-    else:
-        return value.cast(ibis_dtypes.float64)

bigframes/core/blocks.py

@@ -571,12 +571,16 @@ def aggregate_all_and_pivot(
         *,
         value_col_id: str = "values",
         dropna: bool = True,
+        dtype=pd.Float64Dtype(),
     ) -> Block:
         aggregations = [(col_id, operation, col_id) for col_id in self.value_columns]
         result_expr = self.expr.aggregate(
             aggregations, dropna=dropna
         ).transpose_single_row(
-            labels=self.column_labels, index_col_id="index", value_col_id=value_col_id
+            labels=self.column_labels,
+            index_col_id="index",
+            value_col_id=value_col_id,
+            dtype=dtype,
         )
         return Block(result_expr, index_columns=["index"], column_labels=[None])
 

bigframes/dataframe.py

@@ -892,6 +892,30 @@ def dropna(self) -> DataFrame:
 
         return DataFrame(block)
 
+    def any(
+        self,
+        *,
+        bool_only: bool = False,
+    ) -> bigframes.series.Series:
+        if not bool_only:
+            frame = self._raise_on_non_boolean("any")
+        else:
+            frame = self._drop_non_bool()
+        block = frame._block.aggregate_all_and_pivot(
+            agg_ops.any_op, dtype=pd.BooleanDtype()
+        )
+        return bigframes.series.Series(block.select_column("values"))
+
+    def all(self, *, bool_only: bool = False) -> bigframes.series.Series:
+        if not bool_only:
+            frame = self._raise_on_non_boolean("all")
+        else:
+            frame = self._drop_non_bool()
+        block = frame._block.aggregate_all_and_pivot(
+            agg_ops.all_op, dtype=pd.BooleanDtype()
+        )
+        return bigframes.series.Series(block.select_column("values"))
+
     def sum(self, *, numeric_only: bool = False) -> bigframes.series.Series:
         if not numeric_only:
             frame = self._raise_on_non_numeric("sum")
@@ -940,6 +964,16 @@ def max(self, *, numeric_only: bool = False) -> bigframes.series.Series:
         block = frame._block.aggregate_all_and_pivot(agg_ops.max_op)
         return bigframes.series.Series(block.select_column("values"))
 
+    def prod(self, *, numeric_only: bool = False) -> bigframes.series.Series:
+        if not numeric_only:
+            frame = self._raise_on_non_numeric("prod")
+        else:
+            frame = self._drop_non_numeric()
+        block = frame._block.aggregate_all_and_pivot(agg_ops.product_op)
+        return bigframes.series.Series(block.select_column("values"))
+
+    product = prod
+
     def count(self, *, numeric_only: bool = False) -> bigframes.series.Series:
         if not numeric_only:
             frame = self
@@ -960,6 +994,14 @@ def _drop_non_numeric(self) -> DataFrame:
         ]
         return DataFrame(self._block.drop_columns(non_numeric_cols))
 
+    def _drop_non_bool(self) -> DataFrame:
+        non_bool_cols = [
+            col_id
+            for col_id, dtype in zip(self._block.value_columns, self._block.dtypes)
+            if dtype not in bigframes.dtypes.BOOL_BIGFRAMES_TYPES
+        ]
+        return DataFrame(self._block.drop_columns(non_bool_cols))
+
     def _raise_on_non_numeric(self, op: str):
         if not all(
             dtype in bigframes.dtypes.NUMERIC_BIGFRAMES_TYPES
@@ -970,6 +1012,16 @@ def _raise_on_non_numeric(self, op: str):
             )
         return self
 
+    def _raise_on_non_boolean(self, op: str):
+        if not all(
+            dtype in bigframes.dtypes.BOOL_BIGFRAMES_TYPES
+            for dtype in self._block.dtypes
+        ):
+            raise NotImplementedError(
+                f"'{op}' does not support non-bool columns. Set 'bool_only'=True to ignore non-bool columns"
+            )
+        return self
+
     def merge(
         self,
         right: DataFrame,

bigframes/dtypes.py

@@ -61,6 +61,8 @@
     ibis_dtypes.Timestamp,
 ]
 
+BOOL_BIGFRAMES_TYPES = [pd.BooleanDtype()]
+
 # Several operations are restricted to these types.
 NUMERIC_BIGFRAMES_TYPES = [pd.BooleanDtype(), pd.Float64Dtype(), pd.Int64Dtype()]
 
@@ -296,4 +298,7 @@ def cast_ibis_value(value: ibis_types.Value, to_type: IbisDtype) -> ibis_types.V
     if value.type() == ibis_dtypes.bool and to_type == ibis_dtypes.string:
         return typing.cast(ibis_types.StringValue, value.cast(to_type)).capitalize()
 
+    if value.type() == ibis_dtypes.bool and to_type == ibis_dtypes.float64:
+        return value.cast(ibis_dtypes.int64).cast(ibis_dtypes.float64)
+
     raise TypeError(f"Unsupported cast {value.type()} to {to_type}")

tests/system/small/test_dataframe.py

@@ -1305,6 +1305,39 @@ def test_dataframe_aggregates(scalars_df_index, scalars_pandas_df_index, op):
     pd.testing.assert_series_equal(pd_series, bf_result, check_index_type=False)
 
 
+@pytest.mark.parametrize(
+    ("op"),
+    [
+        (lambda x: x.all(bool_only=True)),
+        (lambda x: x.any(bool_only=True)),
+    ],
+    ids=["all", "any"],
+)
+def test_dataframe_bool_aggregates(scalars_df_index, scalars_pandas_df_index, op):
+    # Pandas will drop nullable 'boolean' dtype so we convert first to bool, then cast back later
+    scalars_pandas_df_index = scalars_pandas_df_index.assign(
+        bool_col=scalars_pandas_df_index.bool_col.fillna(False).astype("bool")
+    )
+    bf_series = op(scalars_df_index)
+    pd_series = op(scalars_pandas_df_index).astype("boolean")
+    bf_result = bf_series.compute()
+
+    # Pandas has object index type
+    pd.testing.assert_series_equal(pd_series, bf_result, check_index_type=False)
+
+
+def test_dataframe_prod(scalars_df_index, scalars_pandas_df_index):
+    col_names = ["int64_too", "float64_col"]
+    bf_series = scalars_df_index[col_names].prod()
+    pd_series = scalars_pandas_df_index[col_names].prod()
+    bf_result = bf_series.compute()
+
+    # Pandas may produce narrower numeric types, but bigframes always produces Float64
+    pd_series = pd_series.astype("Float64")
+    # Pandas has object index type
+    pd.testing.assert_series_equal(pd_series, bf_result, check_index_type=False)
+
+
 @pytest.mark.parametrize(
     ("frac", "n", "random_state"),
     [

third_party/bigframes_vendored/pandas/core/frame.py

@@ -874,6 +874,53 @@ def merge(
     # ----------------------------------------------------------------------
     # ndarray-like stats methods
 
+    def any(self, *, bool_only: bool = False):
+        """
+        Return whether any element is True, potentially over an axis.
+
+        Returns False unless there is at least one element within a series or
+        along a Dataframe axis that is True or equivalent (e.g. non-zero or
+        non-empty).
+
+        Args:
+            bool_only:
+                Include only boolean columns.
+
+        Returns:
+            Series
+        """
+        raise NotImplementedError("abstract method")
+
+    def all(self, *, bool_only: bool = False):
+        """
+        Return whether all elements are True, potentially over an axis.
+
+        Returns True unless there at least one element within a series or
+        along a Dataframe axis that is False or equivalent (e.g. zero or
+        empty).
+
+        Args:
+            bool_only:
+                Include only boolean columns.
+
+        Returns:
+            Series
+        """
+        raise NotImplementedError("abstract method")
+
+    def prod(self, *, numeric_only: bool = False):
+        """
+        Return the product of the values over the requested axis.
+
+        Args:
+            numeric_only:
+                Include only float, int, boolean columns.
+
+        Returns:
+            Series
+        """
+        raise NotImplementedError("abstract method")
+
     def min(self, *, numeric_only: bool = False):
         """Return the minimum of the values over the requested axis.