Merge branch 'main' into migrate-euclidean-dist

Author: jialuoo

.github/workflows/docs-deploy.yml

@@ -0,0 +1,57 @@
+name: Deploy docs to GitHub Pages
+
+on:
+  # Runs on pushes targeting the default branch
+  # TODO(tswast): Update this to only be releases once we confirm it's working.
+  push:
+    branches: ["main"]
+
+  # Allows you to run this workflow manually from the Actions tab
+  workflow_dispatch:
+
+# Sets permissions of the GITHUB_TOKEN to allow deployment to GitHub Pages
+permissions:
+  contents: read
+  pages: write
+  id-token: write
+
+# Allow only one concurrent deployment, skipping runs queued between the run in-progress and latest queued.
+# However, do NOT cancel in-progress runs as we want to allow these production deployments to complete.
+concurrency:
+  group: "pages"
+  cancel-in-progress: false
+
+jobs:
+  # Build job
+  build:
+    runs-on: ubuntu-latest
+    steps:
+    - name: Checkout
+      uses: actions/checkout@v4
+    - name: Setup Python
+      uses: actions/setup-python@v5
+      with:
+        python-version: "3.13"
+    - name: Install nox
+      run: |
+        python -m pip install --upgrade setuptools pip wheel
+        python -m pip install nox
+    - name: Run docs
+      run: |
+        nox -s docs
+    - name: Upload artifact
+      uses: actions/upload-pages-artifact@v3
+      with:
+        path: docs/_build/html/
+
+  # Deployment job
+  deploy:
+    environment:
+      name: github-pages
+      url: ${{ steps.deployment.outputs.page_url }}
+    runs-on: ubuntu-latest
+    needs: build
+    steps:
+      - name: Deploy to GitHub Pages
+        id: deployment
+        uses: actions/deploy-pages@v4

.github/workflows/docs.yml

@@ -2,6 +2,9 @@ on:
   pull_request:
     branches:
       - main
+  push:
+    branches:
+      - main
 name: docs
 jobs:
   docs:
@@ -12,7 +15,7 @@ jobs:
     - name: Setup Python
       uses: actions/setup-python@v5
       with:
-        python-version: "3.10"
+        python-version: "3.13"
     - name: Install nox
       run: |
         python -m pip install --upgrade setuptools pip wheel

.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:

README.rst

@@ -1,3 +1,5 @@
+:orphan:
+
 BigQuery DataFrames (BigFrames)
 ===============================
 

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

docs/conf.py

@@ -24,9 +24,11 @@
 # All configuration values have a default; values that are commented out
 # serve to show the default.
 
+from __future__ import annotations
+
 import os
-import shlex
 import sys
+from typing import Any
 
 # If extensions (or modules to document with autodoc) are in another directory,
 # add these directories to sys.path here. If the directory is relative to the
@@ -56,7 +58,7 @@
     "sphinx.ext.napoleon",
     "sphinx.ext.todo",
     "sphinx.ext.viewcode",
-    "recommonmark",
+    "myst_parser",
 ]
 
 # autodoc/autosummary flags
@@ -98,7 +100,7 @@
 #
 # This is also used if you do content translation via gettext catalogs.
 # Usually you set "language" from the command line for these cases.
-language = None
+language = "en-US"
 
 # There are two options for replacing |today|: either, you set today to some
 # non-false value, then it is used:
@@ -148,19 +150,27 @@
 
 # The theme to use for HTML and HTML Help pages.  See the documentation for
 # a list of builtin themes.
-html_theme = "alabaster"
+html_theme = "pydata_sphinx_theme"
 
 # Theme options are theme-specific and customize the look and feel of a theme
 # further.  For a list of options available for each theme, see the
 # documentation.
+# https://pydata-sphinx-theme.readthedocs.io/en/stable/user_guide/layout.html#references
 html_theme_options = {
-    "description": "BigQuery DataFrames provides DataFrame APIs on the BigQuery engine",
-    "github_user": "googleapis",
-    "github_repo": "python-bigquery-dataframes",
-    "github_banner": True,
-    "font_family": "'Roboto', Georgia, sans",
-    "head_font_family": "'Roboto', Georgia, serif",
-    "code_font_family": "'Roboto Mono', 'Consolas', monospace",
+    "github_url": "https://github.com/googleapis/python-bigquery-dataframes",
+    "logo": {
+        "text": "BigQuery DataFrames (BigFrames)",
+    },
+    "external_links": [
+        {
+            "name": "Getting started",
+            "url": "https://docs.cloud.google.com/bigquery/docs/dataframes-quickstart",
+        },
+        {
+            "name": "User guide",
+            "url": "https://docs.cloud.google.com/bigquery/docs/bigquery-dataframes-introduction",
+        },
+    ],
 }
 
 # Add any paths that contain custom themes here, relative to this directory.
@@ -264,7 +274,7 @@
 
 # -- Options for LaTeX output ---------------------------------------------
 
-latex_elements = {
+latex_elements: dict[str, Any] = {
     # The paper size ('letterpaper' or 'a4paper').
     #'papersize': 'letterpaper',
     # The font size ('10pt', '11pt' or '12pt').
@@ -282,7 +292,7 @@
     (
         root_doc,
         "bigframes.tex",
-        "bigframes Documentation",
+        "BigQuery DataFrames (BigFrames)",
         author,
         "manual",
     )
@@ -317,7 +327,7 @@
     (
         root_doc,
         "bigframes",
-        "bigframes Documentation",
+        "BigQuery DataFrames (BigFrames)",
         [author],
         1,
     )
@@ -336,7 +346,7 @@
     (
         root_doc,
         "bigframes",
-        "bigframes Documentation",
+        "BigQuery DataFrames (BigFrames)",
         author,
         "bigframes",
         "bigframes Library",