Skip to content

Commit 92a2377

Browse files
feat: Add row numbering local pushdown in hybrid execution (#1932)
1 parent b47a124 commit 92a2377

File tree

2 files changed

+34
-0
lines changed

2 files changed

+34
-0
lines changed

bigframes/session/polars_executor.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -40,6 +40,7 @@
4040
nodes.ConcatNode,
4141
nodes.JoinNode,
4242
nodes.InNode,
43+
nodes.PromoteOffsetsNode,
4344
)
4445

4546
_COMPATIBLE_SCALAR_OPS = (
Lines changed: 33 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,33 @@
1+
# Copyright 2025 Google LLC
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
import pytest
16+
17+
from bigframes.core import array_value
18+
from bigframes.session import polars_executor
19+
from bigframes.testing.engine_utils import assert_equivalence_execution
20+
21+
pytest.importorskip("polars")
22+
23+
# Polars used as reference as its fast and local. Generally though, prefer gbq engine where they disagree.
24+
REFERENCE_ENGINE = polars_executor.PolarsExecutor()
25+
26+
27+
@pytest.mark.parametrize("engine", ["polars", "bq", "bq-sqlglot"], indirect=True)
28+
def test_engines_with_offsets(
29+
scalars_array_value: array_value.ArrayValue,
30+
engine,
31+
):
32+
result, _ = scalars_array_value.promote_offsets()
33+
assert_equivalence_execution(result.node, REFERENCE_ENGINE, engine)

0 commit comments

Comments
 (0)