From cf814f731c20cd0dc1d09e1f924b17c32e4b1d02 Mon Sep 17 00:00:00 2001
From: ritchie <ritchie46@gmail.com>
Date: Mon, 23 Dec 2024 15:53:49 +0100
Subject: [PATCH] fix(python): Fix nullable object in map_elements

---
 crates/polars-python/src/map/series.rs        | 13 ++++-------
 py-polars/tests/unit/datatypes/test_object.py | 22 +++++++++++++++++++
 2 files changed, 26 insertions(+), 9 deletions(-)

diff --git a/crates/polars-python/src/map/series.rs b/crates/polars-python/src/map/series.rs
index c29d29a876f0..273567f9f70a 100644
--- a/crates/polars-python/src/map/series.rs
+++ b/crates/polars-python/src/map/series.rs
@@ -224,15 +224,10 @@ where
     S: FromPyObject<'py>,
 {
     let out = call_lambda(py, lambda, in_val)?;
-    match out.extract::<S>() {
-        Ok(s) => Ok(Some(s)),
-        Err(e) => {
-            if out.is_none() {
-                Ok(None)
-            } else {
-                Err(e)
-            }
-        },
+    if out.is_none() {
+        Ok(None)
+    } else {
+        out.extract::<S>().map(Some)
     }
 }
 
diff --git a/py-polars/tests/unit/datatypes/test_object.py b/py-polars/tests/unit/datatypes/test_object.py
index c8e024ec80dd..d540b217989a 100644
--- a/py-polars/tests/unit/datatypes/test_object.py
+++ b/py-polars/tests/unit/datatypes/test_object.py
@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import io
 from pathlib import Path
 from uuid import uuid4
@@ -75,6 +77,26 @@ def test_nullable_object_13538() -> None:
     }
 
 
+def test_nullable_object_17936() -> None:
+    class Custom:
+        value: int
+
+        def __init__(self, value: int) -> None:
+            self.value = value
+
+    def mapper(value: int) -> Custom | None:
+        if value == 2:
+            return None
+        return Custom(value)
+
+    df = pl.DataFrame({"a": [1, 2, 3]})
+
+    assert df.select(
+        pl.col("a").map_elements(mapper, return_dtype=pl.Object).alias("with_dtype"),
+        pl.col("a").map_elements(mapper).alias("without_dtype"),
+    ).null_count().row(0) == (1, 1)
+
+
 def test_empty_sort() -> None:
     df = pl.DataFrame(
         data=[