Merge branch 'main' into sycai_ai_doc_fix

Author: sycai

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

@@ -151,6 +151,23 @@ def _(
     )
 
 
+@UNARY_OP_REGISTRATION.register(agg_ops.DiffOp)
+def _(
+    op: agg_ops.DiffOp,
+    column: typed_expr.TypedExpr,
+    window: typing.Optional[window_spec.WindowSpec] = None,
+) -> sge.Expression:
+    shift_op_impl = UNARY_OP_REGISTRATION[agg_ops.ShiftOp(0)]
+    shifted = shift_op_impl(agg_ops.ShiftOp(op.periods), column, window)
+    if column.dtype in (dtypes.BOOL_DTYPE, dtypes.INT_DTYPE, dtypes.FLOAT_DTYPE):
+        if column.dtype == dtypes.BOOL_DTYPE:
+            return sge.NEQ(this=column.expr, expression=shifted)
+        else:
+            return sge.Sub(this=column.expr, expression=shifted)
+    else:
+        raise TypeError(f"Cannot perform diff on type {column.dtype}")
+
+
 @UNARY_OP_REGISTRATION.register(agg_ops.MaxOp)
 def _(
     op: agg_ops.MaxOp,
@@ -240,6 +257,27 @@ def _(
     return apply_window_if_present(sge.func("COUNT", sge.convert(1)), window)
 
 
+@UNARY_OP_REGISTRATION.register(agg_ops.ShiftOp)
+def _(
+    op: agg_ops.ShiftOp,
+    column: typed_expr.TypedExpr,
+    window: typing.Optional[window_spec.WindowSpec] = None,
+) -> sge.Expression:
+    if op.periods == 0:  # No-op
+        return column.expr
+    if op.periods > 0:
+        return apply_window_if_present(
+            sge.func("LAG", column.expr, sge.convert(op.periods)),
+            window,
+            include_framing_clauses=False,
+        )
+    return apply_window_if_present(
+        sge.func("LEAD", column.expr, sge.convert(-op.periods)),
+        window,
+        include_framing_clauses=False,
+    )
+
+
 @UNARY_OP_REGISTRATION.register(agg_ops.SumOp)
 def _(
     op: agg_ops.SumOp,

notebooks/generative_ai/ai_functions.ipynb

@@ -56,7 +56,7 @@
    "id": "aee05821",
    "metadata": {},
    "source": [
-    "This notebook provides a brief introduction to how to use BigFrames AI functions"
+    "This notebook provides a brief introduction to AI functions in BigQuery Dataframes."
    ]
   },
   {
@@ -145,7 +145,7 @@
    "id": "b606c51f",
    "metadata": {},
    "source": [
-    "You can also include additional model parameters into your function call, as long as they satisfy the structure of `generateContent` [request body format](https://cloud.google.com/vertex-ai/docs/reference/rest/v1/projects.locations.endpoints/generateContent#request-body). In the next example, you use `maxOutputTokens` to limite the length of the generated content."
+    "You can also include additional model parameters into your function call, as long as they conform to the structure of `generateContent` [request body format](https://cloud.google.com/vertex-ai/docs/reference/rest/v1/projects.locations.endpoints/generateContent#request-body). In the next example, you use `maxOutputTokens` to limit the length of the generated content."
    ]
   },
   {
@@ -186,7 +186,7 @@
    "source": [
     "The answers are cut short as expected.\n",
     "\n",
-    "In addition to `ai.generate`, you can use `ai.generate_bool`, `ai.generate_int`, and `ai.generate_double` for other type of outputs."
+    "In addition to `ai.generate`, you can use `ai.generate_bool`, `ai.generate_int`, and `ai.generate_double` for other output types."
    ]
   },
   {
@@ -196,7 +196,7 @@
    "source": [
     "## ai.if_\n",
     "\n",
-    "`ai.if_` generates a series of booleans, unlike `ai.generate_bool` where you get a series of structs. It's a handy tool for filtering your data. not only because it directly returns a boolean, but also because it provides more optimization during data processing. Here is an example of using `ai.if_`:"
+    "`ai.if_` generates a series of booleans. It's a handy tool for joining and filtering your data, not only because it directly returns boolean values, but also because it provides more optimization during data processing. Here is an example of using `ai.if_`:"
    ]
   },
   {
@@ -284,7 +284,7 @@
    "id": "63b5a59f",
    "metadata": {},
    "source": [
-    "`ai.score` ranks your input based on the prompt. You can then sort your data based on their ranks. For example:"
+    "`ai.score` ranks your input based on the prompt and assigns a double value (i.e. a score) to each item. You can then sort your data based on their scores. For example:"
    ]
   },
   {
@@ -460,7 +460,7 @@
    "id": "9e4037bc",
    "metadata": {},
    "source": [
-    "Note that this function can only return the values that are present in your provided categories. If your categories do not cover all cases, your will get wrong answers:"
+    "Note that this function can only return the values that are provided in the `categories` argument. If your categories do not cover all cases, your may get wrong answers:"
    ]
   },
   {

tests/unit/core/compile/sqlglot/aggregations/snapshots/test_unary_compiler/test_diff/diff_bool.sql

@@ -0,0 +1,13 @@
+WITH `bfcte_0` AS (
+  SELECT
+    `bool_col` AS `bfcol_0`
+  FROM `bigframes-dev`.`sqlglot_test`.`scalar_types`
+), `bfcte_1` AS (
+  SELECT
+    *,
+    `bfcol_0` <> LAG(`bfcol_0`, 1) OVER (ORDER BY `bfcol_0` IS NULL ASC NULLS LAST, `bfcol_0` ASC NULLS LAST) AS `bfcol_1`
+  FROM `bfcte_0`
+)
+SELECT
+  `bfcol_1` AS `diff_bool`
+FROM `bfcte_1`

tests/unit/core/compile/sqlglot/aggregations/snapshots/test_unary_compiler/test_diff/diff_int.sql

@@ -0,0 +1,13 @@
+WITH `bfcte_0` AS (
+  SELECT
+    `int64_col` AS `bfcol_0`
+  FROM `bigframes-dev`.`sqlglot_test`.`scalar_types`
+), `bfcte_1` AS (
+  SELECT
+    *,
+    `bfcol_0` - LAG(`bfcol_0`, 1) OVER (ORDER BY `bfcol_0` IS NULL ASC NULLS LAST, `bfcol_0` ASC NULLS LAST) AS `bfcol_1`
+  FROM `bfcte_0`
+)
+SELECT
+  `bfcol_1` AS `diff_int`
+FROM `bfcte_1`

tests/unit/core/compile/sqlglot/aggregations/snapshots/test_unary_compiler/test_shift/lag.sql

@@ -0,0 +1,13 @@
+WITH `bfcte_0` AS (
+  SELECT
+    `int64_col` AS `bfcol_0`
+  FROM `bigframes-dev`.`sqlglot_test`.`scalar_types`
+), `bfcte_1` AS (
+  SELECT
+    *,
+    LAG(`bfcol_0`, 1) OVER (ORDER BY `bfcol_0` IS NULL ASC NULLS LAST, `bfcol_0` ASC NULLS LAST) AS `bfcol_1`
+  FROM `bfcte_0`
+)
+SELECT
+  `bfcol_1` AS `lag`
+FROM `bfcte_1`

tests/unit/core/compile/sqlglot/aggregations/snapshots/test_unary_compiler/test_shift/lead.sql

@@ -0,0 +1,13 @@
+WITH `bfcte_0` AS (
+  SELECT
+    `int64_col` AS `bfcol_0`
+  FROM `bigframes-dev`.`sqlglot_test`.`scalar_types`
+), `bfcte_1` AS (
+  SELECT
+    *,
+    LEAD(`bfcol_0`, 1) OVER (ORDER BY `bfcol_0` IS NULL ASC NULLS LAST, `bfcol_0` ASC NULLS LAST) AS `bfcol_1`
+  FROM `bfcte_0`
+)
+SELECT
+  `bfcol_1` AS `lead`
+FROM `bfcte_1`

tests/unit/core/compile/sqlglot/aggregations/snapshots/test_unary_compiler/test_shift/noop.sql

@@ -0,0 +1,13 @@
+WITH `bfcte_0` AS (
+  SELECT
+    `int64_col` AS `bfcol_0`
+  FROM `bigframes-dev`.`sqlglot_test`.`scalar_types`
+), `bfcte_1` AS (
+  SELECT
+    *,
+    `bfcol_0` AS `bfcol_1`
+  FROM `bfcte_0`
+)
+SELECT
+  `bfcol_1` AS `noop`
+FROM `bfcte_1`

tests/unit/core/compile/sqlglot/aggregations/test_unary_compiler.py

@@ -127,6 +127,28 @@ def test_dense_rank(scalar_types_df: bpd.DataFrame, snapshot):
     snapshot.assert_match(sql, "out.sql")
 
 
+def test_diff(scalar_types_df: bpd.DataFrame, snapshot):
+    # Test integer
+    int_col = "int64_col"
+    bf_df_int = scalar_types_df[[int_col]]
+    window = window_spec.WindowSpec(ordering=(ordering.ascending_over(int_col),))
+    int_op = agg_exprs.UnaryAggregation(
+        agg_ops.DiffOp(periods=1), expression.deref(int_col)
+    )
+    int_sql = _apply_unary_window_op(bf_df_int, int_op, window, "diff_int")
+    snapshot.assert_match(int_sql, "diff_int.sql")
+
+    # Test boolean
+    bool_col = "bool_col"
+    bf_df_bool = scalar_types_df[[bool_col]]
+    window = window_spec.WindowSpec(ordering=(ordering.ascending_over(bool_col),))
+    bool_op = agg_exprs.UnaryAggregation(
+        agg_ops.DiffOp(periods=1), expression.deref(bool_col)
+    )
+    bool_sql = _apply_unary_window_op(bf_df_bool, bool_op, window, "diff_bool")
+    snapshot.assert_match(bool_sql, "diff_bool.sql")
+
+
 def test_first(scalar_types_df: bpd.DataFrame, snapshot):
     if sys.version_info < (3, 12):
         pytest.skip(
@@ -271,6 +293,33 @@ def test_rank(scalar_types_df: bpd.DataFrame, snapshot):
     snapshot.assert_match(sql, "out.sql")
 
 
+def test_shift(scalar_types_df: bpd.DataFrame, snapshot):
+    col_name = "int64_col"
+    bf_df = scalar_types_df[[col_name]]
+    window = window_spec.WindowSpec(ordering=(ordering.ascending_over(col_name),))
+
+    # Test lag
+    lag_op = agg_exprs.UnaryAggregation(
+        agg_ops.ShiftOp(periods=1), expression.deref(col_name)
+    )
+    lag_sql = _apply_unary_window_op(bf_df, lag_op, window, "lag")
+    snapshot.assert_match(lag_sql, "lag.sql")
+
+    # Test lead
+    lead_op = agg_exprs.UnaryAggregation(
+        agg_ops.ShiftOp(periods=-1), expression.deref(col_name)
+    )
+    lead_sql = _apply_unary_window_op(bf_df, lead_op, window, "lead")
+    snapshot.assert_match(lead_sql, "lead.sql")
+
+    # Test no-op
+    noop_op = agg_exprs.UnaryAggregation(
+        agg_ops.ShiftOp(periods=0), expression.deref(col_name)
+    )
+    noop_sql = _apply_unary_window_op(bf_df, noop_op, window, "noop")
+    snapshot.assert_match(noop_sql, "noop.sql")
+
+
 def test_sum(scalar_types_df: bpd.DataFrame, snapshot):
     bf_df = scalar_types_df[["int64_col", "bool_col"]]
     agg_ops_map = {

tests/unit/session/test_read_gbq_query.py

@@ -35,3 +35,4 @@ def test_read_gbq_query_sets_destination_table():
 
     assert query == "SELECT 'my-test-query';"
     assert config.destination is not None
+    session.close()