Skip to content

Commit 7a83224

Browse files
test: Add slice op cross validation testing (#1837)
1 parent c53c86f commit 7a83224

File tree

2 files changed

+57
-0
lines changed

2 files changed

+57
-0
lines changed

bigframes/session/polars_executor.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,7 @@
3131
nodes.OrderByNode,
3232
nodes.ReversedNode,
3333
nodes.SelectionNode,
34+
nodes.SliceNode,
3435
)
3536

3637
_COMPATIBLE_SCALAR_OPS = ()
Lines changed: 56 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,56 @@
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, nodes
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"], indirect=True)
28+
@pytest.mark.parametrize(
29+
("start", "stop", "step"),
30+
[
31+
(1, None, None),
32+
(None, 4, None),
33+
(None, None, 2),
34+
(None, 50_000_000_000, 1),
35+
(5, 4, None),
36+
(3, None, 2),
37+
(1, 7, 2),
38+
(1, 7, 50_000_000_000),
39+
(-1, -7, -2),
40+
(None, -7, -2),
41+
(-1, None, -2),
42+
(-7, -1, 2),
43+
(-7, -1, None),
44+
(-7, 7, None),
45+
(7, -7, -2),
46+
],
47+
)
48+
def test_engines_slice(
49+
scalars_array_value: array_value.ArrayValue,
50+
engine,
51+
start,
52+
stop,
53+
step,
54+
):
55+
node = nodes.SliceNode(scalars_array_value.node, start, stop, step)
56+
assert_equivalence_execution(node, REFERENCE_ENGINE, engine)

0 commit comments

Comments
 (0)