Merge remote-tracking branch 'origin/main' into feat-inplace-param-for-drop

Author: tswast

bigframes/core/compile/sqlglot/__init__.py

@@ -14,6 +14,7 @@
 from __future__ import annotations
 
 from bigframes.core.compile.sqlglot.compiler import SQLGlotCompiler
+import bigframes.core.compile.sqlglot.expressions.ai_ops  # noqa: F401
 import bigframes.core.compile.sqlglot.expressions.array_ops  # noqa: F401
 import bigframes.core.compile.sqlglot.expressions.blob_ops  # noqa: F401
 import bigframes.core.compile.sqlglot.expressions.comparison_ops  # noqa: F401

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

@@ -0,0 +1,65 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+import sqlglot.expressions as sge
+
+from bigframes import operations as ops
+from bigframes.core.compile.sqlglot import scalar_compiler
+from bigframes.core.compile.sqlglot.expressions.typed_expr import TypedExpr
+
+register_nary_op = scalar_compiler.scalar_op_compiler.register_nary_op
+
+
+@register_nary_op(ops.AIGenerateBool, pass_op=True)
+def _(*exprs: TypedExpr, op: ops.AIGenerateBool) -> sge.Expression:
+
+    prompt: list[str | sge.Expression] = []
+    column_ref_idx = 0
+
+    for elem in op.prompt_context:
+        if elem is None:
+            prompt.append(exprs[column_ref_idx].expr)
+        else:
+            prompt.append(sge.Literal.string(elem))
+
+    args = [sge.Kwarg(this="prompt", expression=sge.Tuple(expressions=prompt))]
+
+    args.append(
+        sge.Kwarg(this="connection_id", expression=sge.Literal.string(op.connection_id))
+    )
+
+    if op.endpoint is not None:
+        args.append(
+            sge.Kwarg(this="endpoint", expression=sge.Literal.string(op.endpoint))
+        )
+
+    args.append(
+        sge.Kwarg(
+            this="request_type", expression=sge.Literal.string(op.request_type.upper())
+        )
+    )
+
+    if op.model_params is not None:
+        args.append(
+            sge.Kwarg(
+                this="model_params",
+                # sge.JSON requires a newer SQLGlot version than 23.6.3.
+                # PARSE_JSON won't work as the function requires a JSON literal.
+                expression=sge.JSON(this=sge.Literal.string(op.model_params)),
+            )
+        )
+
+    return sge.func("AI.GENERATE_BOOL", *args)

bigframes/core/indexes/base.py

@@ -740,6 +740,9 @@ def to_numpy(self, dtype=None, *, allow_large_results=None, **kwargs) -> np.ndar
 
     __array__ = to_numpy
 
+    def to_list(self, *, allow_large_results: Optional[bool] = None) -> list:
+        return self.to_pandas(allow_large_results=allow_large_results).to_list()
+
     def __len__(self):
         return self.shape[0]
 

tests/system/small/operations/test_strings.py

@@ -236,7 +236,20 @@ def test_reverse(scalars_dfs):
 
 
 @pytest.mark.parametrize(
-    ["start", "stop"], [(0, 1), (3, 5), (100, 101), (None, 1), (0, 12), (0, None)]
+    ["start", "stop"],
+    [
+        (0, 1),
+        (3, 5),
+        (100, 101),
+        (None, 1),
+        (0, 12),
+        (0, None),
+        (None, -1),
+        (-1, None),
+        (-5, -1),
+        (1, -1),
+        (-10, 10),
+    ],
 )
 def test_slice(scalars_dfs, start, stop):
     scalars_df, scalars_pandas_df = scalars_dfs

tests/system/small/test_index.py

@@ -638,6 +638,12 @@ def test_index_item_with_empty(session):
         bf_idx_empty.item()
 
 
+def test_index_to_list(scalars_df_index, scalars_pandas_df_index):
+    bf_result = scalars_df_index.index.to_list()
+    pd_result = scalars_pandas_df_index.index.to_list()
+    assert bf_result == pd_result
+
+
 @pytest.mark.parametrize(
     ("key", "value"),
     [

tests/unit/core/compile/sqlglot/expressions/snapshots/test_ai_ops/test_ai_generate_bool/out.sql

@@ -0,0 +1,18 @@
+WITH `bfcte_0` AS (
+  SELECT
+    `string_col` AS `bfcol_0`
+  FROM `bigframes-dev`.`sqlglot_test`.`scalar_types`
+), `bfcte_1` AS (
+  SELECT
+    *,
+    AI.GENERATE_BOOL(
+      prompt => (`bfcol_0`, ' is the same as ', `bfcol_0`),
+      connection_id => 'test_connection_id',
+      endpoint => 'gemini-2.5-flash',
+      request_type => 'SHARED'
+    ) AS `bfcol_1`
+  FROM `bfcte_0`
+)
+SELECT
+  `bfcol_1` AS `result`
+FROM `bfcte_1`

tests/unit/core/compile/sqlglot/expressions/snapshots/test_ai_ops/test_ai_generate_bool_with_model_param/out.sql

@@ -0,0 +1,18 @@
+WITH `bfcte_0` AS (
+  SELECT
+    `string_col` AS `bfcol_0`
+  FROM `bigframes-dev`.`sqlglot_test`.`scalar_types`
+), `bfcte_1` AS (
+  SELECT
+    *,
+    AI.GENERATE_BOOL(
+      prompt => (`bfcol_0`, ' is the same as ', `bfcol_0`),
+      connection_id => 'test_connection_id',
+      request_type => 'SHARED',
+      model_params => JSON '{}'
+    ) AS `bfcol_1`
+  FROM `bfcte_0`
+)
+SELECT
+  `bfcol_1` AS `result`
+FROM `bfcte_1`

tests/unit/core/compile/sqlglot/expressions/test_ai_ops.py

@@ -0,0 +1,67 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import json
+import sys
+
+import pytest
+
+from bigframes import dataframe
+from bigframes import operations as ops
+from bigframes.testing import utils
+
+pytest.importorskip("pytest_snapshot")
+
+
+def test_ai_generate_bool(scalar_types_df: dataframe.DataFrame, snapshot):
+    col_name = "string_col"
+
+    op = ops.AIGenerateBool(
+        prompt_context=(None, " is the same as ", None),
+        connection_id="test_connection_id",
+        endpoint="gemini-2.5-flash",
+        request_type="shared",
+        model_params=None,
+    )
+
+    sql = utils._apply_unary_ops(
+        scalar_types_df, [op.as_expr(col_name, col_name)], ["result"]
+    )
+
+    snapshot.assert_match(sql, "out.sql")
+
+
+def test_ai_generate_bool_with_model_param(
+    scalar_types_df: dataframe.DataFrame, snapshot
+):
+    if sys.version_info < (3, 10):
+        pytest.skip(
+            "Skip test because SQLGLot cannot compile model params to JSON at this env."
+        )
+
+    col_name = "string_col"
+
+    op = ops.AIGenerateBool(
+        prompt_context=(None, " is the same as ", None),
+        connection_id="test_connection_id",
+        endpoint=None,
+        request_type="shared",
+        model_params=json.dumps(dict()),
+    )
+
+    sql = utils._apply_unary_ops(
+        scalar_types_df, [op.as_expr(col_name, col_name)], ["result"]
+    )
+
+    snapshot.assert_match(sql, "out.sql")

tests/unit/test_index.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import pandas as pd
 import pytest
 
 from bigframes.testing import mocks
@@ -38,3 +39,13 @@ def test_index_rename_inplace_returns_none(monkeypatch: pytest.MonkeyPatch):
     # Make sure the linked DataFrame is updated, too.
     assert dataframe.index.name == "my_index_name"
     assert index.name == "my_index_name"
+
+
+def test_index_to_list(monkeypatch: pytest.MonkeyPatch):
+    pd_index = pd.Index([1, 2, 3], name="my_index")
+    df = mocks.create_dataframe(
+        monkeypatch,
+        data={"my_index": [1, 2, 3]},
+    ).set_index("my_index")
+    bf_index = df.index
+    assert bf_index.to_list() == pd_index.to_list()

third_party/bigframes_vendored/ibis/expr/rewrites.py

@@ -206,21 +206,26 @@ def replace_parameter(_, params, **kwargs):
 @replace(p.StringSlice)
 def lower_stringslice(_, **kwargs):
     """Rewrite StringSlice in terms of Substring."""
-    if _.end is None:
-        return ops.Substring(_.arg, start=_.start)
     if _.start is None:
-        return ops.Substring(_.arg, start=0, length=_.end)
-    if (
-        isinstance(_.start, ops.Literal)
-        and isinstance(_.start.value, int)
-        and isinstance(_.end, ops.Literal)
-        and isinstance(_.end.value, int)
-    ):
-        # optimization for constant values
-        length = _.end.value - _.start.value
+        real_start = 0
     else:
-        length = ops.Subtract(_.end, _.start)
-    return ops.Substring(_.arg, start=_.start, length=length)
+        real_start = ops.IfElse(
+            ops.GreaterEqual(_.start, 0),
+            _.start,
+            ops.Greatest((0, ops.Add(ops.StringLength(_.arg), _.start))),
+        )
+
+    if _.end is None:
+        real_end = ops.StringLength(_.arg)
+    else:
+        real_end = ops.IfElse(
+            ops.GreaterEqual(_.end, 0),
+            _.end,
+            ops.Greatest((0, ops.Add(ops.StringLength(_.arg), _.end))),
+        )
+
+    length = ops.Greatest((0, ops.Subtract(real_end, real_start)))
+    return ops.Substring(_.arg, start=real_start, length=length)
 
 
 @replace(p.Analytic)