Merge branch 'main' into main_chelsealin_addwindowtests

Author: chelsea-lin

bigframes/core/compile/polars/compiler.py

@@ -493,9 +493,9 @@ def compile_agg_op(
             if isinstance(op, agg_ops.MedianOp):
                 return pl.median(*inputs)
             if isinstance(op, agg_ops.AllOp):
-                return pl.all(*inputs)
+                return pl.col(inputs).cast(pl.Boolean).all()
             if isinstance(op, agg_ops.AnyOp):
-                return pl.any(*inputs)  # type: ignore
+                return pl.col(inputs).cast(pl.Boolean).any()
             if isinstance(op, agg_ops.NuniqueOp):
                 return pl.col(*inputs).drop_nulls().n_unique()
             if isinstance(op, agg_ops.MinOp):

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

@@ -23,6 +23,7 @@
 import bigframes.core.compile.sqlglot.scalar_compiler as scalar_compiler
 
 register_unary_op = scalar_compiler.scalar_op_compiler.register_unary_op
+register_ternary_op = scalar_compiler.scalar_op_compiler.register_ternary_op
 
 
 @register_unary_op(ops.AsTypeOp, pass_op=True)
@@ -66,6 +67,18 @@ def _(expr: TypedExpr, op: ops.AsTypeOp) -> sge.Expression:
     return _cast(sg_expr, sg_to_type, op.safe)
 
 
+@register_ternary_op(ops.clip_op)
+def _(
+    original: TypedExpr,
+    lower: TypedExpr,
+    upper: TypedExpr,
+) -> sge.Expression:
+    return sge.Greatest(
+        this=sge.Least(this=original.expr, expressions=[upper.expr]),
+        expressions=[lower.expr],
+    )
+
+
 @register_unary_op(ops.hash_op)
 def _(expr: TypedExpr) -> sge.Expression:
     return sge.func("FARM_FINGERPRINT", expr.expr)
@@ -94,6 +107,13 @@ def _(expr: TypedExpr) -> sge.Expression:
     return sge.Not(this=sge.Is(this=expr.expr, expression=sge.Null()))
 
 
+@register_ternary_op(ops.where_op)
+def _(
+    original: TypedExpr, condition: TypedExpr, replacement: TypedExpr
+) -> sge.Expression:
+    return sge.If(this=condition.expr, true=original.expr, false=replacement.expr)
+
+
 # Helper functions
 def _cast_to_json(expr: TypedExpr, op: ops.AsTypeOp) -> sge.Expression:
     from_type = expr.dtype

bigframes/operations/blob.py

@@ -228,9 +228,14 @@ def display(
         df._set_internal_query_job(query_job)
 
         def display_single_url(
-            read_url: str, content_type: Union[str, pd._libs.missing.NAType]
+            read_url: Union[str, pd._libs.missing.NAType],
+            content_type: Union[str, pd._libs.missing.NAType],
         ):
-            if content_type is pd.NA:  # display as raw data or error
+            if pd.isna(read_url):
+                ipy_display.display("<NA>")
+                return
+
+            if pd.isna(content_type):  # display as raw data or error
                 response = requests.get(read_url)
                 ipy_display.display(response.content)
                 return

bigframes/testing/polars_session.py

@@ -94,11 +94,24 @@ def __init__(self):
         self._loader = None  # type: ignore
 
     def read_pandas(self, pandas_dataframe, write_engine="default"):
+        original_input = pandas_dataframe
+
         # override read_pandas to always keep data local-only
-        if isinstance(pandas_dataframe, pandas.Series):
+        if isinstance(pandas_dataframe, (pandas.Series, pandas.Index)):
             pandas_dataframe = pandas_dataframe.to_frame()
+
         local_block = bigframes.core.blocks.Block.from_local(pandas_dataframe, self)
-        return bigframes.dataframe.DataFrame(local_block)
+        bf_df = bigframes.dataframe.DataFrame(local_block)
+
+        if isinstance(original_input, pandas.Series):
+            series = bf_df[bf_df.columns[0]]
+            series.name = original_input.name
+            return series
+
+        if isinstance(original_input, pandas.Index):
+            return bf_df.index
+
+        return bf_df
 
     @property
     def bqclient(self):

noxfile.py

@@ -46,9 +46,7 @@
     "3.11",
 ]
 
-# pytest-retry is not yet compatible with pytest 8.x.
-# https://github.com/str0zzapreti/pytest-retry/issues/32
-PYTEST_VERSION = "pytest<8.0.0dev"
+PYTEST_VERSION = "pytest==8.4.2"
 SPHINX_VERSION = "sphinx==4.5.0"
 LINT_PATHS = [
     "docs",
@@ -91,7 +89,7 @@
 # 3.10 is needed for Windows tests as it is the only version installed in the
 # bigframes-windows container image. For more information, search
 # bigframes/windows-docker, internally.
-SYSTEM_TEST_PYTHON_VERSIONS = ["3.9", "3.10", "3.11", "3.13"]
+SYSTEM_TEST_PYTHON_VERSIONS = ["3.9", "3.10", "3.11", "3.12", "3.13"]
 SYSTEM_TEST_STANDARD_DEPENDENCIES = [
     "jinja2",
     "mock",
@@ -115,7 +113,7 @@
     # Make sure we leave some versions without "extras" so we know those
     # dependencies are actually optional.
     "3.10": ["tests", "scikit-learn", "anywidget"],
-    "3.11": ["tests", "scikit-learn", "polars", "anywidget"],
+    LATEST_FULLY_SUPPORTED_PYTHON: ["tests", "scikit-learn", "polars", "anywidget"],
     "3.13": ["tests", "polars", "anywidget"],
 }
 
@@ -126,8 +124,13 @@
 # Sessions are executed in the order so putting the smaller sessions
 # ahead to fail fast at presubmit running.
 nox.options.sessions = [
+    # Include unit_noextras to ensure at least some unit tests contribute to
+    # coverage.
+    # TODO(tswast): Consider removing this when unit_noextras and cover is run
+    # from GitHub actions.
+    "unit_noextras",
     "system-3.9",  # No extras.
-    "system-3.11",
+    f"system-{LATEST_FULLY_SUPPORTED_PYTHON}",  # All extras.
     "cover",
     # TODO(b/401609005): remove
     "cleanup",

tests/system/large/blob/test_function.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import logging
 import os
 import traceback
 from typing import Generator
@@ -434,6 +435,15 @@ def test_blob_transcribe(
         actual_text = actual[0]["content"]
     else:
         actual_text = actual[0]
+
+    if pd.isna(actual_text) or actual_text == "":
+        # Ensure the tests are robust to flakes in the model, which isn't
+        # particularly useful information for the bigframes team.
+        logging.warning(
+            f"blob_transcribe() model {model_name} verbose={verbose} failure"
+        )
+        return
+
     actual_len = len(actual_text)
 
     relative_length_tolerance = 0.2

tests/system/small/blob/test_io.py

@@ -12,6 +12,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from unittest import mock
+
+import IPython.display
 import pandas as pd
 
 import bigframes
@@ -92,3 +95,33 @@ def test_blob_create_read_gbq_object_table(
     pd.testing.assert_frame_equal(
         pd_blob_df, expected_df, check_dtype=False, check_index_type=False
     )
+
+
+def test_display_images(monkeypatch, images_mm_df: bpd.DataFrame):
+    mock_display = mock.Mock()
+    monkeypatch.setattr(IPython.display, "display", mock_display)
+
+    images_mm_df["blob_col"].blob.display()
+
+    for call in mock_display.call_args_list:
+        args, _ = call
+        arg = args[0]
+        assert isinstance(arg, IPython.display.Image)
+
+
+def test_display_nulls(
+    monkeypatch,
+    bq_connection: str,
+    session: bigframes.Session,
+):
+    uri_series = bpd.Series([None, None, None], dtype="string", session=session)
+    blob_series = uri_series.str.to_blob(connection=bq_connection)
+    mock_display = mock.Mock()
+    monkeypatch.setattr(IPython.display, "display", mock_display)
+
+    blob_series.blob.display()
+
+    for call in mock_display.call_args_list:
+        args, _ = call
+        arg = args[0]
+        assert arg == "<NA>"

tests/system/small/engines/test_generic_ops.py

@@ -22,7 +22,7 @@
 from bigframes.session import polars_executor
 from bigframes.testing.engine_utils import assert_equivalence_execution
 
-pytest.importorskip("polars")
+polars = pytest.importorskip("polars")
 
 # Polars used as reference as its fast and local. Generally though, prefer gbq engine where they disagree.
 REFERENCE_ENGINE = polars_executor.PolarsExecutor()
@@ -54,6 +54,12 @@ def apply_op(
 
 @pytest.mark.parametrize("engine", ["polars", "bq", "bq-sqlglot"], indirect=True)
 def test_engines_astype_int(scalars_array_value: array_value.ArrayValue, engine):
+    polars_version = tuple([int(part) for part in polars.__version__.split(".")])
+    if polars_version >= (1, 34, 0):
+        # TODO(https://github.com/pola-rs/polars/issues/24841): Remove this when
+        # polars fixes Decimal to Int cast.
+        scalars_array_value = scalars_array_value.drop_columns(["numeric_col"])
+
     arr = apply_op(
         scalars_array_value,
         ops.AsTypeOp(to_type=bigframes.dtypes.INT_DTYPE),
@@ -308,7 +314,7 @@ def test_engines_astype_timedelta(scalars_array_value: array_value.ArrayValue, e
     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_where_op(scalars_array_value: array_value.ArrayValue, engine):
     arr, _ = scalars_array_value.compute_values(
         [

tests/unit/core/compile/sqlglot/expressions/snapshots/test_generic_ops/test_clip/out.sql

@@ -0,0 +1,15 @@
+WITH `bfcte_0` AS (
+  SELECT
+    `int64_col` AS `bfcol_0`,
+    `int64_too` AS `bfcol_1`,
+    `rowindex` AS `bfcol_2`
+  FROM `bigframes-dev`.`sqlglot_test`.`scalar_types`
+), `bfcte_1` AS (
+  SELECT
+    *,
+    GREATEST(LEAST(`bfcol_2`, `bfcol_1`), `bfcol_0`) AS `bfcol_3`
+  FROM `bfcte_0`
+)
+SELECT
+  `bfcol_3` AS `result_col`
+FROM `bfcte_1`

tests/unit/core/compile/sqlglot/expressions/snapshots/test_generic_ops/test_where/out.sql

@@ -0,0 +1,15 @@
+WITH `bfcte_0` AS (
+  SELECT
+    `bool_col` AS `bfcol_0`,
+    `int64_col` AS `bfcol_1`,
+    `float64_col` AS `bfcol_2`
+  FROM `bigframes-dev`.`sqlglot_test`.`scalar_types`
+), `bfcte_1` AS (
+  SELECT
+    *,
+    IF(`bfcol_0`, `bfcol_1`, `bfcol_2`) AS `bfcol_3`
+  FROM `bfcte_0`
+)
+SELECT
+  `bfcol_3` AS `result_col`
+FROM `bfcte_1`