This repository was archived by the owner on Apr 1, 2026. It is now read-only.
-
Notifications
You must be signed in to change notification settings - Fork 68
Expand file tree
/
Copy pathsql.py
More file actions
397 lines (338 loc) · 16.3 KB
/
sql.py
File metadata and controls
397 lines (338 loc) · 16.3 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
# Copyright 2023 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.
"""
Generates SQL queries needed for BigQuery DataFrames ML
"""
from typing import Iterable, Literal, Mapping, Optional, Union
import bigframes_vendored.constants as constants
import google.cloud.bigquery
import bigframes.core.compile.googlesql as sql_utils
import bigframes.core.sql as sql_vals
# TODO: Add proper escaping logic from core/compile module
class BaseSqlGenerator:
"""Generate base SQL strings for ML. Model name isn't needed in this class."""
# General methods
def encode_value(self, v: Union[str, int, float, Iterable[str]]) -> str:
"""Encode a parameter value for SQL"""
if isinstance(v, (str, int, float)):
return sql_vals.simple_literal(v)
elif isinstance(v, Iterable):
inner = ", ".join([self.encode_value(x) for x in v])
return f"[{inner}]"
else:
raise ValueError(
f"Unexpected value type {type(v)}. {constants.FEEDBACK_LINK}"
)
def build_parameters(self, **kwargs: Union[str, int, float, Iterable[str]]) -> str:
"""Encode a dict of values into a formatted Iterable of key-value pairs for SQL"""
indent_str = " "
param_strs = [f"{k}={self.encode_value(v)}" for k, v in kwargs.items()]
return "\n" + indent_str + f",\n{indent_str}".join(param_strs)
def build_structs(self, **kwargs: Union[int, float]) -> str:
"""Encode a dict of values into a formatted STRUCT items for SQL"""
indent_str = " "
param_strs = [
f"{sql_vals.simple_literal(v)} AS {sql_utils.identifier(k)}"
for k, v in kwargs.items()
]
return "\n" + indent_str + f",\n{indent_str}".join(param_strs)
def build_expressions(self, *expr_sqls: str) -> str:
"""Encode a Iterable of SQL expressions into a formatted Iterable for SQL"""
indent_str = " "
return "\n" + indent_str + f",\n{indent_str}".join(expr_sqls)
def build_schema(self, **kwargs: str) -> str:
"""Encode a dict of values into a formatted schema type items for SQL"""
indent_str = " "
param_strs = [f"{sql_utils.identifier(k)} {v}" for k, v in kwargs.items()]
return "\n" + indent_str + f",\n{indent_str}".join(param_strs)
def options(self, **kwargs: Union[str, int, float, Iterable[str]]) -> str:
"""Encode the OPTIONS clause for BQML"""
return f"OPTIONS({self.build_parameters(**kwargs)})"
def struct_options(self, **kwargs: Union[int, float]) -> str:
"""Encode a BQ STRUCT as options."""
return f"STRUCT({self.build_structs(**kwargs)})"
def struct_columns(self, columns: Iterable[str]) -> str:
"""Encode a BQ Table columns to a STRUCT."""
columns_str = ", ".join(map(sql_utils.identifier, columns))
return f"STRUCT({columns_str})"
def input(self, **kwargs: str) -> str:
"""Encode a BQML INPUT clause."""
return f"INPUT({self.build_schema(**kwargs)})"
def output(self, **kwargs: str) -> str:
"""Encode a BQML OUTPUT clause."""
return f"OUTPUT({self.build_schema(**kwargs)})"
# Connection
def connection(self, conn_name: str) -> str:
"""Encode the REMOTE WITH CONNECTION clause for BQML. conn_name is of the format <PROJECT_NUMBER/PROJECT_ID>.<REGION>.<CONNECTION_NAME>."""
return f"REMOTE WITH CONNECTION `{conn_name}`"
# Transformers
def transform(self, *expr_sqls: str) -> str:
"""Encode the TRANSFORM clause for BQML"""
return f"TRANSFORM({self.build_expressions(*expr_sqls)})"
def ml_standard_scaler(self, numeric_expr_sql: str, name: str) -> str:
"""Encode ML.STANDARD_SCALER for BQML"""
return f"""ML.STANDARD_SCALER({sql_utils.identifier(numeric_expr_sql)}) OVER() AS {sql_utils.identifier(name)}"""
def ml_max_abs_scaler(self, numeric_expr_sql: str, name: str) -> str:
"""Encode ML.MAX_ABS_SCALER for BQML"""
return f"""ML.MAX_ABS_SCALER({sql_utils.identifier(numeric_expr_sql)}) OVER() AS {sql_utils.identifier(name)}"""
def ml_min_max_scaler(self, numeric_expr_sql: str, name: str) -> str:
"""Encode ML.MIN_MAX_SCALER for BQML"""
return f"""ML.MIN_MAX_SCALER({sql_utils.identifier(numeric_expr_sql)}) OVER() AS {sql_utils.identifier(name)}"""
def ml_imputer(
self,
col_name: str,
strategy: str,
name: str,
) -> str:
"""Encode ML.IMPUTER for BQML"""
return f"""ML.IMPUTER({sql_utils.identifier(col_name)}, '{strategy}') OVER() AS {sql_utils.identifier(name)}"""
def ml_bucketize(
self,
input_id: str,
array_split_points: Iterable[Union[int, float]],
output_id: str,
) -> str:
"""Encode ML.BUCKETIZE for BQML"""
# Use Python value rather than Numpy value to serialization.
points = [
point.item() if hasattr(point, "item") else point
for point in array_split_points
]
return f"""ML.BUCKETIZE({sql_utils.identifier(input_id)}, {points}, FALSE) AS {sql_utils.identifier(output_id)}"""
def ml_quantile_bucketize(
self,
numeric_expr_sql: str,
num_bucket: int,
name: str,
) -> str:
"""Encode ML.QUANTILE_BUCKETIZE for BQML"""
return f"""ML.QUANTILE_BUCKETIZE({sql_utils.identifier(numeric_expr_sql)}, {num_bucket}) OVER() AS {sql_utils.identifier(name)}"""
def ml_one_hot_encoder(
self,
numeric_expr_sql: str,
drop: str,
top_k: int,
frequency_threshold: int,
name: str,
) -> str:
"""Encode ML.ONE_HOT_ENCODER for BQML.
https://cloud.google.com/bigquery/docs/reference/standard-sql/bigqueryml-syntax-one-hot-encoder for params."""
return f"""ML.ONE_HOT_ENCODER({sql_utils.identifier(numeric_expr_sql)}, '{drop}', {top_k}, {frequency_threshold}) OVER() AS {sql_utils.identifier(name)}"""
def ml_label_encoder(
self,
numeric_expr_sql: str,
top_k: int,
frequency_threshold: int,
name: str,
) -> str:
"""Encode ML.LABEL_ENCODER for BQML.
https://cloud.google.com/bigquery/docs/reference/standard-sql/bigqueryml-syntax-label-encoder for params."""
return f"""ML.LABEL_ENCODER({sql_utils.identifier(numeric_expr_sql)}, {top_k}, {frequency_threshold}) OVER() AS {sql_utils.identifier(name)}"""
def ml_polynomial_expand(
self, columns: Iterable[str], degree: int, name: str
) -> str:
"""Encode ML.POLYNOMIAL_EXPAND.
https://cloud.google.com/bigquery/docs/reference/standard-sql/bigqueryml-syntax-polynomial-expand"""
return f"""ML.POLYNOMIAL_EXPAND({self.struct_columns(columns)}, {degree}) AS {sql_utils.identifier(name)}"""
def ml_distance(
self,
col_x: str,
col_y: str,
type: Literal["EUCLIDEAN", "MANHATTAN", "COSINE"],
source_sql: str,
name: str,
) -> str:
"""Encode ML.DISTANCE for BQML.
https://cloud.google.com/bigquery/docs/reference/standard-sql/bigqueryml-syntax-distance"""
return f"""SELECT *, ML.DISTANCE({sql_utils.identifier(col_x)}, {sql_utils.identifier(col_y)}, '{type}') AS {sql_utils.identifier(name)} FROM ({source_sql})"""
class ModelCreationSqlGenerator(BaseSqlGenerator):
"""Sql generator for creating a model entity. Model id is the standalone id without project id and dataset id."""
def _model_id_sql(
self,
model_ref: google.cloud.bigquery.ModelReference,
):
return f"{sql_utils.identifier(model_ref.project)}.{sql_utils.identifier(model_ref.dataset_id)}.{sql_utils.identifier(model_ref.model_id)}"
# Model create and alter
def create_model(
self,
source_sql: str,
model_ref: google.cloud.bigquery.ModelReference,
options: Mapping[str, Union[str, int, float, Iterable[str]]] = {},
transforms: Optional[Iterable[str]] = None,
) -> str:
"""Encode the CREATE OR REPLACE MODEL statement for BQML"""
parts = [f"CREATE OR REPLACE MODEL {self._model_id_sql(model_ref)}"]
if transforms:
parts.append(self.transform(*transforms))
if options:
parts.append(self.options(**options))
parts.append(f"AS {source_sql}")
return "\n".join(parts)
def create_llm_remote_model(
self,
source_sql: str,
connection_name: str,
model_ref: google.cloud.bigquery.ModelReference,
options: Mapping[str, Union[str, int, float, Iterable[str]]] = {},
) -> str:
"""Encode the CREATE OR REPLACE MODEL statement for BQML"""
parts = [f"CREATE OR REPLACE MODEL {self._model_id_sql(model_ref)}"]
parts.append(self.connection(connection_name))
if options:
parts.append(self.options(**options))
parts.append(f"AS {source_sql}")
return "\n".join(parts)
def create_remote_model(
self,
connection_name: str,
model_ref: google.cloud.bigquery.ModelReference,
input: Mapping[str, str] = {},
output: Mapping[str, str] = {},
options: Mapping[str, Union[str, int, float, Iterable[str]]] = {},
) -> str:
"""Encode the CREATE OR REPLACE MODEL statement for BQML remote model."""
parts = [f"CREATE OR REPLACE MODEL {self._model_id_sql(model_ref)}"]
if input:
parts.append(self.input(**input))
if output:
parts.append(self.output(**output))
parts.append(self.connection(connection_name))
if options:
parts.append(self.options(**options))
return "\n".join(parts)
def create_imported_model(
self,
model_ref: google.cloud.bigquery.ModelReference,
options: Mapping[str, Union[str, int, float, Iterable[str]]] = {},
) -> str:
"""Encode the CREATE OR REPLACE MODEL statement for BQML remote model."""
parts = [f"CREATE OR REPLACE MODEL {self._model_id_sql(model_ref)}"]
if options:
parts.append(self.options(**options))
return "\n".join(parts)
def create_xgboost_imported_model(
self,
model_ref: google.cloud.bigquery.ModelReference,
input: Mapping[str, str] = {},
output: Mapping[str, str] = {},
options: Mapping[str, Union[str, int, float, Iterable[str]]] = {},
) -> str:
"""Encode the CREATE OR REPLACE MODEL statement for BQML remote model."""
parts = [f"CREATE OR REPLACE MODEL {self._model_id_sql(model_ref)}"]
if input:
parts.append(self.input(**input))
if output:
parts.append(self.output(**output))
if options:
parts.append(self.options(**options))
return "\n".join(parts)
class ModelManipulationSqlGenerator(BaseSqlGenerator):
"""Sql generator for manipulating a model entity. Model name is the full model path of project_id.dataset_id.model_id."""
def __init__(self, model_ref: google.cloud.bigquery.ModelReference):
self._model_ref = model_ref
def _model_ref_sql(self) -> str:
return f"{sql_utils.identifier(self._model_ref.project)}.{sql_utils.identifier(self._model_ref.dataset_id)}.{sql_utils.identifier(self._model_ref.model_id)}"
# Alter model
def alter_model(
self,
options: Mapping[str, Union[str, int, float, Iterable[str]]] = {},
) -> str:
"""Encode the ALTER MODEL statement for BQML"""
options_sql = self.options(**options)
parts = [f"ALTER MODEL {self._model_ref_sql()}"]
parts.append(f"SET {options_sql}")
return "\n".join(parts)
# ML prediction TVFs
def ml_predict(self, source_sql: str) -> str:
"""Encode ML.PREDICT for BQML"""
return f"""SELECT * FROM ML.PREDICT(MODEL {self._model_ref_sql()},
({source_sql}))"""
def ml_explain_predict(
self, source_sql: str, struct_options: Mapping[str, Union[int, float]]
) -> str:
"""Encode ML.EXPLAIN_PREDICT for BQML"""
struct_options_sql = self.struct_options(**struct_options)
return f"""SELECT * FROM ML.EXPLAIN_PREDICT(MODEL {self._model_ref_sql()},
({source_sql}), {struct_options_sql})"""
def ml_forecast(self, struct_options: Mapping[str, Union[int, float]]) -> str:
"""Encode ML.FORECAST for BQML"""
struct_options_sql = self.struct_options(**struct_options)
return f"""SELECT * FROM ML.FORECAST(MODEL {self._model_ref_sql()},
{struct_options_sql})"""
def ml_explain_forecast(
self, struct_options: Mapping[str, Union[int, float]]
) -> str:
"""Encode ML.EXPLAIN_FORECAST for BQML"""
struct_options_sql = self.struct_options(**struct_options)
return f"""SELECT * FROM ML.EXPLAIN_FORECAST(MODEL {self._model_ref_sql()},
{struct_options_sql})"""
def ml_generate_text(
self, source_sql: str, struct_options: Mapping[str, Union[int, float]]
) -> str:
"""Encode ML.GENERATE_TEXT for BQML"""
struct_options_sql = self.struct_options(**struct_options)
return f"""SELECT * FROM ML.GENERATE_TEXT(MODEL {self._model_ref_sql()},
({source_sql}), {struct_options_sql})"""
def ml_generate_embedding(
self, source_sql: str, struct_options: Mapping[str, Union[int, float]]
) -> str:
"""Encode ML.GENERATE_EMBEDDING for BQML"""
struct_options_sql = self.struct_options(**struct_options)
return f"""SELECT * FROM ML.GENERATE_EMBEDDING(MODEL {self._model_ref_sql()},
({source_sql}), {struct_options_sql})"""
def ml_detect_anomalies(
self, source_sql: str, struct_options: Mapping[str, Union[int, float]]
) -> str:
"""Encode ML.DETECT_ANOMALIES for BQML"""
struct_options_sql = self.struct_options(**struct_options)
return f"""SELECT * FROM ML.DETECT_ANOMALIES(MODEL {self._model_ref_sql()},
{struct_options_sql}, ({source_sql}))"""
# ML evaluation TVFs
def ml_evaluate(self, source_sql: Optional[str] = None) -> str:
"""Encode ML.EVALUATE for BQML"""
if source_sql is None:
return f"""SELECT * FROM ML.EVALUATE(MODEL {self._model_ref_sql()})"""
else:
return f"""SELECT * FROM ML.EVALUATE(MODEL {self._model_ref_sql()},
({source_sql}))"""
def ml_arima_coefficients(self) -> str:
"""Encode ML.ARIMA_COEFFICIENTS for BQML"""
return f"""SELECT * FROM ML.ARIMA_COEFFICIENTS(MODEL {self._model_ref_sql()})"""
# ML evaluation TVFs
def ml_llm_evaluate(self, source_sql: str, task_type: Optional[str] = None) -> str:
"""Encode ML.EVALUATE for BQML"""
# Note: don't need index as evaluate returns a new table
return f"""SELECT * FROM ML.EVALUATE(MODEL {self._model_ref_sql()},
({source_sql}), STRUCT("{task_type}" AS task_type))"""
# ML evaluation TVFs
def ml_arima_evaluate(self, show_all_candidate_models: bool = False) -> str:
"""Encode ML.ARMIA_EVALUATE for BQML"""
return f"""SELECT * FROM ML.ARIMA_EVALUATE(MODEL {self._model_ref_sql()},
STRUCT({show_all_candidate_models} AS show_all_candidate_models))"""
def ml_centroids(self) -> str:
"""Encode ML.CENTROIDS for BQML"""
return f"""SELECT * FROM ML.CENTROIDS(MODEL {self._model_ref_sql()})"""
def ml_principal_components(self) -> str:
"""Encode ML.PRINCIPAL_COMPONENTS for BQML"""
return (
f"""SELECT * FROM ML.PRINCIPAL_COMPONENTS(MODEL {self._model_ref_sql()})"""
)
def ml_principal_component_info(self) -> str:
"""Encode ML.PRINCIPAL_COMPONENT_INFO for BQML"""
return f"""SELECT * FROM ML.PRINCIPAL_COMPONENT_INFO(MODEL {self._model_ref_sql()})"""
# ML transform TVF, that require a transform_only type model
def ml_transform(self, source_sql: str) -> str:
"""Encode ML.TRANSFORM for BQML"""
return f"""SELECT * FROM ML.TRANSFORM(MODEL {self._model_ref_sql()},
({source_sql}))"""