fix: Use exact median implementation by default (#619)

Author: TrevorBergeron

bigframes/core/block_transforms.py

@@ -111,6 +111,7 @@ def quantile(
     columns: Sequence[str],
     qs: Sequence[float],
     grouping_column_ids: Sequence[str] = (),
+    dropna: bool = False,
 ) -> blocks.Block:
     # TODO: handle windowing and more interpolation methods
     window = core.WindowSpec(
@@ -134,7 +135,7 @@ def quantile(
     block, results = block.aggregate(
         grouping_column_ids,
         tuple((col, agg_ops.AnyValueOp()) for col in quantile_cols),
-        dropna=True,
+        dropna=dropna,
     )
     return block.select_columns(results).with_column_labels(labels)
 

bigframes/core/groupby/__init__.py

@@ -113,9 +113,7 @@ def mean(self, numeric_only: bool = False, *args) -> df.DataFrame:
             self._raise_on_non_numeric("mean")
         return self._aggregate_all(agg_ops.mean_op, numeric_only=True)
 
-    def median(
-        self, numeric_only: bool = False, *, exact: bool = False
-    ) -> df.DataFrame:
+    def median(self, numeric_only: bool = False, *, exact: bool = True) -> df.DataFrame:
         if not numeric_only:
             self._raise_on_non_numeric("median")
         if exact:
@@ -138,6 +136,7 @@ def quantile(
             q_cols,
             qs=tuple(q) if multi_q else (q,),  # type: ignore
             grouping_column_ids=self._by_col_ids,
+            dropna=self._dropna,
         )
         result_df = df.DataFrame(result)
         if multi_q:
@@ -491,7 +490,7 @@ def mean(self, *args) -> series.Series:
     def median(
         self,
         *args,
-        exact: bool = False,
+        exact: bool = True,
         **kwargs,
     ) -> series.Series:
         if exact:
@@ -508,6 +507,7 @@ def quantile(
             (self._value_column,),
             qs=tuple(q) if multi_q else (q,),  # type: ignore
             grouping_column_ids=self._by_col_ids,
+            dropna=self._dropna,
         )
         if multi_q:
             return series.Series(result.stack())

bigframes/dataframe.py

@@ -1995,18 +1995,16 @@ def mean(
         return bigframes.series.Series(block.select_column("values"))
 
     def median(
-        self, *, numeric_only: bool = False, exact: bool = False
+        self, *, numeric_only: bool = False, exact: bool = True
     ) -> bigframes.series.Series:
-        if exact:
-            raise NotImplementedError(
-                f"Only approximate median is supported. {constants.FEEDBACK_LINK}"
-            )
         if not numeric_only:
             frame = self._raise_on_non_numeric("median")
         else:
             frame = self._drop_non_numeric()
         if exact:
-            return self.quantile()
+            result = frame.quantile()
+            result.name = None
+            return result
         else:
             block = frame._block.aggregate_all_and_stack(agg_ops.median_op)
             return bigframes.series.Series(block.select_column("values"))

bigframes/series.py

@@ -966,7 +966,7 @@ def mode(self) -> Series:
     def mean(self) -> float:
         return typing.cast(float, self._apply_aggregation(agg_ops.mean_op))
 
-    def median(self, *, exact: bool = False) -> float:
+    def median(self, *, exact: bool = True) -> float:
         if exact:
             return typing.cast(float, self.quantile(0.5))
         else:

tests/system/small/test_series.py

@@ -1345,10 +1345,9 @@ def test_numeric_literal(scalars_dfs):
     scalars_df, _ = scalars_dfs
     col_name = "numeric_col"
     assert scalars_df[col_name].dtype == pd.ArrowDtype(pa.decimal128(38, 9))
-    bf_result = scalars_df[col_name] - scalars_df[col_name].median()
+    bf_result = scalars_df[col_name] + 42
     assert bf_result.size == scalars_df[col_name].size
-    # TODO(b/323387826): The precision increased by 1 unexpectedly.
-    # assert bf_result.dtype == pd.ArrowDtype(pa.decimal128(38, 9))
+    assert bf_result.dtype == pd.ArrowDtype(pa.decimal128(38, 9))
 
 
 def test_repr(scalars_dfs):
@@ -1523,12 +1522,32 @@ def test_groupby_mean(scalars_dfs):
     )
 
 
-def test_groupby_median(scalars_dfs):
+def test_groupby_median_exact(scalars_dfs):
     scalars_df, scalars_pandas_df = scalars_dfs
     col_name = "int64_too"
-    bf_series = (
+    bf_result = (
         scalars_df[col_name].groupby(scalars_df["string_col"], dropna=False).median()
     )
+    pd_result = (
+        scalars_pandas_df[col_name]
+        .groupby(scalars_pandas_df["string_col"], dropna=False)
+        .median()
+    )
+
+    assert_series_equal(
+        pd_result,
+        bf_result.to_pandas(),
+    )
+
+
+def test_groupby_median_inexact(scalars_dfs):
+    scalars_df, scalars_pandas_df = scalars_dfs
+    col_name = "int64_too"
+    bf_series = (
+        scalars_df[col_name]
+        .groupby(scalars_df["string_col"], dropna=False)
+        .median(exact=False)
+    )
     pd_max = (
         scalars_pandas_df[col_name]
         .groupby(scalars_pandas_df["string_col"], dropna=False)

third_party/bigframes_vendored/pandas/core/frame.py

@@ -4481,7 +4481,7 @@ def mean(self, axis=0, *, numeric_only: bool = False):
         """
         raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE)
 
-    def median(self, *, numeric_only: bool = False, exact: bool = False):
+    def median(self, *, numeric_only: bool = False, exact: bool = True):
         """Return the median of the values over colunms.
 
         **Examples:**
@@ -4500,15 +4500,15 @@ def median(self, *, numeric_only: bool = False, exact: bool = False):
         Finding the median value of each column.
 
             >>> df.median()
-            A    1
-            B    2
-            dtype: Int64
+            A    2.0
+            B    3.0
+            dtype: Float64
 
         Args:
             numeric_only (bool. default False):
                 Default False. Include only float, int, boolean columns.
-            exact (bool. default False):
-                Default False. Get the exact median instead of an approximate
+            exact (bool. default True):
+                Default True. Get the exact median instead of an approximate
                 one.
 
         Returns:

third_party/bigframes_vendored/pandas/core/groupby/__init__.py

@@ -68,17 +68,16 @@ def median(
         self,
         numeric_only: bool = False,
         *,
-        exact: bool = False,
+        exact: bool = True,
     ):
         """
         Compute median of groups, excluding missing values.
 
         Args:
             numeric_only (bool, default False):
                 Include only float, int, boolean columns.
-            exact (bool, default False):
-                Calculate the exact median instead of an approximation. Note:
-                    ``exact=True`` is not supported.
+            exact (bool, default True):
+                Calculate the exact median instead of an approximation.
 
         Returns:
             pandas.Series or pandas.DataFrame: Median of groups.

third_party/bigframes_vendored/pandas/core/series.py

@@ -3150,13 +3150,13 @@ def mean(self):
         """
         raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE)
 
-    def median(self, *, exact: bool = False):
+    def median(self, *, exact: bool = True):
         """Return the median of the values over the requested axis.
 
         Args:
-            exact (bool. default False):
-                Default False. Get the exact median instead of an approximate
-                one. Note: ``exact=True`` not yet supported.
+            exact (bool. default True):
+                Default True. Get the exact median instead of an approximate
+                one.
 
         Returns:
             scalar: Scalar.