cleanup recursion limit stuff

Author: TrevorBergeron

bigframes/core/blocks.py

@@ -140,6 +140,7 @@ def __init__(
         column_labels: typing.Union[pd.Index, typing.Iterable[Label]],
         index_labels: typing.Union[pd.Index, typing.Iterable[Label], None] = None,
         *,
+        value_columns: Optional[Iterable[str]] = None,
         transpose_cache: Optional[Block] = None,
     ):
         """Construct a block object, will create default index if no index columns specified."""
@@ -158,7 +159,13 @@ def __init__(
             if index_labels
             else tuple([None for _ in index_columns])
         )
-        self._expr = self._normalize_expression(expr, self._index_columns)
+        if value_columns is None:
+            value_columns = [
+                col_id for col_id in expr.column_ids if col_id not in index_columns
+            ]
+        self._expr = self._normalize_expression(
+            expr, self._index_columns, value_columns
+        )
         # Use pandas index to more easily replicate column indexing, especially for hierarchical column index
         self._column_labels = (
             column_labels.copy()
@@ -1114,13 +1121,15 @@ def project_exprs(
         labels: Union[Sequence[Label], pd.Index],
         drop=False,
     ) -> Block:
-        new_array, _ = self.expr.compute_values(exprs)
+        new_array, new_cols = self.expr.compute_values(exprs)
         if drop:
             new_array = new_array.drop_columns(self.value_columns)
 
+        new_val_cols = new_cols if drop else (*self.value_columns, *new_cols)
         return Block(
             new_array,
             index_columns=self.index_columns,
+            value_columns=new_val_cols,
             column_labels=labels
             if drop
             else self.column_labels.append(pd.Index(labels)),
@@ -1542,17 +1551,13 @@ def _get_labels_for_columns(self, column_ids: typing.Sequence[str]) -> pd.Index:
     def _normalize_expression(
         self,
         expr: core.ArrayValue,
-        index_columns: typing.Sequence[str],
-        assert_value_size: typing.Optional[int] = None,
+        index_columns: Iterable[str],
+        value_columns: Iterable[str],
     ):
         """Normalizes expression by moving index columns to left."""
-        value_columns = [
-            col_id for col_id in expr.column_ids if col_id not in index_columns
-        ]
-        if (assert_value_size is not None) and (
-            len(value_columns) != assert_value_size
-        ):
-            raise ValueError("Unexpected number of value columns.")
+        normalized_ids = (*index_columns, *value_columns)
+        if tuple(expr.column_ids) == normalized_ids:
+            return expr
         return expr.select_columns([*index_columns, *value_columns])
 
     def grouped_head(

bigframes/pandas/__init__.py

@@ -20,7 +20,6 @@
 import datetime
 import inspect
 import sys
-import traceback
 import typing
 from typing import Literal, Optional, Sequence, Union
 
@@ -370,19 +369,6 @@ def reset_session():
 # https://github.com/python/cpython/issues/112282
 sys.setrecursionlimit(max(10000000, sys.getrecursionlimit()))
 
-
-original_setrecursionlimit = sys.setrecursionlimit
-print(f"recursion limit set to {sys.getrecursionlimit()}")
-
-
-def patched_setrecursionlimit(n):
-    print(f"\n[DEBUG] Recursion limit being set to {n} by:")
-    traceback.print_stack()
-    original_setrecursionlimit(n)
-
-
-sys.setrecursionlimit = patched_setrecursionlimit
-
 if resource is not None:
     soft_limit, hard_limit = resource.getrlimit(resource.RLIMIT_STACK)
     if soft_limit < hard_limit or hard_limit == resource.RLIM_INFINITY:

tests/system/small/test_dataframe.py

@@ -5755,30 +5755,10 @@ def test_df_dot_operator_series(
 
 
 def test_recursion_limit(scalars_df_index):
-    import sys
-
-    print(f"doing recursion test, recursion limit set to {sys.getrecursionlimit()}")
     scalars_df_index = scalars_df_index[["int64_too", "int64_col", "float64_col"]]
     for i in range(400):
         scalars_df_index = scalars_df_index + 4
-    try:
-        scalars_df_index.to_pandas()
-    except Exception:
-
-        try:
-            import resource
-        except ImportError:
-            # resource is only available on Unix-like systems.
-            # https://docs.python.org/3/library/resource.html
-            resource = None  # type: ignore
-        print(f"recursion limit: {sys.getrecursionlimit()}")
-        if resource is not None:
-            soft_limit, hard_limit = resource.getrlimit(resource.RLIMIT_STACK)
-            print(f"stack limits: {soft_limit}, {hard_limit}")
-        else:
-            print("resource module not available")
-        raise
-    assert False
+    scalars_df_index.to_pandas()
 
 
 @pytest.mark.skipif(

tests/unit/test_dataframe_polars.py

@@ -4453,27 +4453,7 @@ def test_dataframe_explode_xfail(col_names):
 
 
 def test_recursion_limit_unit(scalars_df_index):
-    import sys
-
-    print(f"doing recursion test, recursion limit set to {sys.getrecursionlimit()}")
     scalars_df_index = scalars_df_index[["int64_too", "int64_col", "float64_col"]]
     for i in range(400):
         scalars_df_index = scalars_df_index + 4
-    try:
-        scalars_df_index.to_pandas()
-    except Exception:
-
-        try:
-            import resource
-        except ImportError:
-            # resource is only available on Unix-like systems.
-            # https://docs.python.org/3/library/resource.html
-            resource = None  # type: ignore
-        print(f"recursion limit: {sys.getrecursionlimit()}")
-        if resource is not None:
-            soft_limit, hard_limit = resource.getrlimit(resource.RLIMIT_STACK)
-            print(f"stack limits: {soft_limit}, {hard_limit}")
-        else:
-            print("resource module not available")
-        raise
-    assert False
+    scalars_df_index.to_pandas()