test: Add dataframe unit test suite (#1751)

Author: TrevorBergeron

bigframes/core/compile/polars/compiler.py

@@ -112,3 +112,23 @@ def get_global_session():
 
 def with_default_session(func: Callable[..., _T], *args, **kwargs) -> _T:
     return func(get_global_session(), *args, **kwargs)
+
+
+class _GlobalSessionContext:
+    """
+    Context manager for testing that sets global session.
+    """
+
+    def __init__(self, session: bigframes.session.Session):
+        self._session = session
+
+    def __enter__(self):
+        global _global_session, _global_session_lock
+        with _global_session_lock:
+            self._previous_session = _global_session
+            _global_session = self._session
+
+    def __exit__(self, *exc_details):
+        global _global_session, _global_session_lock
+        with _global_session_lock:
+            _global_session = self._previous_session

bigframes/core/global_session.py

@@ -24,7 +24,7 @@
 )
 from bigframes.core.rewrite.slices import pull_out_limit, pull_up_limits, rewrite_slice
 from bigframes.core.rewrite.timedeltas import rewrite_timedelta_expressions
-from bigframes.core.rewrite.windows import rewrite_range_rolling
+from bigframes.core.rewrite.windows import pull_out_window_order, rewrite_range_rolling
 
 __all__ = [
     "legacy_join_as_projection",
@@ -41,4 +41,5 @@
     "bake_order",
     "try_reduce_to_local_scan",
     "fold_row_counts",
+    "pull_out_window_order",
 ]

bigframes/core/rewrite/__init__.py

@@ -17,7 +17,7 @@
 import dataclasses
 
 from bigframes import operations as ops
-from bigframes.core import nodes
+from bigframes.core import guid, identifiers, nodes, ordering
 
 
 def rewrite_range_rolling(node: nodes.BigFrameNode) -> nodes.BigFrameNode:
@@ -43,3 +43,34 @@ def rewrite_range_rolling(node: nodes.BigFrameNode) -> nodes.BigFrameNode:
         node,
         window_spec=dataclasses.replace(node.window_spec, ordering=(new_ordering,)),
     )
+
+
+def pull_out_window_order(root: nodes.BigFrameNode) -> nodes.BigFrameNode:
+    return root.bottom_up(rewrite_window_node)
+
+
+def rewrite_window_node(node: nodes.BigFrameNode) -> nodes.BigFrameNode:
+    if not isinstance(node, nodes.WindowOpNode):
+        return node
+    if len(node.window_spec.ordering) == 0:
+        return node
+    else:
+        offsets_id = guid.generate_guid()
+        w_offsets = nodes.PromoteOffsetsNode(
+            node.child, identifiers.ColumnId(offsets_id)
+        )
+        sorted_child = nodes.OrderByNode(w_offsets, node.window_spec.ordering)
+        new_window_node = dataclasses.replace(
+            node,
+            child=sorted_child,
+            window_spec=node.window_spec.without_order(force=True),
+        )
+        w_resetted_order = nodes.OrderByNode(
+            new_window_node,
+            by=(ordering.ascending_over(identifiers.ColumnId(offsets_id)),),
+            is_total_order=True,
+        )
+        w_offsets_dropped = nodes.SelectionNode(
+            w_resetted_order, tuple(nodes.AliasedRef.identity(id) for id in node.ids)
+        )
+        return w_offsets_dropped

bigframes/core/rewrite/windows.py

@@ -234,7 +234,9 @@ def is_row_bounded(self):
         This is relevant for determining whether the window requires a total order
         to calculate deterministically.
         """
-        return isinstance(self.bounds, RowsWindowBounds)
+        return isinstance(self.bounds, RowsWindowBounds) and (
+            (self.bounds.start is not None) or (self.bounds.end is not None)
+        )
 
     @property
     def is_range_bounded(self):
@@ -254,7 +256,9 @@ def is_unbounded(self):
         This is relevant for determining whether the window requires a total order
         to calculate deterministically.
         """
-        return self.bounds is None
+        return self.bounds is None or (
+            self.bounds.start is None and self.bounds.end is None
+        )
 
     @property
     def all_referenced_columns(self) -> Set[ids.ColumnId]:
@@ -266,9 +270,9 @@ def all_referenced_columns(self) -> Set[ids.ColumnId]:
         )
         return set(itertools.chain((i.id for i in self.grouping_keys), ordering_vars))
 
-    def without_order(self) -> WindowSpec:
+    def without_order(self, force: bool = False) -> WindowSpec:
         """Removes ordering clause if ordering isn't required to define bounds."""
-        if self.is_row_bounded:
+        if self.is_row_bounded and not force:
             raise ValueError("Cannot remove order from row-bounded window")
         return replace(self, ordering=())
 

bigframes/core/window_spec.py

@@ -439,7 +439,6 @@ def output_type(self, *input_types: dtypes.ExpressionType) -> dtypes.ExpressionT
         return dtypes.INT_DTYPE
 
 
-# TODO: Convert to NullaryWindowOp
 @dataclasses.dataclass(frozen=True)
 class RankOp(UnaryWindowOp):
     name: ClassVar[str] = "rank"
@@ -456,7 +455,6 @@ def implicitly_inherits_order(self):
         return False
 
 
-# TODO: Convert to NullaryWindowOp
 @dataclasses.dataclass(frozen=True)
 class DenseRankOp(UnaryWindowOp):
     @property

bigframes/operations/aggregations.py

@@ -20,12 +20,9 @@
 import polars
 
 import bigframes
-import bigframes.clients
 import bigframes.core.blocks
 import bigframes.core.compile.polars
-import bigframes.core.ordering
 import bigframes.dataframe
-import bigframes.session.clients
 import bigframes.session.executor
 import bigframes.session.metrics
 
@@ -35,6 +32,26 @@
 class TestExecutor(bigframes.session.executor.Executor):
     compiler = bigframes.core.compile.polars.PolarsCompiler()
 
+    def peek(
+        self,
+        array_value: bigframes.core.ArrayValue,
+        n_rows: int,
+        use_explicit_destination: Optional[bool] = False,
+    ):
+        """
+        A 'peek' efficiently accesses a small number of rows in the dataframe.
+        """
+        lazy_frame: polars.LazyFrame = self.compiler.compile(array_value)
+        pa_table = lazy_frame.collect().limit(n_rows).to_arrow()
+        # Currently, pyarrow types might not quite be exactly the ones in the bigframes schema.
+        # Nullability may be different, and might use large versions of list, string datatypes.
+        return bigframes.session.executor.ExecuteResult(
+            arrow_batches=pa_table.to_batches(),
+            schema=array_value.schema,
+            total_bytes=pa_table.nbytes,
+            total_rows=pa_table.num_rows,
+        )
+
     def execute(
         self,
         array_value: bigframes.core.ArrayValue,
@@ -58,6 +75,14 @@ def execute(
             total_rows=pa_table.num_rows,
         )
 
+    def cached(
+        self,
+        array_value: bigframes.core.ArrayValue,
+        *,
+        config,
+    ) -> None:
+        return
+
 
 class TestSession(bigframes.session.Session):
     def __init__(self):
@@ -92,3 +117,8 @@ def read_pandas(self, pandas_dataframe, write_engine="default"):
             pandas_dataframe = pandas_dataframe.to_frame()
         local_block = bigframes.core.blocks.Block.from_local(pandas_dataframe, self)
         return bigframes.dataframe.DataFrame(local_block)
+
+    @property
+    def bqclient(self):
+        # prevents logger from trying to call bq upon any errors
+        return None

bigframes/testing/polars_session.py

@@ -79,7 +79,7 @@
 UNIT_TEST_DEPENDENCIES: List[str] = []
 UNIT_TEST_EXTRAS: List[str] = ["tests"]
 UNIT_TEST_EXTRAS_BY_PYTHON: Dict[str, List[str]] = {
-    "3.12": ["polars", "scikit-learn"],
+    "3.12": ["tests", "polars", "scikit-learn"],
 }
 
 # 3.10 is needed for Windows tests as it is the only version installed in the
@@ -202,14 +202,11 @@ def install_unittest_dependencies(session, install_test_extra, *constraints):
     if UNIT_TEST_LOCAL_DEPENDENCIES:
         session.install(*UNIT_TEST_LOCAL_DEPENDENCIES, *constraints)
 
-    if install_test_extra and UNIT_TEST_EXTRAS_BY_PYTHON:
-        extras = UNIT_TEST_EXTRAS_BY_PYTHON.get(session.python, [])
-    if install_test_extra and UNIT_TEST_EXTRAS:
-        extras = UNIT_TEST_EXTRAS
-    else:
-        extras = []
-
-    if extras:
+    if install_test_extra:
+        if session.python in UNIT_TEST_EXTRAS_BY_PYTHON:
+            extras = UNIT_TEST_EXTRAS_BY_PYTHON[session.python]
+        else:
+            extras = UNIT_TEST_EXTRAS
         session.install("-e", f".[{','.join(extras)}]", *constraints)
     else:
         session.install("-e", ".", *constraints)