polars_quickstart_dataframe_filter_groupby_lazy_api.py

1import polars as pl
2from datetime import datetime
3
4# 1. Create a DataFrame
5df = pl.DataFrame(
6    {
7        "integer": [1, 2, 3, 4, 5],
8        "date": [
9            datetime(2025, 1, 1),
10            datetime(2025, 1, 2),
11            datetime(2025, 1, 3),
12            datetime(2025, 1, 4),
13            datetime(2025, 1, 5),
14        ],
15        "float": [4.0, 5.0, 6.0, 7.0, 8.0],
16        "string": ["a", "b", "c", "d", "e"],
17    }
18)
19
20print("--- Initial DataFrame ---")
21print(df)
22
23# 2. Expressions & Contexts (Filtering and Selecting)
24out = df.select(
25    pl.col("integer"),
26    pl.col("string"),
27).filter(
28    pl.col("integer") > 2
29)
30
31print("\n--- Filtered and Selected ---")
32print(out)
33
34# 3. Grouping and Aggregating
35out = df.group_by("string").agg(
36    pl.col("float").sum().alias("float_sum"),
37)
38
39print("\n--- Grouped and Aggregated ---")
40print(out)
41
42# 4. Lazy API (Combining operations for optimization)
43lazy_query = (
44    df.lazy()
45    .filter(pl.col("integer") > 2)
46    .select(pl.col("integer"), pl.col("float") * 2)
47)
48
49# Execute the query
50out_lazy = lazy_query.collect()
51
52print("\n--- Lazy Query Result ---")
53print(out_lazy)