pola-rs / polars

Dataframes powered by a multithreaded, vectorized query engine, written in Rust
https://docs.pola.rs
Other
29.19k stars 1.84k forks source link

JoinArgs.suffix doesn't work when using LazyFrame.join #13424

Open keithyin opened 8 months ago

keithyin commented 8 months ago

Checks

Reproducible example

let df = df!(
    "keys" => &["a","b", "c"],
    "values" => &[10, 7, 1],
).unwrap()
.lazy();

let df2 = df!(
    "keys" => &["a","b", "c"],
    "values" => &[10, 7, 1],
).unwrap()
.lazy();

let result = df.join(df2, [col("keys")], [col("keys")], 
    JoinArgs { 
        how: JoinType::Inner, 
        validation: JoinValidation::default(),
        suffix: Some("_2".to_string()),
        slice: None, join_nulls: false});

println!("{:?}", result.collect().unwrap().head(Some(3)));

Log output

No response

Issue description

the result dataframe is

┌──────┬────────┬──────────────┐ │ keys ┆ values ┆ values_right │ │ --- ┆ --- ┆ --- │ │ str ┆ i32 ┆ i32 │ ╞══════╪════════╪══════════════╡ │ a ┆ 10 ┆ 10 │ │ b ┆ 7 ┆ 7 │ │ c ┆ 1 ┆ 1 │ └──────┴────────┴──────────────┘

Expected behavior

the expected output should be

┌──────┬────────┬──────────────┐ │ keys ┆ values ┆ values_2 │ │ --- ┆ --- ┆ --- │ │ str ┆ i32 ┆ i32 │ ╞══════╪════════╪══════════════╡ │ a ┆ 10 ┆ 10 │ │ b ┆ 7 ┆ 7 │ │ c ┆ 1 ┆ 1 │ └──────┴────────┴──────────────┘ i found that the code in https://github.com/pola-rs/polars/blob/rs-0.36.2/crates/polars-lazy/src/frame/mod.rs#L1231 just ignore the suffix of join_args. I made some changes shown below, it gives me the expected result. is it a valid solution?

self.join_builder()
          .with(other)
          .left_on(left_on)
          .right_on(right_on)
          .how(args.how)
          .finish()

to

self.join_builder()
          .with(other)
          .left_on(left_on)
          .right_on(right_on)
          .suffix(args.suffix()) // new line
          .how(args.how)
          .finish()

Installed versions

polars = {version = "0.36.2", features = ["lazy", "rank", "strings", "lazy_regex", "concat_str", "dtype-struct"]}

ritchie46 commented 8 months ago

Ai.. Can you make A PR?