Merge branch 'main' into b329865893-groupby-iter

Author: tswast

bigframes/core/array_value.py

@@ -480,6 +480,14 @@ def relational_join(
         type: typing.Literal["inner", "outer", "left", "right", "cross"] = "inner",
         propogate_order: Optional[bool] = None,
     ) -> typing.Tuple[ArrayValue, typing.Tuple[dict[str, str], dict[str, str]]]:
+        for lcol, rcol in conditions:
+            ltype = self.get_column_type(lcol)
+            rtype = other.get_column_type(rcol)
+            if not bigframes.dtypes.can_compare(ltype, rtype):
+                raise TypeError(
+                    f"Cannot join with non-comparable join key types: {ltype}, {rtype}"
+                )
+
         l_mapping = {  # Identity mapping, only rename right side
             lcol.name: lcol.name for lcol in self.node.ids
         }

bigframes/core/blocks.py

@@ -1786,15 +1786,19 @@ def pivot(
         else:
             return result_block.with_column_labels(columns_values)
 
-    def stack(self, how="left", levels: int = 1):
+    def stack(
+        self, how="left", levels: int = 1, *, override_labels: Optional[pd.Index] = None
+    ):
         """Unpivot last column axis level into row axis"""
         if levels == 0:
             return self
 
         # These are the values that will be turned into rows
 
         col_labels, row_labels = utils.split_index(self.column_labels, levels=levels)
-        row_labels = row_labels.drop_duplicates()
+        row_labels = (
+            row_labels.drop_duplicates() if override_labels is None else override_labels
+        )
 
         if col_labels is None:
             result_index: pd.Index = pd.Index([None])

bigframes/core/compile/api.py

@@ -15,19 +15,18 @@
 
 from typing import TYPE_CHECKING
 
-from bigframes.core import rewrite
-from bigframes.core.compile.ibis_compiler import ibis_compiler
-
 if TYPE_CHECKING:
     import bigframes.core.nodes
 
 
 def test_only_ibis_inferred_schema(node: bigframes.core.nodes.BigFrameNode):
     """Use only for testing paths to ensure ibis inferred schema does not diverge from bigframes inferred schema."""
+    from bigframes.core.compile.ibis_compiler import ibis_compiler
+    import bigframes.core.rewrite
     import bigframes.core.schema
 
     node = ibis_compiler._replace_unsupported_ops(node)
-    node = rewrite.bake_order(node)
+    node = bigframes.core.rewrite.bake_order(node)
     ir = ibis_compiler.compile_node(node)
     items = tuple(
         bigframes.core.schema.SchemaItem(name, ir.get_column_type(ibis_id))

bigframes/core/groupby/dataframe_group_by.py

@@ -153,6 +153,20 @@ def head(self, n: int = 5) -> df.DataFrame:
             )
         )
 
+    def describe(self, include: None | Literal["all"] = None):
+        from bigframes.pandas.core.methods import describe
+
+        return df.DataFrame(
+            describe._describe(
+                self._block,
+                self._selected_cols,
+                include,
+                as_index=self._as_index,
+                by_col_ids=self._by_col_ids,
+                dropna=self._dropna,
+            )
+        )
+
     def __iter__(self) -> Iterable[Tuple[blocks.Label, df.DataFrame]]:
         for group_keys, filtered_block in group_by.block_groupby_iter(
             self._block,

bigframes/core/groupby/series_group_by.py

@@ -81,6 +81,20 @@ def head(self, n: int = 5) -> series.Series:
             )
         )
 
+    def describe(self, include: None | Literal["all"] = None):
+        from bigframes.pandas.core.methods import describe
+
+        return df.DataFrame(
+            describe._describe(
+                self._block,
+                columns=[self._value_column],
+                include=include,
+                as_index=True,
+                by_col_ids=self._by_col_ids,
+                dropna=self._dropna,
+            )
+        ).droplevel(level=0, axis=1)
+
     def __iter__(self) -> Iterable[Tuple[blocks.Label, series.Series]]:
         for group_keys, filtered_block in group_by.block_groupby_iter(
             self._block,

bigframes/core/rewrite/implicit_align.py

@@ -18,12 +18,8 @@
 from typing import cast, Optional, Sequence, Set, Tuple
 
 import bigframes.core.expression
-import bigframes.core.guid
 import bigframes.core.identifiers
-import bigframes.core.join_def
 import bigframes.core.nodes
-import bigframes.core.window_spec
-import bigframes.operations.aggregations
 
 # Combination of selects and additive nodes can be merged as an explicit keyless "row join"
 ALIGNABLE_NODES = (

bigframes/dtypes.py

@@ -358,6 +358,11 @@ def is_comparable(type_: ExpressionType) -> bool:
     return (type_ is not None) and is_orderable(type_)
 
 
+def can_compare(type1: ExpressionType, type2: ExpressionType) -> bool:
+    coerced_type = coerce_to_common(type1, type2)
+    return is_comparable(coerced_type)
+
+
 def get_struct_fields(type_: ExpressionType) -> dict[str, Dtype]:
     assert isinstance(type_, pd.ArrowDtype)
     assert isinstance(type_.pyarrow_dtype, pa.StructType)

bigframes/operations/aggregations.py

@@ -251,12 +251,7 @@ def name(self):
     def output_type(self, *input_types: dtypes.ExpressionType) -> dtypes.ExpressionType:
         if not dtypes.is_orderable(input_types[0]):
             raise TypeError(f"Type {input_types[0]} is not orderable")
-        if pd.api.types.is_bool_dtype(input_types[0]) or pd.api.types.is_integer_dtype(
-            input_types[0]
-        ):
-            return dtypes.FLOAT_DTYPE
-        else:
-            return input_types[0]
+        return input_types[0]
 
 
 @dataclasses.dataclass(frozen=True)

bigframes/operations/type.py

@@ -174,15 +174,7 @@ class CoerceCommon(BinaryTypeSignature):
     def output_type(
         self, left_type: ExpressionType, right_type: ExpressionType
     ) -> ExpressionType:
-        try:
-            return bigframes.dtypes.coerce_to_common(left_type, right_type)
-        except TypeError:
-            pass
-        if bigframes.dtypes.can_coerce(left_type, right_type):
-            return right_type
-        if bigframes.dtypes.can_coerce(right_type, left_type):
-            return left_type
-        raise TypeError(f"Cannot coerce {left_type} and {right_type} to a common type.")
+        return bigframes.dtypes.coerce_to_common(left_type, right_type)
 
 
 @dataclasses.dataclass
@@ -192,8 +184,7 @@ class Comparison(BinaryTypeSignature):
     def output_type(
         self, left_type: ExpressionType, right_type: ExpressionType
     ) -> ExpressionType:
-        common_type = CoerceCommon().output_type(left_type, right_type)
-        if not bigframes.dtypes.is_comparable(common_type):
+        if not bigframes.dtypes.can_compare(left_type, right_type):
             raise TypeError(f"Types {left_type} and {right_type} are not comparable")
         return bigframes.dtypes.BOOL_DTYPE
 

bigframes/pandas/core/methods/describe.py

@@ -16,8 +16,15 @@
 
 import typing
 
+import pandas as pd
+
 from bigframes import dataframe, dtypes, series
-from bigframes.core.reshape import api as rs
+from bigframes.core import agg_expressions, blocks
+from bigframes.operations import aggregations
+
+_DEFAULT_DTYPES = (
+    dtypes.NUMERIC_BIGFRAMES_TYPES_RESTRICTIVE + dtypes.TEMPORAL_NUMERIC_BIGFRAMES_TYPES
+)
 
 
 def describe(
@@ -30,100 +37,88 @@ def describe(
     elif not isinstance(input, dataframe.DataFrame):
         raise TypeError(f"Unsupported type: {type(input)}")
 
+    block = input._block
+
+    describe_block = _describe(block, columns=block.value_columns, include=include)
+    # we override default stack behavior, because we want very specific ordering
+    stack_cols = pd.Index(
+        [
+            "count",
+            "nunique",
+            "top",
+            "freq",
+            "mean",
+            "std",
+            "min",
+            "25%",
+            "50%",
+            "75%",
+            "max",
+        ]
+    ).intersection(describe_block.column_labels.get_level_values(-1))
+    describe_block = describe_block.stack(override_labels=stack_cols)
+
+    return dataframe.DataFrame(describe_block).droplevel(level=0)
+
+
+def _describe(
+    block: blocks.Block,
+    columns: typing.Sequence[str],
+    include: None | typing.Literal["all"] = None,
+    *,
+    as_index: bool = True,
+    by_col_ids: typing.Sequence[str] = [],
+    dropna: bool = False,
+) -> blocks.Block:
+    stats: list[agg_expressions.Aggregation] = []
+    column_labels: list[typing.Hashable] = []
+
+    # include=None behaves like include='all' if no numeric columns present
     if include is None:
-        numeric_df = _select_dtypes(
-            input,
-            dtypes.NUMERIC_BIGFRAMES_TYPES_RESTRICTIVE
-            + dtypes.TEMPORAL_NUMERIC_BIGFRAMES_TYPES,
-        )
-        if len(numeric_df.columns) == 0:
-            # Describe eligible non-numeric columns
-            return _describe_non_numeric(input)
-
-        # Otherwise, only describe numeric columns
-        return _describe_numeric(input)
-
-    elif include == "all":
-        numeric_result = _describe_numeric(input)
-        non_numeric_result = _describe_non_numeric(input)
-
-        if len(numeric_result.columns) == 0:
-            return non_numeric_result
-        elif len(non_numeric_result.columns) == 0:
-            return numeric_result
-        else:
-            # Use reindex after join to preserve the original column order.
-            return rs.concat(
-                [non_numeric_result, numeric_result], axis=1
-            )._reindex_columns(input.columns)
-
-    else:
-        raise ValueError(f"Unsupported include type: {include}")
-
-
-def _describe_numeric(df: dataframe.DataFrame) -> dataframe.DataFrame:
-    number_df_result = typing.cast(
-        dataframe.DataFrame,
-        _select_dtypes(df, dtypes.NUMERIC_BIGFRAMES_TYPES_RESTRICTIVE).agg(
-            [
-                "count",
-                "mean",
-                "std",
-                "min",
-                "25%",
-                "50%",
-                "75%",
-                "max",
-            ]
-        ),
-    )
-    temporal_df_result = typing.cast(
-        dataframe.DataFrame,
-        _select_dtypes(df, dtypes.TEMPORAL_NUMERIC_BIGFRAMES_TYPES).agg(["count"]),
+        if not any(
+            block.expr.get_column_type(col) in _DEFAULT_DTYPES for col in columns
+        ):
+            include = "all"
+
+    for col_id in columns:
+        label = block.col_id_to_label[col_id]
+        dtype = block.expr.get_column_type(col_id)
+        if include != "all" and dtype not in _DEFAULT_DTYPES:
+            continue
+        agg_ops = _get_aggs_for_dtype(dtype)
+        stats.extend(op.as_expr(col_id) for op in agg_ops)
+        label_tuple = (label,) if block.column_labels.nlevels == 1 else label
+        column_labels.extend((*label_tuple, op.name) for op in agg_ops)  # type: ignore
+
+    agg_block, _ = block.aggregate(
+        by_column_ids=by_col_ids,
+        aggregations=stats,
+        dropna=dropna,
+        column_labels=pd.Index(column_labels, name=(*block.column_labels.names, None)),
     )
-
-    if len(number_df_result.columns) == 0:
-        return temporal_df_result
-    elif len(temporal_df_result.columns) == 0:
-        return number_df_result
+    return agg_block if as_index else agg_block.reset_index(drop=False)
+
+
+def _get_aggs_for_dtype(dtype) -> list[aggregations.UnaryAggregateOp]:
+    if dtype in dtypes.NUMERIC_BIGFRAMES_TYPES_RESTRICTIVE:
+        return [
+            aggregations.count_op,
+            aggregations.mean_op,
+            aggregations.std_op,
+            aggregations.min_op,
+            aggregations.ApproxQuartilesOp(1),
+            aggregations.ApproxQuartilesOp(2),
+            aggregations.ApproxQuartilesOp(3),
+            aggregations.max_op,
+        ]
+    elif dtype in dtypes.TEMPORAL_NUMERIC_BIGFRAMES_TYPES:
+        return [aggregations.count_op]
+    elif dtype in [
+        dtypes.STRING_DTYPE,
+        dtypes.BOOL_DTYPE,
+        dtypes.BYTES_DTYPE,
+        dtypes.TIME_DTYPE,
+    ]:
+        return [aggregations.count_op, aggregations.nunique_op]
     else:
-        import bigframes.core.reshape.api as rs
-
-        original_columns = _select_dtypes(
-            df,
-            dtypes.NUMERIC_BIGFRAMES_TYPES_RESTRICTIVE
-            + dtypes.TEMPORAL_NUMERIC_BIGFRAMES_TYPES,
-        ).columns
-
-        # Use reindex after join to preserve the original column order.
-        return rs.concat(
-            [number_df_result, temporal_df_result],
-            axis=1,
-        )._reindex_columns(original_columns)
-
-
-def _describe_non_numeric(df: dataframe.DataFrame) -> dataframe.DataFrame:
-    return typing.cast(
-        dataframe.DataFrame,
-        _select_dtypes(
-            df,
-            [
-                dtypes.STRING_DTYPE,
-                dtypes.BOOL_DTYPE,
-                dtypes.BYTES_DTYPE,
-                dtypes.TIME_DTYPE,
-            ],
-        ).agg(["count", "nunique"]),
-    )
-
-
-def _select_dtypes(
-    df: dataframe.DataFrame, dtypes: typing.Sequence[dtypes.Dtype]
-) -> dataframe.DataFrame:
-    """Selects columns without considering inheritance relationships."""
-    columns = [
-        col_id
-        for col_id, dtype in zip(df._block.value_columns, df._block.dtypes)
-        if dtype in dtypes
-    ]
-    return dataframe.DataFrame(df._block.select_columns(columns))
+        return []