feat: add support for hparam_range and hparam_candidates to bigframes.bigquery.create_model (#16640)

Fixes internal issue b/501171054
 🦕

---------

Co-authored-by: google-labs-jules[bot] <161369871+google-labs-jules[bot]@users.noreply.github.com>
Co-authored-by: tswast <247555+tswast@users.noreply.github.com>
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>

Author: 3 people

packages/bigframes/bigframes/bigquery/__init__.py

@@ -87,7 +87,11 @@
     to_json,
     to_json_string,
 )
-from bigframes.bigquery._operations.mathematical import rand
+from bigframes.bigquery._operations.mathematical import (
+    hparam_candidates,
+    hparam_range,
+    rand,
+)
 from bigframes.bigquery._operations.search import create_vector_index, vector_search
 from bigframes.bigquery._operations.sql import sql_scalar
 from bigframes.bigquery._operations.struct import struct
@@ -130,6 +134,8 @@
     to_json,
     to_json_string,
     # mathematical ops
+    hparam_candidates,
+    hparam_range,
     rand,
     # search ops
     create_vector_index,
@@ -187,6 +193,8 @@
     "to_json",
     "to_json_string",
     # mathematical ops
+    "hparam_candidates",
+    "hparam_range",
     "rand",
     # search ops
     "create_vector_index",

packages/bigframes/bigframes/bigquery/_operations/mathematical.py

@@ -14,6 +14,8 @@
 
 from __future__ import annotations
 
+from typing import Sequence
+
 import bigframes.core.col
 import bigframes.core.expression
 from bigframes import dtypes
@@ -51,3 +53,71 @@ def rand() -> bigframes.core.col.Expression:
         is_deterministic=False,
     )
     return bigframes.core.col.Expression(bigframes.core.expression.OpExpression(op, ()))
+
+
+def hparam_range(min: float, max: float) -> bigframes.core.col.Expression:
+    """
+    Defines the minimum and maximum bounds of the search space of continuous
+    values for a hyperparameter.
+
+    **Examples:**
+
+        >>> import bigframes.pandas as bpd
+        >>> import bigframes.bigquery as bbq
+        >>> # Specify a range of values for a hyperparameter.
+        >>> learn_rate = bbq.hparam_range(0.0001, 1.0)
+
+    Args:
+        min (float or int):
+            The minimum bound of the search space.
+        max (float or int):
+            The maximum bound of the search space.
+
+    Returns:
+        bigframes.pandas.api.typing.Expression:
+            An expression that can be used in model options.
+    """
+    min_expr = bigframes.core.expression.const(min)
+    max_expr = bigframes.core.expression.const(max)
+
+    op = ops.SqlScalarOp(
+        _output_type=dtypes.FLOAT_DTYPE,
+        sql_template="HPARAM_RANGE({0}, {1})",
+        is_deterministic=True,
+    )
+    return bigframes.core.col.Expression(
+        bigframes.core.expression.OpExpression(op, (min_expr, max_expr))
+    )
+
+
+def hparam_candidates(
+    candidates: Sequence[float | str],
+) -> bigframes.core.col.Expression:
+    """
+    Specifies the set of discrete values for the hyperparameter.
+
+    **Examples:**
+
+        >>> import bigframes.pandas as bpd
+        >>> import bigframes.bigquery as bbq
+        >>> # Specify a set of values for a hyperparameter.
+        >>> optimizer = bbq.hparam_candidates(['ADAGRAD', 'SGD', 'FTRL'])
+
+    Args:
+        candidates (Sequence[float | str]):
+            The set of discrete values for the hyperparameter.
+
+    Returns:
+        bigframes.pandas.api.typing.Expression:
+            An expression that can be used in model options.
+    """
+    candidates_expr = bigframes.core.expression.const(tuple(candidates))
+
+    op = ops.SqlScalarOp(
+        _output_type=dtypes.STRING_DTYPE,
+        sql_template="HPARAM_CANDIDATES({0})",
+        is_deterministic=True,
+    )
+    return bigframes.core.col.Expression(
+        bigframes.core.expression.OpExpression(op, (candidates_expr,))
+    )

packages/bigframes/tests/unit/core/sql/snapshots/test_ml/test_create_model_expression_option/create_model_expression_option.sql

@@ -0,0 +1,3 @@
+CREATE MODEL `my_model`
+OPTIONS(model_type = 'LINEAR_REG', learn_rate = HPARAM_RANGE(0.0001, 1.0), optimizer = HPARAM_CANDIDATES(['ADAGRAD', 'SGD']))
+AS SELECT * FROM t

packages/bigframes/tests/unit/core/sql/snapshots/test_ml/test_create_model_hparam_tuning/create_model_hparam_tuning.sql

@@ -14,6 +14,7 @@
 
 import pytest
 
+import bigframes.bigquery as bbq
 import bigframes.core.col as col
 import bigframes.core.expression as ex
 import bigframes.core.sql.ml
@@ -101,24 +102,17 @@ def test_create_model_list_option(snapshot):
     snapshot.assert_match(sql, "create_model_list_option.sql")
 
 
-def test_create_model_expression_option(snapshot):
-    # An expression that calls a function on a literal value
-    # e.g. 0.1 * 10
-    literal_expr = ex.ScalarConstantExpression(0.1, dtypes.FLOAT_DTYPE)
-    multiplier_expr = ex.ScalarConstantExpression(10, dtypes.INT_DTYPE)
-    math_expr = col.Expression(
-        ex.OpExpression(op=numeric_ops.mul_op, inputs=(literal_expr, multiplier_expr))
-    )
-
+def test_create_model_hparam_tuning(snapshot):
     sql = bigframes.core.sql.ml.create_model_ddl(
         model_name="my_model",
         options={
-            "l2_reg": math_expr,
-            "booster_type": "gbtree",
+            "model_type": "LINEAR_REG",
+            "learn_rate": bbq.hparam_range(0.0001, 1.0),
+            "optimizer": bbq.hparam_candidates(["ADAGRAD", "SGD"]),
         },
         training_data="SELECT * FROM t",
     )
-    snapshot.assert_match(sql, "create_model_expression_option.sql")
+    snapshot.assert_match(sql, "create_model_hparam_tuning.sql")
 
 
 def test_evaluate_model_basic(snapshot):