feat: make all and any compatible with integer columns on Polars session (#2154)

* docs: remove import bigframes.pandas as bpd boilerplate from many samples

Also, fixes several constructors that didn't take a session for
compatibility with multi-session applications.

* fix docs

* fix unit tests

* skip sklearn test

* fix snapshot

* plumb through session for from_tuples and from_arrays

* add from_frame

* make sure polars session isnt skipped on Kokoro

* fix apply doctest

* make doctest conftest available everywhere

* add python version flexibility for to_dict

* disambiguate explicit names

* disambiguate explicit name none versus no name

* fix for column name comparison in pandas bin op

* avoid setting column labels in special case of Series(block)

* revert doctest changes

* revert doctest changes

* revert df docstrings

* add polars series unit tests

* restore a test

* Revert "restore a test"

This reverts commit 765b678.

* skip null

* skip unsupported tests

* revert more docs changes

* revert more docs

* revert more docs

* fix unit tests python 3.13

* add test to reproduce name error

* revert new session methods

* fix TestSession read_pandas for Series

* revert more unnecessary changes

* even more

* add unit_noextras to improve code coverage

* run system tests on latest fully supported

* system-3.12 not found

* cap polars version

* hide progress bar

* relax polars upper pin

Author: tswast

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/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"],
 }
 
@@ -132,7 +130,7 @@
     # 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/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),

tests/unit/test_local_engine.py

@@ -42,6 +42,14 @@ def small_inline_frame() -> pd.DataFrame:
     return df
 
 
+def test_polars_local_engine_series(polars_session: bigframes.Session):
+    bf_series = bpd.Series([1, 2, 3], session=polars_session)
+    pd_series = pd.Series([1, 2, 3], dtype=bf_series.dtype)
+    bf_result = bf_series.to_pandas()
+    pd_result = pd_series
+    pandas.testing.assert_series_equal(bf_result, pd_result, check_index_type=False)
+
+
 def test_polars_local_engine_add(
     small_inline_frame: pd.DataFrame, polars_session: bigframes.Session
 ):