add logical type example

pola-rs · Oct 17, 2023 · fca88bb · fca88bb
1 parent f3548ec
commit fca88bb
Show file tree

Hide file tree

Showing 6 changed files with 45 additions and 6 deletions.
diff --git a/Cargo.toml b/Cargo.toml
@@ -8,8 +8,14 @@ members = [
 ]
 
 [workspace.dependencies]
-polars = { git = "https://github.com/pola-rs/polars", rev = "5d48cc800bc9c71fe6d4ff97b96d7fed4601793b", version = "0.33.2", default-features = false }
-polars-core = { git = "https://github.com/pola-rs/polars", rev = "5d48cc800bc9c71fe6d4ff97b96d7fed4601793b", version = "0.33.2", default-features = false }
-polars-ffi = { git = "https://github.com/pola-rs/polars", rev = "5d48cc800bc9c71fe6d4ff97b96d7fed4601793b", version = "0.33.2", default-features = false }
-polars-plan = { git = "https://github.com/pola-rs/polars", rev = "5d48cc800bc9c71fe6d4ff97b96d7fed4601793b", version = "0.33.2", default-feautres = false }
-polars-lazy = { git = "https://github.com/pola-rs/polars", rev = "5d48cc800bc9c71fe6d4ff97b96d7fed4601793b", version = "0.33.2", default-features = false }
+#polars = { git = "https://github.com/pola-rs/polars", rev = "5d48cc800bc9c71fe6d4ff97b96d7fed4601793b", version = "0.33.2", default-features = false }
+#polars-core = { git = "https://github.com/pola-rs/polars", rev = "5d48cc800bc9c71fe6d4ff97b96d7fed4601793b", version = "0.33.2", default-features = false }
+#polars-ffi = { git = "https://github.com/pola-rs/polars", rev = "5d48cc800bc9c71fe6d4ff97b96d7fed4601793b", version = "0.33.2", default-features = false }
+#polars-plan = { git = "https://github.com/pola-rs/polars", rev = "5d48cc800bc9c71fe6d4ff97b96d7fed4601793b", version = "0.33.2", default-feautres = false }
+#polars-lazy = { git = "https://github.com/pola-rs/polars", rev = "5d48cc800bc9c71fe6d4ff97b96d7fed4601793b", version = "0.33.2", default-features = false }
+
+polars = { path = "../polars/crates/polars", version = "0.33.2", default-features = false }
+polars-core = { path = "../polars/crates/polars-core", version = "0.33.2", default-features = false }
+polars-ffi = { path = "../polars/crates/polars-ffi", version = "0.33.2", default-features = false }
+polars-plan = { path = "../polars/crates/polars-plan", version = "0.33.2", default-feautres = false }
+polars-lazy = { path = "../polars/crates/polars-lazy", version = "0.33.2", default-features = false }
diff --git a/example/derive_expression/expression_lib/Cargo.toml b/example/derive_expression/expression_lib/Cargo.toml
@@ -10,7 +10,7 @@ crate-type = ["cdylib"]
 
 [dependencies]
 jemallocator = { version = "0.5", features = ["disable_initial_exec_tls"] }
-polars = { workspace = true, features = ["fmt"], default-features = false }
+polars = { workspace = true, features = ["fmt", "dtype-date"], default-features = false }
 polars-plan = { workspace = true, default-features = false }
 pyo3 = { version = "0.20.0", features = ["extension-module"] }
 pyo3-polars = { version = "*", path = "../../../pyo3-polars", features = ["derive"] }

diff --git a/example/derive_expression/expression_lib/expression_lib/__init__.py b/example/derive_expression/expression_lib/expression_lib/__init__.py
@@ -76,3 +76,16 @@ def haversine(
             is_elementwise=True,
             cast_to_supertypes=True,
         )
+
+@pl.api.register_expr_namespace("date_util")
+class DateUtil:
+    def __init__(self, expr: pl.Expr):
+        self._expr = expr
+
+
+    def is_leap_year(self) -> pl.Expr:
+        return self._expr._register_plugin(
+            lib=lib,
+            symbol="is_leap_year",
+            is_elementwise=True,
+        )
diff --git a/example/derive_expression/expression_lib/src/expressions.rs b/example/derive_expression/expression_lib/src/expressions.rs
@@ -90,3 +90,17 @@ fn append_kwargs(input: &[Series], kwargs: Option<MyKwargs>) -> PolarsResult<Ser
         })
         .into_series())
 }
+
+#[polars_expr(output_type=Boolean)]
+fn is_leap_year(input: &[Series], _kwargs: Option<DefaultKwargs>) -> PolarsResult<Series> {
+    let input = &input[0];
+    let ca = input.date()?;
+
+    let out: BooleanChunked = ca.as_date_iter().map(|opt_dt| {
+        opt_dt.map(|dt| {
+            dt.leap_year()
+        })
+    }).collect_ca(ca.name());
+
+    Ok(out.into_series())
+}
diff --git a/example/derive_expression/expression_lib/src/lib.rs b/example/derive_expression/expression_lib/src/lib.rs
@@ -1,6 +1,9 @@
 mod distances;
 mod expressions;
 
+#[cfg(target_os = "linux")]
+use jemallocator::Jemalloc;
+
 #[global_allocator]
 #[cfg(target_os = "linux")]
 static ALLOC: Jemalloc = Jemalloc;
diff --git a/example/derive_expression/run.py b/example/derive_expression/run.py
@@ -1,10 +1,12 @@
 import polars as pl
 from expression_lib import Language, Distance
+from datetime import date
 
 df = pl.DataFrame(
     {
         "names": ["Richard", "Alice", "Bob"],
         "moons": ["full", "half", "red"],
+        "dates": [date(2023, 1, 1), date(2024, 1, 1), date(2025, 1, 1)],
         "dist_a": [[12, 32, 1], [], [1, -2]],
         "dist_b": [[-12, 1], [43], [876, -45, 9]],
         "floats": [5.6, -1245.8, 242.224],
@@ -18,6 +20,7 @@
     hamming_dist=pl.col("names").dist.hamming_distance("pig_latin"),
     jaccard_sim=pl.col("dist_a").dist.jaccard_similarity("dist_b"),
     haversine=pl.col("floats").dist.haversine("floats", "floats", "floats", "floats"),
+    leap_year=pl.col("dates").date_util.is_leap_year(),
     appended_args=pl.col("names").language.append_args(
         float_arg=11.234,
         integer_arg=93,