From 8491f510766b4e9347f94dcc6c04162e0b5c2f5d Mon Sep 17 00:00:00 2001 From: Trevor Bergeron Date: Wed, 23 Jul 2025 00:05:13 +0000 Subject: [PATCH] feat: Add row numbering local pushdown in hybrid execution --- bigframes/session/polars_executor.py | 1 + tests/system/small/engines/test_windowing.py | 33 ++++++++++++++++++++ 2 files changed, 34 insertions(+) create mode 100644 tests/system/small/engines/test_windowing.py diff --git a/bigframes/session/polars_executor.py b/bigframes/session/polars_executor.py index 9b2346a7ed..2c04a0016b 100644 --- a/bigframes/session/polars_executor.py +++ b/bigframes/session/polars_executor.py @@ -40,6 +40,7 @@ nodes.ConcatNode, nodes.JoinNode, nodes.InNode, + nodes.PromoteOffsetsNode, ) _COMPATIBLE_SCALAR_OPS = ( diff --git a/tests/system/small/engines/test_windowing.py b/tests/system/small/engines/test_windowing.py new file mode 100644 index 0000000000..f4c2b61e6f --- /dev/null +++ b/tests/system/small/engines/test_windowing.py @@ -0,0 +1,33 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import pytest + +from bigframes.core import array_value +from bigframes.session import polars_executor +from bigframes.testing.engine_utils import assert_equivalence_execution + +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() + + +@pytest.mark.parametrize("engine", ["polars", "bq", "bq-sqlglot"], indirect=True) +def test_engines_with_offsets( + scalars_array_value: array_value.ArrayValue, + engine, +): + result, _ = scalars_array_value.promote_offsets() + assert_equivalence_execution(result.node, REFERENCE_ENGINE, engine)