From 79f7a605b7c3e65605ff7409a4834524eb083d53 Mon Sep 17 00:00:00 2001 From: ritchie Date: Wed, 4 Oct 2023 13:39:05 +0200 Subject: [PATCH] fix: literal hash --- crates/polars-plan/src/logical_plan/lit.rs | 24 +++++++++++++++---- py-polars/tests/unit/test_cse.py | 27 ++++++++++++++++++++++ 2 files changed, 47 insertions(+), 4 deletions(-) diff --git a/crates/polars-plan/src/logical_plan/lit.rs b/crates/polars-plan/src/logical_plan/lit.rs index 6a7fb87ad6f..0937a06ddab 100644 --- a/crates/polars-plan/src/logical_plan/lit.rs +++ b/crates/polars-plan/src/logical_plan/lit.rs @@ -312,10 +312,26 @@ pub fn lit(t: L) -> Expr { impl Hash for LiteralValue { fn hash(&self, state: &mut H) { - if let Some(v) = self.to_anyvalue() { - v.hash_impl(state, true) - } else { - 0.hash(state) + std::mem::discriminant(self).hash(state); + match self { + LiteralValue::Series(s) => { + s.dtype().hash(state); + s.len().hash(state); + }, + LiteralValue::Range { + low, + high, + data_type, + } => { + low.hash(state); + high.hash(state); + data_type.hash(state) + }, + _ => { + if let Some(v) = self.to_anyvalue() { + v.hash_impl(state, true) + } + }, } } } diff --git a/py-polars/tests/unit/test_cse.py b/py-polars/tests/unit/test_cse.py index e1efcdecdf4..f68d1eb5ce4 100644 --- a/py-polars/tests/unit/test_cse.py +++ b/py-polars/tests/unit/test_cse.py @@ -481,3 +481,30 @@ def test_no_cse_in_with_context() -> None: ], "label": [0, 1, 1], } + + +def test_cse_is_in_11489() -> None: + df = pl.DataFrame( + {"cond": [1, 2, 3, 2, 1], "x": [1.0, 0.20, 3.0, 4.0, 0.50]} + ).lazy() + any_cond = ( + pl.when(pl.col("cond").is_in([2, 3])) + .then(True) + .when(pl.col("cond").is_in([1])) + .then(False) + .otherwise(None) + .alias("any_cond") + ) + val = ( + pl.when(any_cond) + .then(1.0) + .when(~any_cond) + .then(0.0) + .otherwise(None) + .alias("val") + ) + assert df.select("cond", any_cond, val).collect().to_dict(False) == { + "cond": [1, 2, 3, 2, 1], + "any_cond": [False, True, True, True, False], + "val": [0.0, 1.0, 1.0, 1.0, 0.0], + }