Merge remote-tracking branch 'origin/main' into tswast-doctest-boilerplate

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

@@ -97,7 +97,7 @@ 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)

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