Repository: pola-rs/polars (language: python)

Issue:
Type coercion of inequality joins propagates wrong types to literals in expressions referencing columns of right table
### Checks

- [X] I have checked that this issue has not already been reported.
- [X] I have confirmed this bug exists on the [latest version](https://pypi.org/project/polars/) of Polars.

### Reproducible example

```python
import polars as pl

left = pl.LazyFrame({"a": pl.Series([1, 2, 3], dtype=pl.Float64)})

right = pl.LazyFrame({"a": pl.Series([1, 2, 3], dtype=pl.Int64), "b": pl.Series([1, 2, 3], dtype=pl.Int8)})

q_bad_1 = left.join_where(right, pl.col("a") <= pl.col("a_right") * 2)

q_bad_2 = left.join_where(right, pl.col("a") <= pl.col("b") * 2)

print(q_bad_1.explain())

print(q_bad_2.explain())
```


### Log output

_No response_

### Issue description

Output of those prints is:

```
IEJOIN JOIN:
LEFT PLAN ON: [col("a")]
  DF ["a"]; PROJECT */1 COLUMNS; SELECTION: None
RIGHT PLAN ON: [[(col("a")) * (2.0)]]
  DF ["a", "b"]; PROJECT */2 COLUMNS; SELECTION: None
END IEJOIN JOIN
```

And
```
IEJOIN JOIN:
LEFT PLAN ON: [col("a")]
  DF ["a"]; PROJECT */1 COLUMNS; SELECTION: None
RIGHT PLAN ON: [[(col("b")) * (dyn int: 2)]]
  DF ["a", "b"]; PROJECT */2 COLUMNS; SELECTION: None
END IEJOIN JOIN
```

Notice how in the first case, the right plan expression has a float literal, where it should be an int literal. In the second case, the right plan expression hasn't managed to determine the type of the literal (it's still a `dyn int`).

I did some printf-debugging:

```diff
diff --git a/crates/polars-plan/src/plans/conversion/type_coercion/binary.rs b/crates/polars-plan/src/plans/conversion/type_coercion/binary.rs
index 24f65b3465..9c95826416 100644
--- a/crates/polars-plan/src/plans/conversion/type_coercion/binary.rs
+++ b/crates/polars-plan/src/plans/conversion/type_coercion/binary.rs
@@ -124,6 +124,9 @@ pub(super) fn process_binary(
     node_right: Node,
 ) -> PolarsResult<Option<AExpr>> {
     let input_schema = get_schema(lp_arena, lp_node);
+    println!("{:?}", input_schema);
+    println!("{:?}", expr_arena.get(node_left));
+    println!("{:?}", expr_arena.get(node_right));
     let (left, type_left): (&AExpr, DataType) =
         unpack!(get_aexpr_and_type(expr_arena, node_left, &input_schema));
     let (right, type_right): (&AExpr, DataType) =
```

The output from the two queries is:

```
Schema:
name: a, field: Float64

Column("a")
Literal(dyn int: 2)
Schema:
name: a, field: Float64

Column("a")
Cast { expr: Node(2), dtype: Float64, options: NonStrict }
```

and
```
Schema:
name: a, field: Float64

Column("b")
Literal(dyn int: 2)
```

So it looks like `run_coercion` for join nodes is always sending down the schema for the left table, but then I got a bit lost.

### Expected behavior

I expect the correct dtypes to be propagated to literals in both cases. I stumbled into this implementing conditional joins in the GPU engine, I trust the dtypes that polars has given me, but in these cases, I get mismatches and then libcudf doesn't like it. I can work around the bugs by doing my own type propagation, but it feels like this is an actual bug, even though it is probably quite benign.

### Installed versions

<details>

```
--------Version info---------
Polars:              1.12.0
Index type:          UInt32
Platform:            Linux-6.8.0-48-generic-x86_64-with-glibc2.35
Python:              3.12.7 | packaged by conda-forge | (main, Oct  4 2024, 16:05:46) [GCC 13.3.0]
LTS CPU:             False

----Optional dependencies----
adbc_driver_manager  1.2.0
altair               5.4.1
cloudpickle          3.1.0
connectorx           0.3.3
deltalake            0.21.0
fastexcel            0.12.0
fsspec               2024.10.0
gevent               24.10.3
great_tables         0.13.0
matplotlib           3.9.2
nest_asyncio         1.6.0
numpy                2.0.2
openpyxl             3.1.5
pandas               2.2.3
pyarrow              17.0.0
pydantic             2.9.2
pyiceberg            <not installed>
sqlalchemy           2.0.36
torch                2.4.1.post302
xlsx2csv             0.8.3
xlsxwriter           3.2.0
```

</details>



Discussion / hints:
Ah, I think this is generic to joins (not just join_where)...


Fix patch:
diff --git a/crates/polars-plan/src/plans/conversion/join.rs b/crates/polars-plan/src/plans/conversion/join.rs
index 5222ca73cda1..8a10db7de177 100644
--- a/crates/polars-plan/src/plans/conversion/join.rs
+++ b/crates/polars-plan/src/plans/conversion/join.rs
@@ -109,9 +109,31 @@ pub fn resolve_join(
 
     ctxt.conversion_optimizer
         .fill_scratch(&left_on, ctxt.expr_arena);
+    ctxt.conversion_optimizer
+        .coerce_types(ctxt.expr_arena, ctxt.lp_arena, input_left)
+        .map_err(|e| e.context("'join' failed".into()))?;
     ctxt.conversion_optimizer
         .fill_scratch(&right_on, ctxt.expr_arena);
-
+    ctxt.conversion_optimizer
+        .coerce_types(ctxt.expr_arena, ctxt.lp_arena, input_right)
+        .map_err(|e| e.context("'join' failed".into()))?;
+
+    let get_dtype = |expr: &ExprIR, schema: &SchemaRef| {
+        ctxt.expr_arena
+            .get(expr.node())
+            .get_type(schema, Context::Default, ctxt.expr_arena)
+            .ok()
+            .unwrap()
+    };
+    for (lnode, rnode) in left_on.iter().zip(right_on.iter()) {
+        let ltype = get_dtype(lnode, &schema_left);
+        let rtype = get_dtype(rnode, &schema_right);
+        polars_ensure!(
+            ltype == rtype,
+            SchemaMismatch: "datatypes of join keys don't match - `{}`: {} on left does not match `{}`: {} on right",
+            lnode.output_name(), ltype, rnode.output_name(), rtype
+        )
+    }
     // Every expression must be elementwise so that we are
     // guaranteed the keys for a join are all the same length.
     let all_elementwise =
@@ -128,7 +150,7 @@ pub fn resolve_join(
         right_on,
         options,
     };
-    run_conversion(lp, ctxt, "join")
+    Ok(ctxt.lp_arena.add(lp))
 }
 
 #[cfg(feature = "iejoin")]


Test patch:
diff --git a/py-polars/tests/unit/operations/test_join.py b/py-polars/tests/unit/operations/test_join.py
index 87ad4dfc5a9b..c65be5ad61c0 100644
--- a/py-polars/tests/unit/operations/test_join.py
+++ b/py-polars/tests/unit/operations/test_join.py
@@ -14,6 +14,7 @@
     ComputeError,
     DuplicateError,
     InvalidOperationError,
+    SchemaError,
 )
 from polars.testing import assert_frame_equal, assert_series_equal
 
@@ -1099,3 +1100,16 @@ def test_left_join_slice_pushdown_19405(set_sorted: bool) -> None:
 
     q = left.join(right, on="k", how="left").head(5)
     assert_frame_equal(q.collect(), pl.DataFrame({"k": [1, 1, 1, 1, 2]}))
+
+
+def test_join_key_type_coercion_19597() -> None:
+    left = pl.LazyFrame({"a": pl.Series([1, 2, 3], dtype=pl.Float64)})
+    right = pl.LazyFrame({"a": pl.Series([1, 2, 3], dtype=pl.Int64)})
+
+    with pytest.raises(SchemaError, match="datatypes of join keys don't match"):
+        left.join(right, left_on=pl.col("a"), right_on=pl.col("a")).collect_schema()
+
+    with pytest.raises(SchemaError, match="datatypes of join keys don't match"):
+        left.join(
+            right, left_on=pl.col("a") * 2, right_on=pl.col("a") * 2
+        ).collect_schema()
diff --git a/py-polars/tests/unit/operations/test_join_asof.py b/py-polars/tests/unit/operations/test_join_asof.py
index 21b7db4ebb90..4ab2a9fe8962 100644
--- a/py-polars/tests/unit/operations/test_join_asof.py
+++ b/py-polars/tests/unit/operations/test_join_asof.py
@@ -141,7 +141,7 @@ def test_join_asof_mismatched_dtypes() -> None:
     )
 
     with pytest.raises(
-        pl.exceptions.ComputeError, match="datatypes of join keys don't match"
+        pl.exceptions.SchemaError, match="datatypes of join keys don't match"
     ):
         df1.join_asof(df2, on="a", strategy="forward")