Skip to content

Commit a9621ee

Browse files
committed
Minor refactor
1 parent c3211db commit a9621ee

1 file changed

Lines changed: 3 additions & 8 deletions

File tree

examples/datafusion-ffi-example/python/tests/_test_physical_optimizer_rule.py

Lines changed: 3 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -22,21 +22,16 @@
2222
from datafusion_ffi_example import MyPhysicalOptimizerRule
2323

2424

25-
def _setup_session_with_rule() -> tuple[SessionContext, MyPhysicalOptimizerRule]:
25+
def test_ffi_physical_optimizer_rule_runs_during_planning():
26+
"""A rule supplied via physical_optimizer_rules is invoked while the
27+
physical plan is built, and the query still returns correct results."""
2628
rule = MyPhysicalOptimizerRule()
2729
ctx = SessionContext(physical_optimizer_rules=[rule])
2830
batch = pa.RecordBatch.from_arrays(
2931
[pa.array([1, 2, 3])],
3032
names=["a"],
3133
)
3234
ctx.register_record_batches("t", [[batch]])
33-
return ctx, rule
34-
35-
36-
def test_ffi_physical_optimizer_rule_runs_during_planning():
37-
"""A rule supplied via physical_optimizer_rules is invoked while the
38-
physical plan is built, and the query still returns correct results."""
39-
ctx, rule = _setup_session_with_rule()
4035

4136
before = rule.optimize_calls()
4237
result = ctx.sql("SELECT a FROM t").collect()

0 commit comments

Comments
 (0)