diff --git a/crates/polars-sql/src/context.rs b/crates/polars-sql/src/context.rs index ec8f892720a01..4830af98d4f4e 100644 --- a/crates/polars-sql/src/context.rs +++ b/crates/polars-sql/src/context.rs @@ -336,7 +336,7 @@ impl SQLContext { let (l_name, mut lf) = self.get_table(&tbl_expr.relation)?; if !tbl_expr.joins.is_empty() { for tbl in &tbl_expr.joins { - let (r_name, rf) = self.get_table(&tbl.relation)?; + let (r_name, mut rf) = self.get_table(&tbl.relation)?; let left_schema = lf.schema()?; let right_schema = rf.schema()?; @@ -548,9 +548,9 @@ impl SQLContext { lf = lf.with_columns(projections); lf = self.process_order_by(lf, &query.order_by)?; - column_names.retain(|&name| !retained_names.contains(name)); lf.drop(column_names) + } else if contains_wildcard_exclude { let mut dropped_names = Vec::with_capacity(projections.len()); let exclude_expr = projections.iter().find(|expr| { diff --git a/crates/polars-sql/tests/statements.rs b/crates/polars-sql/tests/statements.rs index a5bed3cf3635a..42a19b32256f4 100644 --- a/crates/polars-sql/tests/statements.rs +++ b/crates/polars-sql/tests/statements.rs @@ -265,20 +265,21 @@ fn test_compound_join_different_column_names() { .unwrap(); let mut ctx = SQLContext::new(); - ctx.register("df1", df1.lazy()); - ctx.register("df2", df2.lazy()); + ctx.register("lf1", df1.lazy()); + ctx.register("lf2", df2.lazy()); let sql = r#" - SELECT df1.a, df2.b, df2.c - FROM df1 INNER JOIN df2 - ON df1.a = df2.b AND df1.a = df2.c + SELECT lf1.a, lf2.b, lf2.c + FROM lf1 INNER JOIN lf2 + -- note: uses "lf1.a" for *both* constraint arms + ON lf1.a = lf2.b AND lf1.a = lf2.c ORDER BY a "#; let actual = ctx.execute(sql).unwrap().collect().unwrap(); let expected = df! { - "a" => [2, 3], - "b" => [2, 3], - "c" => [8, 9], + "a" => [3, 5], + "b" => [3, 5], + "c" => [3, 5], } .unwrap();