Merge branch 'main' into more_local_numerics

TrevorBergeron · web-flow · commit 91c32a9d76cf · 2025-08-11T13:02:10.000-07:00
diff --git a/GEMINI.md b/GEMINI.md
@@ -45,6 +45,8 @@ We use `nox` to instrument our tests.
   nox -r -s lint
   ```
 
+- When writing tests, use the idiomatic "pytest" style.
+
 ## Documentation
 
 If a method or property is implementing the same interface as a third-party
diff --git a/bigframes/core/compile/sqlglot/expressions/binary_compiler.py b/bigframes/core/compile/sqlglot/expressions/binary_compiler.py
@@ -14,6 +14,7 @@
 
 from __future__ import annotations
 
+import bigframes_vendored.constants as constants
 import sqlglot.expressions as sge
 
 from bigframes import dtypes
@@ -35,8 +36,83 @@ def _(op, left: TypedExpr, right: TypedExpr) -> sge.Expression:
         # String addition
         return sge.Concat(expressions=[left.expr, right.expr])
 
-    # Numerical addition
-    return sge.Add(this=left.expr, expression=right.expr)
+    if dtypes.is_numeric(left.dtype) and dtypes.is_numeric(right.dtype):
+        left_expr = left.expr
+        if left.dtype == dtypes.BOOL_DTYPE:
+            left_expr = sge.Cast(this=left_expr, to="INT64")
+        right_expr = right.expr
+        if right.dtype == dtypes.BOOL_DTYPE:
+            right_expr = sge.Cast(this=right_expr, to="INT64")
+        return sge.Add(this=left_expr, expression=right_expr)
+
+    if (
+        dtypes.is_time_or_date_like(left.dtype)
+        and right.dtype == dtypes.TIMEDELTA_DTYPE
+    ):
+        left_expr = left.expr
+        if left.dtype == dtypes.DATE_DTYPE:
+            left_expr = sge.Cast(this=left_expr, to="DATETIME")
+        return sge.TimestampAdd(
+            this=left_expr, expression=right.expr, unit=sge.Var(this="MICROSECOND")
+        )
+    if (
+        dtypes.is_time_or_date_like(right.dtype)
+        and left.dtype == dtypes.TIMEDELTA_DTYPE
+    ):
+        right_expr = right.expr
+        if right.dtype == dtypes.DATE_DTYPE:
+            right_expr = sge.Cast(this=right_expr, to="DATETIME")
+        return sge.TimestampAdd(
+            this=right_expr, expression=left.expr, unit=sge.Var(this="MICROSECOND")
+        )
+    if left.dtype == dtypes.TIMEDELTA_DTYPE and right.dtype == dtypes.TIMEDELTA_DTYPE:
+        return sge.Add(this=left.expr, expression=right.expr)
+
+    raise TypeError(
+        f"Cannot add type {left.dtype} and {right.dtype}. {constants.FEEDBACK_LINK}"
+    )
+
+
+@BINARY_OP_REGISTRATION.register(ops.sub_op)
+def _(op, left: TypedExpr, right: TypedExpr) -> sge.Expression:
+    if dtypes.is_numeric(left.dtype) and dtypes.is_numeric(right.dtype):
+        left_expr = left.expr
+        if left.dtype == dtypes.BOOL_DTYPE:
+            left_expr = sge.Cast(this=left_expr, to="INT64")
+        right_expr = right.expr
+        if right.dtype == dtypes.BOOL_DTYPE:
+            right_expr = sge.Cast(this=right_expr, to="INT64")
+        return sge.Sub(this=left_expr, expression=right_expr)
+
+    if (
+        dtypes.is_time_or_date_like(left.dtype)
+        and right.dtype == dtypes.TIMEDELTA_DTYPE
+    ):
+        left_expr = left.expr
+        if left.dtype == dtypes.DATE_DTYPE:
+            left_expr = sge.Cast(this=left_expr, to="DATETIME")
+        return sge.TimestampSub(
+            this=left_expr, expression=right.expr, unit=sge.Var(this="MICROSECOND")
+        )
+    if dtypes.is_time_or_date_like(left.dtype) and dtypes.is_time_or_date_like(
+        right.dtype
+    ):
+        left_expr = left.expr
+        if left.dtype == dtypes.DATE_DTYPE:
+            left_expr = sge.Cast(this=left_expr, to="DATETIME")
+        right_expr = right.expr
+        if right.dtype == dtypes.DATE_DTYPE:
+            right_expr = sge.Cast(this=right_expr, to="DATETIME")
+        return sge.TimestampDiff(
+            this=left_expr, expression=right_expr, unit=sge.Var(this="MICROSECOND")
+        )
+
+    if left.dtype == dtypes.TIMEDELTA_DTYPE and right.dtype == dtypes.TIMEDELTA_DTYPE:
+        return sge.Sub(this=left.expr, expression=right.expr)
+
+    raise TypeError(
+        f"Cannot subtract type {left.dtype} and {right.dtype}. {constants.FEEDBACK_LINK}"
+    )
 
 
 @BINARY_OP_REGISTRATION.register(ops.ge_op)
diff --git a/bigframes/dtypes.py b/bigframes/dtypes.py
@@ -289,6 +289,10 @@ def is_time_like(type_: ExpressionType) -> bool:
     return type_ in (DATETIME_DTYPE, TIMESTAMP_DTYPE, TIME_DTYPE)
 
 
+def is_time_or_date_like(type_: ExpressionType) -> bool:
+    return type_ in (DATE_DTYPE, DATETIME_DTYPE, TIME_DTYPE, TIMESTAMP_DTYPE)
+
+
 def is_geo_like(type_: ExpressionType) -> bool:
     return type_ in (GEO_DTYPE,)
 
diff --git a/bigframes/functions/_function_client.py b/bigframes/functions/_function_client.py
@@ -245,7 +245,7 @@ def provision_bq_managed_function(
 
         # Augment user package requirements with any internal package
         # requirements.
-        packages = _utils._get_updated_package_requirements(
+        packages = _utils.get_updated_package_requirements(
             packages, is_row_processor, capture_references, ignore_package_version=True
         )
         if packages:
@@ -258,7 +258,7 @@ def provision_bq_managed_function(
         bq_function_name = name
         if not bq_function_name:
             # Compute a unique hash representing the user code.
-            function_hash = _utils._get_hash(func, packages)
+            function_hash = _utils.get_hash(func, packages)
             bq_function_name = _utils.get_bigframes_function_name(
                 function_hash,
                 session_id,
@@ -539,12 +539,12 @@ def provision_bq_remote_function(
         """Provision a BigQuery remote function."""
         # Augment user package requirements with any internal package
         # requirements
-        package_requirements = _utils._get_updated_package_requirements(
+        package_requirements = _utils.get_updated_package_requirements(
             package_requirements, is_row_processor
         )
 
         # Compute a unique hash representing the user code
-        function_hash = _utils._get_hash(def_, package_requirements)
+        function_hash = _utils.get_hash(def_, package_requirements)
 
         # If reuse of any existing function with the same name (indicated by the
         # same hash of its source code) is not intended, then attach a unique
diff --git a/bigframes/functions/_function_session.py b/bigframes/functions/_function_session.py
@@ -597,7 +597,7 @@ def wrapper(func):
                 bqrf_metadata = _utils.get_bigframes_metadata(
                     python_output_type=py_sig.return_annotation
                 )
-                post_process_routine = _utils._build_unnest_post_routine(
+                post_process_routine = _utils.build_unnest_post_routine(
                     py_sig.return_annotation
                 )
                 py_sig = py_sig.replace(return_annotation=str)
diff --git a/bigframes/functions/_utils.py b/bigframes/functions/_utils.py
@@ -63,7 +63,7 @@ def get_remote_function_locations(bq_location):
     return bq_location, cloud_function_region
 
 
-def _get_updated_package_requirements(
+def get_updated_package_requirements(
     package_requirements=None,
     is_row_processor=False,
     capture_references=True,
@@ -105,7 +105,7 @@ def _get_updated_package_requirements(
     return requirements
 
 
-def _clean_up_by_session_id(
+def clean_up_by_session_id(
     bqclient: bigquery.Client,
     gcfclient: functions_v2.FunctionServiceClient,
     dataset: bigquery.DatasetReference,
@@ -169,7 +169,7 @@ def _clean_up_by_session_id(
             pass
 
 
-def _get_hash(def_, package_requirements=None):
+def get_hash(def_, package_requirements=None):
     "Get hash (32 digits alphanumeric) of a function."
     # There is a known cell-id sensitivity of the cloudpickle serialization in
     # notebooks https://github.com/cloudpipe/cloudpickle/issues/538. Because of
@@ -279,7 +279,7 @@ def get_python_version(is_compat: bool = False) -> str:
     return f"python{major}{minor}" if is_compat else f"python-{major}.{minor}"
 
 
-def _build_unnest_post_routine(py_list_type: type[list]):
+def build_unnest_post_routine(py_list_type: type[list]):
     sdk_type = function_typing.sdk_array_output_type_from_python_type(py_list_type)
     assert sdk_type.array_element_type is not None
     inner_sdk_type = sdk_type.array_element_type
diff --git a/bigframes/functions/function.py b/bigframes/functions/function.py
@@ -90,7 +90,7 @@ def _try_import_routine(
         return BigqueryCallableRoutine(
             udf_def,
             session,
-            post_routine=_utils._build_unnest_post_routine(override_type),
+            post_routine=_utils.build_unnest_post_routine(override_type),
         )
     return BigqueryCallableRoutine(udf_def, session, is_managed=not is_remote)
 
@@ -107,7 +107,7 @@ def _try_import_row_routine(
         return BigqueryCallableRowRoutine(
             udf_def,
             session,
-            post_routine=_utils._build_unnest_post_routine(override_type),
+            post_routine=_utils.build_unnest_post_routine(override_type),
         )
     return BigqueryCallableRowRoutine(udf_def, session, is_managed=not is_remote)
 
diff --git a/bigframes/pandas/__init__.py b/bigframes/pandas/__init__.py
@@ -293,7 +293,7 @@ def clean_up_by_session_id(
         session.bqclient, dataset, session_id
     )
 
-    bff_utils._clean_up_by_session_id(
+    bff_utils.clean_up_by_session_id(
         session.bqclient, session.cloudfunctionsclient, dataset, session_id
     )
 
diff --git a/bigframes/session/metrics.py b/bigframes/session/metrics.py
@@ -42,9 +42,9 @@ def count_job_stats(
             assert row_iterator is not None
 
             # TODO(tswast): Pass None after making benchmark publishing robust to missing data.
-            bytes_processed = getattr(row_iterator, "total_bytes_processed", 0)
-            query_char_count = len(getattr(row_iterator, "query", ""))
-            slot_millis = getattr(row_iterator, "slot_millis", 0)
+            bytes_processed = getattr(row_iterator, "total_bytes_processed", 0) or 0
+            query_char_count = len(getattr(row_iterator, "query", "") or "")
+            slot_millis = getattr(row_iterator, "slot_millis", 0) or 0
             exec_seconds = 0.0
 
             self.execution_count += 1
@@ -63,10 +63,10 @@ def count_job_stats(
         elif (stats := get_performance_stats(query_job)) is not None:
             query_char_count, bytes_processed, slot_millis, exec_seconds = stats
             self.execution_count += 1
-            self.query_char_count += query_char_count
-            self.bytes_processed += bytes_processed
-            self.slot_millis += slot_millis
-            self.execution_secs += exec_seconds
+            self.query_char_count += query_char_count or 0
+            self.bytes_processed += bytes_processed or 0
+            self.slot_millis += slot_millis or 0
+            self.execution_secs += exec_seconds or 0
             write_stats_to_disk(
                 query_char_count=query_char_count,
                 bytes_processed=bytes_processed,
diff --git a/bigframes/testing/utils.py b/bigframes/testing/utils.py
@@ -440,11 +440,11 @@ def get_function_name(func, package_requirements=None, is_row_processor=False):
     """Get a bigframes function name for testing given a udf."""
     # Augment user package requirements with any internal package
     # requirements.
-    package_requirements = bff_utils._get_updated_package_requirements(
+    package_requirements = bff_utils.get_updated_package_requirements(
         package_requirements, is_row_processor
     )
 
     # Compute a unique hash representing the user code.
-    function_hash = bff_utils._get_hash(func, package_requirements)
+    function_hash = bff_utils.get_hash(func, package_requirements)
 
     return f"bigframes_{function_hash}"
diff --git a/tests/system/large/functions/test_remote_function.py b/tests/system/large/functions/test_remote_function.py
@@ -527,8 +527,8 @@ def add_one(x):
         add_one_uniq, add_one_uniq_dir = make_uniq_udf(add_one)
 
         # Expected cloud function name for the unique udf
-        package_requirements = bff_utils._get_updated_package_requirements()
-        add_one_uniq_hash = bff_utils._get_hash(add_one_uniq, package_requirements)
+        package_requirements = bff_utils.get_updated_package_requirements()
+        add_one_uniq_hash = bff_utils.get_hash(add_one_uniq, package_requirements)
         add_one_uniq_cf_name = bff_utils.get_cloud_function_name(
             add_one_uniq_hash, session.session_id
         )
diff --git a/tests/system/small/engines/test_numeric_ops.py b/tests/system/small/engines/test_numeric_ops.py
@@ -53,7 +53,7 @@ def apply_op_pairwise(
     return new_arr
 
 
-@pytest.mark.parametrize("engine", ["polars", "bq"], indirect=True)
+@pytest.mark.parametrize("engine", ["polars", "bq", "bq-sqlglot"], indirect=True)
 def test_engines_project_add(
     scalars_array_value: array_value.ArrayValue,
     engine,
@@ -62,7 +62,7 @@ def test_engines_project_add(
     assert_equivalence_execution(arr.node, REFERENCE_ENGINE, engine)
 
 
-@pytest.mark.parametrize("engine", ["polars", "bq"], indirect=True)
+@pytest.mark.parametrize("engine", ["polars", "bq", "bq-sqlglot"], indirect=True)
 def test_engines_project_sub(
     scalars_array_value: array_value.ArrayValue,
     engine,
diff --git a/tests/unit/core/compile/sqlglot/expressions/snapshots/test_binary_compiler/test_add_numeric/out.sql b/tests/unit/core/compile/sqlglot/expressions/snapshots/test_binary_compiler/test_add_numeric/out.sql
@@ -1,13 +1,54 @@
 WITH `bfcte_0` AS (
   SELECT
-    `int64_col` AS `bfcol_0`
+    `bool_col` AS `bfcol_0`,
+    `int64_col` AS `bfcol_1`,
+    `rowindex` AS `bfcol_2`
   FROM `bigframes-dev`.`sqlglot_test`.`scalar_types`
 ), `bfcte_1` AS (
   SELECT
     *,
-    `bfcol_0` + `bfcol_0` AS `bfcol_1`
+    `bfcol_2` AS `bfcol_6`,
+    `bfcol_1` AS `bfcol_7`,
+    `bfcol_0` AS `bfcol_8`,
+    `bfcol_1` + `bfcol_1` AS `bfcol_9`
   FROM `bfcte_0`
+), `bfcte_2` AS (
+  SELECT
+    *,
+    `bfcol_6` AS `bfcol_14`,
+    `bfcol_7` AS `bfcol_15`,
+    `bfcol_8` AS `bfcol_16`,
+    `bfcol_9` AS `bfcol_17`,
+    `bfcol_7` + 1 AS `bfcol_18`
+  FROM `bfcte_1`
+), `bfcte_3` AS (
+  SELECT
+    *,
+    `bfcol_14` AS `bfcol_24`,
+    `bfcol_15` AS `bfcol_25`,
+    `bfcol_16` AS `bfcol_26`,
+    `bfcol_17` AS `bfcol_27`,
+    `bfcol_18` AS `bfcol_28`,
+    `bfcol_15` + CAST(`bfcol_16` AS INT64) AS `bfcol_29`
+  FROM `bfcte_2`
+), `bfcte_4` AS (
+  SELECT
+    *,
+    `bfcol_24` AS `bfcol_36`,
+    `bfcol_25` AS `bfcol_37`,
+    `bfcol_26` AS `bfcol_38`,
+    `bfcol_27` AS `bfcol_39`,
+    `bfcol_28` AS `bfcol_40`,
+    `bfcol_29` AS `bfcol_41`,
+    CAST(`bfcol_26` AS INT64) + `bfcol_25` AS `bfcol_42`
+  FROM `bfcte_3`
 )
 SELECT
-  `bfcol_1` AS `int64_col`
-FROM `bfcte_1`
+  `bfcol_36` AS `rowindex`,
+  `bfcol_37` AS `int64_col`,
+  `bfcol_38` AS `bool_col`,
+  `bfcol_39` AS `int_add_int`,
+  `bfcol_40` AS `int_add_1`,
+  `bfcol_41` AS `int_add_bool`,
+  `bfcol_42` AS `bool_add_int`
+FROM `bfcte_4`
diff --git a/tests/unit/core/compile/sqlglot/expressions/snapshots/test_binary_compiler/test_add_numeric_w_scalar/out.sql b/tests/unit/core/compile/sqlglot/expressions/snapshots/test_binary_compiler/test_add_numeric_w_scalar/out.sql
diff --git a/tests/unit/core/compile/sqlglot/expressions/snapshots/test_binary_compiler/test_add_timedelta/out.sql b/tests/unit/core/compile/sqlglot/expressions/snapshots/test_binary_compiler/test_add_timedelta/out.sql
diff --git a/tests/unit/core/compile/sqlglot/expressions/snapshots/test_binary_compiler/test_sub_numeric/out.sql b/tests/unit/core/compile/sqlglot/expressions/snapshots/test_binary_compiler/test_sub_numeric/out.sql
diff --git a/tests/unit/core/compile/sqlglot/expressions/snapshots/test_binary_compiler/test_sub_timedelta/out.sql b/tests/unit/core/compile/sqlglot/expressions/snapshots/test_binary_compiler/test_sub_timedelta/out.sql
diff --git a/tests/unit/core/compile/sqlglot/expressions/test_binary_compiler.py b/tests/unit/core/compile/sqlglot/expressions/test_binary_compiler.py
diff --git a/tests/unit/session/test_metrics.py b/tests/unit/session/test_metrics.py

Original file line number	Diff line number	Diff line change
`@@ -597,7 +597,7 @@ def wrapper(func):`
`597`	`597`	`bqrf_metadata = _utils.get_bigframes_metadata(`
`598`	`598`	`python_output_type=py_sig.return_annotation`
`599`	`599`	`)`
`600`		`- post_process_routine = _utils._build_unnest_post_routine(`
	`600`	`+ post_process_routine = _utils.build_unnest_post_routine(`
`601`	`601`	`py_sig.return_annotation`
`602`	`602`	`)`
`603`	`603`	`py_sig = py_sig.replace(return_annotation=str)`
Original file line number	Diff line number	Diff line change
`@@ -293,7 +293,7 @@ def clean_up_by_session_id(`
`293`	`293`	`session.bqclient, dataset, session_id`
`294`	`294`	`)`
`295`	`295`
`296`		`- bff_utils._clean_up_by_session_id(`
	`296`	`+ bff_utils.clean_up_by_session_id(`
`297`	`297`	`session.bqclient, session.cloudfunctionsclient, dataset, session_id`
`298`	`298`	`)`
`299`	`299`