Skip to content

Commit b7e5a06

Browse files
sycaisycai
committed
checkpoint: add logic for rendering signature template
1 parent eceea95 commit b7e5a06

7 files changed

Lines changed: 565 additions & 0 deletions

File tree

‎packages/bigframes/scripts/generate_bigframes_bigquery/__init__.py‎

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,13 @@
1+
# Copyright 2026 Google LLC
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.

‎packages/bigframes/scripts/generate_bigframes_bigquery/__main__.py‎

Lines changed: 42 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,42 @@
1+
#!/usr/bin/env -S uv run --active --script
2+
#
3+
# /// script
4+
# dependencies = [
5+
# "jinja2",
6+
# "pyyaml",
7+
# "ruff==0.14.14",
8+
# ]
9+
# ///
10+
#
11+
# Copyright 2026 Google LLC
12+
#
13+
# Licensed under the Apache License, Version 2.0 (the "License");
14+
# you may not use this file except in compliance with the License.
15+
# You may obtain a copy of the License at
16+
#
17+
# http://www.apache.org/licenses/LICENSE-2.0
18+
#
19+
# Unless required by applicable law or agreed to in writing, software
20+
# distributed under the License is distributed on an "AS IS" BASIS,
21+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
22+
# See the License for the specific language governing permissions and
23+
# limitations under the License.
24+
25+
import pprint
26+
27+
import constants
28+
import yaml_parser
29+
import file_generator
30+
31+
32+
def main():
33+
modules = []
34+
35+
for yaml_file in sorted(constants.DATA_DIR.glob("**/*.yaml")):
36+
modules.append(yaml_parser.parse_yaml(yaml_file))
37+
38+
file_generator.generate(modules)
39+
40+
41+
if __name__ == "__main__":
42+
main()

‎packages/bigframes/scripts/generate_bigframes_bigquery/constants.py‎

Lines changed: 155 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,155 @@
1+
# Copyright 2026 Google LLC
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
import pathlib
16+
import jinja2
17+
18+
SCRIPTS_DIRECTORY = pathlib.Path(__file__).parent.parent.absolute()
19+
PACKAGE_ROOT = SCRIPTS_DIRECTORY.parent
20+
CODE_ROOT = PACKAGE_ROOT / "bigframes"
21+
SCRIPT_PATH_RELATIVE = pathlib.Path(__file__).relative_to(PACKAGE_ROOT)
22+
23+
# Directory containing the YAML files
24+
DATA_DIR = SCRIPTS_DIRECTORY / "data" / "sql-functions"
25+
# Directory where the generated Python files will be placed
26+
OUTPUT_DIR = CODE_ROOT / "operations" / "googlesql"
27+
# Directory where the generated test files will be placed
28+
TEST_OUTPUT_DIR = PACKAGE_ROOT / "tests" / "unit" / "bigquery" / "generated"
29+
# Directory containing the Jinja2 templates
30+
TEMPLATE_DIR = SCRIPTS_DIRECTORY / "templates"
31+
32+
PYTHON_BUILTINS = {
33+
"abs",
34+
"all",
35+
"any",
36+
"ascii",
37+
"bin",
38+
"bool",
39+
"breakpoint",
40+
"bytearray",
41+
"bytes",
42+
"callable",
43+
"chr",
44+
"classmethod",
45+
"compile",
46+
"complex",
47+
"delattr",
48+
"dict",
49+
"dir",
50+
"divmod",
51+
"enumerate",
52+
"eval",
53+
"exec",
54+
"filter",
55+
"float",
56+
"format",
57+
"frozenset",
58+
"getattr",
59+
"globals",
60+
"hasattr",
61+
"hash",
62+
"help",
63+
"hex",
64+
"id",
65+
"input",
66+
"int",
67+
"isinstance",
68+
"issubclass",
69+
"iter",
70+
"len",
71+
"list",
72+
"locals",
73+
"map",
74+
"max",
75+
"memoryview",
76+
"min",
77+
"next",
78+
"object",
79+
"oct",
80+
"open",
81+
"ord",
82+
"pow",
83+
"print",
84+
"property",
85+
"range",
86+
"repr",
87+
"reversed",
88+
"round",
89+
"set",
90+
"setattr",
91+
"slice",
92+
"sorted",
93+
"staticmethod",
94+
"str",
95+
"sum",
96+
"super",
97+
"tuple",
98+
"type",
99+
"vars",
100+
"zip",
101+
}
102+
103+
DTYPE_MAP = {
104+
"binary": "dtypes.BYTES_DTYPE",
105+
"string": "dtypes.STRING_DTYPE",
106+
"int64": "dtypes.INT_DTYPE",
107+
"i64": "dtypes.INT_DTYPE",
108+
"float64": "dtypes.FLOAT_DTYPE",
109+
"fp64": "dtypes.FLOAT_DTYPE",
110+
"bool": "dtypes.BOOL_DTYPE",
111+
"boolean": "dtypes.BOOL_DTYPE",
112+
"geography": "dtypes.GEO_DTYPE",
113+
"json": "dtypes.JSON_DTYPE",
114+
"date": "dtypes.DATE_DTYPE",
115+
"time": "dtypes.TIME_DTYPE",
116+
"datetime": "dtypes.DATETIME_DTYPE",
117+
"timestamp": "dtypes.TIMESTAMP_DTYPE",
118+
"decimal<38,9>": "dtypes.NUMERIC_DTYPE",
119+
"decimal<76,38>": "dtypes.BIGNUMERIC_DTYPE",
120+
}
121+
122+
123+
TEMPLATES: dict[str, jinja2.Template] # Lazily-loaded global variable
124+
_templates = None
125+
126+
127+
def _load_templates() -> dict[str, jinja2.Template]:
128+
env = jinja2.Environment(
129+
loader=jinja2.FileSystemLoader(TEMPLATE_DIR),
130+
trim_blocks=True,
131+
lstrip_blocks=True,
132+
)
133+
return {
134+
"operation": env.get_template("operation.py.j2"),
135+
"test_operation": env.get_template("test_operation.py.j2"),
136+
"license": env.get_template("license.py.j2"),
137+
"signature_def": env.get_template("signature_def.py.j2"),
138+
"core_series_accessor": env.get_template("core_series_accessor.py.j2"),
139+
"bigframes_series_accessor": env.get_template(
140+
"bigframes_series_accessor.py.j2"
141+
),
142+
"pandas_series_accessor": env.get_template("pandas_series_accessor.py.j2"),
143+
}
144+
145+
146+
def __getattr__(name: str):
147+
global _templates
148+
149+
if name != "TEMPLATES":
150+
raise AttributeError(f"module '{__name__}' has no attribute '{name}'")
151+
152+
if _templates is None:
153+
_templates = _load_templates()
154+
155+
return _templates

‎packages/bigframes/scripts/generate_bigframes_bigquery/data_models.py‎

Lines changed: 95 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,95 @@
1+
# Copyright 2026 Google LLC
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
16+
import dataclasses
17+
import pathlib
18+
19+
import constants
20+
from typing import Any
21+
22+
23+
@dataclasses.dataclass
24+
class BQFuncArg:
25+
name: str
26+
value: str # The type of the arg
27+
optional: bool
28+
keyword_only: bool
29+
30+
31+
@dataclasses.dataclass
32+
class BQFuncImpl:
33+
args: list[BQFuncArg]
34+
return_type: str
35+
36+
@property
37+
def uses_any1(self) -> bool:
38+
if "any1" in self.return_type:
39+
return True
40+
41+
return any("any1" in arg.value for arg in self.args)
42+
43+
def to_dict(self) -> dict[str, Any]:
44+
result = dataclasses.asdict(self)
45+
46+
result["uses_any1"] = self.uses_any1
47+
48+
# We cannot use "return" as a field name, but we need to use it
49+
# as a key for template rendering
50+
result["return"] = self.return_type
51+
del result["return_type"]
52+
53+
return result
54+
55+
56+
@dataclasses.dataclass
57+
class BQFunc:
58+
name: str
59+
description: str
60+
impls: list[BQFuncImpl]
61+
62+
63+
@dataclasses.dataclass
64+
class BQModule:
65+
module_path: pathlib.Path
66+
functions: list[BQFunc]
67+
68+
@property
69+
def name(self) -> str:
70+
return self.module_path.name
71+
72+
@property
73+
def is_global_namespace(self) -> bool:
74+
return "global_namespace" in self.module_path.parts
75+
76+
@property
77+
def output_file(self):
78+
return constants.OUTPUT_DIR.joinpath(self.module_path).with_suffix(".py")
79+
80+
81+
@dataclasses.dataclass
82+
class BigFramesOp:
83+
internal_name: str
84+
sql_name: str
85+
arg_specs: str
86+
signature: str
87+
signature_definition: str
88+
89+
90+
@dataclasses.dataclass
91+
class BigFramesFunc:
92+
name: str
93+
description: str
94+
args: list[str]
95+
series_accessor_arg: list[str]

‎packages/bigframes/scripts/generate_bigframes_bigquery/file_generator.py‎

Lines changed: 44 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,44 @@
1+
# Copyright 2026 Google LLC
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
import constants
16+
import jinja2
17+
import data_models
18+
import pprint
19+
20+
import template_renderer
21+
22+
23+
def _generate_sql_operator_def():
24+
pass
25+
26+
27+
def _generate_accesor():
28+
pass
29+
30+
31+
def _generate_tests():
32+
pass
33+
34+
35+
def generate(bq_modules: list[data_models.BQModule]):
36+
37+
for bq_module in bq_modules:
38+
for bq_func in bq_module.functions:
39+
pprint.pp(template_renderer.render_signature_def(bq_func, bq_module))
40+
41+
# Write to file
42+
43+
# Ruff format
44+
pass

0 commit comments

Comments
 (0)