| FazBrowse GitHub Viewer | Trending | | Home |
| Tools: [Download Repo ZIP] [Original HTTPS Page] |
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
@@ -18,8 +18,10 @@ | |||
| 18 | 18 | ||
| 19 | 19 | import bigframes_vendored.pandas.core.col as pd_col | |
| 20 | 20 | ||
| 21 | + from bigframes.core import agg_expressions, window_spec | ||
| 21 | 22 | import bigframes.core.expression as bf_expression | |
| 22 | 23 | import bigframes.operations as bf_ops | |
| 24 | + import bigframes.operations.aggregations as agg_ops | ||
| 23 | 25 | ||
| 24 | 26 | ||
| 25 | 27 | # Not to be confused with the Expression class in `bigframes.core.expressions` | |
@@ -33,6 +35,15 @@ class Expression: | |||
| 33 | 35 | def _apply_unary(self, op: bf_ops.UnaryOp) -> Expression: | |
| 34 | 36 | return Expression(op.as_expr(self._value)) | |
| 35 | 37 | ||
| 38 | + def _apply_unary_agg(self, op: agg_ops.UnaryAggregateOp) -> Expression: | ||
| 39 | + # We probably shouldn't need to windowize here, but block apis expect pre-windowized expressions | ||
| 40 | + # Later on, we will probably have col expressions in windowed context, so will need to defer windowization | ||
| 41 | + # instead of automatically applying the default unbound window | ||
| 42 | + agg_expr = op.as_expr(self._value) | ||
| 43 | + return Expression( | ||
| 44 | + agg_expressions.WindowExpression(agg_expr, window_spec.unbound()) | ||
| 45 | + ) | ||
| 46 | + | ||
| 36 | 47 | def _apply_binary(self, other: Any, op: bf_ops.BinaryOp, reverse: bool = False): | |
| 37 | 48 | if isinstance(other, Expression): | |
| 38 | 49 | other_value = other._value | |
@@ -118,6 +129,24 @@ def __rxor__(self, other: Any) -> Expression: | |||
| 118 | 129 | def __invert__(self) -> Expression: | |
| 119 | 130 | return self._apply_unary(bf_ops.invert_op) | |
| 120 | 131 | ||
| 132 | + def sum(self) -> Expression: | ||
| 133 | + return self._apply_unary_agg(agg_ops.sum_op) | ||
| 134 | + | ||
| 135 | + def mean(self) -> Expression: | ||
| 136 | + return self._apply_unary_agg(agg_ops.mean_op) | ||
| 137 | + | ||
| 138 | + def var(self) -> Expression: | ||
| 139 | + return self._apply_unary_agg(agg_ops.var_op) | ||
| 140 | + | ||
| 141 | + def std(self) -> Expression: | ||
| 142 | + return self._apply_unary_agg(agg_ops.std_op) | ||
| 143 | + | ||
| 144 | + def min(self) -> Expression: | ||
| 145 | + return self._apply_unary_agg(agg_ops.min_op) | ||
| 146 | + | ||
| 147 | + def max(self) -> Expression: | ||
| 148 | + return self._apply_unary_agg(agg_ops.max_op) | ||
| 149 | + | ||
| 121 | 150 | ||
| 122 | 151 | def col(col_name: Hashable) -> Expression: | |
| 123 | 152 | return Expression(bf_expression.free_var(col_name)) | |
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
@@ -100,6 +100,57 @@ def test_pd_col_unary_operators(scalars_dfs, op): | |||
| 100 | 100 | assert_frame_equal(bf_result, pd_result) | |
| 101 | 101 | ||
| 102 | 102 | ||
| 103 | + @pytest.mark.parametrize( | ||
| 104 | + ("op"), | ||
| 105 | + [ | ||
| 106 | + (lambda x: x.sum()), | ||
| 107 | + (lambda x: x.mean()), | ||
| 108 | + (lambda x: x.min()), | ||
| 109 | + (lambda x: x.max()), | ||
| 110 | + (lambda x: x.std()), | ||
| 111 | + (lambda x: x.var()), | ||
| 112 | + ], | ||
| 113 | + ids=[ | ||
| 114 | + "sum", | ||
| 115 | + "mean", | ||
| 116 | + "min", | ||
| 117 | + "max", | ||
| 118 | + "std", | ||
| 119 | + "var", | ||
| 120 | + ], | ||
| 121 | + ) | ||
| 122 | + def test_pd_col_aggregate_op(scalars_dfs, op): | ||
| 123 | + scalars_df, scalars_pandas_df = scalars_dfs | ||
| 124 | + bf_kwargs = { | ||
| 125 | + "result": op(bpd.col("float64_col")), | ||
| 126 | + } | ||
| 127 | + pd_kwargs = { | ||
| 128 | + "result": op(pd.col("float64_col")), # type: ignore | ||
| 129 | + } | ||
| 130 | + df = scalars_df.assign(**bf_kwargs) | ||
| 131 | + | ||
| 132 | + bf_result = df.to_pandas() | ||
| 133 | + pd_result = scalars_pandas_df.assign(**pd_kwargs) | ||
| 134 | + | ||
| 135 | + assert_frame_equal(bf_result, pd_result) | ||
| 136 | + | ||
| 137 | + | ||
| 138 | + def test_pd_col_aggregate_of_aggregate(scalars_dfs): | ||
| 139 | + scalars_df, scalars_pandas_df = scalars_dfs | ||
| 140 | + bf_kwargs = { | ||
| 141 | + "result": (bpd.col("int64_col") - bpd.col("int64_col").mean()).mean(), | ||
| 142 | + } | ||
| 143 | + pd_kwargs = { | ||
| 144 | + "result": (pd.col("int64_col") - pd.col("int64_col").mean()).mean(), # type: ignore | ||
| 145 | + } | ||
| 146 | + df = scalars_df.assign(**bf_kwargs) | ||
| 147 | + | ||
| 148 | + bf_result = df.to_pandas() | ||
| 149 | + pd_result = scalars_pandas_df.assign(**pd_kwargs) | ||
| 150 | + | ||
| 151 | + assert_frame_equal(bf_result, pd_result) | ||
| 152 | + | ||
| 153 | + | ||
| 103 | 154 | @pytest.mark.parametrize( | |
| 104 | 155 | ("op",), | |
| 105 | 156 | [ | |
| Back | FazBrowse Home | New Git URL |
0 commit comments