Skip to content

Commit

Permalink
fix(python): Fix nullable object in map_elements
Browse files Browse the repository at this point in the history
  • Loading branch information
ritchie46 committed Dec 23, 2024
1 parent 5294cb5 commit cf814f7
Show file tree
Hide file tree
Showing 2 changed files with 26 additions and 9 deletions.
13 changes: 4 additions & 9 deletions crates/polars-python/src/map/series.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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)
}
}

Expand Down
22 changes: 22 additions & 0 deletions py-polars/tests/unit/datatypes/test_object.py
Original file line number Diff line number Diff line change
@@ -1,3 +1,5 @@
from __future__ import annotations

import io
from pathlib import Path
from uuid import uuid4
Expand Down Expand Up @@ -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=[
Expand Down

0 comments on commit cf814f7

Please sign in to comment.