Merge remote-tracking branch 'github/main' into crosstab

Author: TrevorBergeron

.github/workflows/docs.yml

@@ -2,6 +2,9 @@ on:
   pull_request:
     branches:
       - main
+  push:
+    branches:
+      - main
 name: docs
 jobs:
   docs:

.github/workflows/lint.yml

@@ -2,6 +2,9 @@ on:
   pull_request:
     branches:
       - main
+  push:
+    branches:
+      - main
 name: lint
 jobs:
   lint:

.github/workflows/mypy.yml

@@ -2,6 +2,9 @@ on:
   pull_request:
     branches:
       - main
+  push:
+    branches:
+      - main
 name: mypy
 jobs:
   mypy:

.github/workflows/unittest.yml

@@ -2,6 +2,9 @@ on:
   pull_request:
     branches:
       - main
+  push:
+    branches:
+      - main
 name: unittest
 jobs:
   unit:

bigframes/core/compile/sqlglot/expressions/geo_ops.py

@@ -21,6 +21,7 @@
 import bigframes.core.compile.sqlglot.scalar_compiler as scalar_compiler
 
 register_unary_op = scalar_compiler.scalar_op_compiler.register_unary_op
+register_binary_op = scalar_compiler.scalar_op_compiler.register_binary_op
 
 
 @register_unary_op(ops.geo_area_op)
@@ -108,3 +109,8 @@ def _(expr: TypedExpr) -> sge.Expression:
 @register_unary_op(ops.geo_y_op)
 def _(expr: TypedExpr) -> sge.Expression:
     return sge.func("SAFE.ST_Y", expr.expr)
+
+
+@register_binary_op(ops.geo_st_difference_op)
+def _(left: TypedExpr, right: TypedExpr) -> sge.Expression:
+    return sge.func("ST_DIFFERENCE", left.expr, right.expr)

bigframes/core/compile/sqlglot/expressions/numeric_ops.py

@@ -77,6 +77,13 @@ def _(expr: TypedExpr) -> sge.Expression:
     return sge.func("ASINH", expr.expr)
 
 
+@register_binary_op(ops.arctan2_op)
+def _(left: TypedExpr, right: TypedExpr) -> sge.Expression:
+    left_expr = _coerce_bool_to_int(left)
+    right_expr = _coerce_bool_to_int(right)
+    return sge.func("ATAN2", left_expr, right_expr)
+
+
 @register_unary_op(ops.arctan_op)
 def _(expr: TypedExpr) -> sge.Expression:
     return sge.func("ATAN", expr.expr)
@@ -118,6 +125,18 @@ def _(expr: TypedExpr) -> sge.Expression:
     )
 
 
+@register_binary_op(ops.cosine_distance_op)
+def _(left: TypedExpr, right: TypedExpr) -> sge.Expression:
+    return sge.Anonymous(
+        this="ML.DISTANCE",
+        expressions=[
+            left.expr,
+            right.expr,
+            sge.Literal.string("COSINE"),
+        ],
+    )
+
+
 @register_unary_op(ops.exp_op)
 def _(expr: TypedExpr) -> sge.Expression:
     return sge.Case(

bigframes/core/local_data.py

@@ -253,9 +253,16 @@ def _(
         value_generator = iter_array(
             array.flatten(), bigframes.dtypes.get_array_inner_type(dtype)
         )
-        for (start, end) in _pairwise(array.offsets):
-            arr_size = end.as_py() - start.as_py()
-            yield list(itertools.islice(value_generator, arr_size))
+        offset_generator = iter_array(array.offsets, bigframes.dtypes.INT_DTYPE)
+
+        start_offset = None
+        end_offset = None
+        for offset in offset_generator:
+            start_offset = end_offset
+            end_offset = offset
+            if start_offset is not None:
+                arr_size = end_offset - start_offset
+                yield list(itertools.islice(value_generator, arr_size))
 
     @iter_array.register
     def _(
@@ -267,8 +274,15 @@ def _(
             sub_generators[field_name] = iter_array(array.field(field_name), dtype)
 
         keys = list(sub_generators.keys())
-        for row_values in zip(*sub_generators.values()):
-            yield {key: value for key, value in zip(keys, row_values)}
+        is_null_generator = iter_array(array.is_null(), bigframes.dtypes.BOOL_DTYPE)
+
+        for values in zip(is_null_generator, *sub_generators.values()):
+            is_row_null = values[0]
+            row_values = values[1:]
+            if not is_row_null:
+                yield {key: value for key, value in zip(keys, row_values)}
+            else:
+                yield None
 
     for batch in table.to_batches():
         sub_generators: dict[str, Generator[Any, None, None]] = {}
@@ -491,16 +505,3 @@ def _schema_durations_to_ints(schema: pa.Schema) -> pa.Schema:
     return pa.schema(
         pa.field(field.name, _durations_to_ints(field.type)) for field in schema
     )
-
-
-def _pairwise(iterable):
-    do_yield = False
-    a = None
-    b = None
-    for item in iterable:
-        a = b
-        b = item
-        if do_yield:
-            yield (a, b)
-        else:
-            do_yield = True

owlbot.py

@@ -58,8 +58,9 @@
         ".kokoro/build.sh",
         ".kokoro/continuous/common.cfg",
         ".kokoro/presubmit/common.cfg",
-        # Temporary workaround to update docs job to use python 3.10
         ".github/workflows/docs.yml",
+        ".github/workflows/lint.yml",
+        ".github/workflows/unittest.yml",
     ],
 )
 

tests/system/small/engines/test_read_local.py

@@ -88,8 +88,9 @@ def test_engines_read_local_w_zero_row_source(
     assert_equivalence_execution(local_node, REFERENCE_ENGINE, engine)
 
 
-# TODO: Fix sqlglot impl
-@pytest.mark.parametrize("engine", ["polars", "bq", "pyarrow"], indirect=True)
+@pytest.mark.parametrize(
+    "engine", ["polars", "bq", "pyarrow", "bq-sqlglot"], indirect=True
+)
 def test_engines_read_local_w_nested_source(
     fake_session: bigframes.Session,
     nested_data_source: local_data.ManagedArrowTable,

tests/unit/core/compile/sqlglot/expressions/snapshots/test_geo_ops/test_geo_st_difference/out.sql

@@ -0,0 +1,13 @@
+WITH `bfcte_0` AS (
+  SELECT
+    `geography_col` AS `bfcol_0`
+  FROM `bigframes-dev`.`sqlglot_test`.`scalar_types`
+), `bfcte_1` AS (
+  SELECT
+    *,
+    ST_DIFFERENCE(`bfcol_0`, `bfcol_0`) AS `bfcol_1`
+  FROM `bfcte_0`
+)
+SELECT
+  `bfcol_1` AS `geography_col`
+FROM `bfcte_1`