Merge branch 'main' into tswast-st-regionstats

Author: tswast

bigframes/core/blocks.py

@@ -1996,6 +1996,31 @@ def _generate_resample_label(
             Literal["epoch", "start", "start_day", "end", "end_day"],
         ] = "start_day",
     ) -> Block:
+        if not isinstance(rule, str):
+            raise NotImplementedError(
+                f"Only offset strings are currently supported for rule, but got {repr(rule)}. {constants.FEEDBACK_LINK}"
+            )
+
+        if rule in ("ME", "YE", "QE", "BME", "BA", "BQE", "W"):
+            raise NotImplementedError(
+                f"Offset strings 'ME', 'YE', 'QE', 'BME', 'BA', 'BQE', 'W' are not currently supported for rule, but got {repr(rule)}. {constants.FEEDBACK_LINK}"
+            )
+
+        if closed == "right":
+            raise NotImplementedError(
+                f"Only closed='left' is currently supported. {constants.FEEDBACK_LINK}",
+            )
+
+        if label == "right":
+            raise NotImplementedError(
+                f"Only label='left' is currently supported. {constants.FEEDBACK_LINK}",
+            )
+
+        if origin not in ("epoch", "start", "start_day"):
+            raise NotImplementedError(
+                f"Only origin='epoch', 'start', 'start_day' are currently supported, but got {repr(origin)}. {constants.FEEDBACK_LINK}"
+            )
+
         # Validate and resolve the index or column to use for grouping
         if on is None:
             if len(self.index_columns) == 0:

bigframes/core/compile/sqlglot/aggregations/op_registration.py

@@ -52,5 +52,5 @@ def arg_checker(*args, **kwargs):
     def __getitem__(self, op: str | agg_ops.WindowOp) -> CompilationFunc:
         key = op if isinstance(op, type) else type(op)
         if str(key) not in self._registered_ops:
-            raise ValueError(f"{key} is already not registered")
+            raise ValueError(f"{key} is not registered")
         return self._registered_ops[str(key)]

bigframes/core/compile/sqlglot/aggregations/unary_compiler.py

@@ -239,6 +239,20 @@ def _(
     return apply_window_if_present(sge.func("MIN", column.expr), window)
 
 
+@UNARY_OP_REGISTRATION.register(agg_ops.PopVarOp)
+def _(
+    op: agg_ops.PopVarOp,
+    column: typed_expr.TypedExpr,
+    window: typing.Optional[window_spec.WindowSpec] = None,
+) -> sge.Expression:
+    expr = column.expr
+    if column.dtype == dtypes.BOOL_DTYPE:
+        expr = sge.Cast(this=expr, to="INT64")
+
+    expr = sge.func("VAR_POP", expr)
+    return apply_window_if_present(expr, window)
+
+
 @UNARY_OP_REGISTRATION.register(agg_ops.QuantileOp)
 def _(
     op: agg_ops.QuantileOp,
@@ -278,6 +292,22 @@ def _(
     return apply_window_if_present(sge.func("COUNT", sge.convert(1)), window)
 
 
+@UNARY_OP_REGISTRATION.register(agg_ops.StdOp)
+def _(
+    op: agg_ops.StdOp,
+    column: typed_expr.TypedExpr,
+    window: typing.Optional[window_spec.WindowSpec] = None,
+) -> sge.Expression:
+    expr = column.expr
+    if column.dtype == dtypes.BOOL_DTYPE:
+        expr = sge.Cast(this=expr, to="INT64")
+
+    expr = sge.func("STDDEV", expr)
+    if op.should_floor_result or column.dtype == dtypes.TIMEDELTA_DTYPE:
+        expr = sge.Cast(this=sge.func("FLOOR", expr), to="INT64")
+    return apply_window_if_present(expr, window)
+
+
 @UNARY_OP_REGISTRATION.register(agg_ops.ShiftOp)
 def _(
     op: agg_ops.ShiftOp,
@@ -331,3 +361,17 @@ def _(
         expression=shifted,
         unit=sge.Identifier(this="MICROSECOND"),
     )
+
+
+@UNARY_OP_REGISTRATION.register(agg_ops.VarOp)
+def _(
+    op: agg_ops.VarOp,
+    column: typed_expr.TypedExpr,
+    window: typing.Optional[window_spec.WindowSpec] = None,
+) -> sge.Expression:
+    expr = column.expr
+    if column.dtype == dtypes.BOOL_DTYPE:
+        expr = sge.Cast(this=expr, to="INT64")
+
+    expr = sge.func("VAR_SAMP", expr)
+    return apply_window_if_present(expr, window)

bigframes/core/compile/sqlglot/expressions/comparison_ops.py

@@ -109,6 +109,11 @@ def _(left: TypedExpr, right: TypedExpr) -> sge.Expression:
     return sge.LTE(this=left_expr, expression=right_expr)
 
 
+@register_binary_op(ops.maximum_op)
+def _(left: TypedExpr, right: TypedExpr) -> sge.Expression:
+    return sge.Greatest(expressions=[left.expr, right.expr])
+
+
 @register_binary_op(ops.minimum_op)
 def _(left: TypedExpr, right: TypedExpr) -> sge.Expression:
     return sge.Least(this=left.expr, expressions=right.expr)

bigframes/core/compile/sqlglot/expressions/generic_ops.py

@@ -24,6 +24,7 @@
 import bigframes.core.compile.sqlglot.scalar_compiler as scalar_compiler
 
 register_unary_op = scalar_compiler.scalar_op_compiler.register_unary_op
+register_binary_op = scalar_compiler.scalar_op_compiler.register_binary_op
 register_nary_op = scalar_compiler.scalar_op_compiler.register_nary_op
 register_ternary_op = scalar_compiler.scalar_op_compiler.register_ternary_op
 
@@ -159,6 +160,13 @@ def _(*cases_and_outputs: TypedExpr) -> sge.Expression:
     )
 
 
+@register_binary_op(ops.coalesce_op)
+def _(left: TypedExpr, right: TypedExpr) -> sge.Expression:
+    if left.expr == right.expr:
+        return left.expr
+    return sge.Coalesce(this=left.expr, expressions=[right.expr])
+
+
 @register_nary_op(ops.RowKey)
 def _(*values: TypedExpr) -> sge.Expression:
     # All inputs into hash must be non-null or resulting hash will be null

bigframes/dataframe.py

@@ -4182,10 +4182,12 @@ def _split(
         return [DataFrame(block) for block in blocks]
 
     @validations.requires_ordering()
-    def _resample(
+    def resample(
         self,
         rule: str,
         *,
+        closed: Optional[Literal["right", "left"]] = None,
+        label: Optional[Literal["right", "left"]] = None,
         on: blocks.Label = None,
         level: Optional[LevelsType] = None,
         origin: Union[
@@ -4195,64 +4197,10 @@ def _resample(
             Literal["epoch", "start", "start_day", "end", "end_day"],
         ] = "start_day",
     ) -> bigframes.core.groupby.DataFrameGroupBy:
-        """Internal function to support resample. Resample time-series data.
-
-        **Examples:**
-
-        >>> import bigframes.pandas as bpd
-        >>> data = {
-        ...     "timestamp_col": pd.date_range(
-        ...         start="2021-01-01 13:00:00", periods=30, freq="1s"
-        ...     ),
-        ...     "int64_col": range(30),
-        ...     "int64_too": range(10, 40),
-        ... }
-
-        Resample on a DataFrame with index:
-
-        >>> df = bpd.DataFrame(data).set_index("timestamp_col")
-        >>> df._resample(rule="7s").min()
-                             int64_col  int64_too
-        2021-01-01 12:59:55          0         10
-        2021-01-01 13:00:02          2         12
-        2021-01-01 13:00:09          9         19
-        2021-01-01 13:00:16         16         26
-        2021-01-01 13:00:23         23         33
-        <BLANKLINE>
-        [5 rows x 2 columns]
-
-        Resample with column and origin set to 'start':
-
-        >>> df = bpd.DataFrame(data)
-        >>> df._resample(rule="7s", on = "timestamp_col", origin="start").min()
-                             int64_col  int64_too
-        2021-01-01 13:00:00          0         10
-        2021-01-01 13:00:07          7         17
-        2021-01-01 13:00:14         14         24
-        2021-01-01 13:00:21         21         31
-        2021-01-01 13:00:28         28         38
-        <BLANKLINE>
-        [5 rows x 2 columns]
-
-        Args:
-            rule (str):
-                The offset string representing target conversion.
-            on (str, default None):
-                For a DataFrame, column to use instead of index for resampling. Column
-                must be datetime-like.
-            level (str or int, default None):
-                For a MultiIndex, level (name or number) to use for resampling.
-                level must be datetime-like.
-            origin(str, default 'start_day'):
-                The timestamp on which to adjust the grouping. Must be one of the following:
-                'epoch': origin is 1970-01-01
-                'start': origin is the first value of the timeseries
-                'start_day': origin is the first day at midnight of the timeseries
-        Returns:
-            DataFrameGroupBy: DataFrameGroupBy object.
-        """
         block = self._block._generate_resample_label(
             rule=rule,
+            closed=closed,
+            label=label,
             on=on,
             level=level,
             origin=origin,

bigframes/series.py

@@ -2505,7 +2505,7 @@ def explode(self, *, ignore_index: Optional[bool] = False) -> Series:
         )
 
     @validations.requires_ordering()
-    def _resample(
+    def resample(
         self,
         rule: str,
         *,
@@ -2519,43 +2519,6 @@ def _resample(
             Literal["epoch", "start", "start_day", "end", "end_day"],
         ] = "start_day",
     ) -> bigframes.core.groupby.SeriesGroupBy:
-        """Internal function to support resample. Resample time-series data.
-
-        **Examples:**
-
-        >>> import bigframes.pandas as bpd
-        >>> data = {
-        ...     "timestamp_col": pd.date_range(
-        ...         start="2021-01-01 13:00:00", periods=30, freq="1s"
-        ...     ),
-        ...     "int64_col": range(30),
-        ... }
-        >>> s = bpd.DataFrame(data).set_index("timestamp_col")
-        >>> s._resample(rule="7s", origin="epoch").min()
-                             int64_col
-        2021-01-01 12:59:56          0
-        2021-01-01 13:00:03          3
-        2021-01-01 13:00:10         10
-        2021-01-01 13:00:17         17
-        2021-01-01 13:00:24         24
-        <BLANKLINE>
-        [5 rows x 1 columns]
-
-
-        Args:
-            rule (str):
-                The offset string representing target conversion.
-            level (str or int, default None):
-                For a MultiIndex, level (name or number) to use for resampling.
-                level must be datetime-like.
-            origin(str, default 'start_day'):
-                The timestamp on which to adjust the grouping. Must be one of the following:
-                'epoch': origin is 1970-01-01
-                'start': origin is the first value of the timeseries
-                'start_day': origin is the first day at midnight of the timeseries
-        Returns:
-            SeriesGroupBy: SeriesGroupBy object.
-        """
         block = self._block._generate_resample_label(
             rule=rule,
             closed=closed,

tests/system/small/engines/test_aggregation.py

@@ -111,7 +111,7 @@ def test_engines_unary_aggregates(
     assert_equivalence_execution(node, REFERENCE_ENGINE, engine)
 
 
-@pytest.mark.parametrize("engine", ["polars", "bq"], indirect=True)
+@pytest.mark.parametrize("engine", ["polars", "bq", "bq-sqlglot"], indirect=True)
 @pytest.mark.parametrize(
     "op",
     [agg_ops.std_op, agg_ops.var_op, agg_ops.PopVarOp()],

tests/system/small/engines/test_generic_ops.py

@@ -329,7 +329,7 @@ def test_engines_where_op(scalars_array_value: array_value.ArrayValue, engine):
     assert_equivalence_execution(arr.node, REFERENCE_ENGINE, engine)
 
 
-@pytest.mark.parametrize("engine", ["polars", "bq"], indirect=True)
+@pytest.mark.parametrize("engine", ["polars", "bq", "bq-sqlglot"], indirect=True)
 def test_engines_coalesce_op(scalars_array_value: array_value.ArrayValue, engine):
     arr, _ = scalars_array_value.compute_values(
         [

tests/system/small/test_dataframe.py

@@ -5915,21 +5915,15 @@ def test_dataframe_explode_xfail(col_names):
         pytest.param("datetime_col", "5M", "epoch"),
         pytest.param("datetime_col", "3Q", "start_day"),
         pytest.param("datetime_col", "3YE", "start"),
-        pytest.param(
-            "int64_col", "100D", "start", marks=pytest.mark.xfail(raises=TypeError)
-        ),
-        pytest.param(
-            "datetime_col", "100D", "end", marks=pytest.mark.xfail(raises=ValueError)
-        ),
     ],
 )
-def test__resample_with_column(
+def test_resample_with_column(
     scalars_df_index, scalars_pandas_df_index, on, rule, origin
 ):
     # TODO: supply a reason why this isn't compatible with pandas 1.x
     pytest.importorskip("pandas", minversion="2.0.0")
     bf_result = (
-        scalars_df_index._resample(rule=rule, on=on, origin=origin)[
+        scalars_df_index.resample(rule=rule, on=on, origin=origin)[
             ["int64_col", "int64_too"]
         ]
         .max()
@@ -5943,30 +5937,54 @@ def test__resample_with_column(
     )
 
 
+@pytest.mark.parametrize("index_col", ["timestamp_col", "datetime_col"])
+@pytest.mark.parametrize(
+    ("index_append", "level"),
+    [(True, 1), (False, None), (False, 0)],
+)
 @pytest.mark.parametrize(
-    ("append", "level", "col", "rule"),
+    "rule",
     [
-        pytest.param(False, None, "timestamp_col", "100d"),
-        pytest.param(True, 1, "timestamp_col", "1200h"),
-        pytest.param(False, None, "datetime_col", "100d"),
+        # TODO(tswast): support timedeltas and dataoffsets.
+        # TODO(tswast): support bins that default to "right".
+        "100d",
+        "1200h",
     ],
 )
-def test__resample_with_index(
-    scalars_df_index, scalars_pandas_df_index, append, level, col, rule
+# TODO(tswast): support "right"
+@pytest.mark.parametrize("closed", ["left", None])
+# TODO(tswast): support "right"
+@pytest.mark.parametrize("label", ["left", None])
+@pytest.mark.parametrize(
+    "origin",
+    ["epoch", "start", "start_day"],  # TODO(tswast): support end, end_day.
+)
+def test_resample_with_index(
+    scalars_df_index,
+    scalars_pandas_df_index,
+    index_append,
+    level,
+    index_col,
+    rule,
+    closed,
+    origin,
+    label,
 ):
     # TODO: supply a reason why this isn't compatible with pandas 1.x
     pytest.importorskip("pandas", minversion="2.0.0")
-    scalars_df_index = scalars_df_index.set_index(col, append=append)
-    scalars_pandas_df_index = scalars_pandas_df_index.set_index(col, append=append)
+    scalars_df_index = scalars_df_index.set_index(index_col, append=index_append)
+    scalars_pandas_df_index = scalars_pandas_df_index.set_index(
+        index_col, append=index_append
+    )
     bf_result = (
         scalars_df_index[["int64_col", "int64_too"]]
-        ._resample(rule=rule, level=level)
+        .resample(rule=rule, level=level, closed=closed, origin=origin, label=label)
         .min()
         .to_pandas()
     )
     pd_result = (
         scalars_pandas_df_index[["int64_col", "int64_too"]]
-        .resample(rule=rule, level=level)
+        .resample(rule=rule, level=level, closed=closed, origin=origin, label=label)
         .min()
     )
     assert_pandas_df_equal(bf_result, pd_result)
@@ -6010,15 +6028,15 @@ def test__resample_with_index(
         ),
     ],
 )
-def test__resample_start_time(rule, origin, data):
+def test_resample_start_time(rule, origin, data):
     # TODO: supply a reason why this isn't compatible with pandas 1.x
     pytest.importorskip("pandas", minversion="2.0.0")
     col = "timestamp_col"
     scalars_df_index = bpd.DataFrame(data).set_index(col)
     scalars_pandas_df_index = pd.DataFrame(data).set_index(col)
     scalars_pandas_df_index.index.name = None
 
-    bf_result = scalars_df_index._resample(rule=rule, origin=origin).min().to_pandas()
+    bf_result = scalars_df_index.resample(rule=rule, origin=origin).min().to_pandas()
 
     pd_result = scalars_pandas_df_index.resample(rule=rule, origin=origin).min()