feat: add bigframes.bigquery.aead.* scalar functions (#17168)

Thank you for opening a Pull Request! Before submitting your PR, there
are a few things you can do to make sure it goes smoothly:
- [ ] Make sure to open an issue as a
[bug/issue](https://github.com/googleapis/google-cloud-python/issues)
before writing your code! That way we can discuss the change, evaluate
designs, and agree on the general idea
- [ ] Ensure the tests and linter pass
- [ ] Code coverage does not decrease (if any source code was changed)
- [ ] Appropriate docs were updated (if necessary)

Fixes #<issue_number_goes_here> 🦕

Author: tswast

.pre-commit-config.yaml

@@ -20,13 +20,10 @@ repos:
     hooks:
     -   id: trailing-whitespace
     -   id: end-of-file-fixer
-    -   id: check-yaml
--   repo: https://github.com/psf/black
-    rev: 23.7.0
-    hooks:
-    - id: black
--   repo: https://github.com/pycqa/flake8
-    rev: 6.1.0 # version-scanner: ignore
-    hooks:
-    - id: flake8
-      args: [--config, packages/google-cloud-alloydb/.flake8]
+- repo: https://github.com/astral-sh/ruff-pre-commit
+  # Ruff version.
+  rev: v0.14.14
+  hooks:
+    # Run the linter.
+    - id: ruff-check
+      args: [ --select, I, --fix, --target-version=py310, --line-length=88 ]

packages/bigframes/bigframes/bigquery/__init__.py

@@ -47,7 +47,7 @@
 
 import sys
 
-from bigframes.bigquery import ai, ml, obj
+from bigframes.bigquery import aead, ai, ml, obj
 from bigframes.bigquery._operations.approx_agg import approx_top_count
 from bigframes.bigquery._operations.array import (
     array_agg,
@@ -208,6 +208,7 @@
     # io ops
     "load_data",
     # Modules / SQL namespaces
+    "aead",
     "ai",
     "ml",
     "obj",

packages/bigframes/bigframes/bigquery/_googlesql.py

@@ -0,0 +1,81 @@
+# Copyright 2026 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.
+
+"""Utilities for working with GoogleSqlScalarOps."""
+
+from __future__ import annotations
+
+from typing import Any, Union
+
+import bigframes.core.col
+import bigframes.core.expression as ex
+import bigframes.core.sentinels as sentinels
+import bigframes.series as series
+from bigframes.operations import googlesql
+
+
+def apply_googlesql_scalar_op(
+    op: googlesql.GoogleSqlScalarOp,
+    *args: Any,
+) -> Union[series.Series, bigframes.core.col.Expression]:
+    """Applies a GoogleSQL scalar operator to the given arguments.
+
+    Handles a mix of Series, Expression, and literal inputs.
+
+    Args:
+        op (googlesql.GoogleSqlScalarOp):
+            The operator to apply.
+        *args (Any):
+            The arguments to apply the operator to.
+
+    Returns:
+        bigframes.pandas.Series | bigframes.core.col.Expression:
+            The result of the operation. If any of ``args`` is a Series, returns
+            a Series. Otherwise, returns an Expression.
+    """
+    # Find the first Series to use for alignment
+    first_series = None
+    for arg in args:
+        if isinstance(arg, series.Series):
+            first_series = arg
+            break
+
+    if first_series is not None:
+        processed_args: list[Union[bigframes.core.col.Expression, series.Series]] = []
+        block = first_series._block
+        for arg in args:
+            if isinstance(arg, bigframes.core.col.Expression):
+                block, col_id = block.project_expr(bigframes.core.col._as_bf_expr(arg))
+                processed_args.append(series.Series(block.select_column(col_id)))
+            elif arg is sentinels.Sentinel.ARGUMENT_DEFAULT:
+                processed_args.append(bigframes.core.col.Expression(ex.OmittedArg()))
+            else:
+                processed_args.append(arg)
+
+        # Apply the n-ary op. _apply_nary_op handles alignment of Series and literals.
+        result = first_series._apply_nary_op(op, processed_args, ignore_self=True)
+        result.name = None
+        return result
+
+    # No Series, return an Expression
+    expr_args = []
+    for arg in args:
+        if isinstance(arg, bigframes.core.col.Expression):
+            expr_args.append(bigframes.core.col._as_bf_expr(arg))
+        elif arg is sentinels.Sentinel.ARGUMENT_DEFAULT:
+            expr_args.append(ex.OmittedArg())
+        else:
+            expr_args.append(ex.const(arg))
+
+    return bigframes.core.col.Expression(ex.OpExpression(op, tuple(expr_args)))

packages/bigframes/bigframes/bigquery/_operations/aead.py

@@ -0,0 +1,127 @@
+# Copyright 2026 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.
+#
+# DO NOT MODIFY THIS FILE DIRECTLY.
+# This file was generated from: scripts/data/sql-functions/aead.yaml
+# by the script: scripts/generate_bigframes_bigquery.py
+
+from __future__ import annotations
+
+import datetime
+from typing import Any, Literal, Optional, TypeVar, Union
+
+import bigframes.bigquery._googlesql
+import bigframes.core.col
+import bigframes.core.expression as ex
+import bigframes.core.sentinels as sentinels
+import bigframes.operations as ops
+import bigframes.series as series
+from bigframes import dtypes
+from bigframes.operations import googlesql
+
+T = TypeVar("T", series.Series, bigframes.core.col.Expression)
+
+_DECRYPT_BYTES_OP = googlesql.GoogleSqlScalarOp(
+    "AEAD.DECRYPT_BYTES",
+    args=(googlesql.ArgSpec(), googlesql.ArgSpec(), googlesql.ArgSpec()),
+    signature=lambda *args: dtypes.BYTES_DTYPE,
+)
+_DECRYPT_STRING_OP = googlesql.GoogleSqlScalarOp(
+    "AEAD.DECRYPT_STRING",
+    args=(googlesql.ArgSpec(), googlesql.ArgSpec(), googlesql.ArgSpec()),
+    signature=lambda *args: dtypes.STRING_DTYPE,
+)
+_ENCRYPT_OP = googlesql.GoogleSqlScalarOp(
+    "AEAD.ENCRYPT",
+    args=(googlesql.ArgSpec(), googlesql.ArgSpec(), googlesql.ArgSpec()),
+    signature=lambda *args: dtypes.BYTES_DTYPE,
+)
+
+
+def decrypt_bytes(
+    keyset: Union[
+        T,
+        bigframes.core.col.Expression,
+        Union[Literal[sentinels.Sentinel.ARGUMENT_DEFAULT], bytes, dict],
+    ],
+    ciphertext: Union[
+        T,
+        bigframes.core.col.Expression,
+        Union[Literal[sentinels.Sentinel.ARGUMENT_DEFAULT], bytes],
+    ],
+    additional_data: Union[
+        T,
+        bigframes.core.col.Expression,
+        Union[Literal[sentinels.Sentinel.ARGUMENT_DEFAULT], bytes],
+    ],
+) -> T:
+    """Uses the matching key from keyset to decrypt ciphertext and verifies the integrity of the data using additional_data. Returns an error if decryption or verification fails."""
+    return bigframes.bigquery._googlesql.apply_googlesql_scalar_op(
+        _DECRYPT_BYTES_OP,
+        keyset,
+        ciphertext,
+        additional_data,
+    )  # type: ignore
+
+
+def decrypt_string(
+    keyset: Union[
+        T,
+        bigframes.core.col.Expression,
+        Union[Literal[sentinels.Sentinel.ARGUMENT_DEFAULT], bytes, dict],
+    ],
+    ciphertext: Union[
+        T,
+        bigframes.core.col.Expression,
+        Union[Literal[sentinels.Sentinel.ARGUMENT_DEFAULT], bytes],
+    ],
+    additional_data: Union[
+        T,
+        bigframes.core.col.Expression,
+        Union[Literal[sentinels.Sentinel.ARGUMENT_DEFAULT], str],
+    ],
+) -> T:
+    """Like AEAD.DECRYPT_BYTES, but where additional_data is of type STRING."""
+    return bigframes.bigquery._googlesql.apply_googlesql_scalar_op(
+        _DECRYPT_STRING_OP,
+        keyset,
+        ciphertext,
+        additional_data,
+    )  # type: ignore
+
+
+def encrypt(
+    keyset: Union[
+        T,
+        bigframes.core.col.Expression,
+        Union[Literal[sentinels.Sentinel.ARGUMENT_DEFAULT], bytes, dict],
+    ],
+    plaintext: Union[
+        T,
+        bigframes.core.col.Expression,
+        Union[Literal[sentinels.Sentinel.ARGUMENT_DEFAULT], bytes, str],
+    ],
+    additional_data: Union[
+        T,
+        bigframes.core.col.Expression,
+        Union[Literal[sentinels.Sentinel.ARGUMENT_DEFAULT], bytes, str],
+    ],
+) -> T:
+    """Encrypts plaintext using the primary cryptographic key in keyset. The algorithm of the primary key must be AEAD_AES_GCM_256. Binds the ciphertext to the context defined by additional_data. Returns NULL if any input is NULL."""
+    return bigframes.bigquery._googlesql.apply_googlesql_scalar_op(
+        _ENCRYPT_OP,
+        keyset,
+        plaintext,
+        additional_data,
+    )  # type: ignore

packages/bigframes/bigframes/bigquery/aead.py

@@ -0,0 +1,25 @@
+# Copyright 2026 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.
+
+"""AEAD encryption functions"""
+
+from __future__ import annotations
+
+from bigframes.bigquery._operations.aead import decrypt_bytes, decrypt_string, encrypt
+
+__all__ = [
+    "decrypt_bytes",
+    "decrypt_string",
+    "encrypt",
+]

packages/bigframes/bigframes/core/sentinels.py

@@ -0,0 +1,31 @@
+# Copyright 2026 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.
+
+"""Sentinel values used throughout BigFrames."""
+
+from __future__ import annotations
+
+import enum
+
+
+class Sentinel(enum.Enum):
+    """Default values used throughout BigFrames."""
+
+    """Default value for an optional argument.
+
+    When a parameter is set to this, that parameter is explicitly omitted
+    from the SQL text. This allows for NULL (None in Python) to be explicitly
+    passed in to optional parameters.
+    """
+    ARGUMENT_DEFAULT = enum.auto()

packages/bigframes/pyproject.toml

@@ -1,3 +1,6 @@
 [build-system]
 requires = ["setuptools"]
 build-backend = "setuptools.build_meta"
+
+[tool.ruff.lint.isort]
+known-first-party = ["bigframes"]