From b7e5a06b603b00d6da3274c5602be2c1c0999aaf Mon Sep 17 00:00:00 2001 From: Shenyang Cai Date: Sun, 26 Jul 2026 01:59:40 +0000 Subject: [PATCH 01/27] checkpoint: add logic for rendering signature template --- .../generate_bigframes_bigquery/__init__.py | 13 ++ .../generate_bigframes_bigquery/__main__.py | 42 +++++ .../generate_bigframes_bigquery/constants.py | 155 ++++++++++++++++++ .../data_models.py | 95 +++++++++++ .../file_generator.py | 44 +++++ .../template_renderer.py | 153 +++++++++++++++++ .../yaml_parser.py | 63 +++++++ 7 files changed, 565 insertions(+) create mode 100644 packages/bigframes/scripts/generate_bigframes_bigquery/__init__.py create mode 100644 packages/bigframes/scripts/generate_bigframes_bigquery/__main__.py create mode 100644 packages/bigframes/scripts/generate_bigframes_bigquery/constants.py create mode 100644 packages/bigframes/scripts/generate_bigframes_bigquery/data_models.py create mode 100644 packages/bigframes/scripts/generate_bigframes_bigquery/file_generator.py create mode 100644 packages/bigframes/scripts/generate_bigframes_bigquery/template_renderer.py create mode 100644 packages/bigframes/scripts/generate_bigframes_bigquery/yaml_parser.py diff --git a/packages/bigframes/scripts/generate_bigframes_bigquery/__init__.py b/packages/bigframes/scripts/generate_bigframes_bigquery/__init__.py new file mode 100644 index 000000000000..58d482ea3866 --- /dev/null +++ b/packages/bigframes/scripts/generate_bigframes_bigquery/__init__.py @@ -0,0 +1,13 @@ +# 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. diff --git a/packages/bigframes/scripts/generate_bigframes_bigquery/__main__.py b/packages/bigframes/scripts/generate_bigframes_bigquery/__main__.py new file mode 100644 index 000000000000..74a5073f33e4 --- /dev/null +++ b/packages/bigframes/scripts/generate_bigframes_bigquery/__main__.py @@ -0,0 +1,42 @@ +#!/usr/bin/env -S uv run --active --script +# +# /// script +# dependencies = [ +# "jinja2", +# "pyyaml", +# "ruff==0.14.14", +# ] +# /// +# +# 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. + +import pprint + +import constants +import yaml_parser +import file_generator + + +def main(): + modules = [] + + for yaml_file in sorted(constants.DATA_DIR.glob("**/*.yaml")): + modules.append(yaml_parser.parse_yaml(yaml_file)) + + file_generator.generate(modules) + + +if __name__ == "__main__": + main() diff --git a/packages/bigframes/scripts/generate_bigframes_bigquery/constants.py b/packages/bigframes/scripts/generate_bigframes_bigquery/constants.py new file mode 100644 index 000000000000..6b8e45c84b70 --- /dev/null +++ b/packages/bigframes/scripts/generate_bigframes_bigquery/constants.py @@ -0,0 +1,155 @@ +# 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. + +import pathlib +import jinja2 + +SCRIPTS_DIRECTORY = pathlib.Path(__file__).parent.parent.absolute() +PACKAGE_ROOT = SCRIPTS_DIRECTORY.parent +CODE_ROOT = PACKAGE_ROOT / "bigframes" +SCRIPT_PATH_RELATIVE = pathlib.Path(__file__).relative_to(PACKAGE_ROOT) + +# Directory containing the YAML files +DATA_DIR = SCRIPTS_DIRECTORY / "data" / "sql-functions" +# Directory where the generated Python files will be placed +OUTPUT_DIR = CODE_ROOT / "operations" / "googlesql" +# Directory where the generated test files will be placed +TEST_OUTPUT_DIR = PACKAGE_ROOT / "tests" / "unit" / "bigquery" / "generated" +# Directory containing the Jinja2 templates +TEMPLATE_DIR = SCRIPTS_DIRECTORY / "templates" + +PYTHON_BUILTINS = { + "abs", + "all", + "any", + "ascii", + "bin", + "bool", + "breakpoint", + "bytearray", + "bytes", + "callable", + "chr", + "classmethod", + "compile", + "complex", + "delattr", + "dict", + "dir", + "divmod", + "enumerate", + "eval", + "exec", + "filter", + "float", + "format", + "frozenset", + "getattr", + "globals", + "hasattr", + "hash", + "help", + "hex", + "id", + "input", + "int", + "isinstance", + "issubclass", + "iter", + "len", + "list", + "locals", + "map", + "max", + "memoryview", + "min", + "next", + "object", + "oct", + "open", + "ord", + "pow", + "print", + "property", + "range", + "repr", + "reversed", + "round", + "set", + "setattr", + "slice", + "sorted", + "staticmethod", + "str", + "sum", + "super", + "tuple", + "type", + "vars", + "zip", +} + +DTYPE_MAP = { + "binary": "dtypes.BYTES_DTYPE", + "string": "dtypes.STRING_DTYPE", + "int64": "dtypes.INT_DTYPE", + "i64": "dtypes.INT_DTYPE", + "float64": "dtypes.FLOAT_DTYPE", + "fp64": "dtypes.FLOAT_DTYPE", + "bool": "dtypes.BOOL_DTYPE", + "boolean": "dtypes.BOOL_DTYPE", + "geography": "dtypes.GEO_DTYPE", + "json": "dtypes.JSON_DTYPE", + "date": "dtypes.DATE_DTYPE", + "time": "dtypes.TIME_DTYPE", + "datetime": "dtypes.DATETIME_DTYPE", + "timestamp": "dtypes.TIMESTAMP_DTYPE", + "decimal<38,9>": "dtypes.NUMERIC_DTYPE", + "decimal<76,38>": "dtypes.BIGNUMERIC_DTYPE", +} + + +TEMPLATES: dict[str, jinja2.Template] # Lazily-loaded global variable +_templates = None + + +def _load_templates() -> dict[str, jinja2.Template]: + env = jinja2.Environment( + loader=jinja2.FileSystemLoader(TEMPLATE_DIR), + trim_blocks=True, + lstrip_blocks=True, + ) + return { + "operation": env.get_template("operation.py.j2"), + "test_operation": env.get_template("test_operation.py.j2"), + "license": env.get_template("license.py.j2"), + "signature_def": env.get_template("signature_def.py.j2"), + "core_series_accessor": env.get_template("core_series_accessor.py.j2"), + "bigframes_series_accessor": env.get_template( + "bigframes_series_accessor.py.j2" + ), + "pandas_series_accessor": env.get_template("pandas_series_accessor.py.j2"), + } + + +def __getattr__(name: str): + global _templates + + if name != "TEMPLATES": + raise AttributeError(f"module '{__name__}' has no attribute '{name}'") + + if _templates is None: + _templates = _load_templates() + + return _templates diff --git a/packages/bigframes/scripts/generate_bigframes_bigquery/data_models.py b/packages/bigframes/scripts/generate_bigframes_bigquery/data_models.py new file mode 100644 index 000000000000..b5af7036a577 --- /dev/null +++ b/packages/bigframes/scripts/generate_bigframes_bigquery/data_models.py @@ -0,0 +1,95 @@ +# 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. + + +import dataclasses +import pathlib + +import constants +from typing import Any + + +@dataclasses.dataclass +class BQFuncArg: + name: str + value: str # The type of the arg + optional: bool + keyword_only: bool + + +@dataclasses.dataclass +class BQFuncImpl: + args: list[BQFuncArg] + return_type: str + + @property + def uses_any1(self) -> bool: + if "any1" in self.return_type: + return True + + return any("any1" in arg.value for arg in self.args) + + def to_dict(self) -> dict[str, Any]: + result = dataclasses.asdict(self) + + result["uses_any1"] = self.uses_any1 + + # We cannot use "return" as a field name, but we need to use it + # as a key for template rendering + result["return"] = self.return_type + del result["return_type"] + + return result + + +@dataclasses.dataclass +class BQFunc: + name: str + description: str + impls: list[BQFuncImpl] + + +@dataclasses.dataclass +class BQModule: + module_path: pathlib.Path + functions: list[BQFunc] + + @property + def name(self) -> str: + return self.module_path.name + + @property + def is_global_namespace(self) -> bool: + return "global_namespace" in self.module_path.parts + + @property + def output_file(self): + return constants.OUTPUT_DIR.joinpath(self.module_path).with_suffix(".py") + + +@dataclasses.dataclass +class BigFramesOp: + internal_name: str + sql_name: str + arg_specs: str + signature: str + signature_definition: str + + +@dataclasses.dataclass +class BigFramesFunc: + name: str + description: str + args: list[str] + series_accessor_arg: list[str] diff --git a/packages/bigframes/scripts/generate_bigframes_bigquery/file_generator.py b/packages/bigframes/scripts/generate_bigframes_bigquery/file_generator.py new file mode 100644 index 000000000000..46744fe1724a --- /dev/null +++ b/packages/bigframes/scripts/generate_bigframes_bigquery/file_generator.py @@ -0,0 +1,44 @@ +# 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. + +import constants +import jinja2 +import data_models +import pprint + +import template_renderer + + +def _generate_sql_operator_def(): + pass + + +def _generate_accesor(): + pass + + +def _generate_tests(): + pass + + +def generate(bq_modules: list[data_models.BQModule]): + + for bq_module in bq_modules: + for bq_func in bq_module.functions: + pprint.pp(template_renderer.render_signature_def(bq_func, bq_module)) + + # Write to file + + # Ruff format + pass diff --git a/packages/bigframes/scripts/generate_bigframes_bigquery/template_renderer.py b/packages/bigframes/scripts/generate_bigframes_bigquery/template_renderer.py new file mode 100644 index 000000000000..ca71c0441c08 --- /dev/null +++ b/packages/bigframes/scripts/generate_bigframes_bigquery/template_renderer.py @@ -0,0 +1,153 @@ +# Render jinja template with module data parsed from yaml +# 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. + +import re +import pathlib + +import constants +import data_models +import jinja2 + + +def _to_snake_case(name): + # Replace dots with underscores + name = name.replace(".", "_") + # Handle CamelCase to snake_case + name = re.sub(r"(? str: + if yaml_type in constants.DTYPE_MAP: + return constants.DTYPE_MAP[yaml_type] + if yaml_type.startswith("list<") and yaml_type.endswith(">"): + inner = yaml_type[5:-1] + if inner in constants.DTYPE_MAP: + return f"dtypes.list_type({constants.DTYPE_MAP[inner]})" + raise ValueError(f"Not a concrete type: {yaml_type}") + + +def _is_concrete_type(yaml_type: str) -> bool: + try: + _get_concrete_type_expr(yaml_type) + return True + except ValueError: + return False + + +def _validate_types(impls): + for impl in impls: + for arg in impl.args: + val = arg.value + if val == "any1": + continue + if val.startswith("list<") and val.endswith(">"): + inner = val[5:-1] + if inner != "any1" and inner not in constants.DTYPE_MAP: + raise ValueError(f"Unsupported inner type: {inner}") + continue + if val == "struct": + continue + if val not in constants.DTYPE_MAP: + raise ValueError(f"Unsupported type: {val}") + + ret = impl.return_type + if ret == "any1": + continue + if ret.startswith("list<") and ret.endswith(">"): + inner = ret[5:-1] + if inner != "any1" and inner not in constants.DTYPE_MAP: + raise ValueError(f"Unsupported inner type: {inner}") + continue + if ret not in constants.DTYPE_MAP: + raise ValueError(f"Unsupported type: {ret}") + + +def render_signature_def( + bq_func: data_models.BQFunc, + hosting_bq_module: data_models.BQModule, +) -> tuple[str, str | None]: + """ + Returns the signature function name and it's definition. + If the signature function can be inlined, the first return value is the lambda, + and the second value is None. + """ + op_base_name = _to_snake_case(bq_func.name) + + module_name = hosting_bq_module.name + if not hosting_bq_module.is_global_namespace and op_base_name.startswith( + module_name + "_" + ): + op_base_name = op_base_name[len(module_name) + 1 :] + + return_types = {impl.return_type for impl in bq_func.impls} + # Optimization: if all impls return the same concrete type, + # inline the signature function as a lambda + if len(return_types) == 1: + ret_type = next(iter(return_types)) + if _is_concrete_type(ret_type): + sig_expr = f"lambda *args: {_get_concrete_type_expr(ret_type)}" + return sig_expr, None + + _validate_types(bq_func.impls) + + python_name = op_base_name + if python_name in constants.PYTHON_BUILTINS: + python_name = python_name + "_" + sig_func_name = f"_{python_name.upper()}_SIG" + + max_args = max(len(impl.args) for impl in bq_func.impls) + + rendered = constants.TEMPLATES["signature_def"].render( + func_name=sig_func_name, + max_args=max_args, + impls=[impl.to_dict() for impl in bq_func.impls], + sql_name=bq_func.name, + dtype_map=constants.DTYPE_MAP, + ) + + return sig_func_name, rendered + + +def _to_bigframes_op(bq_func: data_models.BQFunc) -> data_models.BigFramesOp: + pass + + +def _to_bigframes_func(bq_func: data_models.BQFunc) -> data_models.BigFramesFunc: + pass + + +def render_operation( + yaml_file: pathlib.Path, bq_module: data_models.BQModule, template: jinja2.Template +) -> str: + ops = [] + functions = [] + + return template.render( + yaml_path=yaml_file.relative_to(constants.PACKAGE_ROOT), + script_path=constants.SCRIPT_PATH_RELATIVE, + ops=ops, + functions=functions, + ) + + +def render_series_accessor() -> str: + pass + + +def render_tests() -> str: + pass diff --git a/packages/bigframes/scripts/generate_bigframes_bigquery/yaml_parser.py b/packages/bigframes/scripts/generate_bigframes_bigquery/yaml_parser.py new file mode 100644 index 000000000000..5f83ce52db4d --- /dev/null +++ b/packages/bigframes/scripts/generate_bigframes_bigquery/yaml_parser.py @@ -0,0 +1,63 @@ +# 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. + + +import pathlib +from typing import Any + +import constants +import data_models +import yaml + + +def _parse_func_arg(arg_data: Any) -> data_models.BQFuncArg: + return data_models.BQFuncArg( + name=arg_data["name"], + value=arg_data["value"], + optional=arg_data["optional"], + keyword_only=arg_data["keyword_only"], + ) + + +def _parse_func_impl(impl_data: Any) -> data_models.BQFuncImpl: + return data_models.BQFuncImpl( + args=[_parse_func_arg(arg) for arg in impl_data["args"]], + return_type=impl_data["return"], + ) + + +def _parse_bq_func(func_data: Any) -> data_models.BQFunc: + return data_models.BQFunc( + name=func_data["name"], + description=func_data["description"], + impls=[_parse_func_impl(impl) for impl in func_data["impls"]], + ) + + +def parse_yaml(yaml_file: pathlib.Path) -> data_models.BQModule: + print(f"Parsing {yaml_file}...") + + with open(yaml_file, "r") as f: + data = yaml.safe_load(f) + + functions = [] + if isinstance(data, dict) and "scalar_functions" in data: + functions = [ + _parse_bq_func(func_data) for func_data in data["scalar_functions"] + ] + + return data_models.BQModule( + module_path=yaml_file.relative_to(constants.DATA_DIR).with_suffix(""), + functions=functions, + ) From cc3288ec0b80f476582cf5c4044a3b656ae8f169 Mon Sep 17 00:00:00 2001 From: Shenyang Cai Date: Sun, 26 Jul 2026 20:19:51 +0000 Subject: [PATCH 02/27] checkpoint: impl operator definition gen --- .../generate_bigframes_bigquery/__main__.py | 4 +- .../generate_bigframes_bigquery/constants.py | 56 ++++++- .../data_models.py | 28 ++-- .../file_generator.py | 65 ++++++-- .../template_renderer.py | 144 ++++++++++++++---- .../yaml_parser.py | 27 +++- 6 files changed, 262 insertions(+), 62 deletions(-) diff --git a/packages/bigframes/scripts/generate_bigframes_bigquery/__main__.py b/packages/bigframes/scripts/generate_bigframes_bigquery/__main__.py index 74a5073f33e4..11df62b2bcfb 100644 --- a/packages/bigframes/scripts/generate_bigframes_bigquery/__main__.py +++ b/packages/bigframes/scripts/generate_bigframes_bigquery/__main__.py @@ -22,11 +22,10 @@ # See the License for the specific language governing permissions and # limitations under the License. -import pprint import constants -import yaml_parser import file_generator +import yaml_parser def main(): @@ -34,6 +33,7 @@ def main(): for yaml_file in sorted(constants.DATA_DIR.glob("**/*.yaml")): modules.append(yaml_parser.parse_yaml(yaml_file)) + break file_generator.generate(modules) diff --git a/packages/bigframes/scripts/generate_bigframes_bigquery/constants.py b/packages/bigframes/scripts/generate_bigframes_bigquery/constants.py index 6b8e45c84b70..18dd2dd6e45c 100644 --- a/packages/bigframes/scripts/generate_bigframes_bigquery/constants.py +++ b/packages/bigframes/scripts/generate_bigframes_bigquery/constants.py @@ -13,12 +13,13 @@ # limitations under the License. import pathlib + import jinja2 SCRIPTS_DIRECTORY = pathlib.Path(__file__).parent.parent.absolute() PACKAGE_ROOT = SCRIPTS_DIRECTORY.parent CODE_ROOT = PACKAGE_ROOT / "bigframes" -SCRIPT_PATH_RELATIVE = pathlib.Path(__file__).relative_to(PACKAGE_ROOT) +SCRIPT_PATH_RELATIVE = pathlib.Path(__file__).relative_to(PACKAGE_ROOT).parent # Directory containing the YAML files DATA_DIR = SCRIPTS_DIRECTORY / "data" / "sql-functions" @@ -119,6 +120,59 @@ "decimal<76,38>": "dtypes.BIGNUMERIC_DTYPE", } +YAML_TYPE_TO_COL = { + "binary": "bytes_col", + "string": "string_col", + "int64": "int64_col", + "i64": "int64_col", + "float64": "float64_col", + "fp64": "float64_col", + "bool": "bool_col", + "boolean": "bool_col", + "geography": "geography_col", + "date": "date_col", + "time": "time_col", + "datetime": "datetime_col", + "timestamp": "timestamp_col", + "decimal<38,9>": "numeric_col", + "decimal<76,38>": "bignumeric_col", +} + +PY_TYPE_MAP = { + "binary": "bytes", + "string": "str", + "int64": "int", + "i64": "int", + "float64": "float", + "fp64": "float", + "bool": "bool", + "boolean": "bool", + "geography": "Any", + "json": "Any", + "date": "datetime.date", + "time": "datetime.time", + "datetime": "datetime.datetime", + "timestamp": "datetime.datetime", + "struct": "dict", + "decimal<38,9>": "decimal.Decimal", + "decimal<76,38>": "decimal.Decimal", + "interval_day": "datetime.timedelta", +} + +RUFF_COMMON_ARGS = [ + "--target-version=py310", + "--line-length=88", +] +RUFF_CHECK_ARGS = [ + "check", + "--select", + "I,F", + "--fix", +] + RUFF_COMMON_ARGS +RUFF_FORMAT_ARGS = [ + "format", +] + RUFF_COMMON_ARGS + TEMPLATES: dict[str, jinja2.Template] # Lazily-loaded global variable _templates = None diff --git a/packages/bigframes/scripts/generate_bigframes_bigquery/data_models.py b/packages/bigframes/scripts/generate_bigframes_bigquery/data_models.py index b5af7036a577..ba2b8d6b536d 100644 --- a/packages/bigframes/scripts/generate_bigframes_bigquery/data_models.py +++ b/packages/bigframes/scripts/generate_bigframes_bigquery/data_models.py @@ -15,9 +15,9 @@ import dataclasses import pathlib +from typing import Any import constants -from typing import Any @dataclasses.dataclass @@ -28,6 +28,14 @@ class BQFuncArg: keyword_only: bool +@dataclasses.dataclass +class BigFramesFuncArg: + name: str + types: set[str] + optional: bool + keyword_only: bool + + @dataclasses.dataclass class BQFuncImpl: args: list[BQFuncArg] @@ -56,26 +64,19 @@ def to_dict(self) -> dict[str, Any]: @dataclasses.dataclass class BQFunc: name: str + op_base_name: str description: str impls: list[BQFuncImpl] @dataclasses.dataclass class BQModule: - module_path: pathlib.Path + yaml_file: pathlib.Path functions: list[BQFunc] @property - def name(self) -> str: - return self.module_path.name - - @property - def is_global_namespace(self) -> bool: - return "global_namespace" in self.module_path.parts - - @property - def output_file(self): - return constants.OUTPUT_DIR.joinpath(self.module_path).with_suffix(".py") + def module_path(self): + return self.yaml_file.relative_to(constants.DATA_DIR).with_suffix("") @dataclasses.dataclass @@ -84,12 +85,13 @@ class BigFramesOp: sql_name: str arg_specs: str signature: str - signature_definition: str + signature_definition: str | None @dataclasses.dataclass class BigFramesFunc: name: str + op_name: str description: str args: list[str] series_accessor_arg: list[str] diff --git a/packages/bigframes/scripts/generate_bigframes_bigquery/file_generator.py b/packages/bigframes/scripts/generate_bigframes_bigquery/file_generator.py index 46744fe1724a..90b177652257 100644 --- a/packages/bigframes/scripts/generate_bigframes_bigquery/file_generator.py +++ b/packages/bigframes/scripts/generate_bigframes_bigquery/file_generator.py @@ -12,33 +12,76 @@ # See the License for the specific language governing permissions and # limitations under the License. +import pathlib +import sys +import subprocess + import constants -import jinja2 import data_models -import pprint - import template_renderer -def _generate_sql_operator_def(): - pass +def _ensure_init_py(directory: pathlib.Path, limit_dir: pathlib.Path): + """Ensures __init__.py exists in the directory and its parents up to limit_dir.""" + curr = directory + while curr != limit_dir and curr != curr.parent: + init_file = curr / "__init__.py" + if not init_file.exists(): + print(f" Creating {init_file}") + content = constants.TEMPLATES["license"].render() + with open(init_file, "w") as f: + f.write(content) + curr = curr.parent + + +def _write_file(content: str, output_file: pathlib.Path, limit_dir: pathlib.Path): + output_file.parent.mkdir(parents=True, exist_ok=True) + _ensure_init_py(output_file.parent, limit_dir) + + with open(output_file, "w") as f: + f.write(content) + print(f" Generated {output_file}") + + +def _run_ruff(): + targets = [ + constants.OUTPUT_DIR, + constants.TEST_OUTPUT_DIR, + ] + + subprocess.run( + [sys.executable, "-m", "ruff"] + constants.RUFF_CHECK_ARGS + targets, + check=True, + ) + + subprocess.run( + [sys.executable, "-m", "ruff"] + constants.RUFF_FORMAT_ARGS + targets, + check=True, + ) + + +def _generate_op_defs(bq_module: data_models.BQModule): + content = template_renderer.render_operation(bq_module) + + output_file = constants.OUTPUT_DIR.joinpath(bq_module.module_path).with_suffix( + ".py" + ) + + _write_file(content, output_file, constants.OUTPUT_DIR.parent) def _generate_accesor(): pass -def _generate_tests(): +def _generate_tests(bq_module: data_models.BQModule): pass def generate(bq_modules: list[data_models.BQModule]): for bq_module in bq_modules: - for bq_func in bq_module.functions: - pprint.pp(template_renderer.render_signature_def(bq_func, bq_module)) - - # Write to file + _generate_op_defs(bq_module) # Ruff format - pass + _run_ruff() diff --git a/packages/bigframes/scripts/generate_bigframes_bigquery/template_renderer.py b/packages/bigframes/scripts/generate_bigframes_bigquery/template_renderer.py index ca71c0441c08..3562fe5f6472 100644 --- a/packages/bigframes/scripts/generate_bigframes_bigquery/template_renderer.py +++ b/packages/bigframes/scripts/generate_bigframes_bigquery/template_renderer.py @@ -13,22 +13,9 @@ # See the License for the specific language governing permissions and # limitations under the License. -import re -import pathlib import constants import data_models -import jinja2 - - -def _to_snake_case(name): - # Replace dots with underscores - name = name.replace(".", "_") - # Handle CamelCase to snake_case - name = re.sub(r"(? str: @@ -79,21 +66,12 @@ def _validate_types(impls): def render_signature_def( bq_func: data_models.BQFunc, - hosting_bq_module: data_models.BQModule, ) -> tuple[str, str | None]: """ Returns the signature function name and it's definition. If the signature function can be inlined, the first return value is the lambda, and the second value is None. """ - op_base_name = _to_snake_case(bq_func.name) - - module_name = hosting_bq_module.name - if not hosting_bq_module.is_global_namespace and op_base_name.startswith( - module_name + "_" - ): - op_base_name = op_base_name[len(module_name) + 1 :] - return_types = {impl.return_type for impl in bq_func.impls} # Optimization: if all impls return the same concrete type, # inline the signature function as a lambda @@ -105,7 +83,7 @@ def render_signature_def( _validate_types(bq_func.impls) - python_name = op_base_name + python_name = bq_func.op_base_name if python_name in constants.PYTHON_BUILTINS: python_name = python_name + "_" sig_func_name = f"_{python_name.upper()}_SIG" @@ -123,31 +101,133 @@ def render_signature_def( return sig_func_name, rendered +def _get_bigframes_func_args( + bq_func: data_models.BQFunc, +) -> list[data_models.BigFramesFuncArg]: + """ + Coalesces arguments from all the implementations of this function, + and return them in the order of appearance in the yaml file + """ + args_by_name = {} + arg_order = [] + arg_appearances = {} + for impl in bq_func.impls: + seen_in_impl = set() + for bq_func_arg in impl.args: + name = bq_func_arg.name + seen_in_impl.add(name) + if name not in args_by_name: + args_by_name[name] = data_models.BigFramesFuncArg( + name=name, + types=set(), + optional=bq_func_arg.optional, + keyword_only=bq_func_arg.keyword_only, + ) + arg_order.append(name) + else: + # If it was marked optional or keyword_only in any previous impl, keep it. + # Or if this impl marks it as optional/keyword_only, update it. + if bq_func_arg.optional: + args_by_name[name].optional = True + if bq_func_arg.keyword_only: + args_by_name[name].keyword_only = True + args_by_name[name].types.add(bq_func_arg.value) + for name in seen_in_impl: + arg_appearances[name] = arg_appearances.get(name, 0) + 1 + + # If an argument is not in all impls, it must be optional overall + num_impls = len(bq_func.impls) + for name, count in arg_appearances.items(): + if count < num_impls: + args_by_name[name].optional = True + + return [args_by_name[name] for name in arg_order] + + def _to_bigframes_op(bq_func: data_models.BQFunc) -> data_models.BigFramesOp: - pass + arg_specs = [] + for bf_func_arg in _get_bigframes_func_args(bq_func): + spec = "googlesql.ArgSpec(" + if bf_func_arg.keyword_only: + spec += f'arg_name="{bf_func_arg.name}", ' + if bf_func_arg.optional: + spec += "optional=True, " + spec = spec.rstrip(", ") + ")" + arg_specs.append(spec) + + arg_specs_str = ", ".join(arg_specs) + if len(arg_specs) == 1: + arg_specs_str += "," + + (signature, signature_definition) = render_signature_def(bq_func) + + return data_models.BigFramesOp( + internal_name=f"_{bq_func.op_base_name.upper()}_OP", + sql_name=bq_func.name.upper(), + arg_specs=arg_specs_str, + signature=signature, + signature_definition=signature_definition, + ) def _to_bigframes_func(bq_func: data_models.BQFunc) -> data_models.BigFramesFunc: - pass + func_args = [] + for arg in _get_bigframes_func_args(bq_func): + types = [constants.PY_TYPE_MAP.get(t, "Any") for t in sorted(arg.types)] + [ + "Literal[sentinels.Sentinel.ARGUMENT_DEFAULT]" + ] + type_hint = ( + "Union[" + ", ".join(sorted(set(types))) + "]" + if len(types) > 1 + else types[0] + ) + func_args.append( + { + "name": arg.name, + "type_hint": type_hint, + "default": ( + "sentinels.Sentinel.ARGUMENT_DEFAULT" if arg.optional else "" + ), + } + ) + + # Clean up default values for mandatory args + # In Python, mandatory args come first. + for arg in func_args: + if not arg.get("default"): + arg.pop("default", None) + + python_name = bq_func.op_base_name + if python_name in constants.PYTHON_BUILTINS: + python_name = python_name + "_" + + return data_models.BigFramesFunc( + name=python_name, + op_name=f"_{bq_func.op_base_name.upper()}_OP", + description=bq_func.description, + args=func_args, + series_accessor_arg=[], + ) def render_operation( - yaml_file: pathlib.Path, bq_module: data_models.BQModule, template: jinja2.Template + bq_module: data_models.BQModule, ) -> str: ops = [] functions = [] - return template.render( - yaml_path=yaml_file.relative_to(constants.PACKAGE_ROOT), + for bq_func in bq_module.functions: + ops.append(_to_bigframes_op(bq_func)) + functions.append(_to_bigframes_func(bq_func)) + + return constants.TEMPLATES["operation"].render( + yaml_path=bq_module.yaml_file.relative_to(constants.PACKAGE_ROOT), script_path=constants.SCRIPT_PATH_RELATIVE, ops=ops, functions=functions, ) -def render_series_accessor() -> str: - pass - +def render_tests(bq_module: data_models.BQModule) -> str: -def render_tests() -> str: pass diff --git a/packages/bigframes/scripts/generate_bigframes_bigquery/yaml_parser.py b/packages/bigframes/scripts/generate_bigframes_bigquery/yaml_parser.py index 5f83ce52db4d..257e17aa3e6b 100644 --- a/packages/bigframes/scripts/generate_bigframes_bigquery/yaml_parser.py +++ b/packages/bigframes/scripts/generate_bigframes_bigquery/yaml_parser.py @@ -14,6 +14,7 @@ import pathlib +import re from typing import Any import constants @@ -21,6 +22,16 @@ import yaml +def _to_snake_case(name): + # Replace dots with underscores + name = name.replace(".", "_") + # Handle CamelCase to snake_case + name = re.sub(r"(? data_models.BQFuncArg: return data_models.BQFuncArg( name=arg_data["name"], @@ -37,9 +48,17 @@ def _parse_func_impl(impl_data: Any) -> data_models.BQFuncImpl: ) -def _parse_bq_func(func_data: Any) -> data_models.BQFunc: +def _parse_bq_func(func_data: Any, module_path: pathlib.Path) -> data_models.BQFunc: + op_base_name = _to_snake_case(func_data["name"]) + + module_name = module_path.name + is_global = "global_namespace" in module_path.parts + if not is_global and op_base_name.startswith(module_name + "_"): + op_base_name = op_base_name[len(module_name) + 1 :] + return data_models.BQFunc( name=func_data["name"], + op_base_name=op_base_name, description=func_data["description"], impls=[_parse_func_impl(impl) for impl in func_data["impls"]], ) @@ -52,12 +71,14 @@ def parse_yaml(yaml_file: pathlib.Path) -> data_models.BQModule: data = yaml.safe_load(f) functions = [] + module_path = yaml_file.relative_to(constants.DATA_DIR).with_suffix("") if isinstance(data, dict) and "scalar_functions" in data: functions = [ - _parse_bq_func(func_data) for func_data in data["scalar_functions"] + _parse_bq_func(func_data, module_path) + for func_data in data["scalar_functions"] ] return data_models.BQModule( - module_path=yaml_file.relative_to(constants.DATA_DIR).with_suffix(""), + yaml_file=yaml_file, functions=functions, ) From 8019f08b7090f67a1f77e68350da31711936c4fd Mon Sep 17 00:00:00 2001 From: Shenyang Cai Date: Sun, 26 Jul 2026 20:44:07 +0000 Subject: [PATCH 03/27] checkout: impl test generation --- .../generate_bigframes_bigquery/data_models.py | 1 + .../file_generator.py | 9 ++++++++- .../template_renderer.py | 16 +++++++++++++++- .../generate_bigframes_bigquery/yaml_parser.py | 10 ++++++---- 4 files changed, 30 insertions(+), 6 deletions(-) diff --git a/packages/bigframes/scripts/generate_bigframes_bigquery/data_models.py b/packages/bigframes/scripts/generate_bigframes_bigquery/data_models.py index ba2b8d6b536d..d247f50d8913 100644 --- a/packages/bigframes/scripts/generate_bigframes_bigquery/data_models.py +++ b/packages/bigframes/scripts/generate_bigframes_bigquery/data_models.py @@ -73,6 +73,7 @@ class BQFunc: class BQModule: yaml_file: pathlib.Path functions: list[BQFunc] + is_global: bool @property def module_path(self): diff --git a/packages/bigframes/scripts/generate_bigframes_bigquery/file_generator.py b/packages/bigframes/scripts/generate_bigframes_bigquery/file_generator.py index 90b177652257..9b698593d6e1 100644 --- a/packages/bigframes/scripts/generate_bigframes_bigquery/file_generator.py +++ b/packages/bigframes/scripts/generate_bigframes_bigquery/file_generator.py @@ -75,13 +75,20 @@ def _generate_accesor(): def _generate_tests(bq_module: data_models.BQModule): - pass + content = template_renderer.render_tests(bq_module) + + output_file = constants.TEST_OUTPUT_DIR.joinpath( + bq_module.module_path.with_name(f"test_{bq_module.module_path.name}") + ).with_suffix(".py") + + _write_file(content, output_file, constants.TEST_OUTPUT_DIR.parent) def generate(bq_modules: list[data_models.BQModule]): for bq_module in bq_modules: _generate_op_defs(bq_module) + _generate_tests(bq_module) # Ruff format _run_ruff() diff --git a/packages/bigframes/scripts/generate_bigframes_bigquery/template_renderer.py b/packages/bigframes/scripts/generate_bigframes_bigquery/template_renderer.py index 3562fe5f6472..b7c0535891a6 100644 --- a/packages/bigframes/scripts/generate_bigframes_bigquery/template_renderer.py +++ b/packages/bigframes/scripts/generate_bigframes_bigquery/template_renderer.py @@ -230,4 +230,18 @@ def render_operation( def render_tests(bq_module: data_models.BQModule) -> str: - pass + import_path = "bigframes.operations.googlesql." + ".".join( + bq_module.module_path.parts + ) + functions = [] + for bq_func in bq_module.functions: + functions.append(_to_bigframes_func(bq_func)) + + return constants.TEMPLATES["test_operation"].render( + yaml_path=bq_module.yaml_file.relative_to(constants.PACKAGE_ROOT), + script_path=constants.SCRIPT_PATH_RELATIVE, + import_path=import_path, + short_name=bq_module.module_path.name, + is_global=bq_module.is_global, + functions=functions, + ) diff --git a/packages/bigframes/scripts/generate_bigframes_bigquery/yaml_parser.py b/packages/bigframes/scripts/generate_bigframes_bigquery/yaml_parser.py index 257e17aa3e6b..4676bf388e34 100644 --- a/packages/bigframes/scripts/generate_bigframes_bigquery/yaml_parser.py +++ b/packages/bigframes/scripts/generate_bigframes_bigquery/yaml_parser.py @@ -48,11 +48,11 @@ def _parse_func_impl(impl_data: Any) -> data_models.BQFuncImpl: ) -def _parse_bq_func(func_data: Any, module_path: pathlib.Path) -> data_models.BQFunc: +def _parse_bq_func( + func_data: Any, module_name: str, is_global: bool +) -> data_models.BQFunc: op_base_name = _to_snake_case(func_data["name"]) - module_name = module_path.name - is_global = "global_namespace" in module_path.parts if not is_global and op_base_name.startswith(module_name + "_"): op_base_name = op_base_name[len(module_name) + 1 :] @@ -72,13 +72,15 @@ def parse_yaml(yaml_file: pathlib.Path) -> data_models.BQModule: functions = [] module_path = yaml_file.relative_to(constants.DATA_DIR).with_suffix("") + is_global = "global_namespace" in module_path.parts if isinstance(data, dict) and "scalar_functions" in data: functions = [ - _parse_bq_func(func_data, module_path) + _parse_bq_func(func_data, module_path.name, is_global) for func_data in data["scalar_functions"] ] return data_models.BQModule( yaml_file=yaml_file, functions=functions, + is_global=is_global, ) From 310a93a908af698914c775ca44069f576da76d81 Mon Sep 17 00:00:00 2001 From: Shenyang Cai Date: Sun, 26 Jul 2026 21:21:47 +0000 Subject: [PATCH 04/27] checkpoint: refactor arg rendering --- .../generate_bigframes_bigquery/__main__.py | 1 - .../data_models.py | 37 +++++++++++++------ .../template_renderer.py | 29 +-------------- 3 files changed, 27 insertions(+), 40 deletions(-) diff --git a/packages/bigframes/scripts/generate_bigframes_bigquery/__main__.py b/packages/bigframes/scripts/generate_bigframes_bigquery/__main__.py index 11df62b2bcfb..af5c0cd8847a 100644 --- a/packages/bigframes/scripts/generate_bigframes_bigquery/__main__.py +++ b/packages/bigframes/scripts/generate_bigframes_bigquery/__main__.py @@ -33,7 +33,6 @@ def main(): for yaml_file in sorted(constants.DATA_DIR.glob("**/*.yaml")): modules.append(yaml_parser.parse_yaml(yaml_file)) - break file_generator.generate(modules) diff --git a/packages/bigframes/scripts/generate_bigframes_bigquery/data_models.py b/packages/bigframes/scripts/generate_bigframes_bigquery/data_models.py index d247f50d8913..ee37dacad2e8 100644 --- a/packages/bigframes/scripts/generate_bigframes_bigquery/data_models.py +++ b/packages/bigframes/scripts/generate_bigframes_bigquery/data_models.py @@ -27,15 +27,6 @@ class BQFuncArg: optional: bool keyword_only: bool - -@dataclasses.dataclass -class BigFramesFuncArg: - name: str - types: set[str] - optional: bool - keyword_only: bool - - @dataclasses.dataclass class BQFuncImpl: args: list[BQFuncArg] @@ -89,10 +80,34 @@ class BigFramesOp: signature_definition: str | None +@dataclasses.dataclass +class BigFramesFuncArg: + name: str + types: set[str] + optional: bool + keyword_only: bool + + @property + def type_hint(self) -> str: + types = [constants.PY_TYPE_MAP.get(t, "Any") for t in sorted(self.types)] + [ + "Literal[sentinels.Sentinel.ARGUMENT_DEFAULT]" + ] + + if len(types) > 1: + return "Union[" + ", ".join(sorted(set(types))) + "]" + + return types[0] + + @property + def default(self) -> str | None: + if self.optional: + return "sentinels.Sentinel.ARGUMENT_DEFAULT" + return None + + @dataclasses.dataclass class BigFramesFunc: name: str op_name: str description: str - args: list[str] - series_accessor_arg: list[str] + args: list[BigFramesFuncArg] diff --git a/packages/bigframes/scripts/generate_bigframes_bigquery/template_renderer.py b/packages/bigframes/scripts/generate_bigframes_bigquery/template_renderer.py index b7c0535891a6..ed93fa1433d2 100644 --- a/packages/bigframes/scripts/generate_bigframes_bigquery/template_renderer.py +++ b/packages/bigframes/scripts/generate_bigframes_bigquery/template_renderer.py @@ -171,32 +171,6 @@ def _to_bigframes_op(bq_func: data_models.BQFunc) -> data_models.BigFramesOp: def _to_bigframes_func(bq_func: data_models.BQFunc) -> data_models.BigFramesFunc: - func_args = [] - for arg in _get_bigframes_func_args(bq_func): - types = [constants.PY_TYPE_MAP.get(t, "Any") for t in sorted(arg.types)] + [ - "Literal[sentinels.Sentinel.ARGUMENT_DEFAULT]" - ] - type_hint = ( - "Union[" + ", ".join(sorted(set(types))) + "]" - if len(types) > 1 - else types[0] - ) - func_args.append( - { - "name": arg.name, - "type_hint": type_hint, - "default": ( - "sentinels.Sentinel.ARGUMENT_DEFAULT" if arg.optional else "" - ), - } - ) - - # Clean up default values for mandatory args - # In Python, mandatory args come first. - for arg in func_args: - if not arg.get("default"): - arg.pop("default", None) - python_name = bq_func.op_base_name if python_name in constants.PYTHON_BUILTINS: python_name = python_name + "_" @@ -205,8 +179,7 @@ def _to_bigframes_func(bq_func: data_models.BQFunc) -> data_models.BigFramesFunc name=python_name, op_name=f"_{bq_func.op_base_name.upper()}_OP", description=bq_func.description, - args=func_args, - series_accessor_arg=[], + args=_get_bigframes_func_args(bq_func), ) From 5ba9eb382fae6dc71a6e7f6ce26ed11701695825 Mon Sep 17 00:00:00 2001 From: Shenyang Cai Date: Mon, 27 Jul 2026 01:58:29 +0000 Subject: [PATCH 05/27] finish refactoring --- .../extensions/bigframes/series_accessor.py | 14 +-- .../extensions/core/series_accessor.py | 18 ++-- .../extensions/pandas/series_accessor.py | 14 +-- .../bigframes/operations/googlesql/aead.py | 2 +- .../global_namespace/aead_encryption.py | 2 +- .../googlesql/global_namespace/array.py | 2 +- .../googlesql/global_namespace/bit.py | 2 +- .../googlesql/global_namespace/conversion.py | 2 +- .../googlesql/global_namespace/date.py | 2 +- .../__init__.py | 0 .../__main__.py | 10 --- .../constants.py | 21 +---- .../data_models.py | 30 ++++++- .../file_generator.py | 34 +++++-- .../template_renderer.py | 88 ++++++++++++++++++- .../yaml_parser.py | 2 +- .../global_namespace/test_aead_encryption.py | 2 +- .../generated/global_namespace/test_array.py | 2 +- .../generated/global_namespace/test_bit.py | 2 +- .../global_namespace/test_conversion.py | 2 +- .../generated/global_namespace/test_date.py | 2 +- .../unit/bigquery/generated/test_aead.py | 2 +- 22 files changed, 179 insertions(+), 76 deletions(-) rename packages/bigframes/scripts/{generate_bigframes_bigquery => bigquery_generator}/__init__.py (100%) rename packages/bigframes/scripts/{generate_bigframes_bigquery => bigquery_generator}/__main__.py (86%) rename packages/bigframes/scripts/{generate_bigframes_bigquery => bigquery_generator}/constants.py (90%) rename packages/bigframes/scripts/{generate_bigframes_bigquery => bigquery_generator}/data_models.py (77%) rename packages/bigframes/scripts/{generate_bigframes_bigquery => bigquery_generator}/file_generator.py (76%) rename packages/bigframes/scripts/{generate_bigframes_bigquery => bigquery_generator}/template_renderer.py (70%) rename packages/bigframes/scripts/{generate_bigframes_bigquery => bigquery_generator}/yaml_parser.py (97%) diff --git a/packages/bigframes/bigframes/extensions/bigframes/series_accessor.py b/packages/bigframes/bigframes/extensions/bigframes/series_accessor.py index 8379e6a145a0..4007df5623a0 100644 --- a/packages/bigframes/bigframes/extensions/bigframes/series_accessor.py +++ b/packages/bigframes/bigframes/extensions/bigframes/series_accessor.py @@ -13,7 +13,7 @@ # limitations under the License. # # DO NOT MODIFY THIS FILE DIRECTLY. -# This file was generated by the script: scripts/generate_bigframes_bigquery.py +# This file was generated by the script: scripts # from __future__ import annotations @@ -44,17 +44,17 @@ def _to_dataframe(self, bf_df: dataframe.DataFrame) -> T: def _to_series(self, bf_series: series.Series) -> S: return cast(S, bf_series) - @property - def ai(self) -> BigframesAiSeriesAccessor[T, S]: - return BigframesAiSeriesAccessor(self._obj) - @property def aead(self) -> BigframesAeadSeriesAccessor[T, S]: return BigframesAeadSeriesAccessor(self._obj) + @property + def ai(self) -> BigframesAiSeriesAccessor[T, S]: + return BigframesAiSeriesAccessor(self._obj) + @log_adapter.class_logger -class BigframesAiSeriesAccessor(core_accessor.AiSeriesAccessor[T, S]): +class BigframesAeadSeriesAccessor(core_accessor.AeadSeriesAccessor[T, S]): def __init__(self, bf_obj: S): super().__init__(bf_obj) @@ -71,7 +71,7 @@ def _to_series(self, bf_series: series.Series) -> S: @log_adapter.class_logger -class BigframesAeadSeriesAccessor(core_accessor.AeadSeriesAccessor[T, S]): +class BigframesAiSeriesAccessor(core_accessor.AiSeriesAccessor[T, S]): def __init__(self, bf_obj: S): super().__init__(bf_obj) diff --git a/packages/bigframes/bigframes/extensions/core/series_accessor.py b/packages/bigframes/bigframes/extensions/core/series_accessor.py index 440e6731aba2..8ff12fbdfc28 100644 --- a/packages/bigframes/bigframes/extensions/core/series_accessor.py +++ b/packages/bigframes/bigframes/extensions/core/series_accessor.py @@ -13,7 +13,7 @@ # limitations under the License. # # DO NOT MODIFY THIS FILE DIRECTLY. -# This file was generated by the script: scripts/generate_bigframes_bigquery.py +# This file was generated by the script: scripts # from __future__ import annotations @@ -44,13 +44,13 @@ class BigQuerySeriesAccessor( @property @abc.abstractmethod - def ai(self) -> AiSeriesAccessor[T, S]: - """Accessor for BigQuery ai functions.""" + def aead(self) -> AeadSeriesAccessor[T, S]: + """Accessor for BigQuery aead functions.""" @property @abc.abstractmethod - def aead(self) -> AeadSeriesAccessor[T, S]: - """Accessor for BigQuery aead functions.""" + def ai(self) -> AiSeriesAccessor[T, S]: + """Accessor for BigQuery ai functions.""" def deterministic_decrypt_bytes( self, @@ -1112,10 +1112,6 @@ def unix_date( return self._to_series(cast(series.Series, result)) -class AiSeriesAccessor(series_tvf_mixins.AITVFMixin[T, S]): - """Series accessor for BigQuery ai functions.""" - - class AeadSeriesAccessor(abstract_series_accessor.AbstractBigQuerySeriesAccessor[T, S]): """Series accessor for BigQuery aead functions.""" @@ -1227,3 +1223,7 @@ def encrypt( additional_data, ) return self._to_series(cast(series.Series, result)) + + +class AiSeriesAccessor(series_tvf_mixins.AITVFMixin[T, S]): + """Series accessor for BigQuery ai functions.""" diff --git a/packages/bigframes/bigframes/extensions/pandas/series_accessor.py b/packages/bigframes/bigframes/extensions/pandas/series_accessor.py index 204f1e0d2cf3..80ebdb661ced 100644 --- a/packages/bigframes/bigframes/extensions/pandas/series_accessor.py +++ b/packages/bigframes/bigframes/extensions/pandas/series_accessor.py @@ -13,7 +13,7 @@ # limitations under the License. # # DO NOT MODIFY THIS FILE DIRECTLY. -# This file was generated by the script: scripts/generate_bigframes_bigquery.py +# This file was generated by the script: scripts # from __future__ import annotations @@ -51,17 +51,17 @@ def _to_dataframe(self, bf_df: dataframe.DataFrame) -> T: def _to_series(self, bf_series: series.Series) -> S: return cast(S, bf_series.to_pandas(ordered=True)) - @property - def ai(self) -> PandasAiSeriesAccessor[T, S]: - return PandasAiSeriesAccessor(self._obj) - @property def aead(self) -> PandasAeadSeriesAccessor[T, S]: return PandasAeadSeriesAccessor(self._obj) + @property + def ai(self) -> PandasAiSeriesAccessor[T, S]: + return PandasAiSeriesAccessor(self._obj) + @log_adapter.class_logger -class PandasAiSeriesAccessor(core_accessor.AiSeriesAccessor[T, S]): +class PandasAeadSeriesAccessor(core_accessor.AeadSeriesAccessor[T, S]): def __init__(self, pandas_obj: S): super().__init__(pandas_obj) @@ -80,7 +80,7 @@ def _to_series(self, bf_series: series.Series) -> S: @log_adapter.class_logger -class PandasAeadSeriesAccessor(core_accessor.AeadSeriesAccessor[T, S]): +class PandasAiSeriesAccessor(core_accessor.AiSeriesAccessor[T, S]): def __init__(self, pandas_obj: S): super().__init__(pandas_obj) diff --git a/packages/bigframes/bigframes/operations/googlesql/aead.py b/packages/bigframes/bigframes/operations/googlesql/aead.py index f719d7d69893..c1b9ecb3618c 100644 --- a/packages/bigframes/bigframes/operations/googlesql/aead.py +++ b/packages/bigframes/bigframes/operations/googlesql/aead.py @@ -14,7 +14,7 @@ # # 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 +# by the script: scripts from __future__ import annotations diff --git a/packages/bigframes/bigframes/operations/googlesql/global_namespace/aead_encryption.py b/packages/bigframes/bigframes/operations/googlesql/global_namespace/aead_encryption.py index 4613ddd7e6df..d7e0f58a6435 100644 --- a/packages/bigframes/bigframes/operations/googlesql/global_namespace/aead_encryption.py +++ b/packages/bigframes/bigframes/operations/googlesql/global_namespace/aead_encryption.py @@ -14,7 +14,7 @@ # # DO NOT MODIFY THIS FILE DIRECTLY. # This file was generated from: scripts/data/sql-functions/global_namespace/aead_encryption.yaml -# by the script: scripts/generate_bigframes_bigquery.py +# by the script: scripts from __future__ import annotations diff --git a/packages/bigframes/bigframes/operations/googlesql/global_namespace/array.py b/packages/bigframes/bigframes/operations/googlesql/global_namespace/array.py index 94adbad1839d..bd3a54021b45 100644 --- a/packages/bigframes/bigframes/operations/googlesql/global_namespace/array.py +++ b/packages/bigframes/bigframes/operations/googlesql/global_namespace/array.py @@ -14,7 +14,7 @@ # # DO NOT MODIFY THIS FILE DIRECTLY. # This file was generated from: scripts/data/sql-functions/global_namespace/array.yaml -# by the script: scripts/generate_bigframes_bigquery.py +# by the script: scripts from __future__ import annotations diff --git a/packages/bigframes/bigframes/operations/googlesql/global_namespace/bit.py b/packages/bigframes/bigframes/operations/googlesql/global_namespace/bit.py index e0c22dfc2990..818d8b178753 100644 --- a/packages/bigframes/bigframes/operations/googlesql/global_namespace/bit.py +++ b/packages/bigframes/bigframes/operations/googlesql/global_namespace/bit.py @@ -14,7 +14,7 @@ # # DO NOT MODIFY THIS FILE DIRECTLY. # This file was generated from: scripts/data/sql-functions/global_namespace/bit.yaml -# by the script: scripts/generate_bigframes_bigquery.py +# by the script: scripts from __future__ import annotations diff --git a/packages/bigframes/bigframes/operations/googlesql/global_namespace/conversion.py b/packages/bigframes/bigframes/operations/googlesql/global_namespace/conversion.py index cea4e45d836b..2543c0a9651d 100644 --- a/packages/bigframes/bigframes/operations/googlesql/global_namespace/conversion.py +++ b/packages/bigframes/bigframes/operations/googlesql/global_namespace/conversion.py @@ -14,7 +14,7 @@ # # DO NOT MODIFY THIS FILE DIRECTLY. # This file was generated from: scripts/data/sql-functions/global_namespace/conversion.yaml -# by the script: scripts/generate_bigframes_bigquery.py +# by the script: scripts from __future__ import annotations diff --git a/packages/bigframes/bigframes/operations/googlesql/global_namespace/date.py b/packages/bigframes/bigframes/operations/googlesql/global_namespace/date.py index b6cfc9722b52..4bc675ad29f1 100644 --- a/packages/bigframes/bigframes/operations/googlesql/global_namespace/date.py +++ b/packages/bigframes/bigframes/operations/googlesql/global_namespace/date.py @@ -14,7 +14,7 @@ # # DO NOT MODIFY THIS FILE DIRECTLY. # This file was generated from: scripts/data/sql-functions/global_namespace/date.yaml -# by the script: scripts/generate_bigframes_bigquery.py +# by the script: scripts from __future__ import annotations diff --git a/packages/bigframes/scripts/generate_bigframes_bigquery/__init__.py b/packages/bigframes/scripts/bigquery_generator/__init__.py similarity index 100% rename from packages/bigframes/scripts/generate_bigframes_bigquery/__init__.py rename to packages/bigframes/scripts/bigquery_generator/__init__.py diff --git a/packages/bigframes/scripts/generate_bigframes_bigquery/__main__.py b/packages/bigframes/scripts/bigquery_generator/__main__.py similarity index 86% rename from packages/bigframes/scripts/generate_bigframes_bigquery/__main__.py rename to packages/bigframes/scripts/bigquery_generator/__main__.py index af5c0cd8847a..25a8af08f03c 100644 --- a/packages/bigframes/scripts/generate_bigframes_bigquery/__main__.py +++ b/packages/bigframes/scripts/bigquery_generator/__main__.py @@ -1,13 +1,3 @@ -#!/usr/bin/env -S uv run --active --script -# -# /// script -# dependencies = [ -# "jinja2", -# "pyyaml", -# "ruff==0.14.14", -# ] -# /// -# # Copyright 2026 Google LLC # # Licensed under the Apache License, Version 2.0 (the "License"); diff --git a/packages/bigframes/scripts/generate_bigframes_bigquery/constants.py b/packages/bigframes/scripts/bigquery_generator/constants.py similarity index 90% rename from packages/bigframes/scripts/generate_bigframes_bigquery/constants.py rename to packages/bigframes/scripts/bigquery_generator/constants.py index 18dd2dd6e45c..600944f55923 100644 --- a/packages/bigframes/scripts/generate_bigframes_bigquery/constants.py +++ b/packages/bigframes/scripts/bigquery_generator/constants.py @@ -19,7 +19,8 @@ SCRIPTS_DIRECTORY = pathlib.Path(__file__).parent.parent.absolute() PACKAGE_ROOT = SCRIPTS_DIRECTORY.parent CODE_ROOT = PACKAGE_ROOT / "bigframes" -SCRIPT_PATH_RELATIVE = pathlib.Path(__file__).relative_to(PACKAGE_ROOT).parent +SCRIPT_PATH_RELATIVE = pathlib.Path(__file__).relative_to(PACKAGE_ROOT).parent.parent + # Directory containing the YAML files DATA_DIR = SCRIPTS_DIRECTORY / "data" / "sql-functions" @@ -120,24 +121,6 @@ "decimal<76,38>": "dtypes.BIGNUMERIC_DTYPE", } -YAML_TYPE_TO_COL = { - "binary": "bytes_col", - "string": "string_col", - "int64": "int64_col", - "i64": "int64_col", - "float64": "float64_col", - "fp64": "float64_col", - "bool": "bool_col", - "boolean": "bool_col", - "geography": "geography_col", - "date": "date_col", - "time": "time_col", - "datetime": "datetime_col", - "timestamp": "timestamp_col", - "decimal<38,9>": "numeric_col", - "decimal<76,38>": "bignumeric_col", -} - PY_TYPE_MAP = { "binary": "bytes", "string": "str", diff --git a/packages/bigframes/scripts/generate_bigframes_bigquery/data_models.py b/packages/bigframes/scripts/bigquery_generator/data_models.py similarity index 77% rename from packages/bigframes/scripts/generate_bigframes_bigquery/data_models.py rename to packages/bigframes/scripts/bigquery_generator/data_models.py index ee37dacad2e8..66d9b4c061fa 100644 --- a/packages/bigframes/scripts/generate_bigframes_bigquery/data_models.py +++ b/packages/bigframes/scripts/bigquery_generator/data_models.py @@ -12,6 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +from __future__ import annotations import dataclasses import pathlib @@ -58,18 +59,29 @@ class BQFunc: op_base_name: str description: str impls: list[BQFuncImpl] + series_accessor_arg: str | None @dataclasses.dataclass class BQModule: yaml_file: pathlib.Path functions: list[BQFunc] - is_global: bool @property - def module_path(self): + def module_path(self) -> pathlib.Path: return self.yaml_file.relative_to(constants.DATA_DIR).with_suffix("") + @property + def namespace(self) -> tuple[str, ...]: + parts = self.module_path.parts + if "global_namespace" in parts: + return tuple() + return parts + + @property + def is_global(self) -> bool: + return "global_namespace" in self.module_path.parts + @dataclasses.dataclass class BigFramesOp: @@ -111,3 +123,17 @@ class BigFramesFunc: op_name: str description: str args: list[BigFramesFuncArg] + series_accessor_arg: str | None + import_module: str | None = None + + +@dataclasses.dataclass +class Accessor: + class_name: str + bigframes_class_name: str + pandas_class_name: str + is_root: bool + description: str + children: list[Accessor] + functions: list[BigFramesFunc] + prop_name: str | None = None diff --git a/packages/bigframes/scripts/generate_bigframes_bigquery/file_generator.py b/packages/bigframes/scripts/bigquery_generator/file_generator.py similarity index 76% rename from packages/bigframes/scripts/generate_bigframes_bigquery/file_generator.py rename to packages/bigframes/scripts/bigquery_generator/file_generator.py index 9b698593d6e1..44f42950084e 100644 --- a/packages/bigframes/scripts/generate_bigframes_bigquery/file_generator.py +++ b/packages/bigframes/scripts/bigquery_generator/file_generator.py @@ -13,8 +13,8 @@ # limitations under the License. import pathlib -import sys import subprocess +import sys import constants import data_models @@ -47,6 +47,7 @@ def _run_ruff(): targets = [ constants.OUTPUT_DIR, constants.TEST_OUTPUT_DIR, + constants.CODE_ROOT / "extensions", ] subprocess.run( @@ -62,6 +63,8 @@ def _run_ruff(): def _generate_op_defs(bq_module: data_models.BQModule): content = template_renderer.render_operation(bq_module) + if not content: + return output_file = constants.OUTPUT_DIR.joinpath(bq_module.module_path).with_suffix( ".py" @@ -70,12 +73,10 @@ def _generate_op_defs(bq_module: data_models.BQModule): _write_file(content, output_file, constants.OUTPUT_DIR.parent) -def _generate_accesor(): - pass - - def _generate_tests(bq_module: data_models.BQModule): content = template_renderer.render_tests(bq_module) + if not content: + return output_file = constants.TEST_OUTPUT_DIR.joinpath( bq_module.module_path.with_name(f"test_{bq_module.module_path.name}") @@ -84,11 +85,34 @@ def _generate_tests(bq_module: data_models.BQModule): _write_file(content, output_file, constants.TEST_OUTPUT_DIR.parent) +def _generate_accesor(bq_modules: list[data_models.BQModule]): + (core_content, pd_content, bf_content) = template_renderer.render_accessor( + bq_modules + ) + + core_output_file = ( + constants.CODE_ROOT / "extensions" / "core" / "series_accessor.py" + ) + _write_file(core_content, core_output_file, constants.CODE_ROOT) + + pd_output_file = ( + constants.CODE_ROOT / "extensions" / "pandas" / "series_accessor.py" + ) + _write_file(pd_content, pd_output_file, constants.CODE_ROOT) + + bf_output_file = ( + constants.CODE_ROOT / "extensions" / "bigframes" / "series_accessor.py" + ) + _write_file(bf_content, bf_output_file, constants.CODE_ROOT) + + def generate(bq_modules: list[data_models.BQModule]): for bq_module in bq_modules: _generate_op_defs(bq_module) _generate_tests(bq_module) + _generate_accesor(bq_modules) + # Ruff format _run_ruff() diff --git a/packages/bigframes/scripts/generate_bigframes_bigquery/template_renderer.py b/packages/bigframes/scripts/bigquery_generator/template_renderer.py similarity index 70% rename from packages/bigframes/scripts/generate_bigframes_bigquery/template_renderer.py rename to packages/bigframes/scripts/bigquery_generator/template_renderer.py index ed93fa1433d2..2e3ceb391b0a 100644 --- a/packages/bigframes/scripts/generate_bigframes_bigquery/template_renderer.py +++ b/packages/bigframes/scripts/bigquery_generator/template_renderer.py @@ -16,6 +16,7 @@ import constants import data_models +import dataclasses def _get_concrete_type_expr(yaml_type: str) -> str: @@ -83,10 +84,7 @@ def render_signature_def( _validate_types(bq_func.impls) - python_name = bq_func.op_base_name - if python_name in constants.PYTHON_BUILTINS: - python_name = python_name + "_" - sig_func_name = f"_{python_name.upper()}_SIG" + sig_func_name = f"_{bq_func.op_base_name.upper()}_SIG" max_args = max(len(impl.args) for impl in bq_func.impls) @@ -180,12 +178,17 @@ def _to_bigframes_func(bq_func: data_models.BQFunc) -> data_models.BigFramesFunc op_name=f"_{bq_func.op_base_name.upper()}_OP", description=bq_func.description, args=_get_bigframes_func_args(bq_func), + series_accessor_arg=bq_func.series_accessor_arg, ) def render_operation( bq_module: data_models.BQModule, ) -> str: + if not bq_module.functions: + # If there are no function definitions, do not generate an empty file. + return "" + ops = [] functions = [] @@ -202,6 +205,9 @@ def render_operation( def render_tests(bq_module: data_models.BQModule) -> str: + if not bq_module.functions: + # If there are no function definitions, do not generate an empty test file. + return "" import_path = "bigframes.operations.googlesql." + ".".join( bq_module.module_path.parts @@ -218,3 +224,77 @@ def render_tests(bq_module: data_models.BQModule) -> str: is_global=bq_module.is_global, functions=functions, ) + + +def _create_accessor_class_name(namespace: tuple[str, ...], prefix: str = "") -> str: + if not namespace: + return f"{prefix}BigQuerySeriesAccessor" + camel_parts = [part.capitalize() for part in namespace] + return f"{prefix}{''.join(camel_parts)}SeriesAccessor" + + +def render_accessor(bq_modules: data_models.BQModule) -> tuple[str, str, str]: + """ + Returns the content for core accessor, pandas accessor and BF accessor + """ + + namespaces = set() + for bq_module in bq_modules: + for i in range(len(bq_module.namespace) + 1): + namespaces.add(bq_module.namespace[:i]) + + sorted_namespaces = sorted(list(namespaces), key=len) + + accessors = [] + accessor_lookup_table = {} + for namespace in sorted_namespaces: + accessor = data_models.Accessor( + class_name=_create_accessor_class_name(namespace), + bigframes_class_name=_create_accessor_class_name( + namespace, prefix="Bigframes" + ), + pandas_class_name=_create_accessor_class_name(namespace, prefix="Pandas"), + is_root=len(namespace) == 0, + description=( + f"Series accessor for BigQuery {'.'.join(namespace)} functions." + if namespace + else "Series accessor for BigQuery functions." + ), + children=[], + functions=[], + ) + accessors.append(accessor) + accessor_lookup_table[namespace] = accessor + + # Establish parent-child relations + if len(namespace) > 0: + accessor.prop_name = namespace[-1] + parent_namespace = namespace[:-1] + accessor_lookup_table[parent_namespace].children.append(accessor) + + # Arrange functions by namespaces + for bq_module in bq_modules: + module_parts = bq_module.module_path.parts + for bq_func in bq_module.functions: + if bq_func.series_accessor_arg is None: + continue + bf_func = _to_bigframes_func(bq_func) + bf_func.import_module = ( + f"bigframes.operations.googlesql.{'.'.join(module_parts)}" + ) + accessor_lookup_table[bq_module.namespace].functions.append(bf_func) + + core_content = constants.TEMPLATES["core_series_accessor"].render( + script_path=constants.SCRIPT_PATH_RELATIVE, + namespaces=accessors, + ) + + pandas_content = constants.TEMPLATES["pandas_series_accessor"].render( + script_path=constants.SCRIPT_PATH_RELATIVE, namespaces=accessors + ) + + bigframes_content = constants.TEMPLATES["bigframes_series_accessor"].render( + script_path=constants.SCRIPT_PATH_RELATIVE, namespaces=accessors + ) + + return core_content, pandas_content, bigframes_content diff --git a/packages/bigframes/scripts/generate_bigframes_bigquery/yaml_parser.py b/packages/bigframes/scripts/bigquery_generator/yaml_parser.py similarity index 97% rename from packages/bigframes/scripts/generate_bigframes_bigquery/yaml_parser.py rename to packages/bigframes/scripts/bigquery_generator/yaml_parser.py index 4676bf388e34..f1bc6e5dfa7e 100644 --- a/packages/bigframes/scripts/generate_bigframes_bigquery/yaml_parser.py +++ b/packages/bigframes/scripts/bigquery_generator/yaml_parser.py @@ -61,6 +61,7 @@ def _parse_bq_func( op_base_name=op_base_name, description=func_data["description"], impls=[_parse_func_impl(impl) for impl in func_data["impls"]], + series_accessor_arg=func_data.get("series_accessor_arg", None), ) @@ -82,5 +83,4 @@ def parse_yaml(yaml_file: pathlib.Path) -> data_models.BQModule: return data_models.BQModule( yaml_file=yaml_file, functions=functions, - is_global=is_global, ) diff --git a/packages/bigframes/tests/unit/bigquery/generated/global_namespace/test_aead_encryption.py b/packages/bigframes/tests/unit/bigquery/generated/global_namespace/test_aead_encryption.py index 818151952ff0..b0251a6b8802 100644 --- a/packages/bigframes/tests/unit/bigquery/generated/global_namespace/test_aead_encryption.py +++ b/packages/bigframes/tests/unit/bigquery/generated/global_namespace/test_aead_encryption.py @@ -14,7 +14,7 @@ # # DO NOT MODIFY THIS FILE DIRECTLY. # This file was generated from: scripts/data/sql-functions/global_namespace/aead_encryption.yaml -# by the script: scripts/generate_bigframes_bigquery.py +# by the script: scripts import bigframes.bigquery as bbq import bigframes.core.col diff --git a/packages/bigframes/tests/unit/bigquery/generated/global_namespace/test_array.py b/packages/bigframes/tests/unit/bigquery/generated/global_namespace/test_array.py index 56b853869024..1e8cdbdf61ad 100644 --- a/packages/bigframes/tests/unit/bigquery/generated/global_namespace/test_array.py +++ b/packages/bigframes/tests/unit/bigquery/generated/global_namespace/test_array.py @@ -14,7 +14,7 @@ # # DO NOT MODIFY THIS FILE DIRECTLY. # This file was generated from: scripts/data/sql-functions/global_namespace/array.yaml -# by the script: scripts/generate_bigframes_bigquery.py +# by the script: scripts import bigframes.bigquery as bbq import bigframes.core.col diff --git a/packages/bigframes/tests/unit/bigquery/generated/global_namespace/test_bit.py b/packages/bigframes/tests/unit/bigquery/generated/global_namespace/test_bit.py index 2cccafc0643d..3311d0458168 100644 --- a/packages/bigframes/tests/unit/bigquery/generated/global_namespace/test_bit.py +++ b/packages/bigframes/tests/unit/bigquery/generated/global_namespace/test_bit.py @@ -14,7 +14,7 @@ # # DO NOT MODIFY THIS FILE DIRECTLY. # This file was generated from: scripts/data/sql-functions/global_namespace/bit.yaml -# by the script: scripts/generate_bigframes_bigquery.py +# by the script: scripts import bigframes.bigquery as bbq import bigframes.core.col diff --git a/packages/bigframes/tests/unit/bigquery/generated/global_namespace/test_conversion.py b/packages/bigframes/tests/unit/bigquery/generated/global_namespace/test_conversion.py index 84dfc02465cc..65aedde7ed55 100644 --- a/packages/bigframes/tests/unit/bigquery/generated/global_namespace/test_conversion.py +++ b/packages/bigframes/tests/unit/bigquery/generated/global_namespace/test_conversion.py @@ -14,7 +14,7 @@ # # DO NOT MODIFY THIS FILE DIRECTLY. # This file was generated from: scripts/data/sql-functions/global_namespace/conversion.yaml -# by the script: scripts/generate_bigframes_bigquery.py +# by the script: scripts import bigframes.bigquery as bbq import bigframes.core.col diff --git a/packages/bigframes/tests/unit/bigquery/generated/global_namespace/test_date.py b/packages/bigframes/tests/unit/bigquery/generated/global_namespace/test_date.py index 6484208584f9..6e32ccd8ea14 100644 --- a/packages/bigframes/tests/unit/bigquery/generated/global_namespace/test_date.py +++ b/packages/bigframes/tests/unit/bigquery/generated/global_namespace/test_date.py @@ -14,7 +14,7 @@ # # DO NOT MODIFY THIS FILE DIRECTLY. # This file was generated from: scripts/data/sql-functions/global_namespace/date.yaml -# by the script: scripts/generate_bigframes_bigquery.py +# by the script: scripts import bigframes.bigquery as bbq import bigframes.core.col diff --git a/packages/bigframes/tests/unit/bigquery/generated/test_aead.py b/packages/bigframes/tests/unit/bigquery/generated/test_aead.py index ce728b41899d..2ecbd8d24f1d 100644 --- a/packages/bigframes/tests/unit/bigquery/generated/test_aead.py +++ b/packages/bigframes/tests/unit/bigquery/generated/test_aead.py @@ -14,7 +14,7 @@ # # 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 +# by the script: scripts import bigframes.bigquery as bbq import bigframes.core.col From ab100cbd23498b3f0885e4777d397fbbf383e88a Mon Sep 17 00:00:00 2001 From: Shenyang Cai Date: Mon, 27 Jul 2026 02:02:54 +0000 Subject: [PATCH 06/27] make old script still runnable --- .../bigquery_generator/template_renderer.py | 2 +- .../scripts/generate_bigframes_bigquery.py | 686 +----------------- 2 files changed, 6 insertions(+), 682 deletions(-) diff --git a/packages/bigframes/scripts/bigquery_generator/template_renderer.py b/packages/bigframes/scripts/bigquery_generator/template_renderer.py index 2e3ceb391b0a..edd52f910766 100644 --- a/packages/bigframes/scripts/bigquery_generator/template_renderer.py +++ b/packages/bigframes/scripts/bigquery_generator/template_renderer.py @@ -243,7 +243,7 @@ def render_accessor(bq_modules: data_models.BQModule) -> tuple[str, str, str]: for i in range(len(bq_module.namespace) + 1): namespaces.add(bq_module.namespace[:i]) - sorted_namespaces = sorted(list(namespaces), key=len) + sorted_namespaces = sorted(list(namespaces), key=lambda ns: (len(ns), ns)) accessors = [] accessor_lookup_table = {} diff --git a/packages/bigframes/scripts/generate_bigframes_bigquery.py b/packages/bigframes/scripts/generate_bigframes_bigquery.py index 0381d168329f..6e19214c4ae3 100755 --- a/packages/bigframes/scripts/generate_bigframes_bigquery.py +++ b/packages/bigframes/scripts/generate_bigframes_bigquery.py @@ -23,689 +23,13 @@ # limitations under the License. import pathlib -import re -import subprocess +import sys -import jinja2 -import yaml - -SCRIPTS_DIRECTORY = pathlib.Path(__file__).parent.absolute() -PACKAGE_ROOT = SCRIPTS_DIRECTORY.parent -CODE_ROOT = PACKAGE_ROOT / "bigframes" -SCRIPT_PATH_RELATIVE = pathlib.Path(__file__).relative_to(PACKAGE_ROOT) - -# Directory containing the YAML files -DATA_DIR = SCRIPTS_DIRECTORY / "data" / "sql-functions" -# Directory where the generated Python files will be placed -OUTPUT_DIR = CODE_ROOT / "operations" / "googlesql" -# Directory where the generated test files will be placed -TEST_OUTPUT_DIR = PACKAGE_ROOT / "tests" / "unit" / "bigquery" / "generated" -# Directory containing the Jinja2 templates -TEMPLATE_DIR = SCRIPTS_DIRECTORY / "templates" - - -RUFF_COMMON_ARGS = [ - "--target-version=py310", - "--line-length=88", -] -RUFF_CHECK_ARGS = [ - "ruff", - "check", - "--select", - "I,F", - "--fix", -] + RUFF_COMMON_ARGS -RUFF_FORMAT_ARGS = [ - "ruff", - "format", -] + RUFF_COMMON_ARGS - - -DTYPE_MAP = { - "binary": "dtypes.BYTES_DTYPE", - "string": "dtypes.STRING_DTYPE", - "int64": "dtypes.INT_DTYPE", - "i64": "dtypes.INT_DTYPE", - "float64": "dtypes.FLOAT_DTYPE", - "fp64": "dtypes.FLOAT_DTYPE", - "bool": "dtypes.BOOL_DTYPE", - "boolean": "dtypes.BOOL_DTYPE", - "geography": "dtypes.GEO_DTYPE", - "json": "dtypes.JSON_DTYPE", - "date": "dtypes.DATE_DTYPE", - "time": "dtypes.TIME_DTYPE", - "datetime": "dtypes.DATETIME_DTYPE", - "timestamp": "dtypes.TIMESTAMP_DTYPE", - "decimal<38,9>": "dtypes.NUMERIC_DTYPE", - "decimal<76,38>": "dtypes.BIGNUMERIC_DTYPE", -} - -PY_TYPE_MAP = { - "binary": "bytes", - "string": "str", - "int64": "int", - "i64": "int", - "float64": "float", - "fp64": "float", - "bool": "bool", - "boolean": "bool", - "geography": "Any", - "json": "Any", - "date": "datetime.date", - "time": "datetime.time", - "datetime": "datetime.datetime", - "timestamp": "datetime.datetime", - "struct": "dict", - "decimal<38,9>": "decimal.Decimal", - "decimal<76,38>": "decimal.Decimal", - "interval_day": "datetime.timedelta", -} - -YAML_TYPE_TO_COL = { - "binary": "bytes_col", - "string": "string_col", - "int64": "int64_col", - "i64": "int64_col", - "float64": "float64_col", - "fp64": "float64_col", - "bool": "bool_col", - "boolean": "bool_col", - "geography": "geography_col", - "date": "date_col", - "time": "time_col", - "datetime": "datetime_col", - "timestamp": "timestamp_col", - "decimal<38,9>": "numeric_col", - "decimal<76,38>": "bignumeric_col", -} - -_PYTHON_BUILTINS = { - "abs", - "all", - "any", - "ascii", - "bin", - "bool", - "breakpoint", - "bytearray", - "bytes", - "callable", - "chr", - "classmethod", - "compile", - "complex", - "delattr", - "dict", - "dir", - "divmod", - "enumerate", - "eval", - "exec", - "filter", - "float", - "format", - "frozenset", - "getattr", - "globals", - "hasattr", - "hash", - "help", - "hex", - "id", - "input", - "int", - "isinstance", - "issubclass", - "iter", - "len", - "list", - "locals", - "map", - "max", - "memoryview", - "min", - "next", - "object", - "oct", - "open", - "ord", - "pow", - "print", - "property", - "range", - "repr", - "reversed", - "round", - "set", - "setattr", - "slice", - "sorted", - "staticmethod", - "str", - "sum", - "super", - "tuple", - "type", - "vars", - "zip", -} - - -def to_snake_case(name): - # Replace dots with underscores - name = name.replace(".", "_") - # Handle CamelCase to snake_case - name = re.sub(r"(?"): - inner = yaml_type[5:-1] - return inner in DTYPE_MAP - return False - - -def _get_concrete_type_expr(yaml_type): - if yaml_type in DTYPE_MAP: - return DTYPE_MAP[yaml_type] - if yaml_type.startswith("list<") and yaml_type.endswith(">"): - inner = yaml_type[5:-1] - return f"dtypes.list_type({DTYPE_MAP[inner]})" - raise ValueError(f"Not a concrete type: {yaml_type}") - - -def _validate_types(impls): - for impl in impls: - for arg in impl["args"]: - val = arg["value"] - if val == "any1": - continue - if val.startswith("list<") and val.endswith(">"): - inner = val[5:-1] - if inner != "any1" and inner not in DTYPE_MAP: - raise ValueError(f"Unsupported inner type: {inner}") - continue - if val == "struct": - continue - if val not in DTYPE_MAP: - raise ValueError(f"Unsupported type: {val}") - - ret = impl["return"] - if ret == "any1": - continue - if ret.startswith("list<") and ret.endswith(">"): - inner = ret[5:-1] - if inner != "any1" and inner not in DTYPE_MAP: - raise ValueError(f"Unsupported inner type: {inner}") - continue - if ret not in DTYPE_MAP: - raise ValueError(f"Unsupported type: {ret}") - - -def _generate_signature_def(python_name, impls, sql_name, template): - for impl in impls: - uses_any1 = False - if "any1" in str(impl["return"]): - uses_any1 = True - for arg in impl["args"]: - if "any1" in str(arg["value"]): - uses_any1 = True - impl["uses_any1"] = uses_any1 - - return_types = {impl["return"] for impl in impls} - - # Optimization: if all impls return the same concrete type, - # we can skip type checking and just return that type. - if len(return_types) == 1: - ret_type = next(iter(return_types)) - if _is_concrete_type(ret_type): - sig_expr = f"lambda *args: {_get_concrete_type_expr(ret_type)}" - return sig_expr, None - - _validate_types(impls) - - func_name = f"_{python_name.upper()}_SIG" - max_args = max(len(impl["args"]) for impl in impls) - - rendered = template.render( - func_name=func_name, - max_args=max_args, - impls=impls, - sql_name=sql_name, - dtype_map=DTYPE_MAP, - ) - - return func_name, rendered - - -def _get_func_args(args_by_name, arg_order): - func_args = [] - for name in arg_order: - arg_info = args_by_name[name] - types = [PY_TYPE_MAP.get(t, "Any") for t in sorted(arg_info["types"])] + [ - "Literal[sentinels.Sentinel.ARGUMENT_DEFAULT]" - ] - type_hint = ( - "Union[" + ", ".join(sorted(set(types))) + "]" - if len(types) > 1 - else types[0] - ) - default = "sentinels.Sentinel.ARGUMENT_DEFAULT" if arg_info["optional"] else "" - func_args.append( - { - "name": name, - "type_hint": type_hint, - "default": default, - } - ) - - # Clean up default values for mandatory args - # In Python, mandatory args come first. - for arg in func_args: - if not arg.get("default"): - arg.pop("default", None) - - return func_args - - -def _get_test_args(args_by_name, arg_order): - test_args = [] - for name in arg_order: - arg_info = args_by_name[name] - some_type = sorted(arg_info["types"])[0] - col_name = YAML_TYPE_TO_COL.get(some_type, "string_col") - test_args.append({"col_name": col_name}) - return test_args - - -def parse_scalar_functions(data, module_name, signature_def_template, is_global=False): - ops_list = [] - functions_list = [] - - if "scalar_functions" not in data: - return ops_list, functions_list - - for func_data in data["scalar_functions"]: - sql_name = func_data["name"] - python_name = to_snake_case(sql_name) - if not is_global and python_name.startswith(module_name + "_"): - python_name = python_name[len(module_name) + 1 :] - - op_base_name = python_name - if python_name in _PYTHON_BUILTINS: - python_name = python_name + "_" - - internal_op_name = f"_{op_base_name.upper()}_OP" - - # Aggregate args across impls - args_by_name, arg_order = _collect_args(func_data["impls"]) - - # Build ArgSpecs - arg_specs = _build_arg_specs(args_by_name, arg_order) - arg_specs_str = ", ".join(arg_specs) - if len(arg_specs) == 1: - arg_specs_str += "," - - # Determine return dtype - sig_name, sig_def = _generate_signature_def( - op_base_name, - func_data["impls"], - sql_name, - signature_def_template, - ) - - ops_list.append( - { - "internal_name": internal_op_name, - "sql_name": sql_name.upper(), - "arg_specs": arg_specs_str, - "signature": sig_name, - "signature_definition": sig_def, - } - ) - - # Function args - func_args = _get_func_args(args_by_name, arg_order) - - # Test args - test_args = _get_test_args(args_by_name, arg_order) - - # Read series_accessor_arg - series_accessor_arg = func_data.get("series_accessor_arg") - - functions_list.append( - { - "name": python_name, - "op_name": internal_op_name, - "description": func_data["description"], - "args": func_args, - "test_args": test_args, - "series_accessor_arg": series_accessor_arg, - } - ) - - return ops_list, functions_list - - -def run_ruff(path: pathlib.Path): - import sys - - subprocess.run( - [sys.executable, "-m", "ruff"] - + RUFF_CHECK_ARGS[1:] - + [ - str(path), - ], - check=True, - ) - - subprocess.run( - [sys.executable, "-m", "ruff"] - + RUFF_FORMAT_ARGS[1:] - + [ - str(path), - ], - check=True, - ) - - -def ensure_init_py(directory: pathlib.Path, limit_dir: pathlib.Path, license_template): - """Ensures __init__.py exists in the directory and its parents up to limit_dir.""" - curr = directory - while curr != limit_dir and curr != curr.parent: - init_file = curr / "__init__.py" - if not init_file.exists(): - print(f" Creating {init_file}") - content = license_template.render() - with open(init_file, "w") as f: - f.write(content) - run_ruff(init_file) - curr = curr.parent - - -def process_yaml_file(yaml_file, templates): - print(f"Processing {yaml_file}...") - with open(yaml_file, "r") as f: - data = yaml.safe_load(f) - - rel_path = yaml_file.relative_to(DATA_DIR) - module_path = rel_path.with_suffix("") - module_name = module_path.name - output_file = OUTPUT_DIR.joinpath(module_path).with_suffix(".py") - - is_global = "global_namespace" in module_path.parts - namespace = get_namespace(yaml_file) - - if not data or not isinstance(data, dict) or "scalar_functions" not in data: - # If the file is empty or has no functions, just create the namespace. - return [{"namespace": namespace}] - - ops_list, functions_list = parse_scalar_functions( - data, - module_name, - templates["signature_def"], - is_global=is_global, - ) - - # Render and write - output_file.parent.mkdir(parents=True, exist_ok=True) - ensure_init_py(output_file.parent, OUTPUT_DIR.parent, templates["license"]) - yaml_file_relative = yaml_file.relative_to(PACKAGE_ROOT) - content = templates["operation"].render( - yaml_path=yaml_file_relative, - script_path=SCRIPT_PATH_RELATIVE, - ops=ops_list, - functions=functions_list, - ) - with open(output_file, "w") as f: - f.write(content) - - run_ruff(output_file) - print(f" Generated {output_file}") - - # Render and write test - import_path = "bigframes.operations.googlesql." + ".".join(module_path.parts) - test_output_file = TEST_OUTPUT_DIR.joinpath( - module_path.with_name(f"test_{module_path.name}") - ).with_suffix(".py") - - test_output_file.parent.mkdir(parents=True, exist_ok=True) - ensure_init_py( - test_output_file.parent, TEST_OUTPUT_DIR.parent, templates["license"] - ) - test_content = templates["test_operation"].render( - yaml_path=yaml_file_relative, - script_path=SCRIPT_PATH_RELATIVE, - import_path=import_path, - short_name=module_path.name, - is_global=is_global, - functions=functions_list, - ) - with open(test_output_file, "w") as f: - f.write(test_content) - - run_ruff(test_output_file) - print(f" Generated {test_output_file}") - - # Collect functions for Series accessor - accessor_functions = [] - for func in functions_list: - if func.get("series_accessor_arg"): - import_module = ( - f"bigframes.operations.googlesql.{'.'.join(module_path.parts)}" - ) - accessor_functions.append( - { - "name": func["name"], - "import_module": import_module, - "namespace": namespace, - "description": func["description"], - "args": func["args"], - "series_accessor_arg": func["series_accessor_arg"], - } - ) - - return accessor_functions - - -def get_namespace(yaml_file: pathlib.Path) -> tuple[str, ...] | None: - rel_path = yaml_file.relative_to(DATA_DIR) - parts = rel_path.with_suffix("").parts - if "global_namespace" in parts: - return None - return parts - - -def get_class_name(ns_tuple: tuple[str, ...], prefix: str = "") -> str: - if not ns_tuple: - return f"{prefix}BigQuerySeriesAccessor" - camel_parts = [part.capitalize() for part in ns_tuple] - return f"{prefix}{''.join(camel_parts)}SeriesAccessor" - - -def generate_series_accessors(functions: list[dict], templates: dict): - print("Generating Series accessors...") - # Find all active namespaces - active_namespaces = set() - for func in functions: - ns = func["namespace"] or () - for i in range(len(ns) + 1): - active_namespaces.add(ns[:i]) - - # Sort namespaces by depth so parents come first - sorted_namespaces = sorted(list(active_namespaces), key=len) - - # Build namespace definitions - ns_defs = [] - ns_by_tuple = {} - for ns in sorted_namespaces: - class_name = get_class_name(ns) - bf_class_name = get_class_name(ns, prefix="Bigframes") - pd_class_name = get_class_name(ns, prefix="Pandas") - - ns_def = { - "ns_tuple": ns, - "class_name": class_name, - "bigframes_class_name": bf_class_name, - "pandas_class_name": pd_class_name, - "is_root": len(ns) == 0, - "description": ( - f"Series accessor for BigQuery {'.'.join(ns)} functions." - if ns - else "Series accessor for BigQuery functions." - ), - "children": [], - "functions": [], - } - ns_defs.append(ns_def) - ns_by_tuple[ns] = ns_def - - # Populate functions - for func in functions: - if "name" in func: - ns = func["namespace"] or () - ns_by_tuple[ns]["functions"].append(func) - - # Populate children properties - for ns in sorted_namespaces: - if len(ns) > 0: - parent_ns = ns[:-1] - parent_def = ns_by_tuple[parent_ns] - child_def = ns_by_tuple[ns] - parent_def["children"].append( - { - "prop_name": ns[-1], - "class_name": child_def["class_name"], - "bigframes_class_name": child_def["bigframes_class_name"], - "pandas_class_name": child_def["pandas_class_name"], - } - ) - - # Render and write core - core_output_file = CODE_ROOT / "extensions" / "core" / "series_accessor.py" - core_output_file.parent.mkdir(parents=True, exist_ok=True) - ensure_init_py(core_output_file.parent, CODE_ROOT, templates["license"]) - core_content = templates["core_series_accessor"].render( - script_path=SCRIPT_PATH_RELATIVE, - namespaces=ns_defs, - ) - with open(core_output_file, "w") as f: - f.write(core_content) - run_ruff(core_output_file) - print(f" Generated {core_output_file}") - - # Render and write bigframes - bf_output_file = CODE_ROOT / "extensions" / "bigframes" / "series_accessor.py" - bf_output_file.parent.mkdir(parents=True, exist_ok=True) - ensure_init_py(bf_output_file.parent, CODE_ROOT, templates["license"]) - bf_content = templates["bigframes_series_accessor"].render( - script_path=SCRIPT_PATH_RELATIVE, - namespaces=ns_defs, - ) - with open(bf_output_file, "w") as f: - f.write(bf_content) - run_ruff(bf_output_file) - print(f" Generated {bf_output_file}") - - # Render and write pandas - pd_output_file = CODE_ROOT / "extensions" / "pandas" / "series_accessor.py" - pd_output_file.parent.mkdir(parents=True, exist_ok=True) - ensure_init_py(pd_output_file.parent, CODE_ROOT, templates["license"]) - pd_content = templates["pandas_series_accessor"].render( - script_path=SCRIPT_PATH_RELATIVE, - namespaces=ns_defs, - ) - with open(pd_output_file, "w") as f: - f.write(pd_content) - run_ruff(pd_output_file) - print(f" Generated {pd_output_file}") - - -def main(): - templates = load_templates() - - all_accessor_functions = [] - for yaml_file in sorted(DATA_DIR.glob("**/*.yaml")): - accessor_funcs = process_yaml_file(yaml_file, templates) - all_accessor_functions.extend(accessor_funcs) - - if all_accessor_functions: - generate_series_accessors(all_accessor_functions, templates) +pkg_dir = pathlib.Path(__file__).parent / "bigquery_generator" +if str(pkg_dir) not in sys.path: + sys.path.insert(0, str(pkg_dir)) +from bigquery_generator.__main__ import main if __name__ == "__main__": main() From f85141a6ab3373290831ad69ac8c54899776f04d Mon Sep 17 00:00:00 2001 From: Shenyang Cai Date: Mon, 27 Jul 2026 02:09:51 +0000 Subject: [PATCH 07/27] fix script reference --- .../bigframes/bigframes/extensions/bigframes/series_accessor.py | 2 +- packages/bigframes/bigframes/extensions/core/series_accessor.py | 2 +- .../bigframes/bigframes/extensions/pandas/series_accessor.py | 2 +- packages/bigframes/bigframes/operations/googlesql/aead.py | 2 +- .../operations/googlesql/global_namespace/aead_encryption.py | 2 +- .../bigframes/operations/googlesql/global_namespace/array.py | 2 +- .../bigframes/operations/googlesql/global_namespace/bit.py | 2 +- .../operations/googlesql/global_namespace/conversion.py | 2 +- .../bigframes/operations/googlesql/global_namespace/date.py | 2 +- packages/bigframes/scripts/bigquery_generator/constants.py | 2 +- .../bigquery/generated/global_namespace/test_aead_encryption.py | 2 +- .../unit/bigquery/generated/global_namespace/test_array.py | 2 +- .../tests/unit/bigquery/generated/global_namespace/test_bit.py | 2 +- .../unit/bigquery/generated/global_namespace/test_conversion.py | 2 +- .../tests/unit/bigquery/generated/global_namespace/test_date.py | 2 +- packages/bigframes/tests/unit/bigquery/generated/test_aead.py | 2 +- 16 files changed, 16 insertions(+), 16 deletions(-) diff --git a/packages/bigframes/bigframes/extensions/bigframes/series_accessor.py b/packages/bigframes/bigframes/extensions/bigframes/series_accessor.py index 4007df5623a0..ebc2df2eb7ca 100644 --- a/packages/bigframes/bigframes/extensions/bigframes/series_accessor.py +++ b/packages/bigframes/bigframes/extensions/bigframes/series_accessor.py @@ -13,7 +13,7 @@ # limitations under the License. # # DO NOT MODIFY THIS FILE DIRECTLY. -# This file was generated by the script: scripts +# This file was generated by the script: scripts/bigquery_generator # from __future__ import annotations diff --git a/packages/bigframes/bigframes/extensions/core/series_accessor.py b/packages/bigframes/bigframes/extensions/core/series_accessor.py index 8ff12fbdfc28..2356d22f6c88 100644 --- a/packages/bigframes/bigframes/extensions/core/series_accessor.py +++ b/packages/bigframes/bigframes/extensions/core/series_accessor.py @@ -13,7 +13,7 @@ # limitations under the License. # # DO NOT MODIFY THIS FILE DIRECTLY. -# This file was generated by the script: scripts +# This file was generated by the script: scripts/bigquery_generator # from __future__ import annotations diff --git a/packages/bigframes/bigframes/extensions/pandas/series_accessor.py b/packages/bigframes/bigframes/extensions/pandas/series_accessor.py index 80ebdb661ced..21fd13e2f538 100644 --- a/packages/bigframes/bigframes/extensions/pandas/series_accessor.py +++ b/packages/bigframes/bigframes/extensions/pandas/series_accessor.py @@ -13,7 +13,7 @@ # limitations under the License. # # DO NOT MODIFY THIS FILE DIRECTLY. -# This file was generated by the script: scripts +# This file was generated by the script: scripts/bigquery_generator # from __future__ import annotations diff --git a/packages/bigframes/bigframes/operations/googlesql/aead.py b/packages/bigframes/bigframes/operations/googlesql/aead.py index c1b9ecb3618c..25922f9bc031 100644 --- a/packages/bigframes/bigframes/operations/googlesql/aead.py +++ b/packages/bigframes/bigframes/operations/googlesql/aead.py @@ -14,7 +14,7 @@ # # DO NOT MODIFY THIS FILE DIRECTLY. # This file was generated from: scripts/data/sql-functions/aead.yaml -# by the script: scripts +# by the script: scripts/bigquery_generator from __future__ import annotations diff --git a/packages/bigframes/bigframes/operations/googlesql/global_namespace/aead_encryption.py b/packages/bigframes/bigframes/operations/googlesql/global_namespace/aead_encryption.py index d7e0f58a6435..33e0913da70b 100644 --- a/packages/bigframes/bigframes/operations/googlesql/global_namespace/aead_encryption.py +++ b/packages/bigframes/bigframes/operations/googlesql/global_namespace/aead_encryption.py @@ -14,7 +14,7 @@ # # DO NOT MODIFY THIS FILE DIRECTLY. # This file was generated from: scripts/data/sql-functions/global_namespace/aead_encryption.yaml -# by the script: scripts +# by the script: scripts/bigquery_generator from __future__ import annotations diff --git a/packages/bigframes/bigframes/operations/googlesql/global_namespace/array.py b/packages/bigframes/bigframes/operations/googlesql/global_namespace/array.py index bd3a54021b45..ee3319a9bfdd 100644 --- a/packages/bigframes/bigframes/operations/googlesql/global_namespace/array.py +++ b/packages/bigframes/bigframes/operations/googlesql/global_namespace/array.py @@ -14,7 +14,7 @@ # # DO NOT MODIFY THIS FILE DIRECTLY. # This file was generated from: scripts/data/sql-functions/global_namespace/array.yaml -# by the script: scripts +# by the script: scripts/bigquery_generator from __future__ import annotations diff --git a/packages/bigframes/bigframes/operations/googlesql/global_namespace/bit.py b/packages/bigframes/bigframes/operations/googlesql/global_namespace/bit.py index 818d8b178753..a480ed78b411 100644 --- a/packages/bigframes/bigframes/operations/googlesql/global_namespace/bit.py +++ b/packages/bigframes/bigframes/operations/googlesql/global_namespace/bit.py @@ -14,7 +14,7 @@ # # DO NOT MODIFY THIS FILE DIRECTLY. # This file was generated from: scripts/data/sql-functions/global_namespace/bit.yaml -# by the script: scripts +# by the script: scripts/bigquery_generator from __future__ import annotations diff --git a/packages/bigframes/bigframes/operations/googlesql/global_namespace/conversion.py b/packages/bigframes/bigframes/operations/googlesql/global_namespace/conversion.py index 2543c0a9651d..5d2ae897a423 100644 --- a/packages/bigframes/bigframes/operations/googlesql/global_namespace/conversion.py +++ b/packages/bigframes/bigframes/operations/googlesql/global_namespace/conversion.py @@ -14,7 +14,7 @@ # # DO NOT MODIFY THIS FILE DIRECTLY. # This file was generated from: scripts/data/sql-functions/global_namespace/conversion.yaml -# by the script: scripts +# by the script: scripts/bigquery_generator from __future__ import annotations diff --git a/packages/bigframes/bigframes/operations/googlesql/global_namespace/date.py b/packages/bigframes/bigframes/operations/googlesql/global_namespace/date.py index 4bc675ad29f1..b9a5e920efcc 100644 --- a/packages/bigframes/bigframes/operations/googlesql/global_namespace/date.py +++ b/packages/bigframes/bigframes/operations/googlesql/global_namespace/date.py @@ -14,7 +14,7 @@ # # DO NOT MODIFY THIS FILE DIRECTLY. # This file was generated from: scripts/data/sql-functions/global_namespace/date.yaml -# by the script: scripts +# by the script: scripts/bigquery_generator from __future__ import annotations diff --git a/packages/bigframes/scripts/bigquery_generator/constants.py b/packages/bigframes/scripts/bigquery_generator/constants.py index 600944f55923..62e2ed1a6183 100644 --- a/packages/bigframes/scripts/bigquery_generator/constants.py +++ b/packages/bigframes/scripts/bigquery_generator/constants.py @@ -19,7 +19,7 @@ SCRIPTS_DIRECTORY = pathlib.Path(__file__).parent.parent.absolute() PACKAGE_ROOT = SCRIPTS_DIRECTORY.parent CODE_ROOT = PACKAGE_ROOT / "bigframes" -SCRIPT_PATH_RELATIVE = pathlib.Path(__file__).relative_to(PACKAGE_ROOT).parent.parent +SCRIPT_PATH_RELATIVE = pathlib.Path(__file__).relative_to(PACKAGE_ROOT).parent # Directory containing the YAML files diff --git a/packages/bigframes/tests/unit/bigquery/generated/global_namespace/test_aead_encryption.py b/packages/bigframes/tests/unit/bigquery/generated/global_namespace/test_aead_encryption.py index b0251a6b8802..85f66acbf498 100644 --- a/packages/bigframes/tests/unit/bigquery/generated/global_namespace/test_aead_encryption.py +++ b/packages/bigframes/tests/unit/bigquery/generated/global_namespace/test_aead_encryption.py @@ -14,7 +14,7 @@ # # DO NOT MODIFY THIS FILE DIRECTLY. # This file was generated from: scripts/data/sql-functions/global_namespace/aead_encryption.yaml -# by the script: scripts +# by the script: scripts/bigquery_generator import bigframes.bigquery as bbq import bigframes.core.col diff --git a/packages/bigframes/tests/unit/bigquery/generated/global_namespace/test_array.py b/packages/bigframes/tests/unit/bigquery/generated/global_namespace/test_array.py index 1e8cdbdf61ad..efc80f4feafb 100644 --- a/packages/bigframes/tests/unit/bigquery/generated/global_namespace/test_array.py +++ b/packages/bigframes/tests/unit/bigquery/generated/global_namespace/test_array.py @@ -14,7 +14,7 @@ # # DO NOT MODIFY THIS FILE DIRECTLY. # This file was generated from: scripts/data/sql-functions/global_namespace/array.yaml -# by the script: scripts +# by the script: scripts/bigquery_generator import bigframes.bigquery as bbq import bigframes.core.col diff --git a/packages/bigframes/tests/unit/bigquery/generated/global_namespace/test_bit.py b/packages/bigframes/tests/unit/bigquery/generated/global_namespace/test_bit.py index 3311d0458168..c0c7e1ec2ab3 100644 --- a/packages/bigframes/tests/unit/bigquery/generated/global_namespace/test_bit.py +++ b/packages/bigframes/tests/unit/bigquery/generated/global_namespace/test_bit.py @@ -14,7 +14,7 @@ # # DO NOT MODIFY THIS FILE DIRECTLY. # This file was generated from: scripts/data/sql-functions/global_namespace/bit.yaml -# by the script: scripts +# by the script: scripts/bigquery_generator import bigframes.bigquery as bbq import bigframes.core.col diff --git a/packages/bigframes/tests/unit/bigquery/generated/global_namespace/test_conversion.py b/packages/bigframes/tests/unit/bigquery/generated/global_namespace/test_conversion.py index 65aedde7ed55..b5a8a819115a 100644 --- a/packages/bigframes/tests/unit/bigquery/generated/global_namespace/test_conversion.py +++ b/packages/bigframes/tests/unit/bigquery/generated/global_namespace/test_conversion.py @@ -14,7 +14,7 @@ # # DO NOT MODIFY THIS FILE DIRECTLY. # This file was generated from: scripts/data/sql-functions/global_namespace/conversion.yaml -# by the script: scripts +# by the script: scripts/bigquery_generator import bigframes.bigquery as bbq import bigframes.core.col diff --git a/packages/bigframes/tests/unit/bigquery/generated/global_namespace/test_date.py b/packages/bigframes/tests/unit/bigquery/generated/global_namespace/test_date.py index 6e32ccd8ea14..b0228804c6a9 100644 --- a/packages/bigframes/tests/unit/bigquery/generated/global_namespace/test_date.py +++ b/packages/bigframes/tests/unit/bigquery/generated/global_namespace/test_date.py @@ -14,7 +14,7 @@ # # DO NOT MODIFY THIS FILE DIRECTLY. # This file was generated from: scripts/data/sql-functions/global_namespace/date.yaml -# by the script: scripts +# by the script: scripts/bigquery_generator import bigframes.bigquery as bbq import bigframes.core.col diff --git a/packages/bigframes/tests/unit/bigquery/generated/test_aead.py b/packages/bigframes/tests/unit/bigquery/generated/test_aead.py index 2ecbd8d24f1d..407cd22af6f6 100644 --- a/packages/bigframes/tests/unit/bigquery/generated/test_aead.py +++ b/packages/bigframes/tests/unit/bigquery/generated/test_aead.py @@ -14,7 +14,7 @@ # # DO NOT MODIFY THIS FILE DIRECTLY. # This file was generated from: scripts/data/sql-functions/aead.yaml -# by the script: scripts +# by the script: scripts/bigquery_generator import bigframes.bigquery as bbq import bigframes.core.col From 0e6d04a3cb29ab0c83e9d3d69df496451919a4a0 Mon Sep 17 00:00:00 2001 From: Shenyang Cai Date: Mon, 27 Jul 2026 02:11:00 +0000 Subject: [PATCH 08/27] fix lint --- .../bigframes/scripts/bigquery_generator/template_renderer.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/packages/bigframes/scripts/bigquery_generator/template_renderer.py b/packages/bigframes/scripts/bigquery_generator/template_renderer.py index edd52f910766..a0401f6c6fd0 100644 --- a/packages/bigframes/scripts/bigquery_generator/template_renderer.py +++ b/packages/bigframes/scripts/bigquery_generator/template_renderer.py @@ -14,9 +14,9 @@ # limitations under the License. + import constants import data_models -import dataclasses def _get_concrete_type_expr(yaml_type: str) -> str: From 9353f3543a813829de99b72ceaf8494c10b060b7 Mon Sep 17 00:00:00 2001 From: Shenyang Cai Date: Sun, 26 Jul 2026 19:12:43 -0700 Subject: [PATCH 09/27] Update packages/bigframes/scripts/bigquery_generator/__main__.py Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> --- packages/bigframes/scripts/bigquery_generator/__main__.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/packages/bigframes/scripts/bigquery_generator/__main__.py b/packages/bigframes/scripts/bigquery_generator/__main__.py index 25a8af08f03c..57194358dda2 100644 --- a/packages/bigframes/scripts/bigquery_generator/__main__.py +++ b/packages/bigframes/scripts/bigquery_generator/__main__.py @@ -13,9 +13,9 @@ # limitations under the License. -import constants -import file_generator -import yaml_parser +from . import constants +from . import file_generator +from . import yaml_parser def main(): From 78041aacb9c55d00dfdc48dd1021c2a2234a6ccd Mon Sep 17 00:00:00 2001 From: Shenyang Cai Date: Sun, 26 Jul 2026 19:12:53 -0700 Subject: [PATCH 10/27] Update packages/bigframes/scripts/bigquery_generator/file_generator.py Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> --- .../bigframes/scripts/bigquery_generator/file_generator.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/packages/bigframes/scripts/bigquery_generator/file_generator.py b/packages/bigframes/scripts/bigquery_generator/file_generator.py index 44f42950084e..a607a4c29a7a 100644 --- a/packages/bigframes/scripts/bigquery_generator/file_generator.py +++ b/packages/bigframes/scripts/bigquery_generator/file_generator.py @@ -16,9 +16,9 @@ import subprocess import sys -import constants -import data_models -import template_renderer +from . import constants +from . import data_models +from . import template_renderer def _ensure_init_py(directory: pathlib.Path, limit_dir: pathlib.Path): From 4fd77b8e85b76e8add0756d096f56579ddd452dc Mon Sep 17 00:00:00 2001 From: Shenyang Cai Date: Sun, 26 Jul 2026 19:13:01 -0700 Subject: [PATCH 11/27] Update packages/bigframes/scripts/bigquery_generator/file_generator.py Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> --- packages/bigframes/scripts/bigquery_generator/file_generator.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/packages/bigframes/scripts/bigquery_generator/file_generator.py b/packages/bigframes/scripts/bigquery_generator/file_generator.py index a607a4c29a7a..839e90986d7b 100644 --- a/packages/bigframes/scripts/bigquery_generator/file_generator.py +++ b/packages/bigframes/scripts/bigquery_generator/file_generator.py @@ -29,7 +29,7 @@ def _ensure_init_py(directory: pathlib.Path, limit_dir: pathlib.Path): if not init_file.exists(): print(f" Creating {init_file}") content = constants.TEMPLATES["license"].render() - with open(init_file, "w") as f: + with open(init_file, "w", encoding="utf-8") as f: f.write(content) curr = curr.parent From aa24eb91b4fa8d9328c61555d6a449dad2471243 Mon Sep 17 00:00:00 2001 From: Shenyang Cai Date: Sun, 26 Jul 2026 19:13:09 -0700 Subject: [PATCH 12/27] Update packages/bigframes/scripts/bigquery_generator/file_generator.py Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> --- packages/bigframes/scripts/bigquery_generator/file_generator.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/packages/bigframes/scripts/bigquery_generator/file_generator.py b/packages/bigframes/scripts/bigquery_generator/file_generator.py index 839e90986d7b..504bcd40d6f7 100644 --- a/packages/bigframes/scripts/bigquery_generator/file_generator.py +++ b/packages/bigframes/scripts/bigquery_generator/file_generator.py @@ -38,7 +38,7 @@ def _write_file(content: str, output_file: pathlib.Path, limit_dir: pathlib.Path output_file.parent.mkdir(parents=True, exist_ok=True) _ensure_init_py(output_file.parent, limit_dir) - with open(output_file, "w") as f: + with open(output_file, "w", encoding="utf-8") as f: f.write(content) print(f" Generated {output_file}") From 393010bc6354bb411d088bdc66ae5a49bce9b1cd Mon Sep 17 00:00:00 2001 From: Shenyang Cai Date: Sun, 26 Jul 2026 19:13:29 -0700 Subject: [PATCH 13/27] Update packages/bigframes/scripts/bigquery_generator/template_renderer.py Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> --- .../bigframes/scripts/bigquery_generator/template_renderer.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/packages/bigframes/scripts/bigquery_generator/template_renderer.py b/packages/bigframes/scripts/bigquery_generator/template_renderer.py index a0401f6c6fd0..d0ac8db2f887 100644 --- a/packages/bigframes/scripts/bigquery_generator/template_renderer.py +++ b/packages/bigframes/scripts/bigquery_generator/template_renderer.py @@ -233,7 +233,7 @@ def _create_accessor_class_name(namespace: tuple[str, ...], prefix: str = "") -> return f"{prefix}{''.join(camel_parts)}SeriesAccessor" -def render_accessor(bq_modules: data_models.BQModule) -> tuple[str, str, str]: +def render_accessor(bq_modules: list[data_models.BQModule]) -> tuple[str, str, str]: """ Returns the content for core accessor, pandas accessor and BF accessor """ From f1fea760c003e0cdf14892ae248931588e08e71e Mon Sep 17 00:00:00 2001 From: Shenyang Cai Date: Sun, 26 Jul 2026 19:13:38 -0700 Subject: [PATCH 14/27] Update packages/bigframes/scripts/bigquery_generator/yaml_parser.py Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> --- packages/bigframes/scripts/bigquery_generator/yaml_parser.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/packages/bigframes/scripts/bigquery_generator/yaml_parser.py b/packages/bigframes/scripts/bigquery_generator/yaml_parser.py index f1bc6e5dfa7e..2489406641c0 100644 --- a/packages/bigframes/scripts/bigquery_generator/yaml_parser.py +++ b/packages/bigframes/scripts/bigquery_generator/yaml_parser.py @@ -17,8 +17,8 @@ import re from typing import Any -import constants -import data_models +from . import constants +from . import data_models import yaml From 70e9091ac371aa3a80319ba35f6b89f9c067787e Mon Sep 17 00:00:00 2001 From: Shenyang Cai Date: Sun, 26 Jul 2026 19:13:47 -0700 Subject: [PATCH 15/27] Update packages/bigframes/scripts/bigquery_generator/yaml_parser.py Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> --- packages/bigframes/scripts/bigquery_generator/yaml_parser.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/packages/bigframes/scripts/bigquery_generator/yaml_parser.py b/packages/bigframes/scripts/bigquery_generator/yaml_parser.py index 2489406641c0..e2c0e846de13 100644 --- a/packages/bigframes/scripts/bigquery_generator/yaml_parser.py +++ b/packages/bigframes/scripts/bigquery_generator/yaml_parser.py @@ -22,7 +22,7 @@ import yaml -def _to_snake_case(name): +def _to_snake_case(name: str) -> str: # Replace dots with underscores name = name.replace(".", "_") # Handle CamelCase to snake_case From e1b8a37e5db5a193c139b1d8a675718200046cc0 Mon Sep 17 00:00:00 2001 From: Shenyang Cai Date: Sun, 26 Jul 2026 19:13:54 -0700 Subject: [PATCH 16/27] Update packages/bigframes/scripts/bigquery_generator/yaml_parser.py Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> --- packages/bigframes/scripts/bigquery_generator/yaml_parser.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/packages/bigframes/scripts/bigquery_generator/yaml_parser.py b/packages/bigframes/scripts/bigquery_generator/yaml_parser.py index e2c0e846de13..263b60648894 100644 --- a/packages/bigframes/scripts/bigquery_generator/yaml_parser.py +++ b/packages/bigframes/scripts/bigquery_generator/yaml_parser.py @@ -68,7 +68,7 @@ def _parse_bq_func( def parse_yaml(yaml_file: pathlib.Path) -> data_models.BQModule: print(f"Parsing {yaml_file}...") - with open(yaml_file, "r") as f: + with open(yaml_file, "r", encoding="utf-8") as f: data = yaml.safe_load(f) functions = [] From aa1b51b3a81ed6c3805621f400b815b04596ecbf Mon Sep 17 00:00:00 2001 From: Shenyang Cai Date: Sun, 26 Jul 2026 19:14:18 -0700 Subject: [PATCH 17/27] Update packages/bigframes/scripts/generate_bigframes_bigquery.py Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> --- packages/bigframes/scripts/generate_bigframes_bigquery.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/packages/bigframes/scripts/generate_bigframes_bigquery.py b/packages/bigframes/scripts/generate_bigframes_bigquery.py index 6e19214c4ae3..5882155eb6b5 100755 --- a/packages/bigframes/scripts/generate_bigframes_bigquery.py +++ b/packages/bigframes/scripts/generate_bigframes_bigquery.py @@ -25,9 +25,9 @@ import pathlib import sys -pkg_dir = pathlib.Path(__file__).parent / "bigquery_generator" -if str(pkg_dir) not in sys.path: - sys.path.insert(0, str(pkg_dir)) +scripts_dir = pathlib.Path(__file__).parent +if str(scripts_dir) not in sys.path: + sys.path.insert(0, str(scripts_dir)) from bigquery_generator.__main__ import main From aa74300a866650ca220758b87b021d8bdb4c5d39 Mon Sep 17 00:00:00 2001 From: Shenyang Cai Date: Sun, 26 Jul 2026 19:14:55 -0700 Subject: [PATCH 18/27] Update packages/bigframes/scripts/bigquery_generator/template_renderer.py Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> --- .../bigframes/scripts/bigquery_generator/template_renderer.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/packages/bigframes/scripts/bigquery_generator/template_renderer.py b/packages/bigframes/scripts/bigquery_generator/template_renderer.py index d0ac8db2f887..26d4d7e5ad1b 100644 --- a/packages/bigframes/scripts/bigquery_generator/template_renderer.py +++ b/packages/bigframes/scripts/bigquery_generator/template_renderer.py @@ -15,8 +15,8 @@ -import constants -import data_models +from . import constants +from . import data_models def _get_concrete_type_expr(yaml_type: str) -> str: From dcbd2eae1d275e8daad2bf48d983c62bab2b23b5 Mon Sep 17 00:00:00 2001 From: Shenyang Cai Date: Sun, 26 Jul 2026 19:15:06 -0700 Subject: [PATCH 19/27] Update packages/bigframes/scripts/bigquery_generator/template_renderer.py Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> --- .../bigframes/scripts/bigquery_generator/template_renderer.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/packages/bigframes/scripts/bigquery_generator/template_renderer.py b/packages/bigframes/scripts/bigquery_generator/template_renderer.py index 26d4d7e5ad1b..ab98d3979218 100644 --- a/packages/bigframes/scripts/bigquery_generator/template_renderer.py +++ b/packages/bigframes/scripts/bigquery_generator/template_renderer.py @@ -37,7 +37,7 @@ def _is_concrete_type(yaml_type: str) -> bool: return False -def _validate_types(impls): +def _validate_types(impls: list[data_models.BQFuncImpl]) -> None: for impl in impls: for arg in impl.args: val = arg.value From a048b0c7d424378cbe6e1e7bf5b58a82c997c3d2 Mon Sep 17 00:00:00 2001 From: Shenyang Cai Date: Mon, 27 Jul 2026 02:17:21 +0000 Subject: [PATCH 20/27] fix lint --- packages/bigframes/scripts/bigquery_generator/__main__.py | 4 +--- .../bigframes/scripts/bigquery_generator/file_generator.py | 4 +--- .../bigframes/scripts/bigquery_generator/template_renderer.py | 3 +-- packages/bigframes/scripts/bigquery_generator/yaml_parser.py | 4 ++-- 4 files changed, 5 insertions(+), 10 deletions(-) diff --git a/packages/bigframes/scripts/bigquery_generator/__main__.py b/packages/bigframes/scripts/bigquery_generator/__main__.py index 57194358dda2..2f7ffaacf2d0 100644 --- a/packages/bigframes/scripts/bigquery_generator/__main__.py +++ b/packages/bigframes/scripts/bigquery_generator/__main__.py @@ -13,9 +13,7 @@ # limitations under the License. -from . import constants -from . import file_generator -from . import yaml_parser +from . import constants, file_generator, yaml_parser def main(): diff --git a/packages/bigframes/scripts/bigquery_generator/file_generator.py b/packages/bigframes/scripts/bigquery_generator/file_generator.py index 504bcd40d6f7..9b7853d4d4da 100644 --- a/packages/bigframes/scripts/bigquery_generator/file_generator.py +++ b/packages/bigframes/scripts/bigquery_generator/file_generator.py @@ -16,9 +16,7 @@ import subprocess import sys -from . import constants -from . import data_models -from . import template_renderer +from . import constants, data_models, template_renderer def _ensure_init_py(directory: pathlib.Path, limit_dir: pathlib.Path): diff --git a/packages/bigframes/scripts/bigquery_generator/template_renderer.py b/packages/bigframes/scripts/bigquery_generator/template_renderer.py index ab98d3979218..96dfe64d344a 100644 --- a/packages/bigframes/scripts/bigquery_generator/template_renderer.py +++ b/packages/bigframes/scripts/bigquery_generator/template_renderer.py @@ -15,8 +15,7 @@ -from . import constants -from . import data_models +from . import constants, data_models def _get_concrete_type_expr(yaml_type: str) -> str: diff --git a/packages/bigframes/scripts/bigquery_generator/yaml_parser.py b/packages/bigframes/scripts/bigquery_generator/yaml_parser.py index 263b60648894..277e6431ce94 100644 --- a/packages/bigframes/scripts/bigquery_generator/yaml_parser.py +++ b/packages/bigframes/scripts/bigquery_generator/yaml_parser.py @@ -17,10 +17,10 @@ import re from typing import Any -from . import constants -from . import data_models import yaml +from . import constants, data_models + def _to_snake_case(name: str) -> str: # Replace dots with underscores From e3fb8fd0957f623885727ee50e6580ded4c6c285 Mon Sep 17 00:00:00 2001 From: Shenyang Cai Date: Mon, 27 Jul 2026 02:19:58 +0000 Subject: [PATCH 21/27] fix import --- packages/bigframes/scripts/bigquery_generator/data_models.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/packages/bigframes/scripts/bigquery_generator/data_models.py b/packages/bigframes/scripts/bigquery_generator/data_models.py index 66d9b4c061fa..08e0490358ac 100644 --- a/packages/bigframes/scripts/bigquery_generator/data_models.py +++ b/packages/bigframes/scripts/bigquery_generator/data_models.py @@ -18,7 +18,7 @@ import pathlib from typing import Any -import constants +from . import constants @dataclasses.dataclass From 3049a978515b882e09586558ad26c2d5bb8dc78c Mon Sep 17 00:00:00 2001 From: Shenyang Cai Date: Mon, 27 Jul 2026 02:24:39 +0000 Subject: [PATCH 22/27] fix format --- packages/bigframes/scripts/bigquery_generator/data_models.py | 1 + .../bigframes/scripts/bigquery_generator/template_renderer.py | 1 - 2 files changed, 1 insertion(+), 1 deletion(-) diff --git a/packages/bigframes/scripts/bigquery_generator/data_models.py b/packages/bigframes/scripts/bigquery_generator/data_models.py index 08e0490358ac..614a962fcb72 100644 --- a/packages/bigframes/scripts/bigquery_generator/data_models.py +++ b/packages/bigframes/scripts/bigquery_generator/data_models.py @@ -28,6 +28,7 @@ class BQFuncArg: optional: bool keyword_only: bool + @dataclasses.dataclass class BQFuncImpl: args: list[BQFuncArg] diff --git a/packages/bigframes/scripts/bigquery_generator/template_renderer.py b/packages/bigframes/scripts/bigquery_generator/template_renderer.py index 96dfe64d344a..3dd7fd91b8d5 100644 --- a/packages/bigframes/scripts/bigquery_generator/template_renderer.py +++ b/packages/bigframes/scripts/bigquery_generator/template_renderer.py @@ -14,7 +14,6 @@ # limitations under the License. - from . import constants, data_models From 83b22f0f02ba513139c934125df64e7ed9558169 Mon Sep 17 00:00:00 2001 From: Shenyang Cai Date: Mon, 27 Jul 2026 02:29:34 +0000 Subject: [PATCH 23/27] fix lint --- packages/bigframes/scripts/bigquery_generator/file_generator.py | 1 - 1 file changed, 1 deletion(-) diff --git a/packages/bigframes/scripts/bigquery_generator/file_generator.py b/packages/bigframes/scripts/bigquery_generator/file_generator.py index 9b7853d4d4da..14df5ff98515 100644 --- a/packages/bigframes/scripts/bigquery_generator/file_generator.py +++ b/packages/bigframes/scripts/bigquery_generator/file_generator.py @@ -105,7 +105,6 @@ def _generate_accesor(bq_modules: list[data_models.BQModule]): def generate(bq_modules: list[data_models.BQModule]): - for bq_module in bq_modules: _generate_op_defs(bq_module) _generate_tests(bq_module) From ddc6d4ce39214ebb7450067a3a9027ce97737902 Mon Sep 17 00:00:00 2001 From: Shenyang Cai Date: Tue, 28 Jul 2026 17:46:55 +0000 Subject: [PATCH 24/27] update script path --- .../bigframes/extensions/bigframes/series_accessor.py | 2 +- .../bigframes/bigframes/extensions/core/series_accessor.py | 2 +- .../bigframes/bigframes/extensions/pandas/series_accessor.py | 2 +- packages/bigframes/bigframes/operations/googlesql/aead.py | 2 +- .../operations/googlesql/global_namespace/aead_encryption.py | 2 +- .../bigframes/operations/googlesql/global_namespace/array.py | 2 +- .../bigframes/operations/googlesql/global_namespace/bit.py | 2 +- .../operations/googlesql/global_namespace/conversion.py | 2 +- .../bigframes/operations/googlesql/global_namespace/date.py | 2 +- packages/bigframes/scripts/bigquery_generator/constants.py | 5 ++++- .../generated/global_namespace/test_aead_encryption.py | 2 +- .../unit/bigquery/generated/global_namespace/test_array.py | 2 +- .../unit/bigquery/generated/global_namespace/test_bit.py | 2 +- .../bigquery/generated/global_namespace/test_conversion.py | 2 +- .../unit/bigquery/generated/global_namespace/test_date.py | 2 +- .../bigframes/tests/unit/bigquery/generated/test_aead.py | 2 +- 16 files changed, 19 insertions(+), 16 deletions(-) diff --git a/packages/bigframes/bigframes/extensions/bigframes/series_accessor.py b/packages/bigframes/bigframes/extensions/bigframes/series_accessor.py index ebc2df2eb7ca..c9026595d97e 100644 --- a/packages/bigframes/bigframes/extensions/bigframes/series_accessor.py +++ b/packages/bigframes/bigframes/extensions/bigframes/series_accessor.py @@ -13,7 +13,7 @@ # limitations under the License. # # DO NOT MODIFY THIS FILE DIRECTLY. -# This file was generated by the script: scripts/bigquery_generator +# This file was generated by the script: scripts/generate_bigframes_bigquery.py # from __future__ import annotations diff --git a/packages/bigframes/bigframes/extensions/core/series_accessor.py b/packages/bigframes/bigframes/extensions/core/series_accessor.py index 2356d22f6c88..79a9e511c3c6 100644 --- a/packages/bigframes/bigframes/extensions/core/series_accessor.py +++ b/packages/bigframes/bigframes/extensions/core/series_accessor.py @@ -13,7 +13,7 @@ # limitations under the License. # # DO NOT MODIFY THIS FILE DIRECTLY. -# This file was generated by the script: scripts/bigquery_generator +# This file was generated by the script: scripts/generate_bigframes_bigquery.py # from __future__ import annotations diff --git a/packages/bigframes/bigframes/extensions/pandas/series_accessor.py b/packages/bigframes/bigframes/extensions/pandas/series_accessor.py index 21fd13e2f538..9c33996c4219 100644 --- a/packages/bigframes/bigframes/extensions/pandas/series_accessor.py +++ b/packages/bigframes/bigframes/extensions/pandas/series_accessor.py @@ -13,7 +13,7 @@ # limitations under the License. # # DO NOT MODIFY THIS FILE DIRECTLY. -# This file was generated by the script: scripts/bigquery_generator +# This file was generated by the script: scripts/generate_bigframes_bigquery.py # from __future__ import annotations diff --git a/packages/bigframes/bigframes/operations/googlesql/aead.py b/packages/bigframes/bigframes/operations/googlesql/aead.py index 25922f9bc031..f719d7d69893 100644 --- a/packages/bigframes/bigframes/operations/googlesql/aead.py +++ b/packages/bigframes/bigframes/operations/googlesql/aead.py @@ -14,7 +14,7 @@ # # DO NOT MODIFY THIS FILE DIRECTLY. # This file was generated from: scripts/data/sql-functions/aead.yaml -# by the script: scripts/bigquery_generator +# by the script: scripts/generate_bigframes_bigquery.py from __future__ import annotations diff --git a/packages/bigframes/bigframes/operations/googlesql/global_namespace/aead_encryption.py b/packages/bigframes/bigframes/operations/googlesql/global_namespace/aead_encryption.py index 33e0913da70b..4613ddd7e6df 100644 --- a/packages/bigframes/bigframes/operations/googlesql/global_namespace/aead_encryption.py +++ b/packages/bigframes/bigframes/operations/googlesql/global_namespace/aead_encryption.py @@ -14,7 +14,7 @@ # # DO NOT MODIFY THIS FILE DIRECTLY. # This file was generated from: scripts/data/sql-functions/global_namespace/aead_encryption.yaml -# by the script: scripts/bigquery_generator +# by the script: scripts/generate_bigframes_bigquery.py from __future__ import annotations diff --git a/packages/bigframes/bigframes/operations/googlesql/global_namespace/array.py b/packages/bigframes/bigframes/operations/googlesql/global_namespace/array.py index ee3319a9bfdd..94adbad1839d 100644 --- a/packages/bigframes/bigframes/operations/googlesql/global_namespace/array.py +++ b/packages/bigframes/bigframes/operations/googlesql/global_namespace/array.py @@ -14,7 +14,7 @@ # # DO NOT MODIFY THIS FILE DIRECTLY. # This file was generated from: scripts/data/sql-functions/global_namespace/array.yaml -# by the script: scripts/bigquery_generator +# by the script: scripts/generate_bigframes_bigquery.py from __future__ import annotations diff --git a/packages/bigframes/bigframes/operations/googlesql/global_namespace/bit.py b/packages/bigframes/bigframes/operations/googlesql/global_namespace/bit.py index a480ed78b411..e0c22dfc2990 100644 --- a/packages/bigframes/bigframes/operations/googlesql/global_namespace/bit.py +++ b/packages/bigframes/bigframes/operations/googlesql/global_namespace/bit.py @@ -14,7 +14,7 @@ # # DO NOT MODIFY THIS FILE DIRECTLY. # This file was generated from: scripts/data/sql-functions/global_namespace/bit.yaml -# by the script: scripts/bigquery_generator +# by the script: scripts/generate_bigframes_bigquery.py from __future__ import annotations diff --git a/packages/bigframes/bigframes/operations/googlesql/global_namespace/conversion.py b/packages/bigframes/bigframes/operations/googlesql/global_namespace/conversion.py index 5d2ae897a423..cea4e45d836b 100644 --- a/packages/bigframes/bigframes/operations/googlesql/global_namespace/conversion.py +++ b/packages/bigframes/bigframes/operations/googlesql/global_namespace/conversion.py @@ -14,7 +14,7 @@ # # DO NOT MODIFY THIS FILE DIRECTLY. # This file was generated from: scripts/data/sql-functions/global_namespace/conversion.yaml -# by the script: scripts/bigquery_generator +# by the script: scripts/generate_bigframes_bigquery.py from __future__ import annotations diff --git a/packages/bigframes/bigframes/operations/googlesql/global_namespace/date.py b/packages/bigframes/bigframes/operations/googlesql/global_namespace/date.py index b9a5e920efcc..b6cfc9722b52 100644 --- a/packages/bigframes/bigframes/operations/googlesql/global_namespace/date.py +++ b/packages/bigframes/bigframes/operations/googlesql/global_namespace/date.py @@ -14,7 +14,7 @@ # # DO NOT MODIFY THIS FILE DIRECTLY. # This file was generated from: scripts/data/sql-functions/global_namespace/date.yaml -# by the script: scripts/bigquery_generator +# by the script: scripts/generate_bigframes_bigquery.py from __future__ import annotations diff --git a/packages/bigframes/scripts/bigquery_generator/constants.py b/packages/bigframes/scripts/bigquery_generator/constants.py index 62e2ed1a6183..8fe0e934480a 100644 --- a/packages/bigframes/scripts/bigquery_generator/constants.py +++ b/packages/bigframes/scripts/bigquery_generator/constants.py @@ -19,7 +19,10 @@ SCRIPTS_DIRECTORY = pathlib.Path(__file__).parent.parent.absolute() PACKAGE_ROOT = SCRIPTS_DIRECTORY.parent CODE_ROOT = PACKAGE_ROOT / "bigframes" -SCRIPT_PATH_RELATIVE = pathlib.Path(__file__).relative_to(PACKAGE_ROOT).parent +SCRIPT_PATH_RELATIVE = ( + pathlib.Path(__file__).relative_to(PACKAGE_ROOT).parent.parent + / "generate_bigframes_bigquery.py" +) # Directory containing the YAML files diff --git a/packages/bigframes/tests/unit/bigquery/generated/global_namespace/test_aead_encryption.py b/packages/bigframes/tests/unit/bigquery/generated/global_namespace/test_aead_encryption.py index 85f66acbf498..818151952ff0 100644 --- a/packages/bigframes/tests/unit/bigquery/generated/global_namespace/test_aead_encryption.py +++ b/packages/bigframes/tests/unit/bigquery/generated/global_namespace/test_aead_encryption.py @@ -14,7 +14,7 @@ # # DO NOT MODIFY THIS FILE DIRECTLY. # This file was generated from: scripts/data/sql-functions/global_namespace/aead_encryption.yaml -# by the script: scripts/bigquery_generator +# by the script: scripts/generate_bigframes_bigquery.py import bigframes.bigquery as bbq import bigframes.core.col diff --git a/packages/bigframes/tests/unit/bigquery/generated/global_namespace/test_array.py b/packages/bigframes/tests/unit/bigquery/generated/global_namespace/test_array.py index efc80f4feafb..56b853869024 100644 --- a/packages/bigframes/tests/unit/bigquery/generated/global_namespace/test_array.py +++ b/packages/bigframes/tests/unit/bigquery/generated/global_namespace/test_array.py @@ -14,7 +14,7 @@ # # DO NOT MODIFY THIS FILE DIRECTLY. # This file was generated from: scripts/data/sql-functions/global_namespace/array.yaml -# by the script: scripts/bigquery_generator +# by the script: scripts/generate_bigframes_bigquery.py import bigframes.bigquery as bbq import bigframes.core.col diff --git a/packages/bigframes/tests/unit/bigquery/generated/global_namespace/test_bit.py b/packages/bigframes/tests/unit/bigquery/generated/global_namespace/test_bit.py index c0c7e1ec2ab3..2cccafc0643d 100644 --- a/packages/bigframes/tests/unit/bigquery/generated/global_namespace/test_bit.py +++ b/packages/bigframes/tests/unit/bigquery/generated/global_namespace/test_bit.py @@ -14,7 +14,7 @@ # # DO NOT MODIFY THIS FILE DIRECTLY. # This file was generated from: scripts/data/sql-functions/global_namespace/bit.yaml -# by the script: scripts/bigquery_generator +# by the script: scripts/generate_bigframes_bigquery.py import bigframes.bigquery as bbq import bigframes.core.col diff --git a/packages/bigframes/tests/unit/bigquery/generated/global_namespace/test_conversion.py b/packages/bigframes/tests/unit/bigquery/generated/global_namespace/test_conversion.py index b5a8a819115a..84dfc02465cc 100644 --- a/packages/bigframes/tests/unit/bigquery/generated/global_namespace/test_conversion.py +++ b/packages/bigframes/tests/unit/bigquery/generated/global_namespace/test_conversion.py @@ -14,7 +14,7 @@ # # DO NOT MODIFY THIS FILE DIRECTLY. # This file was generated from: scripts/data/sql-functions/global_namespace/conversion.yaml -# by the script: scripts/bigquery_generator +# by the script: scripts/generate_bigframes_bigquery.py import bigframes.bigquery as bbq import bigframes.core.col diff --git a/packages/bigframes/tests/unit/bigquery/generated/global_namespace/test_date.py b/packages/bigframes/tests/unit/bigquery/generated/global_namespace/test_date.py index b0228804c6a9..6484208584f9 100644 --- a/packages/bigframes/tests/unit/bigquery/generated/global_namespace/test_date.py +++ b/packages/bigframes/tests/unit/bigquery/generated/global_namespace/test_date.py @@ -14,7 +14,7 @@ # # DO NOT MODIFY THIS FILE DIRECTLY. # This file was generated from: scripts/data/sql-functions/global_namespace/date.yaml -# by the script: scripts/bigquery_generator +# by the script: scripts/generate_bigframes_bigquery.py import bigframes.bigquery as bbq import bigframes.core.col diff --git a/packages/bigframes/tests/unit/bigquery/generated/test_aead.py b/packages/bigframes/tests/unit/bigquery/generated/test_aead.py index 407cd22af6f6..ce728b41899d 100644 --- a/packages/bigframes/tests/unit/bigquery/generated/test_aead.py +++ b/packages/bigframes/tests/unit/bigquery/generated/test_aead.py @@ -14,7 +14,7 @@ # # DO NOT MODIFY THIS FILE DIRECTLY. # This file was generated from: scripts/data/sql-functions/aead.yaml -# by the script: scripts/bigquery_generator +# by the script: scripts/generate_bigframes_bigquery.py import bigframes.bigquery as bbq import bigframes.core.col From 8ef98b99a1dda075b0f04937b6da1b7fba7b4d6b Mon Sep 17 00:00:00 2001 From: Shenyang Cai Date: Tue, 28 Jul 2026 22:17:11 +0000 Subject: [PATCH 25/27] remove __main__.py from package --- .../scripts/bigquery_generator/__main__.py | 29 ------------------- .../scripts/generate_bigframes_bigquery.py | 12 +++++++- 2 files changed, 11 insertions(+), 30 deletions(-) delete mode 100644 packages/bigframes/scripts/bigquery_generator/__main__.py diff --git a/packages/bigframes/scripts/bigquery_generator/__main__.py b/packages/bigframes/scripts/bigquery_generator/__main__.py deleted file mode 100644 index 2f7ffaacf2d0..000000000000 --- a/packages/bigframes/scripts/bigquery_generator/__main__.py +++ /dev/null @@ -1,29 +0,0 @@ -# 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. - - -from . import constants, file_generator, yaml_parser - - -def main(): - modules = [] - - for yaml_file in sorted(constants.DATA_DIR.glob("**/*.yaml")): - modules.append(yaml_parser.parse_yaml(yaml_file)) - - file_generator.generate(modules) - - -if __name__ == "__main__": - main() diff --git a/packages/bigframes/scripts/generate_bigframes_bigquery.py b/packages/bigframes/scripts/generate_bigframes_bigquery.py index 5882155eb6b5..c2a3c0ab813a 100755 --- a/packages/bigframes/scripts/generate_bigframes_bigquery.py +++ b/packages/bigframes/scripts/generate_bigframes_bigquery.py @@ -29,7 +29,17 @@ if str(scripts_dir) not in sys.path: sys.path.insert(0, str(scripts_dir)) -from bigquery_generator.__main__ import main +from bigquery_generator import constants, file_generator, yaml_parser + + +def main(): + modules = [] + + for yaml_file in sorted(constants.DATA_DIR.glob("**/*.yaml")): + modules.append(yaml_parser.parse_yaml(yaml_file)) + + file_generator.generate(modules) + if __name__ == "__main__": main() From bd136c611b84b05dd0af3d9aebdf46039c9c6470 Mon Sep 17 00:00:00 2001 From: Shenyang Cai Date: Tue, 28 Jul 2026 22:42:22 +0000 Subject: [PATCH 26/27] update template loading and module-level docstrings --- .../scripts/bigquery_generator/constants.py | 7 ++- .../scripts/bigquery_generator/data_models.py | 54 +++++++++++++++++++ .../bigquery_generator/template_renderer.py | 6 ++- 3 files changed, 62 insertions(+), 5 deletions(-) diff --git a/packages/bigframes/scripts/bigquery_generator/constants.py b/packages/bigframes/scripts/bigquery_generator/constants.py index 8fe0e934480a..9a9a35339b25 100644 --- a/packages/bigframes/scripts/bigquery_generator/constants.py +++ b/packages/bigframes/scripts/bigquery_generator/constants.py @@ -160,10 +160,6 @@ ] + RUFF_COMMON_ARGS -TEMPLATES: dict[str, jinja2.Template] # Lazily-loaded global variable -_templates = None - - def _load_templates() -> dict[str, jinja2.Template]: env = jinja2.Environment( loader=jinja2.FileSystemLoader(TEMPLATE_DIR), @@ -183,6 +179,9 @@ def _load_templates() -> dict[str, jinja2.Template]: } +TEMPLATES: dict[str, jinja2.Template] = _load_templates() + + def __getattr__(name: str): global _templates diff --git a/packages/bigframes/scripts/bigquery_generator/data_models.py b/packages/bigframes/scripts/bigquery_generator/data_models.py index 614a962fcb72..008fa780e7c7 100644 --- a/packages/bigframes/scripts/bigquery_generator/data_models.py +++ b/packages/bigframes/scripts/bigquery_generator/data_models.py @@ -12,6 +12,60 @@ # See the License for the specific language governing permissions and # limitations under the License. +"""Data models for BigQuery code generator. + +`BQ*` models the Substrait YAML extension structure of BigQuery SQL functions, +while `BigFrames*` models the Jinja template inputs. + +BQ* Class Relations: +==================== ++-------------------+ +| BQModule | ++-------------------+ + | + | functions: list[BQFunc] + v ++-------------------+ +| BQFunc | ++-------------------+ + | + | impls: list[BQFuncImpl] + v ++-------------------+ +| BQFuncImpl | ++-------------------+ + | + | args: list[BQFuncArg] + v ++-------------------+ +| BQFuncArg | ++-------------------+ + +BigFrames* Class Relations: +================================= + +--------------------------------+ + | Accessor |<---+ children: list[Accessor] + +--------------------------------+----+ (nested namespace hierarchy) + | + | functions: list[BigFramesFunc] + v + +--------------------------------+ + | BigFramesFunc | + +--------------------------------+ + | + | args: list[BigFramesFuncArg] + v + +--------------------------------+ + | BigFramesFuncArg | + +--------------------------------+ + + ---------------------------------- + + +--------------------------------+ + | BigFramesOp | (Standalone data model for op defs) + +--------------------------------+ +""" + from __future__ import annotations import dataclasses diff --git a/packages/bigframes/scripts/bigquery_generator/template_renderer.py b/packages/bigframes/scripts/bigquery_generator/template_renderer.py index 3dd7fd91b8d5..8bfe1b3c6b3c 100644 --- a/packages/bigframes/scripts/bigquery_generator/template_renderer.py +++ b/packages/bigframes/scripts/bigquery_generator/template_renderer.py @@ -1,4 +1,3 @@ -# Render jinja template with module data parsed from yaml # Copyright 2026 Google LLC # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -14,6 +13,11 @@ # limitations under the License. +""" +Renders jinja template with module data parsed from yaml. +""" + + from . import constants, data_models From 9ea30b81631a9bda16932a6684df5aba01615e60 Mon Sep 17 00:00:00 2001 From: Shenyang Cai Date: Tue, 28 Jul 2026 22:48:34 +0000 Subject: [PATCH 27/27] fix format --- packages/bigframes/scripts/bigquery_generator/data_models.py | 2 +- .../bigframes/scripts/bigquery_generator/template_renderer.py | 1 - 2 files changed, 1 insertion(+), 2 deletions(-) diff --git a/packages/bigframes/scripts/bigquery_generator/data_models.py b/packages/bigframes/scripts/bigquery_generator/data_models.py index 008fa780e7c7..72b02bd026c6 100644 --- a/packages/bigframes/scripts/bigquery_generator/data_models.py +++ b/packages/bigframes/scripts/bigquery_generator/data_models.py @@ -15,7 +15,7 @@ """Data models for BigQuery code generator. `BQ*` models the Substrait YAML extension structure of BigQuery SQL functions, -while `BigFrames*` models the Jinja template inputs. +while `BigFrames*` models the Jinja template outputs. BQ* Class Relations: ==================== diff --git a/packages/bigframes/scripts/bigquery_generator/template_renderer.py b/packages/bigframes/scripts/bigquery_generator/template_renderer.py index 8bfe1b3c6b3c..ee9944e79a13 100644 --- a/packages/bigframes/scripts/bigquery_generator/template_renderer.py +++ b/packages/bigframes/scripts/bigquery_generator/template_renderer.py @@ -17,7 +17,6 @@ Renders jinja template with module data parsed from yaml. """ - from . import constants, data_models