Book a Demo!
CoCalc Logo Icon
StoreFeaturesDocsShareSupportNewsAboutPoliciesSign UpSign In
pola-rs
GitHub Repository: pola-rs/polars
Path: blob/main/pyo3-polars/example/io_plugin/run.py
6939 views
1
import polars as pl
2
from io_plugin import new_bernoulli, new_uniform, scan_random
3
4
5
lf = scan_random(
6
[
7
new_bernoulli("b0.5", p=0.5, seed=1),
8
new_bernoulli("b0.1", p=0.1, seed=2),
9
new_uniform("uniform", low=10, high=100, dtype=pl.Int32, seed=2),
10
]
11
)
12
13
# predicate pushdown
14
assert lf.filter(pl.col("b0.5")).collect()["b0.5"].all()
15
16
# slice pushdown
17
assert lf.head(100).collect().shape == (100, 3)
18
19
# projection pushdown
20
out = lf.select("uniform", "b0.1").collect()
21
assert out.shape == (1000, 2)
22
assert out.columns == ["uniform", "b0.1"]
23
24